Skip to content

pipekit Operators

finitevolx.pipekit wraps the 2-D Cartesian operators as pipekit operators: JSON configuration, Sequential / Graph composition, and Equinox pytree leaves so jit, grad and vmap see through them. Each operator is one Tier-A call on [*batch, Ny, Nx] fields. Requires the pipekit extra: pip install "finitevolx[pipekit]". Walkthrough: Composing finitevolX with pipekit.

Tier-B operator Tier-A counterpart in → out
Difference(axis, src, dst) Difference2D.diff_{axis}_{src}_to_{dst} 1 → 1
Divergence Difference2D.divergence (U, V) → T
Curl Difference2D.curl (U, V) → X
Laplacian Difference2D.laplacian T → T
GradPerp Difference2D.grad_perp X → (U, V)
Interp(src, dst) Interpolation2D.{src}_to_{dst} 1 → 1
Coriolis(f0, beta, y0) coriolis_fn + Coriolis2D (U, V) → (U, V)
FillGhosts(kind) enforce_periodic, zero_boundaries, …, wall_boundaries 1 → 1
ApplyMask(mask, stagger) jnp.where(mask.{h,u,v,xy_corner_strict}, x, 0) 1 → 1
RelativeVorticity Vorticity2D.relative_vorticity (U, V) → X
PotentialVorticity Vorticity2D.potential_vorticity (U, V, T, T) → X
Enstrophy Vorticity2D.enstrophy (U, V) → X
QGPotentialVorticity(f0, beta, y0) QGPotentialVorticity2D T → T
KineticEnergy Energetics2D.kinetic_energy (U, V) → T
AvailablePE(g_prime, H) Energetics2D.available_potential_energy T → T
Strain(component) Strain2D.shear / .tensor / .magnitude_squared (U, V) → X / T
OkuboWeiss Strain2D.okubo_weiss (U, V) → T
TotalEnergy total_energy (T, T) → [*batch]
DomainIntegral / DomainMean area_sum / area_mean T → [*batch]
ExplicitStep(rhs, dt, scheme) euler_step / heun_step / rk3_ssp_step / rk4_step state → state
Pick(index), LinComb(coeffs) — (graph plumbing)

Base classes and helpers

finitevolx.pipekit.FVXOperator

Bases: JaxOperator

Base for finitevolX pipekit operators.

A pipekit_jax.JaxOperator (pipekit Operator + eqx.Module): config fields are eqx.field(static=True), learnable coefficients are leaves with converter=jnp.asarray, and _apply maps fields [*batch, Ny, Nx] to fields.

Source code in finitevolx/_src/pipekit/base.py 
137
138
139
140
141
142
143
144
class FVXOperator(JaxOperator):
    """Base for finitevolX pipekit operators.

    A `pipekit_jax.JaxOperator` (pipekit ``Operator`` + ``eqx.Module``):
    config fields are ``eqx.field(static=True)``, learnable coefficients
    are leaves with ``converter=jnp.asarray``, and ``_apply`` maps fields
    ``[*batch, Ny, Nx]`` to fields.
    """

finitevolx.pipekit.GridOperator

Bases: FVXOperator

An FVXOperator on a 2-D Cartesian C-grid.

Attributes:

Name Type Description
grid GridSpec

(nx_interior, ny_interior, Lx, Ly), static; a list (after a JSON round trip) or a CartesianGrid2D is converted on construction.
Source code in finitevolx/_src/pipekit/base.py 
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
class GridOperator(FVXOperator):
    """An `FVXOperator` on a 2-D Cartesian C-grid.

    Attributes
    ----------
    grid : GridSpec
        ``(nx_interior, ny_interior, Lx, Ly)``, static; a list (after a
        JSON round trip) or a `CartesianGrid2D` is converted on
        construction.
    """

    grid: GridSpec = eqx.field(static=True, converter=as_grid_spec)

    @property
    def cartesian_grid(self) -> CartesianGrid2D:
        """The Tier-A grid, rebuilt from `grid` (cached per spec)."""
        return build_grid(self.grid)

cartesian_grid property

The Tier-A grid, rebuilt from grid (cached per spec).

finitevolx.pipekit.as_grid_spec(spec)

Normalise a grid description to a GridSpec tuple.

Parameters:

Name Type Description Default
spec sequence of 4 numbers or CartesianGrid2D

(nx_interior, ny_interior, Lx, Ly) (a list after a JSON round trip), or a grid to describe.

 required

Returns:

Type Description
GridSpec

(int, int, float, float), hashable, so it can be a static field.

Raises:

Type Description
ValueError

