Skip to content

Diagnostic Quantities

Pointwise and domain-integrated diagnostic quantities on Arakawa C-grids.

Each diagnostic in the table below comes in two forms:

  • Class operators (Layer 3) take the land/ocean mask at construction and zero the dry cells of their output stagger — use these in masked domains.
  • Functional helpers (Layer 2) are mask-free; apply the mask yourself at the call site if you use them directly.

The domain-integrated total_energy / total_enstrophy and vertical_velocity have no class form: pass the reductions fields that are already masked (e.g. from the class operators), and give vertical_velocity a Mask3D as its per-call mask= argument.

See Masks → Operator API for the stagger → mask mapping.

Diagnostic Class method Functional helper Output
Kinetic energy Energetics2D.kinetic_energy kinetic_energy T
Bernoulli potential Energetics2D.bernoulli_potential bernoulli_potential T
Available potential energy Energetics2D.available_potential_energy available_potential_energy T
Relative vorticity Vorticity2D.relative_vorticity relative_vorticity_cgrid X
Shallow-water PV Vorticity2D.potential_vorticity sw_potential_vorticity¹ X
Enstrophy Vorticity2D.enstrophy enstrophy X
Potential enstrophy Vorticity2D.potential_enstrophy potential_enstrophy X
Shear strain Strain2D.shear shear_strain X
Tensor strain Strain2D.tensor tensor_strain T
Strain magnitude² Strain2D.magnitude_squared strain_magnitude_squared T
Okubo–Weiss Strain2D.okubo_weiss okubo_weiss T
QG PV (one layer) QGPotentialVorticity2D.__call__ qg_potential_vorticity T
Stretching term QGPotentialVorticity2D.stretching stretching_term T
QG PV (multilayer) QGPotentialVorticity2D.multilayer potential_vorticity_multilayer T

¹ Not identical where the thickness at a corner is zero: sw_potential_vorticity returns 0 there, while Vorticity2D.potential_vorticity returns a NaN sentinel at wet corners (dry corners are exactly 0 under a mask), so a degenerate layer is flagged rather than hidden.

The functional strain_magnitude_squared / okubo_weiss / enstrophy / potential_enstrophy are pointwise: their inputs must already share a grid point. The class methods do the staggering for you — Strain2D averages the X-point shear and vorticity to T-points, Vorticity2D.potential_enstrophy averages h to X-points.

Class operators

finitevolx.Energetics2D

Bases: Module

Energy diagnostics at T-points on a 2-D Arakawa C-grid.

Class form of :func:~finitevolx.kinetic_energy, :func:~finitevolx.bernoulli_potential and :func:~finitevolx.available_potential_energy. Every method returns a T-point field with a zero ghost ring; when mask is set, inputs are zeroed on dry interior cells first and dry T-cells are zeroed last (via jnp.where(mask.h, ...)); see the module docstring.

Parameters:

Name Type Description Default
grid CartesianGrid2D

The underlying 2-D grid.

 required
mask Mask2D or None

Optional land/ocean mask. None (default) returns exactly the functional forms' output (for :meth:available_potential_energy, on the interior; see there).

 required

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import CartesianGrid2D, Energetics2D
>>> grid = CartesianGrid2D.from_interior(8, 8, 1.0, 1.0)
>>> u = jnp.ones((grid.Ny, grid.Nx))
>>> ke = Energetics2D(grid=grid).kinetic_energy(u, 0.0 * u)
Source code in finitevolx/_src/operators/diagnostic_operators.py 
 87
 88
 89
 90
 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
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
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
186
187
188
189
190
191
192
193
194
195
196
197
198
199
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
226
227
228
229
230
231
class Energetics2D(eqx.Module):
    """Energy diagnostics at T-points on a 2-D Arakawa C-grid.

    Class form of :func:`~finitevolx.kinetic_energy`,
    :func:`~finitevolx.bernoulli_potential` and
    :func:`~finitevolx.available_potential_energy`.  Every method returns a
    T-point field with a zero ghost ring; when ``mask`` is set, inputs are
    zeroed on dry interior cells first and dry T-cells are zeroed last (via
    ``jnp.where(mask.h, ...)``); see the module docstring.

    Parameters
    ----------
    grid : CartesianGrid2D
        The underlying 2-D grid.
    mask : Mask2D or None, optional
        Optional land/ocean mask.  ``None`` (default) returns exactly the
        functional forms' output (for :meth:`available_potential_energy`,
        on the interior; see there).

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import CartesianGrid2D, Energetics2D
    >>> grid = CartesianGrid2D.from_interior(8, 8, 1.0, 1.0)
    >>> u = jnp.ones((grid.Ny, grid.Nx))
    >>> ke = Energetics2D(grid=grid).kinetic_energy(u, 0.0 * u)
    """

    grid: CartesianGrid2D
    mask: Mask2D | None = None

    def _uv(
        self, u: Float[Array, "*batch Ny Nx"], v: Float[Array, "*batch Ny Nx"]
    ) -> tuple[Float[Array, "*batch Ny Nx"], Float[Array, "*batch Ny Nx"]]:
        """Zero velocities on dry interior faces (no-normal-flow)."""
        if self.mask is None:
            return u, v
        return _sanitize(self.mask.u, u), _sanitize(self.mask.v, v)

    def _t(self, field: Float[Array, "*batch Ny Nx"]) -> Float[Array, "*batch Ny Nx"]:
        """Zero a T-field on dry interior cells."""
        return _sanitize(None if self.mask is None else self.mask.h, field)

    def _out(self, out: Float[Array, "*batch Ny Nx"]) -> Float[Array, "*batch Ny Nx"]:
        """Zero dry T-cells: out[j, i] = out[j, i] if mask.h[j, i] else 0."""
        return _where(None if self.mask is None else self.mask.h, out)

    def kinetic_energy(
        self,
        u: Float[Array, "*batch Ny Nx"],
        v: Float[Array, "*batch Ny Nx"],
    ) -> Float[Array, "*batch Ny Nx"]:
        """Kinetic energy at T-points.

        ke[j, i] = 1/2 * (u2_on_T[j, i] + v2_on_T[j, i])

        u2_on_T[j, i] = 1/2 * (u[j, i+1/2]^2 + u[j, i-1/2]^2)
        v2_on_T[j, i] = 1/2 * (v[j+1/2, i]^2 + v[j-1/2, i]^2)

        Parameters
        ----------
        u : Float[Array, "*batch Ny Nx"]
            x-velocity at U-points.
        v : Float[Array, "*batch Ny Nx"]
            y-velocity at V-points.

        Returns
        -------
        Float[Array, "*batch Ny Nx"]
            Kinetic energy at T-points, zero in the ghost ring and, when
            ``self.mask`` is set, at dry T-cells.
        """
        u, v = self._uv(u, v)
        return self._out(kinetic_energy(u, v))

    def bernoulli_potential(
        self,
        h: Float[Array, "*batch Ny Nx"],
        u: Float[Array, "*batch Ny Nx"],
        v: Float[Array, "*batch Ny Nx"],
        gravity: float = GRAVITY,
    ) -> Float[Array, "*batch Ny Nx"]:
        """Bernoulli potential at T-points.

        p[j, i] = ke[j, i] + gravity * h[j, i]

        Parameters
        ----------
        h : Float[Array, "*batch Ny Nx"]
            Layer thickness at T-points.
        u : Float[Array, "*batch Ny Nx"]
            x-velocity at U-points.
        v : Float[Array, "*batch Ny Nx"]
            y-velocity at V-points.
        gravity : float, optional
            Gravitational acceleration.

        Returns
        -------
        Float[Array, "*batch Ny Nx"]
            Bernoulli potential at T-points, zero in the ghost ring and,
            when ``self.mask`` is set, at dry T-cells.
        """
        u, v = self._uv(u, v)
        return self._out(bernoulli_potential(self._t(h), u, v, gravity))

    def available_potential_energy(
        self,
        h: Float[Array, "*batch Ny Nx"],
        H: Float[Array, "*batch Ny Nx"],
        g_prime: float | Float[Array, ""],
    ) -> Float[Array, "*batch Ny Nx"]:
        """Available potential energy at T-points.

        ape[j, i] = 1/2 * g_prime * (h[j, i] - H[j, i])^2

        The functional form is pointwise over the whole array; the class
        writes only the interior ``[1:-1, 1:-1]`` and leaves the ghost ring
        zero, like every other operator.

        Parameters
        ----------
        h : Float[Array, "*batch Ny Nx"]
            Layer thickness at T-points.
        H : Float[Array, "*batch Ny Nx"]
            Reference layer thickness at T-points.
        g_prime : float
            Reduced gravity.

        Returns
        -------
        Float[Array, "*batch Ny Nx"]
            Available potential energy at T-points, zero in the ghost ring
            and, when ``self.mask`` is set, at dry T-cells.
        """
        h, H = self._t(h), self._t(H)
        # Evaluate on the interior only, so ghost values never enter the
        # arithmetic (or its gradient):
        # ape[j, i] = 1/2 * g_prime * (h[j, i] - H[j, i])^2
        #             for 1 <= j <= Ny-2, 1 <= i <= Nx-2
        ape = available_potential_energy(
            h[..., 1:-1, 1:-1], H[..., 1:-1, 1:-1], g_prime
        )
        out = jnp.zeros_like(h, dtype=ape.dtype).at[..., 1:-1, 1:-1].set(ape)
        return self._out(out)

available_potential_energy(h, H, g_prime)

Available potential energy at T-points.

ape[j, i] = 1/2 * g_prime * (h[j, i] - H[j, i])^2

The functional form is pointwise over the whole array; the class writes only the interior [1:-1, 1:-1] and leaves the ghost ring zero, like every other operator.

Parameters:

Name Type Description Default
h Float[Array, '*batch Ny Nx']

Layer thickness at T-points.

 required
H Float[Array, '*batch Ny Nx']

Reference layer thickness at T-points.

 required
g_prime float

Reduced gravity.

 required

Returns:

Type Description
Float[Array, '*batch Ny Nx']

