Skip to content

Vorticity Operators

Relative vorticity and Jacobian operators on Arakawa C-grids.

finitevolx.Vorticity2D

Bases: Module

2-D vorticity and PV-flux operators.

Parameters:

Name Type Description Default
grid CartesianGrid2D

The underlying 2-D grid.

 required
mask Mask2D or None

Optional land/ocean mask. When provided, both the internal Difference2D and Interpolation2D are constructed with the same mask, so every method's output inherits the correct post-compute zero via the stagger-matched mask field.

:meth:potential_vorticity is the only method that needs explicit NaN-sanitisation: under a mask, the denominator h_on_q is zero at every dry X-corner, which would trigger the "degenerate layer thickness" NaN sentinel. We preserve that sentinel for wet corners (genuine numerical bugs) and force dry corners back to exact zero — see the method body.

 None
Source code in finitevolx/_src/operators/vorticity.py 
 30
 31
 32
 33
 34
 35
 36
 37
 38
 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
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
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
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
class Vorticity2D(eqx.Module):
    """2-D vorticity and PV-flux operators.

    Parameters
    ----------
    grid : CartesianGrid2D
        The underlying 2-D grid.
    mask : Mask2D or None, optional
        Optional land/ocean mask.  When provided, both the internal
        ``Difference2D`` and ``Interpolation2D`` are constructed with
        the same mask, so every method's output inherits the correct
        post-compute zero via the stagger-matched mask field.

        :meth:`potential_vorticity` is the only method that needs
        explicit NaN-sanitisation: under a mask, the denominator
        ``h_on_q`` is zero at every dry X-corner, which would trigger
        the "degenerate layer thickness" NaN sentinel.  We preserve
        that sentinel for *wet* corners (genuine numerical bugs) and
        force dry corners back to exact zero — see the method body.
    """

    grid: CartesianGrid2D
    mask: Mask2D | None
    diff: Difference2D
    interp: Interpolation2D

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

    def relative_vorticity(
        self,
        u: Float[Array, "*batch Ny Nx"],
        v: Float[Array, "*batch Ny Nx"],
    ) -> Float[Array, "*batch Ny Nx"]:
        """Relative vorticity at X-points (corners).

        zeta[j+1/2, i+1/2] = dv_dx[j+1/2, i+1/2] - du_dy[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"]
            Relative vorticity at X-points.
        """
        return self.diff.curl(u, v)

    def potential_vorticity(
        self,
        u: Float[Array, "*batch Ny Nx"],
        v: Float[Array, "*batch Ny Nx"],
        h: Float[Array, "*batch Ny Nx"],
        f: Float[Array, "*batch Ny Nx"],
    ) -> Float[Array, "*batch Ny Nx"]:
        """Potential vorticity at X-points (corners).

        q[j+1/2, i+1/2] = (zeta[j+1/2, i+1/2] + f_on_q[j+1/2, i+1/2])
                         / h_on_q[j+1/2, i+1/2]

        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.

        Returns
        -------
        Float[Array, "*batch Ny Nx"]
            Potential vorticity at X-points.  Zero-thickness X-corners
            produce a ``NaN`` sentinel at *wet* corners (a numerical
            bug signal); under ``self.mask``, dry corners are forced
            to exact ``0`` instead so the NaN only fires where the
            user would care about it.
        """
        zeta = self.relative_vorticity(u, v)  # zeta at X-points
        f_on_q = self.interp.T_to_X(f)  # f interpolated to X-points
        h_on_q = self.interp.T_to_X(h)  # h interpolated to X-points
        # q[j+1/2, i+1/2] = (zeta + f) / h  at X-points
        num = zeta[..., 1:-1, 1:-1] + f_on_q[..., 1:-1, 1:-1]
        den = h_on_q[..., 1:-1, 1:-1]
        if self.mask is not None:
            # Pass-down masking makes h_on_q exactly 0 at every dry corner.
            # Divide by 1 there instead: the output is zeroed below anyway,
            # and this keeps reverse mode from evaluating 0 * inf at dry
            # corners.  Wet corners keep their real denominator (and the
            # NaN sentinel for zero thickness).
            # den[j, i] = den[j, i] if xy_corner_strict[j+1, i+1] else 1
            den = jnp.where(self.mask.xy_corner_strict[..., 1:-1, 1:-1], den, 1.0)
        pv = jnp.where(den == 0, jnp.nan, num / den)
        out = interior(pv, h)
        if self.mask is not None:
            # Under pass-down masking, h_on_q is zero at dry X-corners, so
            # every dry corner hit the NaN branch above — but that's a
            # mask artefact, not a degenerate-thickness bug.  Restore the
            # post-compute-zero semantic: dry corners exactly 0, wet
            # corners keep any NaN they got.
            out = jnp.where(self.mask.xy_corner_strict, out, 0.0)
        return out

    def enstrophy(
        self,
        u: Float[Array, "*batch Ny Nx"],
        v: Float[Array, "*batch Ny Nx"],
    ) -> Float[Array, "*batch Ny Nx"]:
        """Enstrophy at X-points (corners).

        ens[j+1/2, i+1/2] = 1/2 * zeta[j+1/2, i+1/2]^2

        Class form of ``enstrophy(relative_vorticity_cgrid(u, v, dx, 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"]
            Enstrophy at X-points, zero in the ghost ring and, when
            ``self.mask`` is set, at dry X-corners.  Under a mask ``u`` / ``v``
            are first zeroed on dry interior faces, so ``NaN`` land values
            poison neither the value nor its gradient.
        """
        if self.mask is not None:
            # u[j, i] = 0 on dry interior U-faces; v[j, i] = 0 on dry V-faces
            u = _sanitize(self.mask.u, u)
            v = _sanitize(self.mask.v, v)
        out = enstrophy(self.relative_vorticity(u, v))
        if self.mask is not None:
            # jnp.where, not a multiply: a NaN land velocity makes zeta NaN
            # at dry corners, and NaN * 0 is still NaN.
            out = jnp.where(self.mask.xy_corner_strict, out, 0.0)
        return out

    def potential_enstrophy(
        self,
        u: Float[Array, "*batch Ny Nx"],
        v: Float[Array, "*batch Ny Nx"],
        h: Float[Array, "*batch Ny Nx"],
        f: Float[Array, "*batch Ny Nx"],
    ) -> Float[Array, "*batch Ny Nx"]:
        """Potential enstrophy at X-points (corners).

        pens[j+1/2, i+1/2] = 1/2 * q[j+1/2, i+1/2]^2 * h_on_q[j+1/2, i+1/2]

        h_on_q[j+1/2, i+1/2] = 1/4 * (h[j, i] + h[j, i+1]
                                    + h[j+1, i] + h[j+1, i+1])

        where ``q`` is :meth:`potential_vorticity`.  Class form of
        ``potential_enstrophy(q, h_on_q)``.

        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.

        Returns
        -------
        Float[Array, "*batch Ny Nx"]
            Potential enstrophy at X-points, zero in the ghost ring and,
            when ``self.mask`` is set, at dry X-corners.  Inherits the
            ``NaN`` sentinel of :meth:`potential_vorticity` at wet corners
            with zero thickness.
        """
        if self.mask is not None:
            # Zero dry interior inputs so NaN land values cannot enter the
            # arithmetic (or its gradient); wet corners only read wet cells.
            u = _sanitize(self.mask.u, u)
            v = _sanitize(self.mask.v, v)
            f = _sanitize(self.mask.h, f)
            h = _sanitize(self.mask.h, h)
        q = self.potential_vorticity(u, v, h, f)  # q at X-points
        h_on_q = self.interp.T_to_X(h)  # h interpolated to X-points
        out = potential_enstrophy(q, h_on_q)
        if self.mask is not None:
            # q is already exactly 0 at dry corners; the select keeps that
            # guarantee explicit, mirroring potential_vorticity.
            out = jnp.where(self.mask.xy_corner_strict, out, 0.0)
        return out

    def pv_flux_energy_conserving(
        self,
        q: Float[Array, "*batch Ny Nx"],
        u: Float[Array, "*batch Ny Nx"],
        v: Float[Array, "*batch Ny Nx"],
    ) -> tuple:
        """Energy-conserving PV flux (Sadourny 1975 E-scheme form).

        Multiply q by the velocity at the X-points (corners), where q lives,
        then average the product to the faces:

        qu[j, i+1/2] = 1/2 * ( q[j+1/2, i+1/2] * u_on_q[j+1/2, i+1/2]
                             + q[j-1/2, i+1/2] * u_on_q[j-1/2, i+1/2] )
        qv[j+1/2, i] = 1/2 * ( q[j+1/2, i+1/2] * v_on_q[j+1/2, i+1/2]
                             + q[j+1/2, i-1/2] * v_on_q[j+1/2, i-1/2] )

        with ``u_on_q = U_to_X(u)`` and ``v_on_q = V_to_X(v)``.  The corner
        products are formed on the full arrays, so the first interior row /
        column reads the south / west ghost corners of the caller's
        ghost-filled ``q``, ``u``, ``v`` (periodic, free-slip, open, ...).

        The method multiplies q by whatever face fields it is given, so the
        name refers to the vortex-force use: with the cross-face transport
        (``du = +q·v``, ``dv = -q·u``, as :class:`MomentumAdvection2D` forms
        it) this corner-product form conserves total energy for
        non-divergent flow (Sadourny 1975), while the average-then-multiply
        form (:meth:`pv_flux_enstrophy_conserving`) conserves potential
        enstrophy.  Passed ``(q, u, v)`` directly it is a PV *transport*
        flux (``q·u``, ``q·v``) and neither label applies.

        Parameters
        ----------
        q : Float[Array, "*batch Ny Nx"]
            Potential vorticity at X-points, including the BC-owned south /
            west ghost corners.
        u : Float[Array, "*batch Ny Nx"]
            x-velocity (or x mass flux) at U-points.
        v : Float[Array, "*batch Ny Nx"]
            y-velocity (or y mass flux) at V-points.

        Returns
        -------
        tuple
            (qu at U-points, qv at V-points), zero ghost ring.  Under a
            mask the corner products are zeroed on dry *interior* corners
            and the outputs multiplied by ``mask.u`` / ``mask.v``.
        """
        mask = self.mask
        return _pv_flux_corner_product(
            q,
            u,
            v,
            wet_x=None if mask is None else mask.xy_corner_strict,
            wet_u=None if mask is None else mask.u,
            wet_v=None if mask is None else mask.v,
        )

    def pv_flux_enstrophy_conserving(
        self,
        q: Float[Array, "*batch Ny Nx"],
        u: Float[Array, "*batch Ny Nx"],
        v: Float[Array, "*batch Ny Nx"],
    ) -> tuple:
        """Enstrophy-conserving PV flux (Sadourny 1975 Z-scheme form).

        Average q to the faces, then multiply by the face field:

        qu[j, i+1/2] = q_on_u[j, i+1/2] * u[j, i+1/2]
        qv[j+1/2, i] = q_on_v[j+1/2, i] * v[j+1/2, i]

        with ``q_on_u = X_to_U(q)`` and ``q_on_v = X_to_V(q)``.  As for
        :meth:`pv_flux_energy_conserving`, the name refers to the
        vortex-force use (cross-face transport averaged to each face, as
        :class:`MomentumAdvection2D` does): this form then conserves
        potential enstrophy for non-divergent flow (Sadourny 1975).

        Parameters
        ----------
        q : Float[Array, "*batch Ny Nx"]
            Potential vorticity at X-points.
        u : Float[Array, "*batch Ny Nx"]
            x-velocity (or x mass flux) at U-points.
        v : Float[Array, "*batch Ny Nx"]
            y-velocity (or y mass flux) at V-points.

        Returns
        -------
        tuple
            (qu at U-points, qv at V-points)
        """
        q_on_u = self.interp.X_to_U(q)  # q_on_u[j, i+1/2] = avg in y
        q_on_v = self.interp.X_to_V(q)  # q_on_v[j+1/2, i] = avg in x
        # qu[j, i+1/2] = q_on_u[j, i+1/2] * u[j, i+1/2]
        qu = interior(q_on_u[..., 1:-1, 1:-1] * u[..., 1:-1, 1:-1], u)
        # qv[j+1/2, i] = q_on_v[j+1/2, i] * v[j+1/2, i]
        qv = interior(q_on_v[..., 1:-1, 1:-1] * v[..., 1:-1, 1:-1], v)
        return qu, qv

    def pv_flux_blend(
        self,
        q: Float[Array, "*batch Ny Nx"],
        u: Float[Array, "*batch Ny Nx"],
        v: Float[Array, "*batch Ny Nx"],
        alpha: float = 1.0 / 3.0,
    ) -> tuple:
        """Linear blend of the energy- and enstrophy-conserving PV fluxes.

        flux = alpha * energy_conserving + (1 - alpha) * enstrophy_conserving

        For 0 < alpha < 1 the blend conserves *neither* energy nor
        enstrophy exactly.  It is not the Arakawa-Lamb (1981) scheme (a
        distinct eight-point stencil, not implemented); the old name
        :meth:`pv_flux_arakawa_lamb` is a deprecated alias (#283).

        Parameters
        ----------
        q : Float[Array, "*batch Ny Nx"]
            Potential vorticity at X-points.
        u : Float[Array, "*batch Ny Nx"]
            x-velocity at U-points.
        v : Float[Array, "*batch Ny Nx"]
            y-velocity at V-points.
        alpha : float
            Weight of the energy-conserving form.  Default 1/3.

        Returns
        -------
        tuple
            (qu at U-points, qv at V-points)
        """
        qu_e, qv_e = self.pv_flux_energy_conserving(q, u, v)
        qu_s, qv_s = self.pv_flux_enstrophy_conserving(q, u, v)
        # Weighted blend
        qu = alpha * qu_e + (1.0 - alpha) * qu_s
        qv = alpha * qv_e + (1.0 - alpha) * qv_s
        return qu, qv

    def pv_flux_arakawa_lamb(
        self,
        q: Float[Array, "*batch Ny Nx"],
        u: Float[Array, "*batch Ny Nx"],
        v: Float[Array, "*batch Ny Nx"],
        alpha: float = 1.0 / 3.0,
    ) -> tuple:
        """Deprecated alias of :meth:`pv_flux_blend` (not Arakawa-Lamb 1981)."""
        _warn_arakawa_lamb("Vorticity2D.pv_flux_blend")
        return self.pv_flux_blend(q, u, v, alpha=alpha)