If spec does not have four entries or a count is not positive.
Source code in finitevolx/_src/pipekit/base.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
def as_grid_spec(spec: Sequence[Any] | CartesianGrid2D) -> GridSpec:
    """Normalise a grid description to a `GridSpec` tuple.

    Parameters
    ----------
    spec : sequence of 4 numbers or CartesianGrid2D
        ``(nx_interior, ny_interior, Lx, Ly)`` (a list after a JSON round
        trip), or a grid to describe.

    Returns
    -------
    GridSpec
        ``(int, int, float, float)``, hashable, so it can be a static field.

    Raises
    ------
    ValueError
        If ``spec`` does not have four entries or a count is not positive.
    """
    if isinstance(spec, CartesianGrid2D):
        spec = (spec.Nx - 2, spec.Ny - 2, spec.Lx, spec.Ly)
    if len(spec) != 4:
        raise ValueError(
            f"grid spec must be (nx_interior, ny_interior, Lx, Ly), got {spec!r}"
        )
    nx, ny, lx, ly = spec
    if int(nx) < 1 or int(ny) < 1:
        raise ValueError(f"grid spec needs positive cell counts, got {spec!r}")
    return (int(nx), int(ny), float(lx), float(ly))

finitevolx.pipekit.build_grid(spec) cached

The CartesianGrid2D a GridSpec describes (cached per spec).

Parameters:

Name Type Description Default
spec GridSpec

(nx_interior, ny_interior, Lx, Ly).

 required

Returns:

Type Description
CartesianGrid2D

CartesianGrid2D.from_interior(*spec).
Source code in finitevolx/_src/pipekit/base.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
@functools.lru_cache(maxsize=64)
def build_grid(spec: GridSpec) -> CartesianGrid2D:
    """The `CartesianGrid2D` a `GridSpec` describes (cached per spec).

    Parameters
    ----------
    spec : GridSpec
        ``(nx_interior, ny_interior, Lx, Ly)``.

    Returns
    -------
    CartesianGrid2D
        ``CartesianGrid2D.from_interior(*spec)``.
    """
    nx, ny, lx, ly = as_grid_spec(spec)
    return CartesianGrid2D.from_interior(nx, ny, lx, ly)

finitevolx.pipekit.over_leading(fn, core_ndim=2)

Lift fn on [*core] arrays to [*batch, *core] arrays.

Every input shares the same leading shape; it is flattened to one axis, jax.vmap(fn) maps over it, and every output array gets the leading shape back::

x : [T, E, Ny, Nx] ─reshape─► [T·E, Ny, Nx] ─vmap fn─► [T·E, …] ─reshape─► [T, E, …]

Only needed for Tier-A functions that are still 2-D only; the 2-D Cartesian operators take [*batch, Ny, Nx] natively.

Parameters:

Name Type Description Default
fn callable

Function of positional arrays with core_ndim trailing axes each; it may return an array or a pytree of arrays.

 required
core_ndim int

Number of trailing axes fn consumes (2 for [Ny, Nx]).

 2

Returns:

Type Description
callable

fn lifted over any number of leading axes (zero included).

Raises:

Type Description
ValueError

If the inputs disagree on their leading shape.
Source code in finitevolx/_src/pipekit/base.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
def over_leading(fn: Callable[..., Any], core_ndim: int = 2) -> Callable[..., Any]:
    """Lift ``fn`` on ``[*core]`` arrays to ``[*batch, *core]`` arrays.

    Every input shares the same leading shape; it is flattened to one axis,
    ``jax.vmap(fn)`` maps over it, and every output array gets the leading
    shape back::

        x : [T, E, Ny, Nx] ─reshape─► [T·E, Ny, Nx] ─vmap fn─► [T·E, …] ─reshape─► [T, E, …]

    Only needed for Tier-A functions that are still 2-D only; the 2-D
    Cartesian operators take ``[*batch, Ny, Nx]`` natively.

    Parameters
    ----------
    fn : callable
        Function of positional arrays with ``core_ndim`` trailing axes each;
        it may return an array or a pytree of arrays.
    core_ndim : int
        Number of trailing axes ``fn`` consumes (2 for ``[Ny, Nx]``).

    Returns
    -------
    callable
        ``fn`` lifted over any number of leading axes (zero included).

    Raises
    ------
    ValueError
        If the inputs disagree on their leading shape.
    """

    @functools.wraps(fn)
    def lifted(*xs: Array) -> Any:
        leads = {tuple(x.shape[: x.ndim - core_ndim]) for x in xs}
        if len(leads) != 1:
            raise ValueError(f"inputs have different leading shapes: {sorted(leads)}")
        (lead,) = leads
        if not lead:
            return fn(*xs)
        flat = [x.reshape(-1, *x.shape[x.ndim - core_ndim :]) for x in xs]
        out = jax.vmap(fn)(*flat)
        return jax.tree.map(lambda y: y.reshape(*lead, *y.shape[1:]), out)

    return lifted