Available potential energy at T-points, zero in the ghost ring and, when self.mask is set, at dry T-cells.
Source code in finitevolx/_src/operators/diagnostic_operators.py 
193
194
195
196
197
198
199
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
226
227
228
229
230
231
def available_potential_energy(
    self,
    h: Float[Array, "*batch Ny Nx"],
    H: Float[Array, "*batch Ny Nx"],
    g_prime: float | Float[Array, ""],
) -> Float[Array, "*batch Ny Nx"]:
    """Available potential energy at T-points.

    ape[j, i] = 1/2 * g_prime * (h[j, i] - H[j, i])^2

    The functional form is pointwise over the whole array; the class
    writes only the interior ``[1:-1, 1:-1]`` and leaves the ghost ring
    zero, like every other operator.

    Parameters
    ----------
    h : Float[Array, "*batch Ny Nx"]
        Layer thickness at T-points.
    H : Float[Array, "*batch Ny Nx"]
        Reference layer thickness at T-points.
    g_prime : float
        Reduced gravity.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Available potential energy at T-points, zero in the ghost ring
        and, when ``self.mask`` is set, at dry T-cells.
    """
    h, H = self._t(h), self._t(H)
    # Evaluate on the interior only, so ghost values never enter the
    # arithmetic (or its gradient):
    # ape[j, i] = 1/2 * g_prime * (h[j, i] - H[j, i])^2
    #             for 1 <= j <= Ny-2, 1 <= i <= Nx-2
    ape = available_potential_energy(
        h[..., 1:-1, 1:-1], H[..., 1:-1, 1:-1], g_prime
    )
    out = jnp.zeros_like(h, dtype=ape.dtype).at[..., 1:-1, 1:-1].set(ape)
    return self._out(out)

bernoulli_potential(h, u, v, gravity=GRAVITY)

Bernoulli potential at T-points.

p[j, i] = ke[j, i] + gravity * h[j, i]

Parameters:

Name Type Description Default
h Float[Array, '*batch Ny Nx']

Layer thickness at T-points.

 required
u Float[Array, '*batch Ny Nx']

x-velocity at U-points.

 required
v Float[Array, '*batch Ny Nx']

y-velocity at V-points.

 required
gravity float

Gravitational acceleration.

 GRAVITY

Returns:

Type Description
Float[Array, '*batch Ny Nx']

Bernoulli potential at T-points, zero in the ghost ring and, when self.mask is set, at dry T-cells.
Source code in finitevolx/_src/operators/diagnostic_operators.py 
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
def bernoulli_potential(
    self,
    h: Float[Array, "*batch Ny Nx"],
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
    gravity: float = GRAVITY,
) -> Float[Array, "*batch Ny Nx"]:
    """Bernoulli potential at T-points.

    p[j, i] = ke[j, i] + gravity * h[j, i]

    Parameters
    ----------
    h : Float[Array, "*batch Ny Nx"]
        Layer thickness at T-points.
    u : Float[Array, "*batch Ny Nx"]
        x-velocity at U-points.
    v : Float[Array, "*batch Ny Nx"]
        y-velocity at V-points.
    gravity : float, optional
        Gravitational acceleration.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Bernoulli potential at T-points, zero in the ghost ring and,
        when ``self.mask`` is set, at dry T-cells.
    """
    u, v = self._uv(u, v)
    return self._out(bernoulli_potential(self._t(h), u, v, gravity))

kinetic_energy(u, v)

Kinetic energy at T-points.

ke[j, i] = 1/2 * (u2_on_T[j, i] + v2_on_T[j, i])

u2_on_T[j, i] = 1/2 * (u[j, i+1/2]^2 + u[j, i-1/2]^2) v2_on_T[j, i] = 1/2 * (v[j+1/2, i]^2 + v[j-1/2, i]^2)

Parameters:

Name Type Description Default
u Float[Array, '*batch Ny Nx']

x-velocity at U-points.

 required
v Float[Array, '*batch Ny Nx']

y-velocity at V-points.

 required

Returns:

Type Description
Float[Array, '*batch Ny Nx']

Kinetic energy at T-points, zero in the ghost ring and, when self.mask is set, at dry T-cells.
Source code in finitevolx/_src/operators/diagnostic_operators.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
def kinetic_energy(
    self,
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
) -> Float[Array, "*batch Ny Nx"]:
    """Kinetic energy at T-points.

    ke[j, i] = 1/2 * (u2_on_T[j, i] + v2_on_T[j, i])

    u2_on_T[j, i] = 1/2 * (u[j, i+1/2]^2 + u[j, i-1/2]^2)
    v2_on_T[j, i] = 1/2 * (v[j+1/2, i]^2 + v[j-1/2, i]^2)

    Parameters
    ----------
    u : Float[Array, "*batch Ny Nx"]
        x-velocity at U-points.
    v : Float[Array, "*batch Ny Nx"]
        y-velocity at V-points.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Kinetic energy at T-points, zero in the ghost ring and, when
        ``self.mask`` is set, at dry T-cells.
    """
    u, v = self._uv(u, v)
    return self._out(kinetic_energy(u, v))

finitevolx.Strain2D

Bases: Module

Strain diagnostics on a 2-D Arakawa C-grid.

Class form of :func:~finitevolx.shear_strain, :func:~finitevolx.tensor_strain, :func:~finitevolx.strain_magnitude_squared and :func:~finitevolx.okubo_weiss.

The shear strain lives at X-points and the tensor strain at T-points. :meth:magnitude_squared and :meth:okubo_weiss combine them at T-points: the X-point shear (and vorticity) are averaged to T-points with :meth:Interpolation2D.X_to_T, which reads the south ghost X-row and west ghost X-column. Those X ghosts are BC-owned but hidden inside the method, so they are computed here from the caller's ghost u and v (see :meth:magnitude_squared) -- apply boundary conditions to u and v before calling.

Parameters:

Name Type Description Default
grid CartesianGrid2D

The underlying 2-D grid.

 required
mask Mask2D or None

Optional land/ocean mask. Velocities are zeroed on dry faces first; each method then zeroes the dry cells of its output stagger (mask.xy_corner_strict for :meth:shear, mask.h otherwise). For the T-point combinations the X-point intermediates are masked before averaging (pass-down, Pattern 2 in docs/masks.md), so a coastal T-cell averages in zero from its dry corners. None (default) returns exactly the functional forms' output for :meth:shear and :meth:tensor.

 None

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import CartesianGrid2D, Strain2D
>>> grid = CartesianGrid2D.from_interior(8, 8, 1.0, 1.0)
>>> u = jnp.ones((grid.Ny, grid.Nx))
>>> ow = Strain2D(grid=grid).okubo_weiss(u, 0.0 * u)
Source code in finitevolx/_src/operators/diagnostic_operators.py 
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
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
class Strain2D(eqx.Module):
    """Strain diagnostics on a 2-D Arakawa C-grid.

    Class form of :func:`~finitevolx.shear_strain`,
    :func:`~finitevolx.tensor_strain`,
    :func:`~finitevolx.strain_magnitude_squared` and
    :func:`~finitevolx.okubo_weiss`.

    The shear strain lives at X-points and the tensor strain at T-points.
    :meth:`magnitude_squared` and :meth:`okubo_weiss` combine them at
    **T-points**: the X-point shear (and vorticity) are averaged to T-points
    with :meth:`Interpolation2D.X_to_T`, which reads the south ghost X-row
    and west ghost X-column.  Those X ghosts are BC-owned but hidden inside
    the method, so they are computed here from the caller's ghost ``u`` and
    ``v`` (see :meth:`magnitude_squared`) -- apply boundary conditions to
    ``u`` and ``v`` before calling.

    Parameters
    ----------
    grid : CartesianGrid2D
        The underlying 2-D grid.
    mask : Mask2D or None, optional
        Optional land/ocean mask.  Velocities are zeroed on dry faces
        first; each method then zeroes the dry cells of its output stagger
        (``mask.xy_corner_strict`` for :meth:`shear`, ``mask.h``
        otherwise).  For the T-point combinations the X-point intermediates
        are masked before averaging (pass-down, Pattern 2 in
        ``docs/masks.md``), so a coastal T-cell averages in zero from its
        dry corners.  ``None`` (default) returns exactly the functional
        forms' output for :meth:`shear` and :meth:`tensor`.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import CartesianGrid2D, Strain2D
    >>> grid = CartesianGrid2D.from_interior(8, 8, 1.0, 1.0)
    >>> u = jnp.ones((grid.Ny, grid.Nx))
    >>> ow = Strain2D(grid=grid).okubo_weiss(u, 0.0 * u)
    """

    grid: CartesianGrid2D
    mask: Mask2D | None
    interp: Interpolation2D

    def __init__(
        self,
        grid: CartesianGrid2D,
        mask: Mask2D | None = None,
    ) -> None:
        self.grid = grid
        self.mask = mask
        self.interp = Interpolation2D(grid=grid, mask=mask)

    def _uv(
        self, u: Float[Array, "*batch Ny Nx"], v: Float[Array, "*batch Ny Nx"]
    ) -> tuple[Float[Array, "*batch Ny Nx"], Float[Array, "*batch Ny Nx"]]:
        """Zero velocities on dry interior faces (no-normal-flow)."""
        if self.mask is None:
            return u, v
        return _sanitize(self.mask.u, u), _sanitize(self.mask.v, v)

    def _x_with_ghosts(
        self,
        u: Float[Array, "*batch Ny Nx"],
        v: Float[Array, "*batch Ny Nx"],
        sign: float,
    ) -> Float[Array, "*batch Ny Nx"]:
        """dv/dx + sign * du/dy at X-points, including the BC-owned ghosts.

        sign = +1 gives the shear strain, sign = -1 the relative vorticity:

        x[j+1/2, i+1/2] = (v[j+1/2, i+1] - v[j+1/2, i]) / dx
                        + sign * (u[j+1, i+1/2] - u[j, i+1/2]) / dy

        written for 0 <= j <= Ny-2, 0 <= i <= Nx-2, i.e. the interior plus
        the south ghost X-row (j = 0) and west ghost X-column (i = 0), which
        read the caller's ghost u / v.  The north / east X-ghosts are
        outside the domain and stay zero.  At interior X-points this is the
        same arithmetic as :func:`~finitevolx.shear_strain` /
        :func:`~finitevolx.relative_vorticity_cgrid`.

        Under a mask only the dry *interior* corners are zeroed: the south /
        west ghost corners are BC-owned and keep the values derived from
        the ghost velocities, even where a closed-basin mask marks them dry.
        """
        # dv_dx[j+1/2, i+1/2] = (v[j+1/2, i+1] - v[j+1/2, i]) / dx
        dv_dx = (v[..., :-1, 1:] - v[..., :-1, :-1]) / self.grid.dx
        # du_dy[j+1/2, i+1/2] = (u[j+1, i+1/2] - u[j, i+1/2]) / dy
        du_dy = (u[..., 1:, :-1] - u[..., :-1, :-1]) / self.grid.dy
        x = dv_dx + sign * du_dy
        # out[j, i] = x[j, i] for 0 <= j <= Ny-2, 0 <= i <= Nx-2
        out = jnp.zeros_like(u, dtype=x.dtype).at[..., :-1, :-1].set(x)
        return _sanitize(None if self.mask is None else self.mask.xy_corner_strict, out)

    def shear(
        self,
        u: Float[Array, "*batch Ny Nx"],
        v: Float[Array, "*batch Ny Nx"],
    ) -> Float[Array, "*batch Ny Nx"]:
        """Shear strain at X-points.

        ss[j+1/2, i+1/2] = (v[j+1/2, i+1] - v[j+1/2, i]) / dx
                         + (u[j+1, i+1/2] - u[j, i+1/2]) / dy

        Parameters
        ----------
        u : Float[Array, "*batch Ny Nx"]
            x-velocity at U-points.
        v : Float[Array, "*batch Ny Nx"]
            y-velocity at V-points.

        Returns
        -------
        Float[Array, "*batch Ny Nx"]
            Shear strain at X-points, zero in the ghost ring and, when
            ``self.mask`` is set, at dry X-corners.
        """
        u, v = self._uv(u, v)
        out = shear_strain(u, v, self.grid.dx, self.grid.dy)
        return _where(None if self.mask is None else self.mask.xy_corner_strict, out)

    def tensor(
        self,
        u: Float[Array, "*batch Ny Nx"],
        v: Float[Array, "*batch Ny Nx"],
    ) -> Float[Array, "*batch Ny Nx"]:
        """Tensor (normal) strain at T-points.

        sn[j, i] = (u[j, i+1/2] - u[j, i-1/2]) / dx
                 - (v[j+1/2, i] - v[j-1/2, i]) / dy

        Parameters
        ----------
        u : Float[Array, "*batch Ny Nx"]
            x-velocity at U-points.
        v : Float[Array, "*batch Ny Nx"]
            y-velocity at V-points.

        Returns
        -------
        Float[Array, "*batch Ny Nx"]
            Tensor strain at T-points, zero in the ghost ring and, when
            ``self.mask`` is set, at dry T-cells.
        """
        u, v = self._uv(u, v)
        out = tensor_strain(u, v, self.grid.dx, self.grid.dy)
        return _where(None if self.mask is None else self.mask.h, out)

    def magnitude_squared(
        self,
        u: Float[Array, "*batch Ny Nx"],
        v: Float[Array, "*batch Ny Nx"],
    ) -> Float[Array, "*batch Ny Nx"]:
        """Squared strain magnitude at T-points.

        sigma2[j, i] = sn[j, i]^2 + ss_on_T[j, i]^2

        ss_on_T[j, i] = 1/4 * (ss[j+1/2, i+1/2] + ss[j-1/2, i+1/2]
                             + ss[j+1/2, i-1/2] + ss[j-1/2, i-1/2])

        The first interior row and column read the south / west ghost
        X-points ``ss[1/2, *]`` and ``ss[*, 1/2]``; those are computed from
        the caller's ghost ``u`` / ``v``, so they carry whatever boundary
        condition was applied to the velocities.

        Parameters
        ----------
        u : Float[Array, "*batch Ny Nx"]
            x-velocity at U-points.
        v : Float[Array, "*batch Ny Nx"]
            y-velocity at V-points.

        Returns
        -------
        Float[Array, "*batch Ny Nx"]
            Squared strain magnitude at T-points, zero in the ghost ring
            and, when ``self.mask`` is set, at dry T-cells.
        """
        u, v = self._uv(u, v)
        sn = tensor_strain(u, v, self.grid.dx, self.grid.dy)
        ss_on_T = self.interp.X_to_T(self._x_with_ghosts(u, v, 1.0))
        out = strain_magnitude_squared(sn, ss_on_T)
        return _where(None if self.mask is None else self.mask.h, out)

    def okubo_weiss(
        self,
        u: Float[Array, "*batch Ny Nx"],
        v: Float[Array, "*batch Ny Nx"],
    ) -> Float[Array, "*batch Ny Nx"]:
        """Okubo-Weiss parameter at T-points.

        ow[j, i] = sn[j, i]^2 + ss_on_T[j, i]^2 - omega_on_T[j, i]^2

        where the X-point shear strain and relative vorticity

        omega[j+1/2, i+1/2] = (v[j+1/2, i+1] - v[j+1/2, i]) / dx
                            - (u[j+1, i+1/2] - u[j, i+1/2]) / dy

        are averaged to T-points over their four surrounding corners
        (including the south / west ghost corners), as in
        :meth:`magnitude_squared`.  Positive = strain-dominated, negative =
        vorticity-dominated.

        Parameters
        ----------
        u : Float[Array, "*batch Ny Nx"]
            x-velocity at U-points.
        v : Float[Array, "*batch Ny Nx"]
            y-velocity at V-points.

        Returns
        -------
        Float[Array, "*batch Ny Nx"]
            Okubo-Weiss parameter at T-points, zero in the ghost ring and,
            when ``self.mask`` is set, at dry T-cells.
        """
        u, v = self._uv(u, v)
        sn = tensor_strain(u, v, self.grid.dx, self.grid.dy)
        ss_on_T = self.interp.X_to_T(self._x_with_ghosts(u, v, 1.0))
        omega_on_T = self.interp.X_to_T(self._x_with_ghosts(u, v, -1.0))
        out = okubo_weiss(sn, ss_on_T, omega_on_T)
        return _where(None if self.mask is None else self.mask.h, out)

