Arakawa C-Grid Masks

finitevolX provides three mask classes — Mask1D, Mask2D, Mask3D — that describe land/ocean geometry on an Arakawa C-grid. Every class operator in the package takes an optional mask class attribute of the matching dimension, so masks travel with operators rather than getting threaded through every call. See the Operator API section below for the full pattern and the stagger → mask field mapping.

The rest of this page documents Mask2D — the 2-D class — but the same factories (from_mask, from_dimensions, from_center, from_u_face …) and the same staggered-derivation pipeline are available on Mask1D (no v-face, no corner) and Mask3D (adds a w vertical face and native 3-D classification / stencil capability).

The Arakawa C-grid layout

Mask2D builds all staggered masks from a single cell-centre wet/dry field, following the Arakawa & Lamb (1977) grid layout:

y
^
:           :
w-----v-----w..
|           |
|           |
u     h     u
|           |
|           |
w-----v-----w..   > x

Point	Location	Variable
h	cell centre	tracers, height, pressure
u	east face	zonal velocity
v	north face	meridional velocity
xy_corner	NE corner (lenient)	vorticity
xy_corner_strict	NE corner (strict)	streamfunction

All staggered points use the same-index, positive half-step convention: u[j, i] is the east face of h[j, i], v[j, i] is the north face, and xy_corner[j, i] is the NE corner — matching CartesianGrid2D. See the Arakawa C-grid discretization notes for the underlying convention and a comparison with negative half-step alternatives.

Creating masks

All you need is a binary h-grid mask (True = ocean, False = land). The factory method derives everything else:

import numpy as np
from finitevolx import Mask2D

# Rectangular basin with land boundaries
n = 10
h_mask = np.ones((n, n), dtype=bool)
h_mask[0, :] = h_mask[-1, :] = False
h_mask[:, 0] = h_mask[:, -1] = False

# With an island
h_mask[4:7, 4:7] = False

masks = Mask2D.from_mask(h_mask)

For an all-ocean domain, use the shortcut:

masks = Mask2D.from_dimensions(ny=12, nx=12)

Or construct from any float field at a known grid position, where NaN marks dry cells:

# From a field at cell centres (T-points): SSH, temperature, pressure, …
masks = Mask2D.from_center(ssh_field)

# From a field at u-faces, v-faces, or xy-corners.  Pass `mode=` to
# choose the inversion strategy when the inverse mapping back to the
# h-grid is non-unique:
#   mode='permissive'  → h wet iff *any* surrounding face/corner wet
#   mode='conservative' → h wet iff *all* surrounding face/corner wet
masks = Mask2D.from_u_face(u_field, mode="permissive")
masks = Mask2D.from_v_face(v_field, mode="permissive")
masks = Mask2D.from_corner(vorticity_field, mode="conservative")

Staggered variable locations

Each variable type sits at a different position within a grid cell. The figures below show the actual staggered positions for several domain topologies:

Rectangular basin

Basin with island

Zonal channel

Irregular coastline

Land / coast classification

The mask includes a 4-level classification (0 = land, 1 = coast, 2 = near-coast, 3 = open ocean):

Rectangular basin

Basin with island

Zonal channel

Irregular coastline

Vorticity boundary classification

At xy-corner points, cells are classified based on their relationship to adjacent velocity faces:

• xy_corner_valid — interior: all 4 adjacent velocity faces are wet
• xy_corner_y_wall — on a vertical (y-direction) boundary
• xy_corner_x_wall — on a horizontal (x-direction) boundary
• xy_corner_convex — at convex corners (both boundary types)

Stencil capability and adaptive WENO

Each cell stores how many contiguous wet neighbours it has in each direction via StencilCapability2D. This drives adaptive stencil selection for WENO reconstruction near coastlines:

# Mutually-exclusive masks: largest usable stencil at each point
adaptive = masks.get_adaptive_masks(direction="x", source="h")
# adaptive[2]  → 1st-order upwind only
# adaptive[4]  → WENO3
# adaptive[6]  → WENO5
# adaptive[8]  → WENO7
# adaptive[10] → WENO9

Jupyter notebook

A complete interactive demo is available as a jupytext notebook: notebooks/demo_masks.py

---

Operator API: how operators consume masks

Every class operator in finitevolx takes an optional mask class attribute of the matching dimension. The mask travels with the operator; once constructed, method calls don't need to pass the mask through per call.

import finitevolx as fvx

grid = fvx.CartesianGrid2D.from_interior(64, 64, 1e6, 1e6)
mask = fvx.Mask2D.from_mask(h_mask)