finitevolx.pipekit.testing.check_operator(op, *fields, hlo=True)

Assert that op is a well-behaved finitevolX pipekit operator.

Runs pipekit_jax.testing.check_jax_operator (operator + module, eager ≡ filter_jit, vmap, visible leaves, JSON config round trip), then checks:

  1. Leading axes: on inputs [*lead, Ny, Nx] with lead = (2,) and (2, 3), op equals a loop over the members.
  2. Gradient: ∂/∂x₀ ∑ op(x) is finite.
  3. Lowering (when hlo): the compiled HLO has no gather, scatter or convolution.

Parameters:

Name Type Description Default
op FVXOperator

Operator under test.

 required
*fields Array

Example inputs [Ny, Nx] (non-array inputs are passed through).

 ()
hlo bool

Also check the lowering.

 True

Raises:

Type Description
AssertionError

On the first violated property.
Source code in finitevolx/_src/pipekit/testing.py 
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
def check_operator(op: Any, *fields: Any, hlo: bool = True) -> None:
    """Assert that ``op`` is a well-behaved finitevolX pipekit operator.

    Runs `pipekit_jax.testing.check_jax_operator` (operator + module, eager
    ≡ ``filter_jit``, ``vmap``, visible leaves, JSON config round trip),
    then checks:

    1. **Leading axes**: on inputs ``[*lead, Ny, Nx]`` with ``lead`` = ``(2,)``
       and ``(2, 3)``, ``op`` equals a loop over the members.
    2. **Gradient**: ``∂/∂x₀ ∑ op(x)`` is finite.
    3. **Lowering** (when ``hlo``): the compiled HLO has no ``gather``,
       ``scatter`` or ``convolution``.

    Parameters
    ----------
    op : FVXOperator
        Operator under test.
    *fields : Array
        Example inputs ``[Ny, Nx]`` (non-array inputs are passed through).
    hlo : bool
        Also check the lowering.

    Raises
    ------
    AssertionError
        On the first violated property.
    """
    check_jax_operator(op, *fields)

    for lead in ((2,), (2, 3)):
        members = _member_inputs(fields, lead)
        per = [op(*m) for m in members]
        want = jax.tree.map(
            lambda *ys, lead=lead: jnp.stack(ys).reshape(*lead, *ys[0].shape), *per
        )
        batched = tuple(
            jnp.stack([m[i] for m in members]).reshape(*lead, *f.shape)
            if eqx.is_array(f)
            else f
            for i, f in enumerate(fields)
        )
        got = op(*batched)
        assert jax.tree.structure(got) == jax.tree.structure(want), (
            f"leading axes {lead}: output structure differs"
        )
        for a, b in zip(jax.tree.leaves(got), jax.tree.leaves(want), strict=True):
            # Eager batched ≡ per-member up to rounding (a few ulp).
            tol = 64 * float(jnp.finfo(a.dtype).eps)
            np.testing.assert_allclose(
                a, b, rtol=tol, atol=tol, err_msg=f"leading axes {lead}"
            )

    first = next(i for i, f in enumerate(fields) if eqx.is_inexact_array(f))

    def loss(x: Any) -> Any:
        xs = list(fields)
        xs[first] = x
        return sum(jnp.sum(jnp.real(y)) for y in jax.tree.leaves(op(*xs)))

    grad = jax.grad(loss)(fields[first])
    assert bool(jnp.all(jnp.isfinite(grad))), "gradient w.r.t. the first input"

    if hlo:
        text = compiled_hlo(op, *fields)
        found = [name.rstrip("(") for name in FORBIDDEN_HLO if name in text]
        assert not found, f"compiled HLO contains {found}"

Stencils, interpolation, boundaries and masks

finitevolx.pipekit.Difference

Bases: GridOperator

One-sided difference between staggered points, δ / Δ.

axis="x", src="T", dst="U": δₓh[j, i+½] = (h[j, i+1] − h[j, i]) / Δx.

Attributes:

Name Type Description
grid GridSpec

(nx_interior, ny_interior, Lx, Ly).
axis {'x', 'y'}

Difference direction.
src, dst {'T', 'U', 'V', 'X'}