magnitude_squared(u, v)

Squared strain magnitude at T-points.

sigma2[j, i] = sn[j, i]^2 + ss_on_T[j, i]^2

ss_on_T[j, i] = 1/4 * (ss[j+1/2, i+1/2] + ss[j-1/2, i+1/2] + ss[j+1/2, i-1/2] + ss[j-1/2, i-1/2])

The first interior row and column read the south / west ghost X-points ss[1/2, *] and ss[*, 1/2]; those are computed from the caller's ghost u / v, so they carry whatever boundary condition was applied to the velocities.

Parameters:

Name Type Description Default
u Float[Array, '*batch Ny Nx']

x-velocity at U-points.

 required
v Float[Array, '*batch Ny Nx']

y-velocity at V-points.

 required

Returns:

Type Description
Float[Array, '*batch Ny Nx']

Squared strain magnitude at T-points, zero in the ghost ring and, when self.mask is set, at dry T-cells.
Source code in finitevolx/_src/operators/diagnostic_operators.py 
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
def magnitude_squared(
    self,
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
) -> Float[Array, "*batch Ny Nx"]:
    """Squared strain magnitude at T-points.

    sigma2[j, i] = sn[j, i]^2 + ss_on_T[j, i]^2

    ss_on_T[j, i] = 1/4 * (ss[j+1/2, i+1/2] + ss[j-1/2, i+1/2]
                         + ss[j+1/2, i-1/2] + ss[j-1/2, i-1/2])

    The first interior row and column read the south / west ghost
    X-points ``ss[1/2, *]`` and ``ss[*, 1/2]``; those are computed from
    the caller's ghost ``u`` / ``v``, so they carry whatever boundary
    condition was applied to the velocities.

    Parameters
    ----------
    u : Float[Array, "*batch Ny Nx"]
        x-velocity at U-points.
    v : Float[Array, "*batch Ny Nx"]
        y-velocity at V-points.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Squared strain magnitude at T-points, zero in the ghost ring
        and, when ``self.mask`` is set, at dry T-cells.
    """
    u, v = self._uv(u, v)
    sn = tensor_strain(u, v, self.grid.dx, self.grid.dy)
    ss_on_T = self.interp.X_to_T(self._x_with_ghosts(u, v, 1.0))
    out = strain_magnitude_squared(sn, ss_on_T)
    return _where(None if self.mask is None else self.mask.h, out)

okubo_weiss(u, v)

Okubo-Weiss parameter at T-points.

ow[j, i] = sn[j, i]^2 + ss_on_T[j, i]^2 - omega_on_T[j, i]^2

where the X-point shear strain and relative vorticity

omega[j+1/2, i+1/2] = (v[j+1/2, i+1] - v[j+1/2, i]) / dx - (u[j+1, i+1/2] - u[j, i+1/2]) / dy

are averaged to T-points over their four surrounding corners (including the south / west ghost corners), as in :meth:magnitude_squared. Positive = strain-dominated, negative = vorticity-dominated.

Parameters:

Name Type Description Default
u Float[Array, '*batch Ny Nx']

x-velocity at U-points.

 required
v Float[Array, '*batch Ny Nx']

y-velocity at V-points.

 required

Returns:

Type Description
Float[Array, '*batch Ny Nx']

Okubo-Weiss parameter at T-points, zero in the ghost ring and, when self.mask is set, at dry T-cells.
Source code in finitevolx/_src/operators/diagnostic_operators.py 
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
def okubo_weiss(
    self,
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
) -> Float[Array, "*batch Ny Nx"]:
    """Okubo-Weiss parameter at T-points.

    ow[j, i] = sn[j, i]^2 + ss_on_T[j, i]^2 - omega_on_T[j, i]^2

    where the X-point shear strain and relative vorticity

    omega[j+1/2, i+1/2] = (v[j+1/2, i+1] - v[j+1/2, i]) / dx
                        - (u[j+1, i+1/2] - u[j, i+1/2]) / dy

    are averaged to T-points over their four surrounding corners
    (including the south / west ghost corners), as in
    :meth:`magnitude_squared`.  Positive = strain-dominated, negative =
    vorticity-dominated.

    Parameters
    ----------
    u : Float[Array, "*batch Ny Nx"]
        x-velocity at U-points.
    v : Float[Array, "*batch Ny Nx"]
        y-velocity at V-points.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Okubo-Weiss parameter at T-points, zero in the ghost ring and,
        when ``self.mask`` is set, at dry T-cells.
    """
    u, v = self._uv(u, v)
    sn = tensor_strain(u, v, self.grid.dx, self.grid.dy)
    ss_on_T = self.interp.X_to_T(self._x_with_ghosts(u, v, 1.0))
    omega_on_T = self.interp.X_to_T(self._x_with_ghosts(u, v, -1.0))
    out = okubo_weiss(sn, ss_on_T, omega_on_T)
    return _where(None if self.mask is None else self.mask.h, out)

shear(u, v)

Shear strain at X-points.

ss[j+1/2, i+1/2] = (v[j+1/2, i+1] - v[j+1/2, i]) / dx + (u[j+1, i+1/2] - u[j, i+1/2]) / dy

Parameters:

Name Type Description Default
u Float[Array, '*batch Ny Nx']

x-velocity at U-points.

 required
v Float[Array, '*batch Ny Nx']

y-velocity at V-points.

 required

Returns:

Type Description
Float[Array, '*batch Ny Nx']

Shear strain at X-points, zero in the ghost ring and, when self.mask is set, at dry X-corners.
Source code in finitevolx/_src/operators/diagnostic_operators.py 
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
def shear(
    self,
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
) -> Float[Array, "*batch Ny Nx"]:
    """Shear strain at X-points.

    ss[j+1/2, i+1/2] = (v[j+1/2, i+1] - v[j+1/2, i]) / dx
                     + (u[j+1, i+1/2] - u[j, i+1/2]) / dy

    Parameters
    ----------
    u : Float[Array, "*batch Ny Nx"]
        x-velocity at U-points.
    v : Float[Array, "*batch Ny Nx"]
        y-velocity at V-points.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Shear strain at X-points, zero in the ghost ring and, when
        ``self.mask`` is set, at dry X-corners.
    """
    u, v = self._uv(u, v)
    out = shear_strain(u, v, self.grid.dx, self.grid.dy)
    return _where(None if self.mask is None else self.mask.xy_corner_strict, out)

