Finite-volume operators for JAX on Arakawa C-grids
Documentation • Installation • Quick Start • Examples
finitevolX provides staggered finite-volume building blocks for ocean and atmosphere modelling in JAX. Every operator is a pure function or stateless equinox Module — fully compatible with jax.jit, jax.vmap, and jax.grad.
import finitevolx as fvx
# Build an Arakawa C-grid (64x64 interior + 2-cell ghost ring)
grid = fvx.ArakawaCGrid2D.from_interior(64, 64, Lx=5.12e6, Ly=5.12e6)
# Operators are stateless modules
diff = fvx.Difference2D(grid)
div = fvx.Divergence2D(grid)
vort = fvx.Vorticity2D(grid)
# Compute divergence at T-points from staggered velocities
div_uv = div(u, v)
# Solve the Helmholtz equation (∇² − λ)ψ = q for streamfunction
psi = fvx.solve_helmholtz_dst(q, grid.dx, grid.dy, lambda_=1/Ld**2)
# Time-step with Heun's method (RK2)
state_next = fvx.heun_step(tendency_fn, state, dt)| Component | Description |
|---|---|
ArakawaCGrid1D/2D/3D |
Staggered C-grid containers with T, U, V, X point locations |
Mask1D/2D/3D |
Land/ocean masks with automatic staggered derivation (h, u, v, w, xy_corner, xy_corner_strict) |
| Boundary classification | 4-level land/coast/near-coast/ocean + vorticity boundary categories |
| Stencil capability | Adaptive WENO stencil dispatch at irregular coastlines |
SphericalArakawaCGrid2D/3D |
Spherical coordinate grids |
| Operator | Class | Description |
|---|---|---|
| Difference | Difference1D/2D/3D |
Forward/backward differences, Laplacian, curl, gradient |
| Interpolation | Interpolation1D/2D/3D |
Staggered averaging (T↔U, T↔V, X↔T, etc.) |
| Divergence | Divergence2D/3D |
Backward-difference divergence at T-points |
| Vorticity | Vorticity2D/3D |
Relative vorticity, potential vorticity, PV flux |
| Coriolis | Coriolis2D/3D |
Beta-plane Coriolis tendencies |
| Diffusion | Diffusion2D/3D |
Harmonic and biharmonic diffusion |
| Advection | Advection1D/2D/3D |
Flux-form with upwind, TVD (minmod/van Leer/superbee/MC), WENO3/5/7/9 |
| Jacobian | arakawa_jacobian |
Energy-conserving Arakawa Jacobian |
| Diagnostics | kinetic_energy, enstrophy, okubo_weiss, ... |
Scalar diagnostics |
| Spherical | SphericalDivergence2D, SphericalLaplacian2D, ... |
Operators on the sphere |
Per-face composable BCs with ghost-cell enforcement:
Periodic • Dirichlet • Neumann • Robin • Slip • Sponge • Reflective • Extrapolation • Outflow
| Solver | Use case | Accuracy |
|---|---|---|
| Spectral (DST/DCT/FFT) | Rectangular, constant coefficient | Machine precision |
| Capacitance matrix | Masked domains, simple coastlines | Machine precision |
| Preconditioned CG | Arbitrary masks, tight tolerance | Controllable |
| Multigrid | Variable coefficients, any mask | O(N) per V-cycle |
| MG + CG | Variable coefficients, tight tolerance | Controllable |
Convenience wrappers: streamfunction_from_vorticity, pressure_from_divergence, pv_inversion
Functional steppers (pure functions, no hidden state):
euler_step • heun_step • rk4_step • rk3_ssp_step • ab2_step • ab3_step • leapfrog_raf_step • imex_ssp2_step • split_explicit_step • semi_lagrangian_step
diffrax integration: ForwardEulerDfx, RK2Heun, RK3SSP, RK4Classic, SSP_RK104, IMEX_SSP2, and more — with adaptive stepping, checkpointing, and SaveAt.
multilayer vmap helper • decompose_vertical_modes • layer_to_mode / mode_to_layer transforms • build_coupling_matrix
pip install git+https://github.com/jejjohnson/finitevolXgit clone https://github.com/jejjohnson/finitevolX.git
cd finitevolX
uv sync --all-extrasBuild a grid, compute vorticity, and invert for the streamfunction:
import jax
import jax.numpy as jnp
import finitevolx as fvx
jax.config.update("jax_enable_x64", True)
# 1. Grid
grid = fvx.ArakawaCGrid2D.from_interior(64, 64, Lx=1e6, Ly=1e6)
# 2. Operators
diff = fvx.Difference2D(grid)
vort = fvx.Vorticity2D(grid)
# 3. Velocity field (full grid with ghost ring)
Ny, Nx = grid.Ny, grid.Nx
j, i = jnp.mgrid[:Ny, :Nx]
u = -jnp.sin(jnp.pi * j / Ny) * jnp.cos(jnp.pi * i / Nx)
v = jnp.cos(jnp.pi * j / Ny) * jnp.sin(jnp.pi * i / Nx)
# 4. Relative vorticity at X-points (corners)
zeta = vort.relative_vorticity(u, v)
# 5. Invert for streamfunction: ∇²ψ = ζ
psi = fvx.streamfunction_from_vorticity(
zeta[1:-1, 1:-1], grid.dx, grid.dy
)See the documentation for the full API reference.
The docs/notebooks/ directory contains pedagogical Jupytext notebooks that build models step by step with equations, ASCII diagrams, and inline figures. The scripts/ directory has production simulation scripts that generate Zarr output and animated GIFs.
| Notebook | Description | Key APIs |
|---|---|---|
| Masks | C-grid mask construction, staggered derivation, boundary classification | Mask2D |
| Elliptic Solvers | Spectral, capacitance, CG, multigrid on 4 geometries + inhomogeneous BCs | solve_helmholtz_dst, build_capacitance_solver, solve_cg, build_multigrid_solver |
| Pressure Poisson | Divergence-free projection on the C-grid, DST-I vs DST-II | Divergence2D, Difference2D, solve_poisson_dst |
| Streamfunction Inversion | X-point vs T-point placement, velocity recovery, convergence | streamfunction_from_vorticity, solve_poisson_dst2 |
| Helmholtz Screening | QG PV inversion with screening, JIT/vmap/grad | StaggeredDirichletHelmholtzSolver2D, pv_inversion |
| Notebook | Model | Key APIs |
|---|---|---|
| Linear Shallow Water | Wind-driven double gyre on beta-plane | Difference2D, Interpolation2D, Vorticity2D, heun_step |
| Nonlinear Shallow Water | Full depth continuity + momentum advection | Advection2D, Difference2D, heun_step |
| 1.5-Layer QG | PV advection + Helmholtz inversion | Advection2D, solve_helmholtz_dst, heun_step |
Run the full-resolution simulations with spin-up and Zarr/GIF output:
uv run python scripts/swm_linear.py # Linear shallow-water
uv run python scripts/shallow_water.py # Nonlinear shallow-water
uv run python scripts/qg_1p5_layer.py # 1.5-layer QG| Linear SWM | Nonlinear SWM | 1.5-Layer QG |
|---|---|---|
Full documentation with theory, usage guides, and API reference:
jejjohnson.github.io/finitevolX
| Section | Content |
|---|---|
| C-Grid Discretization | Theory of Arakawa C-grid staggering |
| Operators | Divergence, vorticity, Coriolis, diffusion |
| Advection | TVD and WENO reconstruction theory |
| Boundary Conditions | Per-face BC composition |
| Elliptic Solvers | Spectral, capacitance, CG, multigrid theory |
| Time Integration | Explicit, IMEX, and diffrax-based steppers |
| Solver Comparison | Visual benchmark across geometries |
| API Reference | Auto-generated from docstrings |
Software
- spectraldiffx — Pseudospectral solvers in JAX (the spectral backend for finitevolX's elliptic solvers)
- diffrax — JAX-native ODE/SDE solvers (time integration backend)
- equinox — JAX neural network library (Module system used by all operators)
- PyFVTool — Finite Volume Tool in Python
Algorithms
- Thiry et al, 2023 | MQGeometry — WENO reconstructions for multilayer QG, Arakawa grid masks
- Roullet & Gaillard, 2021 | pyRSW — WENO reconstructions for shallow-water equations
- Gottlieb, Shu & Tadmor, 2001 — Strong Stability-Preserving High-Order Time Discretization Methods
- Ketcheson, 2008 — Highly efficient SSP methods: SSP-RK(10,4)
MIT © J. Emmanuel Johnson