Source and destination stagger; see DIFFERENCES for the valid triples.
Source code in finitevolx/_src/pipekit/stencils.py 
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
class Difference(GridOperator):
    """One-sided difference between staggered points, ``δ / Δ``.

    ``axis="x", src="T", dst="U"``: δₓh[j, i+½] = (h[j, i+1] − h[j, i]) / Δx.

    Attributes
    ----------
    grid : GridSpec
        ``(nx_interior, ny_interior, Lx, Ly)``.
    axis : {"x", "y"}
        Difference direction.
    src, dst : {"T", "U", "V", "X"}
        Source and destination stagger; see `DIFFERENCES` for the valid
        triples.
    """

    axis: str = eqx.field(static=True, kw_only=True)
    src: str = eqx.field(static=True, kw_only=True)
    dst: str = eqx.field(static=True, kw_only=True)

    def __check_init__(self) -> None:
        _check_choice("difference", (self.axis, self.src, self.dst), DIFFERENCES)

    def _apply(self, x: Field) -> Field:
        method = DIFFERENCES[(self.axis, self.src, self.dst)]
        return getattr(Difference2D(grid=self.cartesian_grid), method)(x)

finitevolx.pipekit.Divergence

Bases: GridOperator

∇·(u, v) at T-points: δₓu / Δx + δᵧv / Δy. (U, V) → T.

Source code in finitevolx/_src/pipekit/stencils.py 
102
103
104
105
106
class Divergence(GridOperator):
    """∇·(u, v) at T-points: δₓu / Δx + δᵧv / Δy.  ``(U, V) → T``."""

    def _apply(self, u: Field, v: Field) -> Field:
        return Difference2D(grid=self.cartesian_grid).divergence(u, v)

finitevolx.pipekit.Curl

Bases: GridOperator

Relative vorticity ζ = δₓv / Δx − δᵧu / Δy at X-points. (U, V) → X.

Source code in finitevolx/_src/pipekit/stencils.py 
109
110
111
112
113
class Curl(GridOperator):
    """Relative vorticity ζ = δₓv / Δx − δᵧu / Δy at X-points.  ``(U, V) → X``."""

    def _apply(self, u: Field, v: Field) -> Field:
        return Difference2D(grid=self.cartesian_grid).curl(u, v)

finitevolx.pipekit.Laplacian

Bases: GridOperator

Five-point Laplacian ∇²h at T-points. T → T.

Source code in finitevolx/_src/pipekit/stencils.py 
116
117
118
119
120
class Laplacian(GridOperator):
    """Five-point Laplacian ∇²h at T-points.  ``T → T``."""

    def _apply(self, h: Field) -> Field:
        return Difference2D(grid=self.cartesian_grid).laplacian(h)

finitevolx.pipekit.GradPerp

Bases: GridOperator

Rotated gradient (u, v) = (−∂ᵧψ, ∂ₓψ) of a streamfunction. X → (U, V).

Source code in finitevolx/_src/pipekit/stencils.py 
123
124
125
126
127
class GradPerp(GridOperator):
    """Rotated gradient (u, v) = (−∂ᵧψ, ∂ₓψ) of a streamfunction.  ``X → (U, V)``."""

    def _apply(self, psi: Field) -> tuple[Field, Field]:
        return Difference2D(grid=self.cartesian_grid).grad_perp(psi)

finitevolx.pipekit.Interp

Bases: GridOperator

Average between staggered points (two- or four-point mean).

Attributes:

Name Type Description
grid GridSpec

(nx_interior, ny_interior, Lx, Ly).
src, dst {'T', 'U', 'V', 'X'}

Source and destination stagger; see INTERPOLATIONS.
Source code in finitevolx/_src/pipekit/stencils.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
class Interp(GridOperator):
    """Average between staggered points (two- or four-point mean).

    Attributes
    ----------
    grid : GridSpec
        ``(nx_interior, ny_interior, Lx, Ly)``.
    src, dst : {"T", "U", "V", "X"}
        Source and destination stagger; see `INTERPOLATIONS`.
    """

    src: str = eqx.field(static=True, kw_only=True)
    dst: str = eqx.field(static=True, kw_only=True)

    def __check_init__(self) -> None:
        _check_choice("interpolation", (self.src, self.dst), INTERPOLATIONS)

    def _apply(self, x: Field) -> Field:
        method = INTERPOLATIONS[(self.src, self.dst)]
        return getattr(Interpolation2D(grid=self.cartesian_grid), method)(x)

finitevolx.pipekit.Coriolis

Bases: GridOperator

Coriolis tendencies on a β-plane, (U, V) → (U, V).

f = f₀ + β (y − y₀) at T-points, ghost ring included (y = j Δy), then Coriolis2D(u, v, f): (+f v̄ at U, −f ū at V), energy conserving. f0 and beta are learnable leaves.

Attributes:

Name Type Description
grid GridSpec

(nx_interior, ny_interior, Lx, Ly).
f0 Array

Reference Coriolis parameter.
beta Array