tensor(u, v)

Tensor (normal) strain at T-points.

sn[j, i] = (u[j, i+1/2] - u[j, i-1/2]) / dx - (v[j+1/2, i] - v[j-1/2, i]) / dy

Parameters:

Name Type Description Default
u Float[Array, '*batch Ny Nx']

x-velocity at U-points.

 required
v Float[Array, '*batch Ny Nx']

y-velocity at V-points.

 required

Returns:

Type Description
Float[Array, '*batch Ny Nx']

Tensor strain at T-points, zero in the ghost ring and, when self.mask is set, at dry T-cells.
Source code in finitevolx/_src/operators/diagnostic_operators.py 
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
def tensor(
    self,
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
) -> Float[Array, "*batch Ny Nx"]:
    """Tensor (normal) strain at T-points.

    sn[j, i] = (u[j, i+1/2] - u[j, i-1/2]) / dx
             - (v[j+1/2, i] - v[j-1/2, i]) / dy

    Parameters
    ----------
    u : Float[Array, "*batch Ny Nx"]
        x-velocity at U-points.
    v : Float[Array, "*batch Ny Nx"]
        y-velocity at V-points.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Tensor strain at T-points, zero in the ghost ring and, when
        ``self.mask`` is set, at dry T-cells.
    """
    u, v = self._uv(u, v)
    out = tensor_strain(u, v, self.grid.dx, self.grid.dy)
    return _where(None if self.mask is None else self.mask.h, out)

finitevolx.QGPotentialVorticity2D

Bases: Module

Quasi-geostrophic potential vorticity at T-points.

Class form of :func:~finitevolx.qg_potential_vorticity, :func:~finitevolx.stretching_term and :func:~finitevolx.potential_vorticity_multilayer. Every method returns a T-point field ([Ny, Nx] for one layer, [nl, Ny, Nx] for the multilayer methods); when mask is set, dry T-cells are zeroed via jnp.where(mask.h, ...) (post-compute, Pattern 1 in docs/masks.md), broadcast over the layer axis.

The Laplacian stencil reads psi at the four neighbours of each T-cell. Under a mask, psi is first set to zero on dry interior T-cells -- the usual no-normal-flow condition, and what a masked elliptic solve returns -- so a coastal cell never reads interior land values (NaN or otherwise). The ghost ring is BC-owned and passed through unchanged.

Parameters:

Name Type Description Default
grid CartesianGrid2D

The underlying 2-D grid.

 required
mask Mask2D or None

Optional land/ocean mask. None (default) returns exactly the functional forms' output.

 required

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import CartesianGrid2D, QGPotentialVorticity2D
>>> grid = CartesianGrid2D.from_interior(8, 8, 1.0, 1.0)
>>> psi = jnp.zeros((grid.Ny, grid.Nx))
>>> y = jnp.zeros((grid.Ny, grid.Nx))
>>> q = QGPotentialVorticity2D(grid=grid)(psi, 1e-4, 0.0, y, 0.0)
Source code in finitevolx/_src/operators/diagnostic_operators.py 
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
class QGPotentialVorticity2D(eqx.Module):
    """Quasi-geostrophic potential vorticity at T-points.

    Class form of :func:`~finitevolx.qg_potential_vorticity`,
    :func:`~finitevolx.stretching_term` and
    :func:`~finitevolx.potential_vorticity_multilayer`.  Every method
    returns a T-point field (``[Ny, Nx]`` for one layer, ``[nl, Ny, Nx]``
    for the multilayer methods); when ``mask`` is set, dry T-cells are
    zeroed via ``jnp.where(mask.h, ...)`` (post-compute, Pattern 1 in
    ``docs/masks.md``), broadcast over the layer axis.

    The Laplacian stencil reads ``psi`` at the four neighbours of each
    T-cell.  Under a mask, ``psi`` is first set to zero on dry *interior*
    T-cells -- the usual no-normal-flow condition, and what a masked
    elliptic solve returns -- so a coastal cell never reads interior land
    values (``NaN`` or otherwise).  The ghost ring is BC-owned and passed
    through unchanged.

    Parameters
    ----------
    grid : CartesianGrid2D
        The underlying 2-D grid.
    mask : Mask2D or None, optional
        Optional land/ocean mask.  ``None`` (default) returns exactly the
        functional forms' output.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import CartesianGrid2D, QGPotentialVorticity2D
    >>> grid = CartesianGrid2D.from_interior(8, 8, 1.0, 1.0)
    >>> psi = jnp.zeros((grid.Ny, grid.Nx))
    >>> y = jnp.zeros((grid.Ny, grid.Nx))
    >>> q = QGPotentialVorticity2D(grid=grid)(psi, 1e-4, 0.0, y, 0.0)
    """

    grid: CartesianGrid2D
    mask: Mask2D | None = None

    def _psi(self, psi: Float[Array, "... Ny Nx"]) -> Float[Array, "... Ny Nx"]:
        """Zero psi on dry interior T-cells; the ghost ring is BC-owned."""
        return _sanitize(None if self.mask is None else self.mask.h, psi)

    def _mask_h(self, out: Float[Array, "... Ny Nx"]) -> Float[Array, "... Ny Nx"]:
        """Zero dry T-cells: out[..., j, i] = out[..., j, i] if mask.h[j, i] else 0."""
        return _where(None if self.mask is None else self.mask.h, out)

    def __call__(
        self,
        psi: Float[Array, "*batch Ny Nx"],
        f0: float | Float[Array, ""],
        beta: float | Float[Array, ""],
        y: Float[Array, "*batch Ny Nx"],
        y0: float,
    ) -> Float[Array, "*batch Ny Nx"]:
        """Single-layer QG potential vorticity.

        q[j, i] = lap_psi[j, i] / f0 + beta * (y[j, i] - y0) / f0

        lap_psi[j, i] = (psi[j, i+1] - 2 * psi[j, i] + psi[j, i-1]) / dx^2
                      + (psi[j+1, i] - 2 * psi[j, i] + psi[j-1, i]) / dy^2

        Parameters
        ----------
        psi : Float[Array, "*batch Ny Nx"]
            Streamfunction at T-points.
        f0 : float
            Reference Coriolis parameter.
        beta : float
            Meridional gradient of the Coriolis parameter.
        y : Float[Array, "*batch Ny Nx"]
            Meridional coordinate at T-points.
        y0 : float
            Reference latitude.

        Returns
        -------
        Float[Array, "*batch Ny Nx"]
            QG potential vorticity at T-points, zero in the ghost ring and,
            when ``self.mask`` is set, at dry T-cells.
        """
        psi = self._psi(psi)
        return self._mask_h(
            qg_potential_vorticity(psi, f0, beta, self.grid.dx, self.grid.dy, y, y0)
        )

    def stretching(
        self,
        A: Float[Array, "nl nl"],
        psi: Float[Array, "nl Ny Nx"],
    ) -> Float[Array, "nl Ny Nx"]:
        """Cross-layer stretching term.

        s[k, j, i] = sum_m A[k, m] * psi[m, j, i]

        Parameters
        ----------
        A : Float[Array, "nl nl"]
            Coupling (stretching) matrix.
        psi : Float[Array, "nl Ny Nx"]
            Streamfunction at T-points for all layers.

        Returns
        -------
        Float[Array, "nl Ny Nx"]
            Stretching contribution at T-points, zero in the ghost ring and,
            when ``self.mask`` is set, at dry T-cells of every layer.
        """
        return self._mask_h(stretching_term(A, self._psi(psi)))

    def multilayer(
        self,
        psi: Float[Array, "nl Ny Nx"],
        A: Float[Array, "nl nl"],
        f0: float,
        beta: float,
        y: Float[Array, "Ny Nx"],
        y0: float,
    ) -> Float[Array, "nl Ny Nx"]:
        """Multi-layer QG potential vorticity.

        q[k, j, i] = lap_psi[k, j, i] / f0 + beta * (y[j, i] - y0) / f0
                   - f0 * sum_m A[k, m] * psi[m, j, i]

        ``A`` is the raw :func:`~finitevolx.build_coupling_matrix` output
        (s^2/m^2); ``f0 * A . psi`` is dimensionless, so do not
        pre-multiply ``A`` by ``f0``.

        Parameters
        ----------
        psi : Float[Array, "nl Ny Nx"]
            Streamfunction at T-points for all layers.
        A : Float[Array, "nl nl"]
            Coupling (stretching) matrix, the raw ``build_coupling_matrix``
            output.
        f0 : float
            Reference Coriolis parameter.
        beta : float
            Meridional gradient of the Coriolis parameter.
        y : Float[Array, "Ny Nx"]
            Meridional coordinate at T-points.
        y0 : float
            Reference latitude.

        Returns
        -------
        Float[Array, "nl Ny Nx"]
            QG potential vorticity at T-points, zero in the ghost ring and,
            when ``self.mask`` is set, at dry T-cells of every layer.
        """
        psi = self._psi(psi)
        return self._mask_h(
            potential_vorticity_multilayer(
                psi, A, f0, beta, self.grid.dx, self.grid.dy, y, y0
            )
        )

__call__(psi, f0, beta, y, y0)

Single-layer QG potential vorticity.

q[j, i] = lap_psi[j, i] / f0 + beta * (y[j, i] - y0) / f0

lap_psi[j, i] = (psi[j, i+1] - 2 * psi[j, i] + psi[j, i-1]) / dx^2 + (psi[j+1, i] - 2 * psi[j, i] + psi[j-1, i]) / dy^2

Parameters:

Name Type Description Default
psi Float[Array, '*batch Ny Nx']

Streamfunction at T-points.

 required
f0 float

Reference Coriolis parameter.

 required
beta float

Meridional gradient of the Coriolis parameter.

 required
y Float[Array, '*batch Ny Nx']

Meridional coordinate at T-points.

 required
y0 float

Reference latitude.

 required

Returns:

Type Description
Float[Array, '*batch Ny Nx']