# Mask as a class attribute — set once, consumed on every call.
diff = fvx.Difference2D(grid=grid, mask=mask)
adv  = fvx.Advection2D(grid=grid, mask=mask)
diff_op = fvx.Diffusion2D(grid=grid, mask=mask)

# Unmasked operators are the default — just omit `mask=`.
diff_noop = fvx.Difference2D(grid=grid)   # mask=None

The three-layer rule

finitevolx keeps masks out of the lowest two layers, so stencil and functional code stays coordinate-and-mask agnostic:

Layer	Example	Mask-aware?
Layer 1 — raw stencils	diff_x_fwd, avg_y_bwd, … in stencils.py	No
Layer 2 — functional helpers	_curl_2d, divergence_2d, diffusion_2d, kinetic_energy, enstrophy, _geostrophic_velocity_sphere (private), arakawa_jacobian, …	No
Layer 3 — class operators	Difference2D, Interpolation2D, Divergence2D, Vorticity2D, Coriolis2D, Advection2D, Diffusion2D, MomentumAdvection2D, Energetics2D, Strain2D, QGPotentialVorticity2D, ArakawaJacobian2D, and their 1D / 3D / Spherical* siblings	Yes, via self.mask

The pointwise and stencil diagnostics below have Layer-3 class forms, so in a masked domain you don't need to mask them by hand:

Layer-2 helper	Layer-3 class form
kinetic_energy, bernoulli_potential, available_potential_energy	Energetics2D(grid, mask) methods of the same name
shear_strain, tensor_strain	Strain2D.shear, Strain2D.tensor
strain_magnitude_squared, okubo_weiss	Strain2D.magnitude_squared(u, v), Strain2D.okubo_weiss(u, v)
relative_vorticity_cgrid, sw_potential_vorticity	Vorticity2D.relative_vorticity, Vorticity2D.potential_vorticity
enstrophy, potential_enstrophy	Vorticity2D.enstrophy(u, v), Vorticity2D.potential_enstrophy(u, v, h, f)
qg_potential_vorticity, stretching_term, potential_vorticity_multilayer	QGPotentialVorticity2D.__call__, .stretching, .multilayer
arakawa_jacobian	ArakawaJacobian2D(grid, mask) (full-shape output, zero ghost ring)
_geostrophic_velocity_sphere (private)	SphericalDifference2D.geostrophic_velocity(h, f, gravity)

energetics = fvx.Energetics2D(grid=grid, mask=mask)
ke = energetics.kinetic_energy(u, v)       # dry T-cells exactly 0

strain = fvx.Strain2D(grid=grid, mask=mask)
ow = strain.okubo_weiss(u, v)              # T-points, dry T-cells exactly 0

Energetics2D, Strain2D, QGPotentialVorticity2D and ArakawaJacobian2D zero every field input on the dry interior cells of its stagger before any stencil reads it, and mask outputs with jnp.where, so interior land values stored as NaN (e.g. from Mask2D.from_center) neither leak into wet cells nor survive at dry ones. The ghost ring is BC-owned: it is passed through unchanged and must be finite.

SphericalDifference2D.geostrophic_velocity only guards the division (it sets the face-averaged f to 1 on dry faces); its 4-point stencils read h on neighbouring land cells for wet coastal faces, so land h must be finite there. The same holds for the other existing class operators, which mask outputs but do not sanitize inputs.

Diagnostics without a class form:

• total_energy, total_enstrophy are domain reductions, not stencils. Pass them fields that are already masked (e.g. from the class forms above) so dry cells contribute zero.
• vertical_velocity takes a per-call mask=Mask3D: it zeroes the divergence on dry T-points (mask.h) and its output on dry W-faces (mask.w); the output has the same [Nz, Ny, Nx] shape as mask.w.

If you do call a Layer-2 helper directly and want masked output, apply the mask at the call site: ke = jnp.where(mask.h, kinetic_energy(u, v), 0.0) (a plain * mask.h also works when the inputs are finite on land). For pointwise helpers whose inputs live on different staggers (e.g. okubo_weiss(sn, ss, omega) with sn at T-points and ss, omega at X-points), the class form also does the X → T averaging for you.

Dimension → mask type mapping

Operators take the mask class matching their dimension:

Operator class	Mask type	Example
*1D (Difference1D, Interpolation1D, Advection1D)	Mask1D \| None	Difference1D(grid=g, mask=Mask1D.from_mask(h))
*2D (every 2-D operator)	Mask2D \| None	Advection2D(grid=g, mask=Mask2D.from_mask(h))
*3D (every 3-D operator)	Mask3D \| None	Diffusion3D(grid=g, mask=Mask3D.from_mask(h))

Spherical 2-D / 3-D operators take the Cartesian Mask2D / Mask3D — the mask geometry is coordinate-agnostic at the post-compute-multiply layer. A dedicated SphericalMask2D / SphericalMask3D is a TODO for future work.