Meridional gradient ∂f/∂y.
y0 float or None

Reference latitude coordinate; None uses the mean of the T-point y coordinates (as coriolis_fn).
Source code in finitevolx/_src/pipekit/stencils.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
class Coriolis(GridOperator):
    """Coriolis tendencies on a β-plane, ``(U, V) → (U, V)``.

    f = f₀ + β (y − y₀) at T-points, ghost ring included (y = j Δy), then
    ``Coriolis2D(u, v, f)``: (+f v̄ at U, −f ū at V), energy conserving.
    ``f0`` and ``beta`` are learnable leaves.

    Attributes
    ----------
    grid : GridSpec
        ``(nx_interior, ny_interior, Lx, Ly)``.
    f0 : Array
        Reference Coriolis parameter.
    beta : Array
        Meridional gradient ∂f/∂y.
    y0 : float or None
        Reference latitude coordinate; ``None`` uses the mean of the T-point
        y coordinates (as `coriolis_fn`).
    """

    f0: jax.Array = eqx.field(converter=jnp.asarray, kw_only=True)
    beta: jax.Array = eqx.field(converter=jnp.asarray, default=0.0, kw_only=True)
    y0: float | None = eqx.field(static=True, default=None, kw_only=True)

    def coriolis_parameter(self, dtype: Any = None) -> Float[Array, "Ny Nx"]:
        """f at the T-points of the grid, ghost ring included."""
        g = self.cartesian_grid
        y = jnp.arange(g.Ny, dtype=dtype) * g.dy  # T[j] sits at y = j Δy
        Y = jnp.broadcast_to(y[:, None], (g.Ny, g.Nx))
        return coriolis_fn(Y, f0=self.f0, beta=self.beta, y0=self.y0)

    def _apply(self, u: Field, v: Field) -> tuple[Field, Field]:
        f = self.coriolis_parameter(jnp.result_type(u, v))
        return Coriolis2D(grid=self.cartesian_grid)(u, v, f)

coriolis_parameter(dtype=None)

f at the T-points of the grid, ghost ring included.

Source code in finitevolx/_src/pipekit/stencils.py 
176
177
178
179
180
181
def coriolis_parameter(self, dtype: Any = None) -> Float[Array, "Ny Nx"]:
    """f at the T-points of the grid, ghost ring included."""
    g = self.cartesian_grid
    y = jnp.arange(g.Ny, dtype=dtype) * g.dy  # T[j] sits at y = j Δy
    Y = jnp.broadcast_to(y[:, None], (g.Ny, g.Nx))
    return coriolis_fn(Y, f0=self.f0, beta=self.beta, y0=self.y0)

finitevolx.pipekit.FillGhosts

Bases: FVXOperator

Fill the ghost ring with a boundary condition. x → x.

Attributes:

Name Type Description
kind str

One of GHOST_FILLERS: periodic, zero, zero_gradient, no_flux, no_slip, free_slip or wall.
staggering {'h', 'u', 'v', 'q'}

Only for kind="wall": the C-grid variable (see wall_boundaries).
Source code in finitevolx/_src/pipekit/stencils.py 
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
class FillGhosts(FVXOperator):
    """Fill the ghost ring with a boundary condition.  ``x → x``.

    Attributes
    ----------
    kind : str
        One of `GHOST_FILLERS`: ``periodic``, ``zero``, ``zero_gradient``,
        ``no_flux``, ``no_slip``, ``free_slip`` or ``wall``.
    staggering : {"h", "u", "v", "q"}
        Only for ``kind="wall"``: the C-grid variable (see
        `wall_boundaries`).
    """

    kind: str = eqx.field(static=True)
    staggering: str = eqx.field(static=True, default="h")

    def __check_init__(self) -> None:
        if self.kind not in GHOST_FILLERS:
            raise ValueError(
                f"kind {self.kind!r} is not available; valid: {sorted(GHOST_FILLERS)}"
            )

    def _apply(self, x: Field) -> Field:
        if self.kind == "wall":
            return wall_boundaries(x, staggering=self.staggering)
        return GHOST_FILLERS[self.kind](x)

finitevolx.pipekit.ApplyMask

Bases: FVXOperator

Zero a field on the dry points of its stagger (jnp.where, NaN-safe).

The mask is an array-holding module, so the config has no JSON form: the operator is forbid_in_yaml and persists through serialize_weights / with_weights.

Attributes:

Name Type Description
mask Mask2D

Land / ocean mask.
stagger {'T', 'U', 'V', 'X'}