QG potential vorticity at T-points, zero in the ghost ring and, when self.mask is set, at dry T-cells.
Source code in finitevolx/_src/operators/diagnostic_operators.py 
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
def __call__(
    self,
    psi: Float[Array, "*batch Ny Nx"],
    f0: float | Float[Array, ""],
    beta: float | Float[Array, ""],
    y: Float[Array, "*batch Ny Nx"],
    y0: float,
) -> Float[Array, "*batch Ny Nx"]:
    """Single-layer QG potential vorticity.

    q[j, i] = lap_psi[j, i] / f0 + beta * (y[j, i] - y0) / f0

    lap_psi[j, i] = (psi[j, i+1] - 2 * psi[j, i] + psi[j, i-1]) / dx^2
                  + (psi[j+1, i] - 2 * psi[j, i] + psi[j-1, i]) / dy^2

    Parameters
    ----------
    psi : Float[Array, "*batch Ny Nx"]
        Streamfunction at T-points.
    f0 : float
        Reference Coriolis parameter.
    beta : float
        Meridional gradient of the Coriolis parameter.
    y : Float[Array, "*batch Ny Nx"]
        Meridional coordinate at T-points.
    y0 : float
        Reference latitude.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        QG potential vorticity at T-points, zero in the ghost ring and,
        when ``self.mask`` is set, at dry T-cells.
    """
    psi = self._psi(psi)
    return self._mask_h(
        qg_potential_vorticity(psi, f0, beta, self.grid.dx, self.grid.dy, y, y0)
    )

multilayer(psi, A, f0, beta, y, y0)

Multi-layer QG potential vorticity.

q[k, j, i] = lap_psi[k, j, i] / f0 + beta * (y[j, i] - y0) / f0 - f0 * sum_m A[k, m] * psi[m, j, i]

A is the raw :func:~finitevolx.build_coupling_matrix output (s^2/m^2); f0 * A . psi is dimensionless, so do not pre-multiply A by f0.

Parameters:

Name Type Description Default
psi Float[Array, 'nl Ny Nx']

Streamfunction at T-points for all layers.

 required
A Float[Array, 'nl nl']

Coupling (stretching) matrix, the raw build_coupling_matrix output.

 required
f0 float

Reference Coriolis parameter.

 required
beta float

Meridional gradient of the Coriolis parameter.

 required
y Float[Array, 'Ny Nx']

Meridional coordinate at T-points.

 required
y0 float

Reference latitude.

 required

Returns:

Type Description
Float[Array, 'nl Ny Nx']

QG potential vorticity at T-points, zero in the ghost ring and, when self.mask is set, at dry T-cells of every layer.
Source code in finitevolx/_src/operators/diagnostic_operators.py 
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
def multilayer(
    self,
    psi: Float[Array, "nl Ny Nx"],
    A: Float[Array, "nl nl"],
    f0: float,
    beta: float,
    y: Float[Array, "Ny Nx"],
    y0: float,
) -> Float[Array, "nl Ny Nx"]:
    """Multi-layer QG potential vorticity.

    q[k, j, i] = lap_psi[k, j, i] / f0 + beta * (y[j, i] - y0) / f0
               - f0 * sum_m A[k, m] * psi[m, j, i]

    ``A`` is the raw :func:`~finitevolx.build_coupling_matrix` output
    (s^2/m^2); ``f0 * A . psi`` is dimensionless, so do not
    pre-multiply ``A`` by ``f0``.

    Parameters
    ----------
    psi : Float[Array, "nl Ny Nx"]
        Streamfunction at T-points for all layers.
    A : Float[Array, "nl nl"]
        Coupling (stretching) matrix, the raw ``build_coupling_matrix``
        output.
    f0 : float
        Reference Coriolis parameter.
    beta : float
        Meridional gradient of the Coriolis parameter.
    y : Float[Array, "Ny Nx"]
        Meridional coordinate at T-points.
    y0 : float
        Reference latitude.

    Returns
    -------
    Float[Array, "nl Ny Nx"]
        QG potential vorticity at T-points, zero in the ghost ring and,
        when ``self.mask`` is set, at dry T-cells of every layer.
    """
    psi = self._psi(psi)
    return self._mask_h(
        potential_vorticity_multilayer(
            psi, A, f0, beta, self.grid.dx, self.grid.dy, y, y0
        )
    )

stretching(A, psi)

Cross-layer stretching term.

s[k, j, i] = sum_m A[k, m] * psi[m, j, i]

Parameters:

Name Type Description Default
A Float[Array, 'nl nl']

Coupling (stretching) matrix.

 required
psi Float[Array, 'nl Ny Nx']

Streamfunction at T-points for all layers.

 required

Returns:

Type Description
Float[Array, 'nl Ny Nx']

Stretching contribution at T-points, zero in the ghost ring and, when self.mask is set, at dry T-cells of every layer.
Source code in finitevolx/_src/operators/diagnostic_operators.py 
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
def stretching(
    self,
    A: Float[Array, "nl nl"],
    psi: Float[Array, "nl Ny Nx"],
) -> Float[Array, "nl Ny Nx"]:
    """Cross-layer stretching term.

    s[k, j, i] = sum_m A[k, m] * psi[m, j, i]

    Parameters
    ----------
    A : Float[Array, "nl nl"]
        Coupling (stretching) matrix.
    psi : Float[Array, "nl Ny Nx"]
        Streamfunction at T-points for all layers.

    Returns
    -------
    Float[Array, "nl Ny Nx"]
        Stretching contribution at T-points, zero in the ghost ring and,
        when ``self.mask`` is set, at dry T-cells of every layer.
    """
    return self._mask_h(stretching_term(A, self._psi(psi)))

Enstrophy and potential enstrophy are methods of Vorticity2D (enstrophy, potential_enstrophy), next to the vorticity they are built from.

Functional helpers

Kinetic Energy

finitevolx.kinetic_energy(u, v)

Kinetic energy at T-points (cell centers) on an Arakawa C-grid.

Eq: ke[j, i] = 0.5 * (u²_on_T[j, i] + v²_on_T[j, i])

where u² and v² are averaged from face-points to T-points: u²_on_T[j, i] = 0.5 * (u[j, i+1/2]² + u[j, i-1/2]²) = 0.5 * (u[j, i]² + u[j, i-1]²) v²_on_T[j, i] = 0.5 * (v[j+1/2, i]² + v[j-1/2, i]²) = 0.5 * (v[j, i]² + v[j-1, i]²)

This is a mask-free Layer-2 functional helper per #209. For masked kinetic energy, multiply the result by your T-point mask at the call site: ke = kinetic_energy(u, v) * mask.h.

Args: u (Array): x-velocity at U-points (east faces), shape [Ny, Nx]. v (Array): y-velocity at V-points (north faces), shape [Ny, Nx].

Returns: ke (Array): kinetic energy at T-points, shape [Ny, Nx]. Ghost ring is zero; interior is [1:-1, 1:-1].

Source code in finitevolx/_src/operators/diagnostics.py 
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
def kinetic_energy(
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
) -> Float[Array, "*batch Ny Nx"]:
    """Kinetic energy at T-points (cell centers) on an Arakawa C-grid.

    Eq:
        ke[j, i] = 0.5 * (u²_on_T[j, i] + v²_on_T[j, i])

    where u² and v² are averaged from face-points to T-points:
        u²_on_T[j, i] = 0.5 * (u[j, i+1/2]² + u[j, i-1/2]²)
                       = 0.5 * (u[j, i]² + u[j, i-1]²)
        v²_on_T[j, i] = 0.5 * (v[j+1/2, i]² + v[j-1/2, i]²)
                       = 0.5 * (v[j, i]² + v[j-1, i]²)

    This is a mask-free Layer-2 functional helper per #209.  For masked
    kinetic energy, multiply the result by your T-point mask at the call
    site: ``ke = kinetic_energy(u, v) * mask.h``.

    Args:
        u (Array): x-velocity at U-points (east faces), shape [Ny, Nx].
        v (Array): y-velocity at V-points (north faces), shape [Ny, Nx].

    Returns:
        ke (Array): kinetic energy at T-points, shape [Ny, Nx].
            Ghost ring is zero; interior is [1:-1, 1:-1].
    """
    dtype = jnp.result_type(u, v, 0.0)
    u_float = jnp.asarray(u, dtype=dtype)
    v_float = jnp.asarray(v, dtype=dtype)
    u2 = u_float**2
    v2 = v_float**2
    # u²_on_T[j, i] = 0.5 * (u²[j, i] + u²[j, i-1])  (east + west U-faces)
    # v²_on_T[j, i] = 0.5 * (v²[j, i] + v²[j-1, i])  (north + south V-faces)
    u2_on_T = avg_x_bwd(u2)
    v2_on_T = avg_y_bwd(v2)
    ke_int = 0.5 * (u2_on_T + v2_on_T)
    out = interior(ke_int, u_float)
    return out

Bernoulli Potential

finitevolx.bernoulli_potential(h, u, v, gravity=GRAVITY)

Bernoulli potential at T-points on an Arakawa C-grid.

Eq: p[j, i] = ke[j, i] + g * h[j, i]

where ke is the kinetic energy at T-points.

Args: h (Array): layer thickness at T-points, shape [Ny, Nx]. u (Array): x-velocity at U-points (east faces), shape [Ny, Nx]. v (Array): y-velocity at V-points (north faces), shape [Ny, Nx]. gravity (float): gravitational acceleration. Default = 9.81.

Returns: p (Array): Bernoulli potential at T-points, shape [Ny, Nx]. Ghost ring is zero; interior is [1:-1, 1:-1].

Example: >>> u, v, h = ... >>> p = bernoulli_potential(h=h, u=u, v=v)

Source code in finitevolx/_src/operators/diagnostics.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
def bernoulli_potential(
    h: Float[Array, "*batch Ny Nx"],
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
    gravity: float = GRAVITY,
) -> Float[Array, "*batch Ny Nx"]:
    """Bernoulli potential at T-points on an Arakawa C-grid.

    Eq:
        p[j, i] = ke[j, i] + g * h[j, i]

    where ke is the kinetic energy at T-points.

    Args:
        h (Array): layer thickness at T-points, shape [Ny, Nx].
        u (Array): x-velocity at U-points (east faces), shape [Ny, Nx].
        v (Array): y-velocity at V-points (north faces), shape [Ny, Nx].
        gravity (float): gravitational acceleration. Default = 9.81.

    Returns:
        p (Array): Bernoulli potential at T-points, shape [Ny, Nx].
            Ghost ring is zero; interior is [1:-1, 1:-1].

    Example:
        >>> u, v, h = ...
        >>> p = bernoulli_potential(h=h, u=u, v=v)
    """
    dtype = jnp.result_type(h, u, v, 0.0)
    h_float = jnp.asarray(h, dtype=dtype)
    ke = kinetic_energy(u=u, v=v)
    out = interior(ke[..., 1:-1, 1:-1] + gravity * h_float[..., 1:-1, 1:-1], h_float)
    return out

Relative Vorticity

finitevolx.relative_vorticity_cgrid(u, v, dx, dy)

Relative vorticity at X-points (corners) on an Arakawa C-grid.

zeta[j+1/2, i+1/2] = (v[j+1/2, i+1] - v[j+1/2, i]) / dx - (u[j+1, i+1/2] - u[j, i+1/2]) / dy