Stagger → mask field mapping

The core convention is one operator output → one mask field, chosen by the operator's output stagger. Every operator follows the same mapping:

Output stagger	Mask field	Operators / methods
T (cell centre)	mask.h	divergence, laplacian, U→T / V→T / X→T differences and interpolations, Advection*, Diffusion*, BiharmonicDiffusion*, Energetics2D.*, Strain2D.tensor / .magnitude_squared / .okubo_weiss, QGPotentialVorticity2D.*, ArakawaJacobian2D
U (east face)	mask.u	T→U / X→U / V→U differences and interpolations, grad_perp u-component, Coriolis* du_cor, MomentumAdvection* du_adv, SphericalDifference2D.geostrophic_velocity u_g
V (north face)	mask.v	T→V / X→V / U→V, grad_perp v-component, Coriolis* dv_cor, MomentumAdvection* dv_adv, SphericalDifference2D.geostrophic_velocity v_g
X (NE corner)	mask.xy_corner_strict	curl, T→X / U→X / V→X, relative & potential vorticity, Vorticity2D.enstrophy / .potential_enstrophy, Strain2D.shear, pv_flux_* U/V-stagger outputs where applicable

The X-output mapping uses the strict 4-of-4 corner mask (xy_corner_strict) because a corner-output value is trusted only when all four surrounding T-cells are wet. The lenient xy_corner mask is not currently consumed by any operator — it's available for users building their own.

Four patterns for mask application

Under the hood, operators apply masks in one of four patterns depending on what's safe. You don't need to know which is which for normal usage — the result is always "dry cells at the output stagger are exactly 0" — but the distinctions matter when customising or inheriting:

Pattern 1 — Post-compute multiply. Used by the majority of operators: Difference*, Interpolation*, Vorticity3D, Spherical*, Energetics2D, QGPotentialVorticity2D, ArakawaJacobian2D. Compute the output as if all-ocean, then multiply by the stagger-matched mask field at the end. Cheap and simple. The #206 diagnostic classes (Energetics2D, Strain2D, QGPotentialVorticity2D, ArakawaJacobian2D) use the NaN-safe variant: zero the inputs on dry cells first, then select with jnp.where(mask, out, 0) instead of multiplying.

Pattern 2 — Pass-down into sub-operators. Used by composed operators with no tricky division — Divergence2D, Vorticity2D, MomentumAdvection2D, Strain2D (its T-point combinations mask the X-point shear and vorticity before averaging them X → T, so a coastal T-cell averages in zeros from its dry corners). The outer class wires its mask into every internal sub-operator at __init__ time, so intermediate staggered fields are already zero at dry cells. The outer method trusts the pass-down and typically doesn't need an extra multiply.

• Special case for Vorticity2D.potential_vorticity: the (zeta + f) / h division would produce NaN at every dry X-corner under pass-down (because h_on_q == 0). The method preserves the existing "NaN at wet zero-thickness" sentinel but replaces mask-induced NaNs with exact 0 via a final jnp.where(mask.xy_corner_strict, out, 0). Same fix applies to SphericalVorticity2D.potential_vorticity and Vorticity2D.potential_enstrophy (built on it). SphericalDifference2D.geostrophic_velocity faces the same division (a dry face can have f_on_face = 0): it sets the face-averaged f to 1 on every dry face right before dividing, so the value and its reverse-mode gradient stay finite, then masks with jnp.where. Vorticity2D.potential_vorticity likewise divides by 1 at dry corners, where pass-down masking makes h_on_q exactly 0.

Pattern 3 — Intermediate flux masking. Used by Diffusion2D and Diffusion3D only. Post-compute multiply is not sufficient for diffusion because the divergence at wet T-cells adjacent to land would be contaminated by polluted dry-side face fluxes. Instead the class wrappers apply the mask in three steps:

flux_x *= mask.u        # U-face stage
flux_y *= mask.v        # V-face stage
tendency *= mask.h      # final output