Which wet mask to apply: T → mask.h, U → mask.u, V → mask.v, X → mask.xy_corner_strict.
Source code in finitevolx/_src/pipekit/stencils.py 
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
class ApplyMask(FVXOperator):
    """Zero a field on the dry points of its stagger (``jnp.where``, NaN-safe).

    The mask is an array-holding module, so the config has no JSON form:
    the operator is ``forbid_in_yaml`` and persists through
    ``serialize_weights`` / ``with_weights``.

    Attributes
    ----------
    mask : Mask2D
        Land / ocean mask.
    stagger : {"T", "U", "V", "X"}
        Which wet mask to apply: T → ``mask.h``, U → ``mask.u``, V →
        ``mask.v``, X → ``mask.xy_corner_strict``.
    """

    forbid_in_yaml: ClassVar[bool] = True

    mask: Mask2D
    stagger: str = eqx.field(static=True, default="T")

    def __check_init__(self) -> None:
        if self.stagger not in MASK_ATTRIBUTES:
            raise ValueError(
                f"stagger {self.stagger!r} is not available; valid: {list(STAGGERS)}"
            )

    def _apply(self, x: Field) -> Field:
        wet = getattr(self.mask, MASK_ATTRIBUTES[self.stagger])
        return jnp.where(wet, x, jnp.zeros((), x.dtype))

Diagnostics and reductions

finitevolx.pipekit.RelativeVorticity

Bases: GridOperator

ζ = δₓv / Δx − δᵧu / Δy at X-points. (U, V) → X.

Source code in finitevolx/_src/pipekit/diagnostics.py 
38
39
40
41
42
class RelativeVorticity(GridOperator):
    """ζ = δₓv / Δx − δᵧu / Δy at X-points.  ``(U, V) → X``."""

    def _apply(self, u: Field, v: Field) -> Field:
        return Vorticity2D(grid=self.cartesian_grid).relative_vorticity(u, v)

finitevolx.pipekit.PotentialVorticity

Bases: GridOperator

Shallow-water PV q = (ζ + f) / h at X-points. (U, V, T, T) → X.

Source code in finitevolx/_src/pipekit/diagnostics.py 
45
46
47
48
49
class PotentialVorticity(GridOperator):
    """Shallow-water PV q = (ζ + f) / h at X-points.  ``(U, V, T, T) → X``."""

    def _apply(self, u: Field, v: Field, h: Field, f: Field) -> Field:
        return Vorticity2D(grid=self.cartesian_grid).potential_vorticity(u, v, h, f)

finitevolx.pipekit.Enstrophy

Bases: GridOperator

½ ζ² at X-points. (U, V) → X.

Source code in finitevolx/_src/pipekit/diagnostics.py 
52
53
54
55
56
class Enstrophy(GridOperator):
    """½ ζ² at X-points.  ``(U, V) → X``."""

    def _apply(self, u: Field, v: Field) -> Field:
        return Vorticity2D(grid=self.cartesian_grid).enstrophy(u, v)

finitevolx.pipekit.QGPotentialVorticity

Bases: GridOperator

Single-layer QG PV q = (∇²ψ + β (y − y₀)) / f₀ at T-points. T → T.

y = j Δy at the T-points (ghost ring included).

Attributes:

Name Type Description
grid GridSpec

(nx_interior, ny_interior, Lx, Ly).
f0 Array

Reference Coriolis parameter (learnable).
beta Array

Meridional gradient ∂f/∂y (learnable).
y0 float or None

Reference y; None is the domain centre, (Ny − 1) Δy / 2.
Source code in finitevolx/_src/pipekit/diagnostics.py 
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
class QGPotentialVorticity(GridOperator):
    """Single-layer QG PV q = (∇²ψ + β (y − y₀)) / f₀ at T-points.  ``T → T``.

    y = j Δy at the T-points (ghost ring included).

    Attributes
    ----------
    grid : GridSpec
        ``(nx_interior, ny_interior, Lx, Ly)``.
    f0 : Array
        Reference Coriolis parameter (learnable).
    beta : Array
        Meridional gradient ∂f/∂y (learnable).
    y0 : float or None
        Reference y; ``None`` is the domain centre, ``(Ny − 1) Δy / 2``.
    """

    f0: jax.Array = eqx.field(converter=jnp.asarray, kw_only=True)
    beta: jax.Array = eqx.field(converter=jnp.asarray, default=0.0, kw_only=True)
    y0: float | None = eqx.field(static=True, default=None, kw_only=True)

    def _apply(self, psi: Field) -> Field:
        g = self.cartesian_grid
        y = jnp.arange(g.Ny, dtype=psi.dtype) * g.dy  # T[j] sits at y = j Δy
        Y = jnp.broadcast_to(y[:, None], (g.Ny, g.Nx))
        y0 = (g.Ny - 1) * g.dy / 2 if self.y0 is None else self.y0
        qg = QGPotentialVorticity2D(grid=g)
        return qg(psi, self.f0, self.beta, Y, y0)