This is a standalone functional form of the class-based :meth:Vorticity2D.relative_vorticity. Both share the same underlying implementation (:func:finitevolx._src.operators.difference._curl_2d).

Parameters:

Name Type Description Default
u Float[Array, '*batch Ny Nx']

x-velocity at U-points.

 required
v Float[Array, '*batch Ny Nx']

y-velocity at V-points.

 required
dx float

Grid spacing in x.

 required
dy float

Grid spacing in y.

 required

Returns:

Type Description
Float[Array, '*batch Ny Nx']

Relative vorticity at X-points. Ghost ring is zero.
Source code in finitevolx/_src/operators/diagnostics.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
def relative_vorticity_cgrid(
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
    dx: float,
    dy: float,
) -> Float[Array, "*batch Ny Nx"]:
    """Relative vorticity at X-points (corners) on an Arakawa C-grid.

    zeta[j+1/2, i+1/2] = (v[j+1/2, i+1] - v[j+1/2, i]) / dx
                        - (u[j+1, i+1/2] - u[j, i+1/2]) / dy

    This is a standalone functional form of the class-based
    :meth:`Vorticity2D.relative_vorticity`.  Both share the same
    underlying implementation
    (:func:`finitevolx._src.operators.difference._curl_2d`).

    Parameters
    ----------
    u : Float[Array, "*batch Ny Nx"]
        x-velocity at U-points.
    v : Float[Array, "*batch Ny Nx"]
        y-velocity at V-points.
    dx : float
        Grid spacing in x.
    dy : float
        Grid spacing in y.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Relative vorticity at X-points.  Ghost ring is zero.
    """
    return _curl_2d(u, v, dx, dy)

Potential Vorticity

sw_potential_vorticity is the functional counterpart of Vorticity2D.potential_vorticity (C-grid vorticity and interpolation included); potential_vorticity is the pointwise (omega + f) / h for inputs that already share a grid point.

finitevolx.sw_potential_vorticity(u, v, h, f, dx, dy)

Shallow-water potential vorticity at X-points on an Arakawa C-grid.

q[j+1/2, i+1/2] = (omega[j+1/2, i+1/2] + f_X[j+1/2, i+1/2]) / h_X[j+1/2, i+1/2]

Computes relative vorticity from (u, v) at X-points, interpolates f and h from T-points to X-points, then forms the PV ratio.

Parameters:

Name Type Description Default
u Float[Array, '*batch Ny Nx']

x-velocity at U-points.

 required
v Float[Array, '*batch Ny Nx']

y-velocity at V-points.

 required
h Float[Array, '*batch Ny Nx']

Layer thickness at T-points.

 required
f Float[Array, '*batch Ny Nx']

Coriolis parameter at T-points.

 required
dx float

Grid spacing in x.

 required
dy float

Grid spacing in y.

 required

Returns:

Type Description
Float[Array, '*batch Ny Nx']

Potential vorticity at X-points. Ghost ring is zero. Where h_X == 0 the result is NaN.
Source code in finitevolx/_src/operators/diagnostics.py 
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
def sw_potential_vorticity(
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
    h: Float[Array, "*batch Ny Nx"],
    f: Float[Array, "*batch Ny Nx"],
    dx: float,
    dy: float,
) -> Float[Array, "*batch Ny Nx"]:
    """Shallow-water potential vorticity at X-points on an Arakawa C-grid.

    q[j+1/2, i+1/2] = (omega[j+1/2, i+1/2] + f_X[j+1/2, i+1/2])
                       / h_X[j+1/2, i+1/2]

    Computes relative vorticity from ``(u, v)`` at X-points, interpolates
    ``f`` and ``h`` from T-points to X-points, then forms the PV ratio.

    Parameters
    ----------
    u : Float[Array, "*batch Ny Nx"]
        x-velocity at U-points.
    v : Float[Array, "*batch Ny Nx"]
        y-velocity at V-points.
    h : Float[Array, "*batch Ny Nx"]
        Layer thickness at T-points.
    f : Float[Array, "*batch Ny Nx"]
        Coriolis parameter at T-points.
    dx : float
        Grid spacing in x.
    dy : float
        Grid spacing in y.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Potential vorticity at X-points.  Ghost ring is zero.
        Where ``h_X == 0`` the result is NaN.
    """
    omega = _curl_2d(u, v, dx, dy)
    f_X = _interp_T_to_X(f)
    h_X = _interp_T_to_X(h)
    # Use 0 in ghost ring (where h_X == 0) rather than NaN, matching
    # the ghost-ring-is-zero convention used by all other operators.
    return jnp.where(h_X == 0, 0.0, (omega + f_X) / h_X)

finitevolx.potential_vorticity(omega, f, h)

Potential vorticity (pointwise).

q = (omega + f) / h

All inputs must live on the same grid points (typically X-points after interpolating f and h from T-points).

Parameters:

Name Type Description Default
omega Float[Array, '*batch Ny Nx']

Relative vorticity.

 required
f Float[Array, '*batch Ny Nx']

Coriolis parameter.

 required
h Float[Array, '*batch Ny Nx']

Layer thickness (or depth).

 required

Returns:

Type Description
Float[Array, '*batch Ny Nx']

Potential vorticity. Where h == 0 the result is NaN.
Source code in finitevolx/_src/operators/diagnostics.py 
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
def potential_vorticity(
    omega: Float[Array, "*batch Ny Nx"],
    f: Float[Array, "*batch Ny Nx"],
    h: Float[Array, "*batch Ny Nx"],
) -> Float[Array, "*batch Ny Nx"]:
    """Potential vorticity (pointwise).

    q = (omega + f) / h

    All inputs must live on the same grid points (typically X-points
    after interpolating ``f`` and ``h`` from T-points).

    Parameters
    ----------
    omega : Float[Array, "*batch Ny Nx"]
        Relative vorticity.
    f : Float[Array, "*batch Ny Nx"]
        Coriolis parameter.
    h : Float[Array, "*batch Ny Nx"]
        Layer thickness (or depth).

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Potential vorticity.  Where ``h == 0`` the result is NaN.
    """
    return jnp.where(h == 0, jnp.nan, (omega + f) / h)

QG Potential Vorticity

finitevolx.qg_potential_vorticity(psi, f0, beta, dx, dy, y, y0)

QG potential vorticity at T-points (single layer).

q = nabla^2(psi) / f0 + beta * (y - y0) / f0

This computes the per-layer (barotropic) QG PV for a single [Ny, Nx] streamfunction. For multi-layer QG, lift this function with :func:~finitevolx.multilayer and subtract the cross-layer stretching term separately::

from finitevolx import multilayer, qg_potential_vorticity

# Per-layer PV (vmapped over the layer axis)
q = multilayer(lambda p: qg_potential_vorticity(p, f0, beta, dx, dy, y, y0))(
    psi
)  # psi: [nl, Ny, Nx] -> q: [nl, Ny, Nx]

# Cross-layer stretching (A couples layers, cannot be vmapped);
# A is the raw build_coupling_matrix output, so scale by f0
q = q - f0 * stretching_term(A, psi)

Parameters:

Name Type Description Default
psi Float[Array, '*batch Ny Nx']

Streamfunction at T-points.

 required
f0 float

Reference Coriolis parameter.

 required
beta float

Meridional gradient of the Coriolis parameter.

 required
dx float

Grid spacing in x.

 required
dy float

Grid spacing in y.

 required
y Float[Array, '*batch Ny Nx']

Meridional coordinate at T-points.

 required
y0 float

Reference latitude.

 required

Returns:

Type Description
Float[Array, '*batch Ny Nx']

QG potential vorticity at T-points. Ghost ring is zero; interior is [1:-1, 1:-1].
Source code in finitevolx/_src/operators/diagnostics.py 
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
def qg_potential_vorticity(
    psi: Float[Array, "*batch Ny Nx"],
    f0: float | Float[Array, ""],
    beta: float | Float[Array, ""],
    dx: float,
    dy: float,
    y: Float[Array, "*batch Ny Nx"],
    y0: float,
) -> Float[Array, "*batch Ny Nx"]:
    """QG potential vorticity at T-points (single layer).

    q = nabla^2(psi) / f0 + beta * (y - y0) / f0

    This computes the per-layer (barotropic) QG PV for a single ``[Ny, Nx]``
    streamfunction.  For multi-layer QG, lift this function with
    :func:`~finitevolx.multilayer` and subtract the cross-layer stretching
    term separately::

        from finitevolx import multilayer, qg_potential_vorticity

        # Per-layer PV (vmapped over the layer axis)
        q = multilayer(lambda p: qg_potential_vorticity(p, f0, beta, dx, dy, y, y0))(
            psi
        )  # psi: [nl, Ny, Nx] -> q: [nl, Ny, Nx]

        # Cross-layer stretching (A couples layers, cannot be vmapped);
        # A is the raw build_coupling_matrix output, so scale by f0
        q = q - f0 * stretching_term(A, psi)

    Parameters
    ----------
    psi : Float[Array, "*batch Ny Nx"]
        Streamfunction at T-points.
    f0 : float
        Reference Coriolis parameter.
    beta : float
        Meridional gradient of the Coriolis parameter.
    dx : float
        Grid spacing in x.
    dy : float
        Grid spacing in y.
    y : Float[Array, "*batch Ny Nx"]
        Meridional coordinate at T-points.
    y0 : float
        Reference latitude.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        QG potential vorticity at T-points.
        Ghost ring is zero; interior is ``[1:-1, 1:-1]``.
    """
    # Laplacian at T-points via centred second differences:
    d2x = (diff_x_fwd(psi) - diff_x_bwd(psi)) / dx**2
    d2y = (diff_y_fwd(psi) - diff_y_bwd(psi)) / dy**2
    # q[j,i] = (d2psi/dx2 + d2psi/dy2) / f0 + beta * (y - y0) / f0
    q_int = (d2x + d2y) / f0 + beta * (y[..., 1:-1, 1:-1] - y0) / f0
    out = interior(q_int, psi)
    return out

Stretching Term (Multi-Layer QG)

finitevolx.stretching_term(A, psi)

Cross-layer stretching term for multi-layer QG models.

Computes A . psi — the vertical coupling that mixes information across layers. This is a cross-layer operation and cannot be vmapped; it should be used alongside :func:qg_potential_vorticity (which handles the per-layer Laplacian + beta terms via :func:~finitevolx.multilayer).

Parameters:

Name Type Description Default
A Float[Array, 'nl nl']

Coupling (stretching) matrix. Shape (nl, nl) where nl is the number of layers.

 required
psi Float[Array, 'nl Ny Nx']

Streamfunction at T-points for all layers.

 required

Returns:

Type Description
Float[Array, 'nl Ny Nx']