The shared kernel is _diffusion_2d_impl(h, kappa, dx, dy, mh, mu, mv) in finitevolx._src.diffusion.diffusion; Diffusion3D vmaps it per-z-slice with sliced masks (Pattern B in #209's terminology).

Pattern 4 — The masked harmonic operator applied twice. Used by BiharmonicDiffusion2D / 3D and SphericalBiharmonicDiffusion2D / 3D. The tendency is −κ · Dm(Dm(h)), where Dm is the masked no-flux harmonic operator of Pattern 3 (the inner Diffusion* / SphericalDiffusion* is built with the outer mask). Dry inputs are zeroed with jnp.where first, so land may hold NaN or a fill value. Because the second pass also zeroes the wet→dry fluxes, the dry zeros of the intermediate Dm(h) are never read: the coast is a no-flux wall for both h and ∇²h, exactly like the closed domain walls. The result is conservative (Σ_wet ∂h/∂t = 0) and symmetric on the wet cells, and no land value reaches a wet cell. (Before #292 the inner Laplacian was unmasked and post-multiplied by mask.h, which read land values and was not conservative.)

Pattern A vs Pattern B for 3-D operators

The 3-D wrappers that vmap an inner 2-D operator over z-levels (Vorticity3D, Coriolis3D, Spherical*3D, MomentumAdvection3D) use Pattern A (post-compute): the inner 2-D op is always constructed mask=None and the 3-D wrapper owns the mask, applying * self.mask.<stagger> after the vmap and zero_z_ghosts. This is cleaner than trying to thread per-z mask slices into the vmap'd 2-D op.

Diffusion3D (and BiharmonicDiffusion3D, built on it) is the exception and uses Pattern B (vmap with per-z sliced masks) because its intermediate flux masking makes Pattern A incorrect. See the Diffusion module docstring for the implementation.

Advection hierarchy pre-build

Advection1D / Advection2D / Advection3D are the special case where "mask as field" matters for performance, not just ergonomics. When constructed with a mask, __init__ pre-builds the (2, 4, 6, 8, 10) adaptive stencil hierarchy via mask.get_adaptive_masks(...) — once, up-front — and stores the result on the operator. Every subsequent call then uses the cached hierarchy, narrowed at dispatch time (by the narrow_mask_hierarchy helper) to the subset of tiers the chosen reconstruction method needs:

adv = fvx.Advection2D(grid=grid, mask=mask)  # builds (2,4,6,8,10) hierarchies

# First call uses the cached hierarchies; no get_adaptive_masks recomputation.
dh_dt = adv(h, u, v, method="weno5")

# A second method call on the same instance reuses the cache and
# narrows to (2, 4) on the fly for WENO3.
dh_dt = adv(h, u, v, method="weno3")

This is the main reason the Advection API moved from a per-call mask= kwarg to a class field: masks under JIT now compile once per Advection* instance instead of once per (instance, method) pair.

Changing the mask on an operator

equinox Modules are immutable — op.mask = new_mask will raise. To "change" a mask on an operator, use eqx.tree_at to produce a new operator instance with the field updated:

import equinox as eqx

# Start with one mask ...
adv_summer = fvx.Advection2D(grid=grid, mask=summer_mask)

# ... and produce a new instance with a different mask.  The grid and
# reconstruction module are shared (no copy); only `mask` (and, for
# Advection, the pre-built hierarchies — eqx.tree_at reconstructs the
# __init__ path correctly for all affected fields) is swapped.
adv_winter = eqx.tree_at(lambda o: o.mask, adv_summer, winter_mask)

For composed operators (Divergence2D, Vorticity2D, MomentumAdvection2D, etc.) the outer self.mask is independent of the mask stored on the internal sub-operators, so you may also need to swap those — the simpler alternative is usually to just construct a fresh operator from scratch:

# Cleaner for composed ops
vort_winter = fvx.Vorticity2D(grid=grid, mask=winter_mask)

Migration from the old per-call kwarg API

Some operators used to take mask= as a per-call keyword argument. In the current API, the mask is always a class field. Translation:

# Old (removed)
adv = fvx.Advection2D(grid=grid)
adv(h, u, v, method="weno5", mask=mask)

cor = fvx.Coriolis2D(grid=grid)
cor(u, v, f, mask=mask)

# New
adv = fvx.Advection2D(grid=grid, mask=mask)
adv(h, u, v, method="weno5")

cor = fvx.Coriolis2D(grid=grid, mask=mask)
cor(u, v, f)

Diffusion2D, Diffusion3D, diffusion_2d, Difference2D.grad_perp, kinetic_energy, enstrophy, and potential_enstrophy used to take raw-array mask_h / mask_u / mask_v keyword arguments. Those are all removed:

• Diffusion2D / Diffusion3D take a Mask2D / Mask3D class field instead.
• diffusion_2d (the Layer-2 free function) is mask-free. If you want masked diffusion, use the class operator.
• Difference2D.grad_perp uses the self.mask class field (set on Difference2D at construction).
• kinetic_energy / enstrophy / potential_enstrophy are Layer-2 mask-free. Apply masks at the call site:

ke = fvx.kinetic_energy(u, v) * mask.h
z  = fvx.enstrophy(omega)     * mask.xy_corner_strict
pe = fvx.potential_enstrophy(q, h) * mask.xy_corner_strict