finitevolx.pipekit.KineticEnergy

Bases: GridOperator

½ (ū² + v̄²) at T-points (velocities averaged to T). (U, V) → T.

Source code in finitevolx/_src/pipekit/diagnostics.py 
89
90
91
92
93
class KineticEnergy(GridOperator):
    """½ (ū² + v̄²) at T-points (velocities averaged to T).  ``(U, V) → T``."""

    def _apply(self, u: Field, v: Field) -> Field:
        return Energetics2D(grid=self.cartesian_grid).kinetic_energy(u, v)

finitevolx.pipekit.AvailablePE

Bases: GridOperator

½ g′ (h − H)² at T-points. T → T.

Attributes:

Name Type Description
grid GridSpec

(nx_interior, ny_interior, Lx, Ly).
g_prime Array

Reduced gravity (learnable).
H Array

Reference thickness (learnable), broadcast over the grid.
Source code in finitevolx/_src/pipekit/diagnostics.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
class AvailablePE(GridOperator):
    """½ g′ (h − H)² at T-points.  ``T → T``.

    Attributes
    ----------
    grid : GridSpec
        ``(nx_interior, ny_interior, Lx, Ly)``.
    g_prime : Array
        Reduced gravity (learnable).
    H : Array
        Reference thickness (learnable), broadcast over the grid.
    """

    g_prime: jax.Array = eqx.field(converter=jnp.asarray, kw_only=True)
    H: jax.Array = eqx.field(converter=jnp.asarray, default=0.0, kw_only=True)

    def _apply(self, h: Field) -> Field:
        H = jnp.broadcast_to(self.H, h.shape)
        energetics = Energetics2D(grid=self.cartesian_grid)
        return energetics.available_potential_energy(h, H, self.g_prime)

finitevolx.pipekit.Strain

Bases: GridOperator

A strain component: shear σₛ (X), tensor σₙ (T) or σₛ² + σₙ² (T).

Attributes:

Name Type Description
grid GridSpec

(nx_interior, ny_interior, Lx, Ly).
component {'shear', 'tensor', 'magnitude_squared'}

Which component.
Source code in finitevolx/_src/pipekit/diagnostics.py 
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
class Strain(GridOperator):
    """A strain component: shear σₛ (X), tensor σₙ (T) or σₛ² + σₙ² (T).

    Attributes
    ----------
    grid : GridSpec
        ``(nx_interior, ny_interior, Lx, Ly)``.
    component : {"shear", "tensor", "magnitude_squared"}
        Which component.
    """

    component: str = eqx.field(static=True, kw_only=True)

    def __check_init__(self) -> None:
        if self.component not in STRAIN_COMPONENTS:
            raise ValueError(
                f"component {self.component!r} is not available; "
                f"valid: {sorted(STRAIN_COMPONENTS)}"
            )

    def _apply(self, u: Field, v: Field) -> Field:
        strain = Strain2D(grid=self.cartesian_grid)
        return getattr(strain, STRAIN_COMPONENTS[self.component])(u, v)

finitevolx.pipekit.OkuboWeiss

Bases: GridOperator

W = σₙ² + σₛ² − ζ² at T-points (W < 0: vortex cores). (U, V) → T.

Source code in finitevolx/_src/pipekit/diagnostics.py 
151
152
153
154
155
class OkuboWeiss(GridOperator):
    """W = σₙ² + σₛ² − ζ² at T-points (W < 0: vortex cores).  ``(U, V) → T``."""

    def _apply(self, u: Field, v: Field) -> Field:
        return Strain2D(grid=self.cartesian_grid).okubo_weiss(u, v)

finitevolx.pipekit.TotalEnergy

Bases: GridOperator

E = ∑_interior (KE + APE) Δx Δy, one value per leading index.

(T, T) → [*batch].

Source code in finitevolx/_src/pipekit/diagnostics.py 
158
159
160
161
162
163
164
165
166
class TotalEnergy(GridOperator):
    """E = ∑_interior (KE + APE) Δx Δy, one value per leading index.

    ``(T, T) → [*batch]``.
    """

    def _apply(self, ke: Field, ape: Field) -> Scalar:
        g = self.cartesian_grid
        return total_energy(ke, ape, g.dx, g.dy)

finitevolx.pipekit.DomainIntegral

Bases: GridOperator

∑_interior x ΔA over the T-cells (area_sum). T → [*batch].