enstrophy(u, v)

Enstrophy at X-points (corners).

ens[j+1/2, i+1/2] = 1/2 * zeta[j+1/2, i+1/2]^2

Class form of enstrophy(relative_vorticity_cgrid(u, v, dx, 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']

Enstrophy at X-points, zero in the ghost ring and, when self.mask is set, at dry X-corners. Under a mask u / v are first zeroed on dry interior faces, so NaN land values poison neither the value nor its gradient.
Source code in finitevolx/_src/operators/vorticity.py 
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
def enstrophy(
    self,
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
) -> Float[Array, "*batch Ny Nx"]:
    """Enstrophy at X-points (corners).

    ens[j+1/2, i+1/2] = 1/2 * zeta[j+1/2, i+1/2]^2

    Class form of ``enstrophy(relative_vorticity_cgrid(u, v, dx, 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"]
        Enstrophy at X-points, zero in the ghost ring and, when
        ``self.mask`` is set, at dry X-corners.  Under a mask ``u`` / ``v``
        are first zeroed on dry interior faces, so ``NaN`` land values
        poison neither the value nor its gradient.
    """
    if self.mask is not None:
        # u[j, i] = 0 on dry interior U-faces; v[j, i] = 0 on dry V-faces
        u = _sanitize(self.mask.u, u)
        v = _sanitize(self.mask.v, v)
    out = enstrophy(self.relative_vorticity(u, v))
    if self.mask is not None:
        # jnp.where, not a multiply: a NaN land velocity makes zeta NaN
        # at dry corners, and NaN * 0 is still NaN.
        out = jnp.where(self.mask.xy_corner_strict, out, 0.0)
    return out

potential_enstrophy(u, v, h, f)

Potential enstrophy at X-points (corners).

pens[j+1/2, i+1/2] = 1/2 * q[j+1/2, i+1/2]^2 * h_on_q[j+1/2, i+1/2]

h_on_q[j+1/2, i+1/2] = 1/4 * (h[j, i] + h[j, i+1] + h[j+1, i] + h[j+1, i+1])

where q is :meth:potential_vorticity. Class form of potential_enstrophy(q, h_on_q).

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

Returns:

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

Potential enstrophy at X-points, zero in the ghost ring and, when self.mask is set, at dry X-corners. Inherits the NaN sentinel of :meth:potential_vorticity at wet corners with zero thickness.
Source code in finitevolx/_src/operators/vorticity.py 
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
232
233
234
235
def potential_enstrophy(
    self,
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
    h: Float[Array, "*batch Ny Nx"],
    f: Float[Array, "*batch Ny Nx"],
) -> Float[Array, "*batch Ny Nx"]:
    """Potential enstrophy at X-points (corners).

    pens[j+1/2, i+1/2] = 1/2 * q[j+1/2, i+1/2]^2 * h_on_q[j+1/2, i+1/2]

    h_on_q[j+1/2, i+1/2] = 1/4 * (h[j, i] + h[j, i+1]
                                + h[j+1, i] + h[j+1, i+1])

    where ``q`` is :meth:`potential_vorticity`.  Class form of
    ``potential_enstrophy(q, h_on_q)``.

    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.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Potential enstrophy at X-points, zero in the ghost ring and,
        when ``self.mask`` is set, at dry X-corners.  Inherits the
        ``NaN`` sentinel of :meth:`potential_vorticity` at wet corners
        with zero thickness.
    """
    if self.mask is not None:
        # Zero dry interior inputs so NaN land values cannot enter the
        # arithmetic (or its gradient); wet corners only read wet cells.
        u = _sanitize(self.mask.u, u)
        v = _sanitize(self.mask.v, v)
        f = _sanitize(self.mask.h, f)
        h = _sanitize(self.mask.h, h)
    q = self.potential_vorticity(u, v, h, f)  # q at X-points
    h_on_q = self.interp.T_to_X(h)  # h interpolated to X-points
    out = potential_enstrophy(q, h_on_q)
    if self.mask is not None:
        # q is already exactly 0 at dry corners; the select keeps that
        # guarantee explicit, mirroring potential_vorticity.
        out = jnp.where(self.mask.xy_corner_strict, out, 0.0)
    return out

potential_vorticity(u, v, h, f)

Potential vorticity at X-points (corners).

q[j+1/2, i+1/2] = (zeta[j+1/2, i+1/2] + f_on_q[j+1/2, i+1/2]) / h_on_q[j+1/2, i+1/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
h Float[Array, '*batch Ny Nx']

Layer thickness at T-points.

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

Coriolis parameter at T-points.

 required

Returns:

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

Potential vorticity at X-points. Zero-thickness X-corners produce a NaN sentinel at wet corners (a numerical bug signal); under self.mask, dry corners are forced to exact 0 instead so the NaN only fires where the user would care about it.
Source code in finitevolx/_src/operators/vorticity.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
135
136
137
138
139
140
141
142
143
144
145
146
def potential_vorticity(
    self,
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
    h: Float[Array, "*batch Ny Nx"],
    f: Float[Array, "*batch Ny Nx"],
) -> Float[Array, "*batch Ny Nx"]:
    """Potential vorticity at X-points (corners).

    q[j+1/2, i+1/2] = (zeta[j+1/2, i+1/2] + f_on_q[j+1/2, i+1/2])
                     / h_on_q[j+1/2, i+1/2]

    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.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Potential vorticity at X-points.  Zero-thickness X-corners
        produce a ``NaN`` sentinel at *wet* corners (a numerical
        bug signal); under ``self.mask``, dry corners are forced
        to exact ``0`` instead so the NaN only fires where the
        user would care about it.
    """
    zeta = self.relative_vorticity(u, v)  # zeta at X-points
    f_on_q = self.interp.T_to_X(f)  # f interpolated to X-points
    h_on_q = self.interp.T_to_X(h)  # h interpolated to X-points
    # q[j+1/2, i+1/2] = (zeta + f) / h  at X-points
    num = zeta[..., 1:-1, 1:-1] + f_on_q[..., 1:-1, 1:-1]
    den = h_on_q[..., 1:-1, 1:-1]
    if self.mask is not None:
        # Pass-down masking makes h_on_q exactly 0 at every dry corner.
        # Divide by 1 there instead: the output is zeroed below anyway,
        # and this keeps reverse mode from evaluating 0 * inf at dry
        # corners.  Wet corners keep their real denominator (and the
        # NaN sentinel for zero thickness).
        # den[j, i] = den[j, i] if xy_corner_strict[j+1, i+1] else 1
        den = jnp.where(self.mask.xy_corner_strict[..., 1:-1, 1:-1], den, 1.0)
    pv = jnp.where(den == 0, jnp.nan, num / den)
    out = interior(pv, h)
    if self.mask is not None:
        # Under pass-down masking, h_on_q is zero at dry X-corners, so
        # every dry corner hit the NaN branch above — but that's a
        # mask artefact, not a degenerate-thickness bug.  Restore the
        # post-compute-zero semantic: dry corners exactly 0, wet
        # corners keep any NaN they got.
        out = jnp.where(self.mask.xy_corner_strict, out, 0.0)
    return out

pv_flux_arakawa_lamb(q, u, v, alpha=1.0 / 3.0)

Deprecated alias of :meth:pv_flux_blend (not Arakawa-Lamb 1981).

Source code in finitevolx/_src/operators/vorticity.py 
374
375
376
377
378
379
380
381
382
383
def pv_flux_arakawa_lamb(
    self,
    q: Float[Array, "*batch Ny Nx"],
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
    alpha: float = 1.0 / 3.0,
) -> tuple:
    """Deprecated alias of :meth:`pv_flux_blend` (not Arakawa-Lamb 1981)."""
    _warn_arakawa_lamb("Vorticity2D.pv_flux_blend")
    return self.pv_flux_blend(q, u, v, alpha=alpha)

pv_flux_blend(q, u, v, alpha=1.0 / 3.0)

Linear blend of the energy- and enstrophy-conserving PV fluxes.

flux = alpha * energy_conserving + (1 - alpha) * enstrophy_conserving

For 0 < alpha < 1 the blend conserves neither energy nor enstrophy exactly. It is not the Arakawa-Lamb (1981) scheme (a distinct eight-point stencil, not implemented); the old name :meth:pv_flux_arakawa_lamb is a deprecated alias (#283).

Parameters:

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

Potential vorticity at X-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
alpha float

Weight of the energy-conserving form. Default 1/3.

 1.0 / 3.0

Returns:

Type Description
tuple

(qu at U-points, qv at V-points)
Source code in finitevolx/_src/operators/vorticity.py 
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
def pv_flux_blend(
    self,
    q: Float[Array, "*batch Ny Nx"],
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
    alpha: float = 1.0 / 3.0,
) -> tuple:
    """Linear blend of the energy- and enstrophy-conserving PV fluxes.

    flux = alpha * energy_conserving + (1 - alpha) * enstrophy_conserving

    For 0 < alpha < 1 the blend conserves *neither* energy nor
    enstrophy exactly.  It is not the Arakawa-Lamb (1981) scheme (a
    distinct eight-point stencil, not implemented); the old name
    :meth:`pv_flux_arakawa_lamb` is a deprecated alias (#283).

    Parameters
    ----------
    q : Float[Array, "*batch Ny Nx"]
        Potential vorticity at X-points.
    u : Float[Array, "*batch Ny Nx"]
        x-velocity at U-points.
    v : Float[Array, "*batch Ny Nx"]
        y-velocity at V-points.
    alpha : float
        Weight of the energy-conserving form.  Default 1/3.

    Returns
    -------
    tuple
        (qu at U-points, qv at V-points)
    """
    qu_e, qv_e = self.pv_flux_energy_conserving(q, u, v)
    qu_s, qv_s = self.pv_flux_enstrophy_conserving(q, u, v)
    # Weighted blend
    qu = alpha * qu_e + (1.0 - alpha) * qu_s
    qv = alpha * qv_e + (1.0 - alpha) * qv_s
    return qu, qv

pv_flux_energy_conserving(q, u, v)

Energy-conserving PV flux (Sadourny 1975 E-scheme form).

Multiply q by the velocity at the X-points (corners), where q lives, then average the product to the faces:

qu[j, i+1/2] = 1/2 * ( q[j+1/2, i+1/2] * u_on_q[j+1/2, i+1/2] + q[j-1/2, i+1/2] * u_on_q[j-1/2, i+1/2] ) qv[j+1/2, i] = 1/2 * ( q[j+1/2, i+1/2] * v_on_q[j+1/2, i+1/2] + q[j+1/2, i-1/2] * v_on_q[j+1/2, i-1/2] )

with u_on_q = U_to_X(u) and v_on_q = V_to_X(v). The corner products are formed on the full arrays, so the first interior row / column reads the south / west ghost corners of the caller's ghost-filled q, u, v (periodic, free-slip, open, ...).

The method multiplies q by whatever face fields it is given, so the name refers to the vortex-force use: with the cross-face transport (du = +q·v, dv = -q·u, as :class:MomentumAdvection2D forms it) this corner-product form conserves total energy for non-divergent flow (Sadourny 1975), while the average-then-multiply form (:meth:pv_flux_enstrophy_conserving) conserves potential enstrophy. Passed (q, u, v) directly it is a PV transport flux (q·u, q·v) and neither label applies.

Parameters:

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

Potential vorticity at X-points, including the BC-owned south / west ghost corners.

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

x-velocity (or x mass flux) at U-points.

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

y-velocity (or y mass flux) at V-points.

 required

Returns:

Type Description
tuple

(qu at U-points, qv at V-points), zero ghost ring. Under a mask the corner products are zeroed on dry interior corners and the outputs multiplied by mask.u / mask.v.
Source code in finitevolx/_src/operators/vorticity.py 
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
def pv_flux_energy_conserving(
    self,
    q: Float[Array, "*batch Ny Nx"],
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
) -> tuple:
    """Energy-conserving PV flux (Sadourny 1975 E-scheme form).

    Multiply q by the velocity at the X-points (corners), where q lives,
    then average the product to the faces:

    qu[j, i+1/2] = 1/2 * ( q[j+1/2, i+1/2] * u_on_q[j+1/2, i+1/2]
                         + q[j-1/2, i+1/2] * u_on_q[j-1/2, i+1/2] )
    qv[j+1/2, i] = 1/2 * ( q[j+1/2, i+1/2] * v_on_q[j+1/2, i+1/2]
                         + q[j+1/2, i-1/2] * v_on_q[j+1/2, i-1/2] )

    with ``u_on_q = U_to_X(u)`` and ``v_on_q = V_to_X(v)``.  The corner
    products are formed on the full arrays, so the first interior row /
    column reads the south / west ghost corners of the caller's
    ghost-filled ``q``, ``u``, ``v`` (periodic, free-slip, open, ...).

    The method multiplies q by whatever face fields it is given, so the
    name refers to the vortex-force use: with the cross-face transport
    (``du = +q·v``, ``dv = -q·u``, as :class:`MomentumAdvection2D` forms
    it) this corner-product form conserves total energy for
    non-divergent flow (Sadourny 1975), while the average-then-multiply
    form (:meth:`pv_flux_enstrophy_conserving`) conserves potential
    enstrophy.  Passed ``(q, u, v)`` directly it is a PV *transport*
    flux (``q·u``, ``q·v``) and neither label applies.

    Parameters
    ----------
    q : Float[Array, "*batch Ny Nx"]
        Potential vorticity at X-points, including the BC-owned south /
        west ghost corners.
    u : Float[Array, "*batch Ny Nx"]
        x-velocity (or x mass flux) at U-points.
    v : Float[Array, "*batch Ny Nx"]
        y-velocity (or y mass flux) at V-points.

    Returns
    -------
    tuple
        (qu at U-points, qv at V-points), zero ghost ring.  Under a
        mask the corner products are zeroed on dry *interior* corners
        and the outputs multiplied by ``mask.u`` / ``mask.v``.
    """
    mask = self.mask
    return _pv_flux_corner_product(
        q,
        u,
        v,
        wet_x=None if mask is None else mask.xy_corner_strict,
        wet_u=None if mask is None else mask.u,
        wet_v=None if mask is None else mask.v,
    )

pv_flux_enstrophy_conserving(q, u, v)

Enstrophy-conserving PV flux (Sadourny 1975 Z-scheme form).

Average q to the faces, then multiply by the face field:

qu[j, i+1/2] = q_on_u[j, i+1/2] * u[j, i+1/2] qv[j+1/2, i] = q_on_v[j+1/2, i] * v[j+1/2, i]

with q_on_u = X_to_U(q) and q_on_v = X_to_V(q). As for :meth:pv_flux_energy_conserving, the name refers to the vortex-force use (cross-face transport averaged to each face, as :class:MomentumAdvection2D does): this form then conserves potential enstrophy for non-divergent flow (Sadourny 1975).

Parameters:

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

Potential vorticity at X-points.

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

x-velocity (or x mass flux) at U-points.

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

y-velocity (or y mass flux) at V-points.

 required

Returns:

Type Description
tuple

(qu at U-points, qv at V-points)
Source code in finitevolx/_src/operators/vorticity.py 
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
def pv_flux_enstrophy_conserving(
    self,
    q: Float[Array, "*batch Ny Nx"],
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
) -> tuple:
    """Enstrophy-conserving PV flux (Sadourny 1975 Z-scheme form).

    Average q to the faces, then multiply by the face field:

    qu[j, i+1/2] = q_on_u[j, i+1/2] * u[j, i+1/2]
    qv[j+1/2, i] = q_on_v[j+1/2, i] * v[j+1/2, i]

    with ``q_on_u = X_to_U(q)`` and ``q_on_v = X_to_V(q)``.  As for
    :meth:`pv_flux_energy_conserving`, the name refers to the
    vortex-force use (cross-face transport averaged to each face, as
    :class:`MomentumAdvection2D` does): this form then conserves
    potential enstrophy for non-divergent flow (Sadourny 1975).

    Parameters
    ----------
    q : Float[Array, "*batch Ny Nx"]
        Potential vorticity at X-points.
    u : Float[Array, "*batch Ny Nx"]
        x-velocity (or x mass flux) at U-points.
    v : Float[Array, "*batch Ny Nx"]
        y-velocity (or y mass flux) at V-points.

    Returns
    -------
    tuple
        (qu at U-points, qv at V-points)
    """
    q_on_u = self.interp.X_to_U(q)  # q_on_u[j, i+1/2] = avg in y
    q_on_v = self.interp.X_to_V(q)  # q_on_v[j+1/2, i] = avg in x
    # qu[j, i+1/2] = q_on_u[j, i+1/2] * u[j, i+1/2]
    qu = interior(q_on_u[..., 1:-1, 1:-1] * u[..., 1:-1, 1:-1], u)
    # qv[j+1/2, i] = q_on_v[j+1/2, i] * v[j+1/2, i]
    qv = interior(q_on_v[..., 1:-1, 1:-1] * v[..., 1:-1, 1:-1], v)
    return qu, qv

relative_vorticity(u, v)

Relative vorticity at X-points (corners).

zeta[j+1/2, i+1/2] = dv_dx[j+1/2, i+1/2] - du_dy[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']

Relative vorticity at X-points.
Source code in finitevolx/_src/operators/vorticity.py 
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
def relative_vorticity(
    self,
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
) -> Float[Array, "*batch Ny Nx"]:
    """Relative vorticity at X-points (corners).

    zeta[j+1/2, i+1/2] = dv_dx[j+1/2, i+1/2] - du_dy[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"]
        Relative vorticity at X-points.
    """
    return self.diff.curl(u, v)

finitevolx.Vorticity3D

Bases: Module

3-D vorticity operators (horizontal plane per z-level).

Parameters:

Name Type Description Default
grid CartesianGrid3D
required
mask Mask3D or None

Optional land/ocean mask. When provided, the result of :meth:relative_vorticity is post-multiplied by mask.xy_corner_strict. Applied externally rather than pass-down because the underlying _curl_2d is a free function (no sub-operator to inject a mask into).

 required
Source code in finitevolx/_src/operators/vorticity.py 
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
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
class Vorticity3D(eqx.Module):
    """3-D vorticity operators (horizontal plane per z-level).

    Parameters
    ----------
    grid : CartesianGrid3D
    mask : Mask3D or None, optional
        Optional land/ocean mask.  When provided, the result of
        :meth:`relative_vorticity` is post-multiplied by
        ``mask.xy_corner_strict``.  Applied externally rather than
        pass-down because the underlying ``_curl_2d`` is a free
        function (no sub-operator to inject a mask into).
    """

    grid: CartesianGrid3D
    mask: Mask3D | None = None

    def relative_vorticity(
        self,
        u: Float[Array, "Nz Ny Nx"],
        v: Float[Array, "Nz Ny Nx"],
    ) -> Float[Array, "Nz Ny Nx"]:
        """Relative vorticity at X-points over all z-levels.

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

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

        Returns
        -------
        Float[Array, "Nz Ny Nx"]
            Relative vorticity at X-points.  When ``self.mask`` is
            set, dry X-corners are zero via ``* mask.xy_corner_strict``.
        """
        out = eqx.filter_vmap(
            lambda u_k, v_k: _curl_2d(u_k, v_k, self.grid.dx, self.grid.dy)
        )(u, v)
        # Zero z-ghost slices to match 3D ghost-ring convention.
        out = zero_z_ghosts(out)
        if self.mask is not None:
            out = out * self.mask.xy_corner_strict
        return out

relative_vorticity(u, v)

Relative vorticity at X-points over all z-levels.

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

Parameters:

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

x-velocity at U-points.

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

y-velocity at V-points.

 required

Returns:

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

Relative vorticity at X-points. When self.mask is set, dry X-corners are zero via * mask.xy_corner_strict.
Source code in finitevolx/_src/operators/vorticity.py 
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
def relative_vorticity(
    self,
    u: Float[Array, "Nz Ny Nx"],
    v: Float[Array, "Nz Ny Nx"],
) -> Float[Array, "Nz Ny Nx"]:
    """Relative vorticity at X-points over all z-levels.

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

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

    Returns
    -------
    Float[Array, "Nz Ny Nx"]
        Relative vorticity at X-points.  When ``self.mask`` is
        set, dry X-corners are zero via ``* mask.xy_corner_strict``.
    """
    out = eqx.filter_vmap(
        lambda u_k, v_k: _curl_2d(u_k, v_k, self.grid.dx, self.grid.dy)
    )(u, v)
    # Zero z-ghost slices to match 3D ghost-ring convention.
    out = zero_z_ghosts(out)
    if self.mask is not None:
        out = out * self.mask.xy_corner_strict
    return out

PV flux (functional)

pv_flux_blend is the grid-free form of Vorticity2D.pv_flux_blend. pv_flux_arakawa_lamb is its deprecated alias: the blend is not the Arakawa–Lamb (1981) scheme.

finitevolx.pv_flux_blend(q, u, v, *, alpha=1.0 / 3.0)

Blend of the energy- and enstrophy-conserving PV fluxes (grid-free).

Public functional form of :meth:Vorticity2D.pv_flux_blend:

flux = alpha * energy_conserving + (1 - alpha) * enstrophy_conserving

where the energy-conserving form multiplies q by the velocity at the X-points and averages the product to the faces, and the enstrophy-conserving form averages q to the faces and multiplies there (see :meth:Vorticity2D.pv_flux_energy_conserving / :meth:Vorticity2D.pv_flux_enstrophy_conserving). It reproduces the unmasked :class:Vorticity2D method bit for bit using only the C-grid averaging stencils, so it is reachable without constructing a grid. For masked domains, use :class:Vorticity2D with a mask=.

For 0 < alpha < 1 the blend conserves neither energy nor enstrophy exactly. It is not the Arakawa-Lamb (1981) scheme (a distinct eight-point stencil, not implemented); :func:pv_flux_arakawa_lamb is a deprecated alias (#283).

Parameters:

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

Potential vorticity at X-points (corners), including ghost ring.

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

x-velocity (or x mass flux) at U-points.

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

y-velocity (or y mass flux) at V-points.

 required
alpha float

Weight of the energy-conserving form. Default 1/3.

 1.0 / 3.0

Returns:

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

(qu, qv) — the PV flux at U-points and V-points.
Source code in finitevolx/_src/operators/vorticity.py 
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
def pv_flux_blend(
    q: Float[Array, "*batch Ny Nx"],
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
    *,
    alpha: float = 1.0 / 3.0,
) -> tuple[Float[Array, "*batch Ny Nx"], Float[Array, "*batch Ny Nx"]]:
    """Blend of the energy- and enstrophy-conserving PV fluxes (grid-free).

    Public functional form of :meth:`Vorticity2D.pv_flux_blend`:

        flux = alpha * energy_conserving + (1 - alpha) * enstrophy_conserving

    where the energy-conserving form multiplies q by the velocity at the
    X-points and averages the product to the faces, and the
    enstrophy-conserving form averages q to the faces and multiplies there
    (see :meth:`Vorticity2D.pv_flux_energy_conserving` /
    :meth:`Vorticity2D.pv_flux_enstrophy_conserving`).  It reproduces the
    unmasked :class:`Vorticity2D` method bit for bit using only the C-grid
    averaging stencils, so it is reachable without constructing a grid.
    For masked domains, use :class:`Vorticity2D` with a ``mask=``.

    For 0 < alpha < 1 the blend conserves *neither* energy nor enstrophy
    exactly.  It is not the Arakawa-Lamb (1981) scheme (a distinct
    eight-point stencil, not implemented); :func:`pv_flux_arakawa_lamb` is
    a deprecated alias (#283).

    Parameters
    ----------
    q : Float[Array, "*batch Ny Nx"]
        Potential vorticity at X-points (corners), including ghost ring.
    u : Float[Array, "*batch Ny Nx"]
        x-velocity (or x mass flux) at U-points.
    v : Float[Array, "*batch Ny Nx"]
        y-velocity (or y mass flux) at V-points.
    alpha : float, optional
        Weight of the energy-conserving form.  Default ``1/3``.

    Returns
    -------
    tuple[Float[Array, "*batch Ny Nx"], Float[Array, "*batch Ny Nx"]]
        ``(qu, qv)`` — the PV flux at U-points and V-points.
    """
    # Energy-conserving: multiply at corners (full arrays), then average.
    qu_e, qv_e = _pv_flux_corner_product(q, u, v)

    # Enstrophy-conserving: average q to the faces, then multiply.
    q_on_u = interior(avg_y_bwd(q), q)  # X_to_U
    q_on_v = interior(avg_x_bwd(q), q)  # X_to_V
    qu_s = interior(q_on_u[..., 1:-1, 1:-1] * u[..., 1:-1, 1:-1], u)
    qv_s = interior(q_on_v[..., 1:-1, 1:-1] * v[..., 1:-1, 1:-1], v)

    qu = alpha * qu_e + (1.0 - alpha) * qu_s
    qv = alpha * qv_e + (1.0 - alpha) * qv_s
    return qu, qv

Jacobian

ArakawaJacobian2D is the mask-aware class form: it returns the full [..., Ny, Nx] T-point array with a zero ghost ring. Under a mask, f and g are treated as zero on dry T-cells (so land NaNs cannot reach the stencil) and dry output cells are exactly zero. The functional arakawa_jacobian returns only the interior [..., Ny-2, Nx-2], takes no mask, and reads land values as given.

finitevolx.ArakawaJacobian2D

Bases: Module

Arakawa (1966) Jacobian J(f, g) at T-points, on the full grid.

Class form of :func:arakawa_jacobian. Unlike the function, which returns only the interior (..., Ny-2, Nx-2), the class returns the full (..., Ny, Nx) array with a zero ghost ring, like every other Layer-3 operator:

J_full[..., j, i] = J[..., j-1, i-1] for 1 <= j <= Ny-2, 1 <= i <= Nx-2 J_full[..., j, i] = 0 on the ghost ring

When mask is set, f and g are first zeroed on dry interior T-cells (the stencil reads the eight neighbours of each T-cell, so a coastal cell would otherwise read land values, NaN or not; a streamfunction from a masked elliptic solve is zero there anyway; the ghost ring is BC-owned and passed through), and dry output cells are zeroed last via jnp.where(mask.h, ...) (Pattern 1 in docs/masks.md), broadcast over any leading batch axes.

Masking zeroes the Jacobian at dry cells, so the discrete identities sum(J) = 0 and sum(g * J) = 0 are not guaranteed for sums over a masked domain -- they hold for the unmasked operator with suitable boundary conditions.

Parameters:

Name Type Description Default
grid CartesianGrid2D

The underlying 2-D grid (supplies dx and dy).

 required
mask Mask2D or None

Optional land/ocean mask. None (default) returns the functional form's output padded with a zero ghost ring.

 required

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import ArakawaJacobian2D, CartesianGrid2D
>>> grid = CartesianGrid2D.from_interior(8, 8, 1.0, 1.0)
>>> psi = jnp.zeros((grid.Ny, grid.Nx))
>>> ArakawaJacobian2D(grid=grid)(psi, psi).shape
(10, 10)
Source code in finitevolx/_src/operators/jacobian.py 
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
class ArakawaJacobian2D(eqx.Module):
    """Arakawa (1966) Jacobian J(f, g) at T-points, on the full grid.

    Class form of :func:`arakawa_jacobian`.  Unlike the function, which
    returns only the interior ``(..., Ny-2, Nx-2)``, the class returns the
    full ``(..., Ny, Nx)`` array with a zero ghost ring, like every other
    Layer-3 operator:

    J_full[..., j, i] = J[..., j-1, i-1]   for 1 <= j <= Ny-2, 1 <= i <= Nx-2
    J_full[..., j, i] = 0                  on the ghost ring

    When ``mask`` is set, ``f`` and ``g`` are first zeroed on dry *interior*
    T-cells (the stencil reads the eight neighbours of each T-cell, so a
    coastal cell would otherwise read land values, ``NaN`` or not; a
    streamfunction from a masked elliptic solve is zero there anyway; the
    ghost ring is BC-owned and passed through), and dry output cells
    are zeroed last via ``jnp.where(mask.h, ...)`` (Pattern 1 in
    ``docs/masks.md``), broadcast over any leading batch axes.

    Masking zeroes the Jacobian at dry cells, so the discrete identities
    ``sum(J) = 0`` and ``sum(g * J) = 0`` are **not** guaranteed for sums
    over a masked domain -- they hold for the unmasked operator with
    suitable boundary conditions.

    Parameters
    ----------
    grid : CartesianGrid2D
        The underlying 2-D grid (supplies ``dx`` and ``dy``).
    mask : Mask2D or None, optional
        Optional land/ocean mask.  ``None`` (default) returns the functional
        form's output padded with a zero ghost ring.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import ArakawaJacobian2D, CartesianGrid2D
    >>> grid = CartesianGrid2D.from_interior(8, 8, 1.0, 1.0)
    >>> psi = jnp.zeros((grid.Ny, grid.Nx))
    >>> ArakawaJacobian2D(grid=grid)(psi, psi).shape
    (10, 10)
    """

    grid: CartesianGrid2D
    mask: Mask2D | None = None

    def __call__(
        self,
        f: Float[Array, "... Ny Nx"],
        g: Float[Array, "... Ny Nx"],
    ) -> Float[Array, "... Ny Nx"]:
        """Evaluate J(f, g) at T-points.

        Parameters
        ----------
        f : Float[Array, "... Ny Nx"]
            First scalar field at T-points (e.g. the streamfunction).
        g : Float[Array, "... Ny Nx"]
            Second scalar field at T-points (e.g. the potential vorticity).

        Returns
        -------
        Float[Array, "... Ny Nx"]
            J(f, g) at T-points, zero in the ghost ring and, when
            ``self.mask`` is set, at dry T-cells.
        """
        if self.mask is not None:
            # f[..., j, i] = 0 on dry interior T-cells (same for g)
            f = _sanitize(self.mask.h, f)
            g = _sanitize(self.mask.h, g)
        J = arakawa_jacobian(f, g, self.grid.dx, self.grid.dy)
        shape = jnp.broadcast_shapes(f.shape, g.shape)
        # J_full[..., j, i] = J[..., j-1, i-1]  for 1 <= j <= Ny-2, 1 <= i <= Nx-2
        out = jnp.zeros_like(f, dtype=J.dtype, shape=shape)
        out = out.at[..., 1:-1, 1:-1].set(J)
        if self.mask is not None:
            out = jnp.where(self.mask.h, out, 0.0)
        return out

__call__(f, g)

Evaluate J(f, g) at T-points.

Parameters:

Name Type Description Default
f Float[Array, '... Ny Nx']

First scalar field at T-points (e.g. the streamfunction).

 required
g Float[Array, '... Ny Nx']

Second scalar field at T-points (e.g. the potential vorticity).

 required

Returns:

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

J(f, g) at T-points, zero in the ghost ring and, when self.mask is set, at dry T-cells.
Source code in finitevolx/_src/operators/jacobian.py 
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
def __call__(
    self,
    f: Float[Array, "... Ny Nx"],
    g: Float[Array, "... Ny Nx"],
) -> Float[Array, "... Ny Nx"]:
    """Evaluate J(f, g) at T-points.

    Parameters
    ----------
    f : Float[Array, "... Ny Nx"]
        First scalar field at T-points (e.g. the streamfunction).
    g : Float[Array, "... Ny Nx"]
        Second scalar field at T-points (e.g. the potential vorticity).

    Returns
    -------
    Float[Array, "... Ny Nx"]
        J(f, g) at T-points, zero in the ghost ring and, when
        ``self.mask`` is set, at dry T-cells.
    """
    if self.mask is not None:
        # f[..., j, i] = 0 on dry interior T-cells (same for g)
        f = _sanitize(self.mask.h, f)
        g = _sanitize(self.mask.h, g)
    J = arakawa_jacobian(f, g, self.grid.dx, self.grid.dy)
    shape = jnp.broadcast_shapes(f.shape, g.shape)
    # J_full[..., j, i] = J[..., j-1, i-1]  for 1 <= j <= Ny-2, 1 <= i <= Nx-2
    out = jnp.zeros_like(f, dtype=J.dtype, shape=shape)
    out = out.at[..., 1:-1, 1:-1].set(J)
    if self.mask is not None:
        out = jnp.where(self.mask.h, out, 0.0)
    return out

finitevolx.arakawa_jacobian(f, g, dx, dy)

Arakawa (1966) discretization of J(f, g).

Computes the Jacobian J(f, g) = ∂f/∂x·∂g/∂y − ∂f/∂y·∂g/∂x using the energy- and enstrophy-conserving three-term Arakawa scheme on a collocated grid. The inputs must include a one-point boundary halo on each side so that the returned interior array has shape (..., Ny-2, Nx-2) (i.e. Ny_i = Ny - 2, Nx_i = Nx - 2).

Parameters:

Name Type Description Default
f Float[Array, '... Ny Nx']

First scalar field (including one halo cell on each side).

 required
g Float[Array, '... Ny Nx']

Second scalar field (same shape as f).

 required
dx float

Grid spacing in the x-direction (last array axis).

 required
dy float

Grid spacing in the y-direction (second-to-last array axis).

 required

Returns:

Type Description
Float[Array, '... Ny_i Nx_i']

Jacobian evaluated on the interior grid points, where Ny_i = Ny - 2 and Nx_i = Nx - 2. Boundary points are consumed by the stencil and are not included in the output.
Notes

The Arakawa scheme averages three discrete forms:

  • J⁺⁺ (standard centred form):

.. code-block:: text

 Jpp[j,i] = ( (f[j,i+1] - f[j,i-1]) * (g[j+1,i] - g[j-1,i])
            - (f[j+1,i] - f[j-1,i]) * (g[j,i+1] - g[j,i-1]) ) / (4 dx dy)
  • J⁺× (advective form):

.. code-block:: text

 Jpx[j,i] = ( f[j,i+1] * (g[j+1,i+1] - g[j-1,i+1])
            - f[j,i-1] * (g[j+1,i-1] - g[j-1,i-1])
            - f[j+1,i] * (g[j+1,i+1] - g[j+1,i-1])
            + f[j-1,i] * (g[j-1,i+1] - g[j-1,i-1]) ) / (4 dx dy)
  • J×⁺ (divergence form):

.. code-block:: text

 Jxp[j,i] = ( g[j+1,i] * (f[j+1,i+1] - f[j+1,i-1])
            - g[j-1,i] * (f[j-1,i+1] - f[j-1,i-1])
            - g[j,i+1] * (f[j+1,i+1] - f[j-1,i+1])
            + g[j,i-1] * (f[j+1,i-1] - f[j-1,i-1]) ) / (4 dx dy)

Together: J = (Jpp + Jpx + Jxp) / 3

This triple average conserves energy (∫∫ f·J dA = 0), enstrophy (∫∫ g·J dA = 0), satisfies J(f, f) = 0, and ∫∫ J(f, g) dA = 0 at the discrete level.

The function is JAX-compatible and jit-able. Batch dimensions (...) are supported via standard broadcasting.

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import arakawa_jacobian
>>> Ny, Nx = 12, 10
>>> x = jnp.linspace(0, 1, Nx)
>>> y = jnp.linspace(0, 1, Ny)
>>> dx, dy = x[1] - x[0], y[1] - y[0]
>>> X, Y = jnp.meshgrid(x, y)
>>> J = arakawa_jacobian(X, Y, float(dx), float(dy))
>>> J.shape
(10, 8)
Source code in finitevolx/_src/operators/jacobian.py 
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 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
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
def arakawa_jacobian(
    f: Float[Array, "... Ny Nx"],
    g: Float[Array, "... Ny Nx"],
    dx: float,
    dy: float,
) -> Float[Array, "... Ny_i Nx_i"]:
    """Arakawa (1966) discretization of J(f, g).

    Computes the Jacobian J(f, g) = ∂f/∂x·∂g/∂y − ∂f/∂y·∂g/∂x using the
    energy- and enstrophy-conserving three-term Arakawa scheme on a collocated
    grid.  The inputs must include a one-point boundary halo on each side so
    that the returned interior array has shape ``(..., Ny-2, Nx-2)``
    (i.e. ``Ny_i = Ny - 2``, ``Nx_i = Nx - 2``).

    Parameters
    ----------
    f : Float[Array, "... Ny Nx"]
        First scalar field (including one halo cell on each side).
    g : Float[Array, "... Ny Nx"]
        Second scalar field (same shape as *f*).
    dx : float
        Grid spacing in the x-direction (last array axis).
    dy : float
        Grid spacing in the y-direction (second-to-last array axis).

    Returns
    -------
    Float[Array, "... Ny_i Nx_i"]
        Jacobian evaluated on the interior grid points, where
        ``Ny_i = Ny - 2`` and ``Nx_i = Nx - 2``.
        Boundary points are consumed by the stencil and are not included
        in the output.

    Notes
    -----
    The Arakawa scheme averages three discrete forms:

    * J⁺⁺ (standard centred form):

      .. code-block:: text

         Jpp[j,i] = ( (f[j,i+1] - f[j,i-1]) * (g[j+1,i] - g[j-1,i])
                    - (f[j+1,i] - f[j-1,i]) * (g[j,i+1] - g[j,i-1]) ) / (4 dx dy)

    * J⁺× (advective form):

      .. code-block:: text

         Jpx[j,i] = ( f[j,i+1] * (g[j+1,i+1] - g[j-1,i+1])
                    - f[j,i-1] * (g[j+1,i-1] - g[j-1,i-1])
                    - f[j+1,i] * (g[j+1,i+1] - g[j+1,i-1])
                    + f[j-1,i] * (g[j-1,i+1] - g[j-1,i-1]) ) / (4 dx dy)

    * J×⁺ (divergence form):

      .. code-block:: text

         Jxp[j,i] = ( g[j+1,i] * (f[j+1,i+1] - f[j+1,i-1])
                    - g[j-1,i] * (f[j-1,i+1] - f[j-1,i-1])
                    - g[j,i+1] * (f[j+1,i+1] - f[j-1,i+1])
                    + g[j,i-1] * (f[j+1,i-1] - f[j-1,i-1]) ) / (4 dx dy)

    Together: ``J = (Jpp + Jpx + Jxp) / 3``

    This triple average conserves energy (∫∫ f·J dA = 0), enstrophy
    (∫∫ g·J dA = 0), satisfies J(f, f) = 0, and ∫∫ J(f, g) dA = 0 at the
    discrete level.

    The function is JAX-compatible and ``jit``-able.  Batch dimensions
    (``...``) are supported via standard broadcasting.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import arakawa_jacobian
    >>> Ny, Nx = 12, 10
    >>> x = jnp.linspace(0, 1, Nx)
    >>> y = jnp.linspace(0, 1, Ny)
    >>> dx, dy = x[1] - x[0], y[1] - y[0]
    >>> X, Y = jnp.meshgrid(x, y)
    >>> J = arakawa_jacobian(X, Y, float(dx), float(dy))
    >>> J.shape
    (10, 8)
    """
    # Named 3x3 views, one static slice per offset (see `_Neighbours2D`):
    #   F.e = f[..., 1:-1, 2:]  (f[j, i+1])     F.ne = f[..., 2:, 2:]    (f[j+1, i+1])
    #   F.w = f[..., 1:-1, :-2] (f[j, i-1])     F.nw = f[..., 2:, :-2]   (f[j+1, i-1])
    #   F.n = f[..., 2:, 1:-1]  (f[j+1, i])     F.se = f[..., :-2, 2:]   (f[j-1, i+1])
    #   F.s = f[..., :-2, 1:-1] (f[j-1, i])     F.sw = f[..., :-2, :-2]  (f[j-1, i-1])
    # The three forms below are the same expressions as the original
    # slice-per-term code (bitwise-identical forward); slicing each offset
    # once makes the VJP ~3-4x faster on CPU (one pad per offset).
    F, G = _neighbours(f), _neighbours(g)  # each view [..., Ny-2, Nx-2]

    # J++ (standard centred form)
    # Jpp[j,i] = (df/dx * dg/dy - df/dy * dg/dx)
    # where df/dx ~ (f[j, i+1] - f[j, i-1]) / (2dx), etc.
    Jpp = (F.e - F.w) * (G.n - G.s) - (F.n - F.s) * (G.e - G.w)

    # J+x (advective form)
    # f evaluated at off-centre x-neighbours, g differenced in y at those neighbours
    Jpx = (
        F.e * (G.ne - G.se)
        - F.w * (G.nw - G.sw)
        - F.n * (G.ne - G.nw)
        + F.s * (G.se - G.sw)
    )

    # Jx+ (divergence form)
    # g evaluated at off-centre y-neighbours, f differenced in x at those neighbours
    Jxp = (
        G.n * (F.ne - F.nw)
        - G.s * (F.se - F.sw)
        - G.e * (F.ne - F.se)
        + G.w * (F.nw - F.sw)
    )

    return (Jpp + Jpx + Jxp) / (12.0 * dx * dy)