Stretching contribution, same shape as psi. Ghost ring is zero; interior is [:, 1:-1, 1:-1].
Source code in finitevolx/_src/operators/diagnostics.py 
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
def stretching_term(
    A: Float[Array, "nl nl"],
    psi: Float[Array, "nl Ny Nx"],
) -> Float[Array, "nl Ny Nx"]:
    """Cross-layer stretching term for multi-layer QG models.

    Computes ``A . psi`` — the vertical coupling that mixes information
    across layers.  This is a cross-layer operation and cannot be vmapped;
    it should be used alongside :func:`qg_potential_vorticity` (which
    handles the per-layer Laplacian + beta terms via
    :func:`~finitevolx.multilayer`).

    Parameters
    ----------
    A : Float[Array, "nl nl"]
        Coupling (stretching) matrix.  Shape ``(nl, nl)`` where ``nl``
        is the number of layers.
    psi : Float[Array, "nl Ny Nx"]
        Streamfunction at T-points for all layers.

    Returns
    -------
    Float[Array, "nl Ny Nx"]
        Stretching contribution, same shape as ``psi``.
        Ghost ring is zero; interior is ``[:, 1:-1, 1:-1]``.
    """
    # (A . psi)_k = sum_j A[k,j] * psi_j  for each spatial point
    # Flatten spatial dims, matmul along layer axis, reshape back.
    nl = A.shape[0]
    Ny, Nx = psi.shape[-2], psi.shape[-1]
    psi_flat = psi.reshape(nl, -1)  # [nl, Ny*Nx]
    Ap_flat = A @ psi_flat  # [nl, Ny*Nx]
    Ap = Ap_flat.reshape(nl, Ny, Nx)  # [nl, Ny, Nx]
    out = jnp.zeros_like(psi)
    out = out.at[..., :, 1:-1, 1:-1].set(Ap[..., :, 1:-1, 1:-1])
    return out

Strain

finitevolx.shear_strain(u, v, dx, dy)

Shear strain at X-points (corners) on an Arakawa C-grid.

Ss[j+1/2, i+1/2] = dv/dx + du/dy = (v[j+1/2, i+1] - v[j+1/2, i]) / dx + (u[j+1, i+1/2] - u[j, i+1/2]) / dy

Parameters:

Name Type Description Default
u Float[Array, '*batch Ny Nx']

x-velocity at U-points.

 required
v Float[Array, '*batch Ny Nx']

y-velocity at V-points.

 required
dx float

Grid spacing in x.

 required
dy float

Grid spacing in y.

 required

Returns:

Type Description
Float[Array, '*batch Ny Nx']

Shear strain at X-points. Ghost ring is zero.
Source code in finitevolx/_src/operators/diagnostics.py 
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
def shear_strain(
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
    dx: float,
    dy: float,
) -> Float[Array, "*batch Ny Nx"]:
    """Shear strain at X-points (corners) on an Arakawa C-grid.

    Ss[j+1/2, i+1/2] = dv/dx + du/dy
        = (v[j+1/2, i+1] - v[j+1/2, i]) / dx
        + (u[j+1, i+1/2] - u[j, i+1/2]) / dy

    Parameters
    ----------
    u : Float[Array, "*batch Ny Nx"]
        x-velocity at U-points.
    v : Float[Array, "*batch Ny Nx"]
        y-velocity at V-points.
    dx : float
        Grid spacing in x.
    dy : float
        Grid spacing in y.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Shear strain at X-points.  Ghost ring is zero.
    """
    out = interior(diff_x_fwd(v) / dx + diff_y_fwd(u) / dy, u)
    return out

finitevolx.tensor_strain(u, v, dx, dy)

Tensor (normal) strain at T-points on an Arakawa C-grid.

Sn[j, i] = du/dx - dv/dy = (u[j, i+1/2] - u[j, i-1/2]) / dx - (v[j+1/2, i] - v[j-1/2, i]) / dy

Parameters:

Name Type Description Default
u Float[Array, '*batch Ny Nx']

x-velocity at U-points.

 required
v Float[Array, '*batch Ny Nx']

y-velocity at V-points.

 required
dx float

Grid spacing in x.

 required
dy float

Grid spacing in y.

 required

Returns:

Type Description
Float[Array, '*batch Ny Nx']

Tensor strain at T-points. Ghost ring is zero.
Source code in finitevolx/_src/operators/diagnostics.py 
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
def tensor_strain(
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
    dx: float,
    dy: float,
) -> Float[Array, "*batch Ny Nx"]:
    """Tensor (normal) strain at T-points on an Arakawa C-grid.

    Sn[j, i] = du/dx - dv/dy
        = (u[j, i+1/2] - u[j, i-1/2]) / dx
        - (v[j+1/2, i] - v[j-1/2, i]) / dy

    Parameters
    ----------
    u : Float[Array, "*batch Ny Nx"]
        x-velocity at U-points.
    v : Float[Array, "*batch Ny Nx"]
        y-velocity at V-points.
    dx : float
        Grid spacing in x.
    dy : float
        Grid spacing in y.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Tensor strain at T-points.  Ghost ring is zero.
    """
    out = interior(diff_x_bwd(u) / dx - diff_y_bwd(v) / dy, u)
    return out

finitevolx.strain_magnitude_squared(sn, ss)

Squared strain magnitude (pointwise).

sigma^2 = Sn^2 + Ss^2

Inputs must live on the same grid points. If sn is at T-points and ss at X-points, interpolate one to the other before calling.

Parameters:

Name Type Description Default
sn Float[Array, '*batch Ny Nx']

Tensor (normal) strain.

 required
ss Float[Array, '*batch Ny Nx']

Shear strain.

 required

Returns:

Type Description
Float[Array, '*batch Ny Nx']

Squared strain magnitude.
Source code in finitevolx/_src/operators/diagnostics.py 
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
def strain_magnitude_squared(
    sn: Float[Array, "*batch Ny Nx"],
    ss: Float[Array, "*batch Ny Nx"],
) -> Float[Array, "*batch Ny Nx"]:
    """Squared strain magnitude (pointwise).

    sigma^2 = Sn^2 + Ss^2

    Inputs must live on the same grid points.  If ``sn`` is at T-points
    and ``ss`` at X-points, interpolate one to the other before calling.

    Parameters
    ----------
    sn : Float[Array, "*batch Ny Nx"]
        Tensor (normal) strain.
    ss : Float[Array, "*batch Ny Nx"]
        Shear strain.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Squared strain magnitude.
    """
    return sn**2 + ss**2

Okubo-Weiss Parameter

finitevolx.okubo_weiss(sn, ss, omega)

Okubo-Weiss parameter (pointwise).

OW = Sn^2 + Ss^2 - omega^2

All inputs must live on the same grid points.

Parameters:

Name Type Description Default
sn Float[Array, '*batch Ny Nx']

Tensor (normal) strain.

 required
ss Float[Array, '*batch Ny Nx']

Shear strain.

 required
omega Float[Array, '*batch Ny Nx']

Relative vorticity.

 required

Returns:

Type Description
Float[Array, '*batch Ny Nx']

Okubo-Weiss parameter. Positive = strain-dominated, negative = vorticity-dominated.
Source code in finitevolx/_src/operators/diagnostics.py 
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
def okubo_weiss(
    sn: Float[Array, "*batch Ny Nx"],
    ss: Float[Array, "*batch Ny Nx"],
    omega: Float[Array, "*batch Ny Nx"],
) -> Float[Array, "*batch Ny Nx"]:
    """Okubo-Weiss parameter (pointwise).

    OW = Sn^2 + Ss^2 - omega^2

    All inputs must live on the same grid points.

    Parameters
    ----------
    sn : Float[Array, "*batch Ny Nx"]
        Tensor (normal) strain.
    ss : Float[Array, "*batch Ny Nx"]
        Shear strain.
    omega : Float[Array, "*batch Ny Nx"]
        Relative vorticity.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Okubo-Weiss parameter.  Positive = strain-dominated,
        negative = vorticity-dominated.
    """
    return sn**2 + ss**2 - omega**2

Enstrophy

finitevolx.enstrophy(omega)

Enstrophy (pointwise).

Z = 0.5 * omega^2

Mask-free Layer-2 functional helper per #209. For masked enstrophy, multiply the result by your mask at the call site: z = enstrophy(omega) * mask.xy_corner_strict.

Parameters:

Name Type Description Default
omega Float[Array, '*batch Ny Nx']

Relative vorticity (typically at X-points).

 required

Returns:

Type Description
Float[Array, '*batch Ny Nx']

Enstrophy at the same grid points as omega.
Source code in finitevolx/_src/operators/diagnostics.py 
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
def enstrophy(
    omega: Float[Array, "*batch Ny Nx"],
) -> Float[Array, "*batch Ny Nx"]:
    """Enstrophy (pointwise).

    Z = 0.5 * omega^2

    Mask-free Layer-2 functional helper per #209.  For masked enstrophy,
    multiply the result by your mask at the call site:
    ``z = enstrophy(omega) * mask.xy_corner_strict``.

    Parameters
    ----------
    omega : Float[Array, "*batch Ny Nx"]
        Relative vorticity (typically at X-points).

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Enstrophy at the same grid points as *omega*.
    """
    return 0.5 * omega**2

finitevolx.potential_enstrophy(q, h)

Potential enstrophy (pointwise).

PE = 0.5 * q^2 * h

Inputs must live on the same grid points. Mask-free Layer-2 functional helper per #209; apply masks at the call site.

Parameters:

Name Type Description Default
q Float[Array, '*batch Ny Nx']

Potential vorticity.

 required
h Float[Array, '*batch Ny Nx']

Layer thickness (or depth).

 required

Returns:

Type Description
Float[Array, '*batch Ny Nx']

Potential enstrophy.
Source code in finitevolx/_src/operators/diagnostics.py 
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
def potential_enstrophy(
    q: Float[Array, "*batch Ny Nx"],
    h: Float[Array, "*batch Ny Nx"],
) -> Float[Array, "*batch Ny Nx"]:
    """Potential enstrophy (pointwise).

    PE = 0.5 * q^2 * h

    Inputs must live on the same grid points.  Mask-free Layer-2
    functional helper per #209; apply masks at the call site.

    Parameters
    ----------
    q : Float[Array, "*batch Ny Nx"]
        Potential vorticity.
    h : Float[Array, "*batch Ny Nx"]
        Layer thickness (or depth).

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Potential enstrophy.
    """
    return 0.5 * q**2 * h

Available Potential Energy

finitevolx.available_potential_energy(h, H, g_prime)

Available potential energy at T-points.

APE = 0.5 * g' * (h - H)^2

Parameters:

Name Type Description Default
h Float[Array, '*batch Ny Nx']

Layer thickness at T-points.

 required
H Float[Array, '*batch Ny Nx']

Mean (reference) layer thickness at T-points.

 required
g_prime float

Reduced gravity.

 required

Returns:

Type Description
Float[Array, '*batch Ny Nx']

Available potential energy at T-points.
Source code in finitevolx/_src/operators/diagnostics.py 
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
def available_potential_energy(
    h: Float[Array, "*batch Ny Nx"],
    H: Float[Array, "*batch Ny Nx"],
    g_prime: float | Float[Array, ""],
) -> Float[Array, "*batch Ny Nx"]:
    """Available potential energy at T-points.

    APE = 0.5 * g' * (h - H)^2

    Parameters
    ----------
    h : Float[Array, "*batch Ny Nx"]
        Layer thickness at T-points.
    H : Float[Array, "*batch Ny Nx"]
        Mean (reference) layer thickness at T-points.
    g_prime : float
        Reduced gravity.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Available potential energy at T-points.
    """
    eta = h - H
    return 0.5 * g_prime * eta**2

Domain-Integrated Quantities

finitevolx.total_energy(ke, ape, dx, dy)

Domain-integrated total energy.

E = sum_{interior}(KE + APE) * dx * dy

Parameters:

Name Type Description Default
ke Float[Array, '*batch Ny Nx']

Kinetic energy at T-points.

 required
ape Float[Array, '*batch Ny Nx']

Available potential energy at T-points.

 required
dx float

Grid spacing in x.

 required
dy float

Grid spacing in y.

 required

Returns:

Type Description
Float[Array, '']

Scalar total energy.
Source code in finitevolx/_src/operators/diagnostics.py 
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
def total_energy(
    ke: Float[Array, "*batch Ny Nx"],
    ape: Float[Array, "*batch Ny Nx"],
    dx: float,
    dy: float,
) -> Float[Array, ""]:
    """Domain-integrated total energy.

    E = sum_{interior}(KE + APE) * dx * dy

    Parameters
    ----------
    ke : Float[Array, "*batch Ny Nx"]
        Kinetic energy at T-points.
    ape : Float[Array, "*batch Ny Nx"]
        Available potential energy at T-points.
    dx : float
        Grid spacing in x.
    dy : float
        Grid spacing in y.

    Returns
    -------
    Float[Array, ""]
        Scalar total energy.
    """
    return jnp.sum(ke[..., 1:-1, 1:-1] + ape[..., 1:-1, 1:-1], axis=(-2, -1)) * dx * dy

finitevolx.total_enstrophy(ens, dx, dy)

Domain-integrated enstrophy.

Z = sum_{interior}(0.5 * omega^2) * dx * dy

Parameters:

Name Type Description Default
ens Float[Array, '*batch Ny Nx']

Enstrophy field (e.g. from :func:enstrophy).

 required
dx float

Grid spacing in x.

 required
dy float

Grid spacing in y.

 required

Returns:

Type Description
Float[Array, '']

Scalar total enstrophy.
Source code in finitevolx/_src/operators/diagnostics.py 
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
def total_enstrophy(
    ens: Float[Array, "*batch Ny Nx"],
    dx: float,
    dy: float,
) -> Float[Array, ""]:
    """Domain-integrated enstrophy.

    Z = sum_{interior}(0.5 * omega^2) * dx * dy

    Parameters
    ----------
    ens : Float[Array, "*batch Ny Nx"]
        Enstrophy field (e.g. from :func:`enstrophy`).
    dx : float
        Grid spacing in x.
    dy : float
        Grid spacing in y.

    Returns
    -------
    Float[Array, ""]
        Scalar total enstrophy.
    """
    return jnp.sum(ens[..., 1:-1, 1:-1], axis=(-2, -1)) * dx * dy

Vertical Velocity

finitevolx.vertical_velocity(u, v, dx, dy, dz, mask=None)

Vertical velocity from the continuity equation.

Integrates the horizontal divergence of the physical levels upward from a rigid floor (w = 0):

w[k+1/2] = w[k-1/2] - (du/dx + dv/dy)[k] * dz,   k = 1 .. Nz-2

The result uses the same [Nz, Ny, Nx] convention as Mask3D.w: w[k] is the top face of T[k] (k + 1/2). So w[0] is the floor (the bottom face of the first physical level T[1], zero), w[1:-1] are the interfaces above the physical levels, w[Nz-2] is the surface (zero for a rigid lid when the column conserves mass) and w[Nz-1] is outside the domain (zero). The z-ghost slabs k = 0 and k = Nz-1 of u and v are never read, so BC-filled z-ghosts do not change the answer.

Parameters:

Name Type Description Default
u Float[Array, 'Nz Ny Nx']

x-velocity at U-points, with ghost cells in all three dimensions.

 required
v Float[Array, 'Nz Ny Nx']

y-velocity at V-points, with ghost cells in all three dimensions.

 required
dx float

Grid spacing in x.

 required
dy float

Grid spacing in y.

 required
dz float

Grid spacing in z.

 required
mask Mask3D or Float[Array, 'Nz Ny Nx'] or None

A :class:Mask3D: u / v are zeroed on dry faces with jnp.where before the stencil (so NaN-filled land cannot leak), the divergence on dry T-points (mask.h) before the integration, and the output on dry W-faces (mask.w). mask.w[Nz-2] (the surface) is dry wherever the top ghost slab is dry, so a non-zero rigid-lid residual there is masked out too. A raw float T-point mask (deprecated) only zeroes the divergence.

 None

Returns:

Type Description
Float[Array, 'Nz Ny Nx']

Vertical velocity at W-points, w[k] on the top face of T[k], written on [1:-1, 1:-1, 1:-1] with a zero ghost ring (w[0] is the zero floor).

Examples:

>>> grid = CartesianGrid3D.from_interior(6, 4, 3, Lx=1.0, Ly=1.0, Lz=1.0)
>>> u = jnp.broadcast_to(jnp.arange(grid.Nx) * grid.dx, (grid.Nz, grid.Ny, grid.Nx))
>>> w = vertical_velocity(u, jnp.zeros_like(u), grid.dx, grid.dy, grid.dz)
>>> w.shape, float(w[0, 2, 2]), float(w[1, 2, 2])  # du/dx = 1
((5, 6, 8), 0.0, -0.333...)
Source code in finitevolx/_src/operators/diagnostics.py 
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
def vertical_velocity(
    u: Float[Array, "Nz Ny Nx"],
    v: Float[Array, "Nz Ny Nx"],
    dx: float,
    dy: float,
    dz: float,
    mask: Mask3D | Float[Array, "Nz Ny Nx"] | None = None,
) -> Float[Array, "Nz Ny Nx"]:
    """Vertical velocity from the continuity equation.

    Integrates the horizontal divergence of the physical levels upward from
    a rigid floor (``w = 0``):

        w[k+1/2] = w[k-1/2] - (du/dx + dv/dy)[k] * dz,   k = 1 .. Nz-2

    The result uses the same ``[Nz, Ny, Nx]`` convention as ``Mask3D.w``:
    ``w[k]`` is the **top face** of ``T[k]`` (``k + 1/2``).  So ``w[0]`` is
    the floor (the bottom face of the first physical level ``T[1]``, zero),
    ``w[1:-1]`` are the interfaces above the physical levels, ``w[Nz-2]`` is
    the surface (zero for a rigid lid when the column conserves mass) and
    ``w[Nz-1]`` is outside the domain (zero).  The z-ghost slabs ``k = 0``
    and ``k = Nz-1`` of ``u`` and ``v`` are never read, so BC-filled
    z-ghosts do not change the answer.

    Parameters
    ----------
    u : Float[Array, "Nz Ny Nx"]
        x-velocity at U-points, with ghost cells in all three dimensions.
    v : Float[Array, "Nz Ny Nx"]
        y-velocity at V-points, with ghost cells in all three dimensions.
    dx : float
        Grid spacing in x.
    dy : float
        Grid spacing in y.
    dz : float
        Grid spacing in z.
    mask : Mask3D or Float[Array, "Nz Ny Nx"] or None, optional
        A :class:`Mask3D`: ``u`` / ``v`` are zeroed on dry faces with
        ``jnp.where`` before the stencil (so NaN-filled land cannot leak),
        the divergence on dry T-points (``mask.h``) before the integration,
        and the output on dry W-faces (``mask.w``).  ``mask.w[Nz-2]`` (the
        surface) is dry wherever the top ghost slab is dry, so a non-zero
        rigid-lid residual there is masked out too.  A raw float T-point
        mask (deprecated) only zeroes the divergence.

    Returns
    -------
    Float[Array, "Nz Ny Nx"]
        Vertical velocity at W-points, ``w[k]`` on the top face of ``T[k]``,
        written on ``[1:-1, 1:-1, 1:-1]`` with a zero ghost ring (``w[0]``
        is the zero floor).

    Examples
    --------
    >>> grid = CartesianGrid3D.from_interior(6, 4, 3, Lx=1.0, Ly=1.0, Lz=1.0)
    >>> u = jnp.broadcast_to(jnp.arange(grid.Nx) * grid.dx, (grid.Nz, grid.Ny, grid.Nx))
    >>> w = vertical_velocity(u, jnp.zeros_like(u), grid.dx, grid.dy, grid.dz)
    >>> w.shape, float(w[0, 2, 2]), float(w[1, 2, 2])  # du/dx = 1
    ((5, 6, 8), 0.0, -0.333...)
    """
    if isinstance(mask, Mask3D):
        # zero dry faces before the stencil reads them (NaN-filled land)
        u = jnp.where(mask.u, u, 0.0)
        v = jnp.where(mask.v, v, 0.0)
    # Horizontal divergence at the interior T-points of the physical levels
    # div[k, j, i] = (u[k, j, i+1/2] - u[k, j, i-1/2]) / dx
    #              + (v[k, j+1/2, i] - v[k, j-1/2, i]) / dy,   k = 1 .. Nz-2
    div_h = diff_x_bwd_3d(u) / dx + diff_y_bwd_3d(v) / dy  # [Nz-2, Ny-2, Nx-2]

    if isinstance(mask, Mask3D):
        div_h = jnp.where(mask.h[..., 1:-1, 1:-1, 1:-1], div_h, 0.0)
    elif mask is not None:
        warnings.warn(
            "vertical_velocity(mask=<float array>) is deprecated; pass a "
            "Mask3D (its h zeroes the divergence and its w the output).",
            DeprecationWarning,
            stacklevel=2,
        )
        div_h = div_h * mask[..., 1:-1, 1:-1, 1:-1]

    # w[k+1/2] = -sum_{k'=1..k} div[k'] * dz  ->  w[k] for k = 1 .. Nz-2
    w_int = -jnp.cumsum(div_h, axis=0) * dz  # [Nz-2, Ny-2, Nx-2]
    w = interior(w_int, u)  # [Nz, Ny, Nx]; w[0] = 0 is the floor
    if isinstance(mask, Mask3D):
        w = jnp.where(mask.w, w, 0.0)
    return w