Source code in finitevolx/_src/pipekit/diagnostics.py 
169
170
171
172
173
class DomainIntegral(GridOperator):
    """∑_interior x ΔA over the T-cells (`area_sum`).  ``T → [*batch]``."""

    def _apply(self, x: Field) -> Scalar:
        return over_leading(lambda f: area_sum(f, self.cartesian_grid))(x)

finitevolx.pipekit.DomainMean

Bases: GridOperator

Area-weighted interior mean ⟨x⟩ (area_mean). T → [*batch].

Source code in finitevolx/_src/pipekit/diagnostics.py 
176
177
178
179
180
class DomainMean(GridOperator):
    """Area-weighted interior mean ⟨x⟩ (`area_mean`).  ``T → [*batch]``."""

    def _apply(self, x: Field) -> Scalar:
        return over_leading(lambda f: area_mean(f, self.cartesian_grid))(x)

Time stepping and graph plumbing

finitevolx.pipekit.ExplicitStep

Bases: FVXOperator

One explicit step s ↦ s + Δt Φ(s) of ∂ₜs = F(s), for a dict state.

Attributes:

Name Type Description
rhs Operator

Tendency operator: rhs(**s) returns a dict with the keys of s (typically a pipekit Graph). Its leaves are leaves of the step, so gradients reach them.
dt Array

Time step (a leaf).
scheme {'euler', 'heun', 'rk3_ssp', 'rk4'}

Which SCHEMES entry.
Source code in finitevolx/_src/pipekit/timestep.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
class ExplicitStep(FVXOperator):
    """One explicit step s ↦ s + Δt Φ(s) of ∂ₜs = F(s), for a dict state.

    Attributes
    ----------
    rhs : Operator
        Tendency operator: ``rhs(**s)`` returns a dict with the keys of
        ``s`` (typically a pipekit ``Graph``). Its leaves are leaves of the
        step, so gradients reach them.
    dt : Array
        Time step (a leaf).
    scheme : {"euler", "heun", "rk3_ssp", "rk4"}
        Which `SCHEMES` entry.
    """

    rhs: Operator
    dt: jax.Array = eqx.field(converter=jnp.asarray, kw_only=True)
    scheme: str = eqx.field(static=True, default="rk4", kw_only=True)

    def __check_init__(self) -> None:
        if self.scheme not in SCHEMES:
            raise ValueError(
                f"scheme {self.scheme!r} is not available; valid: {sorted(SCHEMES)}"
            )

    def _apply(self, state: dict[str, Any]) -> dict[str, Any]:
        return SCHEMES[self.scheme](state, lambda s: self.rhs(**s), self.dt)

    def get_config(self) -> dict[str, Any]:
        """Config with ``rhs`` as a nested-operator payload."""
        return {
            "rhs": nested_config(self.rhs),
            "dt": float(self.dt),
            "scheme": self.scheme,
        }

get_config()

Config with rhs as a nested-operator payload.

Source code in finitevolx/_src/pipekit/timestep.py 
67
68
69
70
71
72
73
def get_config(self) -> dict[str, Any]:
    """Config with ``rhs`` as a nested-operator payload."""
    return {
        "rhs": nested_config(self.rhs),
        "dt": float(self.dt),
        "scheme": self.scheme,
    }

finitevolx.pipekit.Pick

Bases: FVXOperator

Select one output of a tuple-valued node: (a, b, …) → [index].

Attributes:

Name Type Description
index int

Position to select.
Source code in finitevolx/_src/pipekit/stencils.py 
264
265
266
267
268
269
270
271
272
273
274
275
276
class Pick(FVXOperator):
    """Select one output of a tuple-valued node: ``(a, b, …) → [index]``.

    Attributes
    ----------
    index : int
        Position to select.
    """

    index: int = eqx.field(static=True)

    def _apply(self, xs: tuple[Any, ...]) -> Any:
        return xs[self.index]

finitevolx.pipekit.LinComb

Bases: FVXOperator

Linear combination y = ∑ᵢ cᵢ xᵢ (coefficients are learnable leaves).

Attributes:

Name Type Description
coeffs Array

[n] coefficients, one per input.
Source code in finitevolx/_src/pipekit/stencils.py 
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
class LinComb(FVXOperator):
    """Linear combination y = ∑ᵢ cᵢ xᵢ (coefficients are learnable leaves).

    Attributes
    ----------
    coeffs : Array
        ``[n]`` coefficients, one per input.
    """

    coeffs: jax.Array = eqx.field(converter=jnp.asarray)

    def _apply(self, *xs: Field) -> Field:
        if len(xs) != self.coeffs.shape[0]:
            raise ValueError(
                f"LinComb has {self.coeffs.shape[0]} coefficients, got {len(xs)} inputs"
            )
        return sum((c * x for c, x in zip(self.coeffs, xs, strict=True)), start=0.0)