Difference Operators

Finite-difference operators for computing spatial derivatives on staggered Arakawa C-grids.

`finitevolx.Difference1D`

Bases: Module

Finite-difference operators on a 1-D Arakawa C-grid.

Parameters:

Name	Type	Description	Default
`grid`	`CartesianGrid1D`	The underlying 1-D grid.	required
`mask`	`Mask1D or None`	Optional land/ocean mask. When provided, every method post-multiplies its output by the mask field matching the output stagger: T-output → `mask.h` U-output → `mask.u` `None` (default) leaves every output untouched.	required

Source code in finitevolx/_src/operators/difference.py

class Difference1D(eqx.Module):
    """Finite-difference operators on a 1-D Arakawa C-grid.

    Parameters
    ----------
    grid : CartesianGrid1D
        The underlying 1-D grid.
    mask : Mask1D or None, optional
        Optional land/ocean mask.  When provided, every method
        post-multiplies its output by the mask field matching the
        output stagger:

        * T-output → ``mask.h``
        * U-output → ``mask.u``

        ``None`` (default) leaves every output untouched.
    """

    grid: CartesianGrid1D
    mask: Mask1D | None = None

    def diff_x_T_to_U(self, h: Float[Array, "Nx"]) -> Float[Array, "Nx"]:
        """Forward difference in x: T-point -> U-point.

        dh_dx[i+1/2] = (h[i+1] - h[i]) / dx

        Parameters
        ----------
        h : Float[Array, "Nx"]
            Scalar field at T-points.

        Returns
        -------
        Float[Array, "Nx"]
            Forward x-difference at U-points, same shape as input.
            When ``self.mask`` is set, the output is zeroed at dry
            U-faces via ``* self.mask.u``.
        """
        out = interior(diff_x_fwd_1d(h) / self.grid.dx, h)
        if self.mask is not None:
            out = out * self.mask.u
        return out

    def diff_x_U_to_T(self, u: Float[Array, "Nx"]) -> Float[Array, "Nx"]:
        """Backward difference in x: U-point -> T-point.

        du_dx[i] = (u[i+1/2] - u[i-1/2]) / dx

        Parameters
        ----------
        u : Float[Array, "Nx"]
            Velocity field at U-points.

        Returns
        -------
        Float[Array, "Nx"]
            Backward x-difference at T-points, same shape as input.
            When ``self.mask`` is set, the output is zeroed at dry
            T-cells via ``* self.mask.h``.
        """
        out = interior(diff_x_bwd_1d(u) / self.grid.dx, u)
        if self.mask is not None:
            out = out * self.mask.h
        return out

    def laplacian(self, h: Float[Array, "Nx"]) -> Float[Array, "Nx"]:
        """Laplacian at T-points.

        nabla2_h[i] = (h[i+1] - 2*h[i] + h[i-1]) / dx^2

        Parameters
        ----------
        h : Float[Array, "Nx"]
            Scalar field at T-points.

        Returns
        -------
        Float[Array, "Nx"]
            Laplacian at T-points, same shape as input.  When
            ``self.mask`` is set, the output is zeroed at dry T-cells
            via ``* self.mask.h``.
        """
        # Each offset is sliced once (see `_Neighbours2D` for why); the
        # arithmetic is unchanged:
        #   (h_e - h_c) - (h_c - h_w)  ≡  diff_x_fwd_1d(h) - diff_x_bwd_1d(h)
        h_c, h_e, h_w = h[1:-1], h[2:], h[:-2]  # h[i], h[i+1], h[i-1]  [Nx-2]
        out = interior(((h_e - h_c) - (h_c - h_w)) / self.grid.dx**2, h)
        if self.mask is not None:
            out = out * self.mask.h
        return out

`diff_x_T_to_U(h)`

Forward difference in x: T-point -> U-point.

dh_dx[i+1/2] = (h[i+1] - h[i]) / dx

Parameters:

Name	Type	Description	Default
`h`	`Float[Array, Nx]`	Scalar field at T-points.	required

Returns:

Type	Description
`Float[Array, Nx]`	Forward x-difference at U-points, same shape as input. When `self.mask` is set, the output is zeroed at dry U-faces via `* self.mask.u`.

Source code in finitevolx/_src/operators/difference.py

def diff_x_T_to_U(self, h: Float[Array, "Nx"]) -> Float[Array, "Nx"]:
    """Forward difference in x: T-point -> U-point.

    dh_dx[i+1/2] = (h[i+1] - h[i]) / dx

    Parameters
    ----------
    h : Float[Array, "Nx"]
        Scalar field at T-points.

    Returns
    -------
    Float[Array, "Nx"]
        Forward x-difference at U-points, same shape as input.
        When ``self.mask`` is set, the output is zeroed at dry
        U-faces via ``* self.mask.u``.
    """
    out = interior(diff_x_fwd_1d(h) / self.grid.dx, h)
    if self.mask is not None:
        out = out * self.mask.u
    return out

`diff_x_U_to_T(u)`

Backward difference in x: U-point -> T-point.

du_dx[i] = (u[i+1/2] - u[i-1/2]) / dx

Parameters:

Name	Type	Description	Default
`u`	`Float[Array, Nx]`	Velocity field at U-points.	required

Returns:

Type	Description
`Float[Array, Nx]`	Backward x-difference at T-points, same shape as input. When `self.mask` is set, the output is zeroed at dry T-cells via `* self.mask.h`.

Source code in finitevolx/_src/operators/difference.py

def diff_x_U_to_T(self, u: Float[Array, "Nx"]) -> Float[Array, "Nx"]:
    """Backward difference in x: U-point -> T-point.

    du_dx[i] = (u[i+1/2] - u[i-1/2]) / dx

    Parameters
    ----------
    u : Float[Array, "Nx"]
        Velocity field at U-points.

    Returns
    -------
    Float[Array, "Nx"]
        Backward x-difference at T-points, same shape as input.
        When ``self.mask`` is set, the output is zeroed at dry
        T-cells via ``* self.mask.h``.
    """
    out = interior(diff_x_bwd_1d(u) / self.grid.dx, u)
    if self.mask is not None:
        out = out * self.mask.h
    return out

`laplacian(h)`

Laplacian at T-points.

nabla2_h[i] = (h[i+1] - 2*h[i] + h[i-1]) / dx^2

Parameters:

Name	Type	Description	Default
`h`	`Float[Array, Nx]`	Scalar field at T-points.	required

Returns:

Type	Description
`Float[Array, Nx]`	Laplacian at T-points, same shape as input. When `self.mask` is set, the output is zeroed at dry T-cells via `* self.mask.h`.

Source code in finitevolx/_src/operators/difference.py

def laplacian(self, h: Float[Array, "Nx"]) -> Float[Array, "Nx"]:
    """Laplacian at T-points.

    nabla2_h[i] = (h[i+1] - 2*h[i] + h[i-1]) / dx^2

    Parameters
    ----------
    h : Float[Array, "Nx"]
        Scalar field at T-points.

    Returns
    -------
    Float[Array, "Nx"]
        Laplacian at T-points, same shape as input.  When
        ``self.mask`` is set, the output is zeroed at dry T-cells
        via ``* self.mask.h``.
    """
    # Each offset is sliced once (see `_Neighbours2D` for why); the
    # arithmetic is unchanged:
    #   (h_e - h_c) - (h_c - h_w)  ≡  diff_x_fwd_1d(h) - diff_x_bwd_1d(h)
    h_c, h_e, h_w = h[1:-1], h[2:], h[:-2]  # h[i], h[i+1], h[i-1]  [Nx-2]
    out = interior(((h_e - h_c) - (h_c - h_w)) / self.grid.dx**2, h)
    if self.mask is not None:
        out = out * self.mask.h
    return out

`finitevolx.Difference2D`

Bases: Module

Finite-difference operators on a 2-D Arakawa C-grid.

Parameters:

Name	Type	Description	Default
`grid`	`CartesianGrid2D`	The underlying 2-D grid.	required
`mask`	`Mask2D or None`	Optional land/ocean mask. When provided, every method post-multiplies its output by the mask field matching the output stagger: T-output → `mask.h` (divergence, laplacian, U→T, V→T) U-output → `mask.u` (T→U, X→U, grad_perp's u component) V-output → `mask.v` (T→V, X→V, grad_perp's v component) X-output → `mask.xy_corner_strict` (curl, U→X, V→X) The strict corner mask is used because an X-point output is trusted only when all four surrounding T-cells are wet. `None` (default) leaves every output untouched.	required

Source code in finitevolx/_src/operators/difference.py

class Difference2D(eqx.Module):
    """Finite-difference operators on a 2-D Arakawa C-grid.

    Parameters
    ----------
    grid : CartesianGrid2D
        The underlying 2-D grid.
    mask : Mask2D or None, optional
        Optional land/ocean mask.  When provided, every method
        post-multiplies its output by the mask field matching the
        output stagger:

        * T-output → ``mask.h`` (divergence, laplacian, U→T, V→T)
        * U-output → ``mask.u`` (T→U, X→U, grad_perp's u component)
        * V-output → ``mask.v`` (T→V, X→V, grad_perp's v component)
        * X-output → ``mask.xy_corner_strict`` (curl, U→X, V→X)

        The strict corner mask is used because an X-point output is
        trusted only when all four surrounding T-cells are wet.
        ``None`` (default) leaves every output untouched.
    """

    grid: CartesianGrid2D
    mask: Mask2D | None = None

    # ------------------------------------------------------------------
    # Forward differences
    # ------------------------------------------------------------------

    def diff_x_T_to_U(
        self, h: Float[Array, "*batch Ny Nx"]
    ) -> Float[Array, "*batch Ny Nx"]:
        """Forward x-difference: T-point -> U-point.

        dh_dx[j, i+1/2] = (h[j, i+1] - h[j, i]) / dx

        Parameters
        ----------
        h : Float[Array, "*batch Ny Nx"]
            Scalar field at T-points.

        Returns
        -------
        Float[Array, "*batch Ny Nx"]
            Forward x-difference at U-points.  When ``self.mask`` is
            set, the output is zeroed at dry U-faces via ``* mask.u``.
        """
        out = interior(diff_x_fwd(h) / self.grid.dx, h)
        if self.mask is not None:
            out = out * self.mask.u
        return out

    def diff_y_T_to_V(
        self, h: Float[Array, "*batch Ny Nx"]
    ) -> Float[Array, "*batch Ny Nx"]:
        """Forward y-difference: T-point -> V-point.

        dh_dy[j+1/2, i] = (h[j+1, i] - h[j, i]) / dy

        Parameters
        ----------
        h : Float[Array, "*batch Ny Nx"]
            Scalar field at T-points.

        Returns
        -------
        Float[Array, "*batch Ny Nx"]
            Forward y-difference at V-points.  When ``self.mask`` is
            set, the output is zeroed at dry V-faces via ``* mask.v``.
        """
        out = interior(diff_y_fwd(h) / self.grid.dy, h)
        if self.mask is not None:
            out = out * self.mask.v
        return out

    def diff_y_U_to_X(
        self, u: Float[Array, "*batch Ny Nx"]
    ) -> Float[Array, "*batch Ny Nx"]:
        """Forward y-difference: U-point -> X-point (corner).

        du_dy[j+1/2, i+1/2] = (u[j+1, i+1/2] - u[j, i+1/2]) / dy

        Parameters
        ----------
        u : Float[Array, "*batch Ny Nx"]
            Velocity field at U-points.

        Returns
        -------
        Float[Array, "*batch Ny Nx"]
            Forward y-difference at X-points.  When ``self.mask`` is
            set, the output is zeroed at dry X-corners via
            ``* mask.xy_corner_strict``.
        """
        out = interior(diff_y_fwd(u) / self.grid.dy, u)
        if self.mask is not None:
            out = out * self.mask.xy_corner_strict
        return out

    def diff_x_V_to_X(
        self, v: Float[Array, "*batch Ny Nx"]
    ) -> Float[Array, "*batch Ny Nx"]:
        """Forward x-difference: V-point -> X-point (corner).

        dv_dx[j+1/2, i+1/2] = (v[j+1/2, i+1] - v[j+1/2, i]) / dx

        Parameters
        ----------
        v : Float[Array, "*batch Ny Nx"]
            Velocity field at V-points.

        Returns
        -------
        Float[Array, "*batch Ny Nx"]
            Forward x-difference at X-points.  When ``self.mask`` is
            set, the output is zeroed at dry X-corners via
            ``* mask.xy_corner_strict``.
        """
        out = interior(diff_x_fwd(v) / self.grid.dx, v)
        if self.mask is not None:
            out = out * self.mask.xy_corner_strict
        return out

    # ------------------------------------------------------------------
    # Backward differences
    # ------------------------------------------------------------------

    def diff_y_X_to_U(
        self, q: Float[Array, "*batch Ny Nx"]
    ) -> Float[Array, "*batch Ny Nx"]:
        """Backward y-difference: X-point -> U-point.

        dq_dy[j, i+1/2] = (q[j+1/2, i+1/2] - q[j-1/2, i+1/2]) / dy

        Parameters
        ----------
        q : Float[Array, "*batch Ny Nx"]
            Scalar field at X-points.

        Returns
        -------
        Float[Array, "*batch Ny Nx"]
            Backward y-difference at U-points.  When ``self.mask`` is
            set, the output is zeroed at dry U-faces via ``* mask.u``.
        """
        out = interior(diff_y_bwd(q) / self.grid.dy, q)
        if self.mask is not None:
            out = out * self.mask.u
        return out

    def diff_x_X_to_V(
        self, q: Float[Array, "*batch Ny Nx"]
    ) -> Float[Array, "*batch Ny Nx"]:
        """Backward x-difference: X-point -> V-point.

        dq_dx[j+1/2, i] = (q[j+1/2, i+1/2] - q[j+1/2, i-1/2]) / dx

        Parameters
        ----------
        q : Float[Array, "*batch Ny Nx"]
            Scalar field at X-points.

        Returns
        -------
        Float[Array, "*batch Ny Nx"]
            Backward x-difference at V-points.  When ``self.mask`` is
            set, the output is zeroed at dry V-faces via ``* mask.v``.
        """
        out = interior(diff_x_bwd(q) / self.grid.dx, q)
        if self.mask is not None:
            out = out * self.mask.v
        return out

    def diff_x_U_to_T(
        self, u: Float[Array, "*batch Ny Nx"]
    ) -> Float[Array, "*batch Ny Nx"]:
        """Backward x-difference: U-point -> T-point.

        du_dx[j, i] = (u[j, i+1/2] - u[j, i-1/2]) / dx

        Parameters
        ----------
        u : Float[Array, "*batch Ny Nx"]
            Velocity field at U-points.

        Returns
        -------
        Float[Array, "*batch Ny Nx"]
            Backward x-difference at T-points.  When ``self.mask`` is
            set, the output is zeroed at dry T-cells via ``* mask.h``.
        """
        out = interior(diff_x_bwd(u) / self.grid.dx, u)
        if self.mask is not None:
            out = out * self.mask.h
        return out

    def diff_y_V_to_T(
        self, v: Float[Array, "*batch Ny Nx"]
    ) -> Float[Array, "*batch Ny Nx"]:
        """Backward y-difference: V-point -> T-point.

        dv_dy[j, i] = (v[j+1/2, i] - v[j-1/2, i]) / dy

        Parameters
        ----------
        v : Float[Array, "*batch Ny Nx"]
            Velocity field at V-points.

        Returns
        -------
        Float[Array, "*batch Ny Nx"]
            Backward y-difference at T-points.  When ``self.mask`` is
            set, the output is zeroed at dry T-cells via ``* mask.h``.
        """
        out = interior(diff_y_bwd(v) / self.grid.dy, v)
        if self.mask is not None:
            out = out * self.mask.h
        return out

    # ------------------------------------------------------------------
    # Compound operators
    # ------------------------------------------------------------------

    def divergence(
        self,
        u: Float[Array, "*batch Ny Nx"],
        v: Float[Array, "*batch Ny Nx"],
    ) -> Float[Array, "*batch Ny Nx"]:
        """Divergence of (u, v) at T-points.

        delta[j, i] = du_dx[j, i] + dv_dy[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"]
            Divergence at T-points.  When ``self.mask`` is set, the
            output is zeroed at dry T-cells via ``* mask.h``.
        """
        out = _divergence_2d(u, v, self.grid.dx, self.grid.dy)
        if self.mask is not None:
            out = out * self.mask.h
        return out

    def curl(
        self,
        u: Float[Array, "*batch Ny Nx"],
        v: Float[Array, "*batch Ny Nx"],
    ) -> Float[Array, "*batch Ny Nx"]:
        """Curl (relative vorticity) of (u, v) 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 (corners).  When
            ``self.mask`` is set, the output is zeroed at dry X-corners
            via ``* mask.xy_corner_strict``.
        """
        out = _curl_2d(u, v, self.grid.dx, self.grid.dy)
        if self.mask is not None:
            out = out * self.mask.xy_corner_strict
        return out

    def laplacian(
        self, h: Float[Array, "*batch Ny Nx"]
    ) -> Float[Array, "*batch Ny Nx"]:
        """Laplacian at T-points.

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

        Parameters
        ----------
        h : Float[Array, "*batch Ny Nx"]
            Scalar field at T-points.

        Returns
        -------
        Float[Array, "*batch Ny Nx"]
            Laplacian at T-points.  When ``self.mask`` is set, the
            output is zeroed at dry T-cells via ``* mask.h``.
        """
        # nabla2_h[j, i] = d^2h/dx^2 + d^2h/dy^2, with each offset sliced
        # once (see `_Neighbours2D`); same arithmetic as
        #   d2x = (diff_x_fwd(h) - diff_x_bwd(h)) / dx**2
        #   d2y = (diff_y_fwd(h) - diff_y_bwd(h)) / dy**2
        nb = _neighbours(h)  # c, e, w, n, s views, each [Ny-2, Nx-2]
        d2x = ((nb.e - nb.c) - (nb.c - nb.w)) / self.grid.dx**2
        d2y = ((nb.n - nb.c) - (nb.c - nb.s)) / self.grid.dy**2
        out = interior(d2x + d2y, h)
        if self.mask is not None:
            out = out * self.mask.h
        return out

    def grad_perp(
        self,
        psi: Float[Array, "*batch Ny Nx"],
    ) -> tuple[Float[Array, "*batch Ny Nx"], Float[Array, "*batch Ny Nx"]]:
        """Perpendicular gradient: T-point streamfunction to geostrophic velocity.

        Maps a streamfunction ψ at T-points to face-centred geostrophic
        velocities (u, v) = (-∂ψ/∂y, ∂ψ/∂x) on an Arakawa C-grid.

        Expanding the T→X bilinear interpolation into the backward X→U / X→V
        differences gives a compact stencil that reads the T-point ghost cells
        of ψ directly:

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

        The resulting (unmasked) velocity field is discretely non-divergent:
        div(u, v) = 0 at T-points. When ``self.mask`` is set, ``u`` is
        multiplied by ``self.mask.u`` and ``v`` by ``self.mask.v``; the
        divergence-free property no longer holds in general in that case.

        Parameters
        ----------
        psi : Float[Array, "*batch Ny Nx"]
            Streamfunction at T-points.

        Returns
        -------
        tuple[Float[Array, "*batch Ny Nx"], Float[Array, "*batch Ny Nx"]]
            (u, v) — zonal velocity at U-points and meridional velocity at
            V-points.

        References
        ----------
        .. [1] louity/MQGeometry ``fd.py`` — ``grad_perp`` function.
           https://github.com/louity/MQGeometry/blob/main/fd.py
        .. [2] louity/qgsw-pytorch ``finite_diff.py`` — ``grad_perp`` function.
           https://github.com/louity/qgsw-pytorch/blob/main/src/finite_diff.py

        Examples
        --------
        >>> from finitevolx import CartesianGrid2D, Difference2D
        >>> grid = CartesianGrid2D.from_interior(8, 8, 1e3, 1e3)
        >>> diff = Difference2D(grid=grid)
        >>> psi = jnp.ones((grid.Ny, grid.Nx))
        >>> u, v = diff.grad_perp(psi)
        """
        # u[j, i+1/2] = -(ψ[j+1,i] + ψ[j+1,i+1] - ψ[j-1,i] - ψ[j-1,i+1]) / (4·dy)
        u = interior(
            -(
                psi[..., 2:, 1:-1]
                + psi[..., 2:, 2:]
                - psi[..., :-2, 1:-1]
                - psi[..., :-2, 2:]
            )
            / (4.0 * self.grid.dy),
            psi,
        )

        # v[j+1/2, i] = (ψ[j,i+1] + ψ[j+1,i+1] - ψ[j,i-1] - ψ[j+1,i-1]) / (4·dx)
        v = interior(
            (
                psi[..., 1:-1, 2:]
                + psi[..., 2:, 2:]
                - psi[..., 1:-1, :-2]
                - psi[..., 2:, :-2]
            )
            / (4.0 * self.grid.dx),
            psi,
        )

        if self.mask is not None:
            u = u * self.mask.u
            v = v * self.mask.v

        return u, v

`curl(u, v)`

Curl (relative vorticity) of (u, v) 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 (corners). When `self.mask` is set, the output is zeroed at dry X-corners via `* mask.xy_corner_strict`.

Source code in finitevolx/_src/operators/difference.py

def curl(
    self,
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
) -> Float[Array, "*batch Ny Nx"]:
    """Curl (relative vorticity) of (u, v) 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 (corners).  When
        ``self.mask`` is set, the output is zeroed at dry X-corners
        via ``* mask.xy_corner_strict``.
    """
    out = _curl_2d(u, v, self.grid.dx, self.grid.dy)
    if self.mask is not None:
        out = out * self.mask.xy_corner_strict
    return out

`diff_x_T_to_U(h)`

Forward x-difference: T-point -> U-point.

dh_dx[j, i+1/2] = (h[j, i+1] - h[j, i]) / dx

Parameters:

Name	Type	Description	Default
`h`	`Float[Array, '*batch Ny Nx']`	Scalar field at T-points.	required

Returns:

Type	Description
`Float[Array, '*batch Ny Nx']`	Forward x-difference at U-points. When `self.mask` is set, the output is zeroed at dry U-faces via `* mask.u`.

Source code in finitevolx/_src/operators/difference.py

def diff_x_T_to_U(
    self, h: Float[Array, "*batch Ny Nx"]
) -> Float[Array, "*batch Ny Nx"]:
    """Forward x-difference: T-point -> U-point.

    dh_dx[j, i+1/2] = (h[j, i+1] - h[j, i]) / dx

    Parameters
    ----------
    h : Float[Array, "*batch Ny Nx"]
        Scalar field at T-points.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Forward x-difference at U-points.  When ``self.mask`` is
        set, the output is zeroed at dry U-faces via ``* mask.u``.
    """
    out = interior(diff_x_fwd(h) / self.grid.dx, h)
    if self.mask is not None:
        out = out * self.mask.u
    return out

`diff_x_U_to_T(u)`

Backward x-difference: U-point -> T-point.

du_dx[j, i] = (u[j, i+1/2] - u[j, i-1/2]) / dx

Parameters:

Name	Type	Description	Default
`u`	`Float[Array, '*batch Ny Nx']`	Velocity field at U-points.	required

Returns:

Type	Description
`Float[Array, '*batch Ny Nx']`	Backward x-difference at T-points. When `self.mask` is set, the output is zeroed at dry T-cells via `* mask.h`.

Source code in finitevolx/_src/operators/difference.py

def diff_x_U_to_T(
    self, u: Float[Array, "*batch Ny Nx"]
) -> Float[Array, "*batch Ny Nx"]:
    """Backward x-difference: U-point -> T-point.

    du_dx[j, i] = (u[j, i+1/2] - u[j, i-1/2]) / dx

    Parameters
    ----------
    u : Float[Array, "*batch Ny Nx"]
        Velocity field at U-points.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Backward x-difference at T-points.  When ``self.mask`` is
        set, the output is zeroed at dry T-cells via ``* mask.h``.
    """
    out = interior(diff_x_bwd(u) / self.grid.dx, u)
    if self.mask is not None:
        out = out * self.mask.h
    return out

`diff_x_V_to_X(v)`

Forward x-difference: V-point -> X-point (corner).

dv_dx[j+1/2, i+1/2] = (v[j+1/2, i+1] - v[j+1/2, i]) / dx

Parameters:

Name	Type	Description	Default
`v`	`Float[Array, '*batch Ny Nx']`	Velocity field at V-points.	required

Returns:

Type	Description
`Float[Array, '*batch Ny Nx']`	Forward x-difference at X-points. When `self.mask` is set, the output is zeroed at dry X-corners via `* mask.xy_corner_strict`.

Source code in finitevolx/_src/operators/difference.py

def diff_x_V_to_X(
    self, v: Float[Array, "*batch Ny Nx"]
) -> Float[Array, "*batch Ny Nx"]:
    """Forward x-difference: V-point -> X-point (corner).

    dv_dx[j+1/2, i+1/2] = (v[j+1/2, i+1] - v[j+1/2, i]) / dx

    Parameters
    ----------
    v : Float[Array, "*batch Ny Nx"]
        Velocity field at V-points.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Forward x-difference at X-points.  When ``self.mask`` is
        set, the output is zeroed at dry X-corners via
        ``* mask.xy_corner_strict``.
    """
    out = interior(diff_x_fwd(v) / self.grid.dx, v)
    if self.mask is not None:
        out = out * self.mask.xy_corner_strict
    return out

`diff_x_X_to_V(q)`

Backward x-difference: X-point -> V-point.

dq_dx[j+1/2, i] = (q[j+1/2, i+1/2] - q[j+1/2, i-1/2]) / dx

Parameters:

Name	Type	Description	Default
`q`	`Float[Array, '*batch Ny Nx']`	Scalar field at X-points.	required

Returns:

Type	Description
`Float[Array, '*batch Ny Nx']`	Backward x-difference at V-points. When `self.mask` is set, the output is zeroed at dry V-faces via `* mask.v`.

Source code in finitevolx/_src/operators/difference.py

def diff_x_X_to_V(
    self, q: Float[Array, "*batch Ny Nx"]
) -> Float[Array, "*batch Ny Nx"]:
    """Backward x-difference: X-point -> V-point.

    dq_dx[j+1/2, i] = (q[j+1/2, i+1/2] - q[j+1/2, i-1/2]) / dx

    Parameters
    ----------
    q : Float[Array, "*batch Ny Nx"]
        Scalar field at X-points.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Backward x-difference at V-points.  When ``self.mask`` is
        set, the output is zeroed at dry V-faces via ``* mask.v``.
    """
    out = interior(diff_x_bwd(q) / self.grid.dx, q)
    if self.mask is not None:
        out = out * self.mask.v
    return out

`diff_y_T_to_V(h)`

Forward y-difference: T-point -> V-point.

dh_dy[j+1/2, i] = (h[j+1, i] - h[j, i]) / dy

Parameters:

Name	Type	Description	Default
`h`	`Float[Array, '*batch Ny Nx']`	Scalar field at T-points.	required

Returns:

Type	Description
`Float[Array, '*batch Ny Nx']`	Forward y-difference at V-points. When `self.mask` is set, the output is zeroed at dry V-faces via `* mask.v`.

Source code in finitevolx/_src/operators/difference.py

def diff_y_T_to_V(
    self, h: Float[Array, "*batch Ny Nx"]
) -> Float[Array, "*batch Ny Nx"]:
    """Forward y-difference: T-point -> V-point.

    dh_dy[j+1/2, i] = (h[j+1, i] - h[j, i]) / dy

    Parameters
    ----------
    h : Float[Array, "*batch Ny Nx"]
        Scalar field at T-points.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Forward y-difference at V-points.  When ``self.mask`` is
        set, the output is zeroed at dry V-faces via ``* mask.v``.
    """
    out = interior(diff_y_fwd(h) / self.grid.dy, h)
    if self.mask is not None:
        out = out * self.mask.v
    return out

`diff_y_U_to_X(u)`

Forward y-difference: U-point -> X-point (corner).

du_dy[j+1/2, i+1/2] = (u[j+1, i+1/2] - u[j, i+1/2]) / dy

Parameters:

Name	Type	Description	Default
`u`	`Float[Array, '*batch Ny Nx']`	Velocity field at U-points.	required

Returns:

Type	Description
`Float[Array, '*batch Ny Nx']`	Forward y-difference at X-points. When `self.mask` is set, the output is zeroed at dry X-corners via `* mask.xy_corner_strict`.

Source code in finitevolx/_src/operators/difference.py

def diff_y_U_to_X(
    self, u: Float[Array, "*batch Ny Nx"]
) -> Float[Array, "*batch Ny Nx"]:
    """Forward y-difference: U-point -> X-point (corner).

    du_dy[j+1/2, i+1/2] = (u[j+1, i+1/2] - u[j, i+1/2]) / dy

    Parameters
    ----------
    u : Float[Array, "*batch Ny Nx"]
        Velocity field at U-points.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Forward y-difference at X-points.  When ``self.mask`` is
        set, the output is zeroed at dry X-corners via
        ``* mask.xy_corner_strict``.
    """
    out = interior(diff_y_fwd(u) / self.grid.dy, u)
    if self.mask is not None:
        out = out * self.mask.xy_corner_strict
    return out

`diff_y_V_to_T(v)`

Backward y-difference: V-point -> T-point.

dv_dy[j, i] = (v[j+1/2, i] - v[j-1/2, i]) / dy

Parameters:

Name	Type	Description	Default
`v`	`Float[Array, '*batch Ny Nx']`	Velocity field at V-points.	required

Returns:

Type	Description
`Float[Array, '*batch Ny Nx']`	Backward y-difference at T-points. When `self.mask` is set, the output is zeroed at dry T-cells via `* mask.h`.

Source code in finitevolx/_src/operators/difference.py

def diff_y_V_to_T(
    self, v: Float[Array, "*batch Ny Nx"]
) -> Float[Array, "*batch Ny Nx"]:
    """Backward y-difference: V-point -> T-point.

    dv_dy[j, i] = (v[j+1/2, i] - v[j-1/2, i]) / dy

    Parameters
    ----------
    v : Float[Array, "*batch Ny Nx"]
        Velocity field at V-points.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Backward y-difference at T-points.  When ``self.mask`` is
        set, the output is zeroed at dry T-cells via ``* mask.h``.
    """
    out = interior(diff_y_bwd(v) / self.grid.dy, v)
    if self.mask is not None:
        out = out * self.mask.h
    return out

`diff_y_X_to_U(q)`

Backward y-difference: X-point -> U-point.

dq_dy[j, i+1/2] = (q[j+1/2, i+1/2] - q[j-1/2, i+1/2]) / dy

Parameters:

Name	Type	Description	Default
`q`	`Float[Array, '*batch Ny Nx']`	Scalar field at X-points.	required

Returns:

Type	Description
`Float[Array, '*batch Ny Nx']`	Backward y-difference at U-points. When `self.mask` is set, the output is zeroed at dry U-faces via `* mask.u`.

Source code in finitevolx/_src/operators/difference.py

def diff_y_X_to_U(
    self, q: Float[Array, "*batch Ny Nx"]
) -> Float[Array, "*batch Ny Nx"]:
    """Backward y-difference: X-point -> U-point.

    dq_dy[j, i+1/2] = (q[j+1/2, i+1/2] - q[j-1/2, i+1/2]) / dy

    Parameters
    ----------
    q : Float[Array, "*batch Ny Nx"]
        Scalar field at X-points.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Backward y-difference at U-points.  When ``self.mask`` is
        set, the output is zeroed at dry U-faces via ``* mask.u``.
    """
    out = interior(diff_y_bwd(q) / self.grid.dy, q)
    if self.mask is not None:
        out = out * self.mask.u
    return out

`divergence(u, v)`

Divergence of (u, v) at T-points.

delta[j, i] = du_dx[j, i] + dv_dy[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']`	Divergence at T-points. When `self.mask` is set, the output is zeroed at dry T-cells via `* mask.h`.

Source code in finitevolx/_src/operators/difference.py

def divergence(
    self,
    u: Float[Array, "*batch Ny Nx"],
    v: Float[Array, "*batch Ny Nx"],
) -> Float[Array, "*batch Ny Nx"]:
    """Divergence of (u, v) at T-points.

    delta[j, i] = du_dx[j, i] + dv_dy[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"]
        Divergence at T-points.  When ``self.mask`` is set, the
        output is zeroed at dry T-cells via ``* mask.h``.
    """
    out = _divergence_2d(u, v, self.grid.dx, self.grid.dy)
    if self.mask is not None:
        out = out * self.mask.h
    return out

`grad_perp(psi)`

Perpendicular gradient: T-point streamfunction to geostrophic velocity.

Maps a streamfunction ψ at T-points to face-centred geostrophic velocities (u, v) = (-∂ψ/∂y, ∂ψ/∂x) on an Arakawa C-grid.

Expanding the T→X bilinear interpolation into the backward X→U / X→V differences gives a compact stencil that reads the T-point ghost cells of ψ directly:

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

The resulting (unmasked) velocity field is discretely non-divergent: div(u, v) = 0 at T-points. When self.mask is set, u is multiplied by self.mask.u and v by self.mask.v; the divergence-free property no longer holds in general in that case.

Parameters:

Name	Type	Description	Default
`psi`	`Float[Array, '*batch Ny Nx']`	Streamfunction at T-points.	required

Returns:

Type	Description
`tuple[Float[Array, 'batch Ny Nx'], Float[Array, 'batch Ny Nx']]`	(u, v) — zonal velocity at U-points and meridional velocity at V-points.

References

.. [1] louity/MQGeometry fd.py — grad_perp function. https://github.com/louity/MQGeometry/blob/main/fd.py .. [2] louity/qgsw-pytorch finite_diff.py — grad_perp function. https://github.com/louity/qgsw-pytorch/blob/main/src/finite_diff.py

Examples:

>>> from finitevolx import CartesianGrid2D, Difference2D
>>> grid = CartesianGrid2D.from_interior(8, 8, 1e3, 1e3)
>>> diff = Difference2D(grid=grid)
>>> psi = jnp.ones((grid.Ny, grid.Nx))
>>> u, v = diff.grad_perp(psi)

Source code in finitevolx/_src/operators/difference.py

def grad_perp(
    self,
    psi: Float[Array, "*batch Ny Nx"],
) -> tuple[Float[Array, "*batch Ny Nx"], Float[Array, "*batch Ny Nx"]]:
    """Perpendicular gradient: T-point streamfunction to geostrophic velocity.

    Maps a streamfunction ψ at T-points to face-centred geostrophic
    velocities (u, v) = (-∂ψ/∂y, ∂ψ/∂x) on an Arakawa C-grid.

    Expanding the T→X bilinear interpolation into the backward X→U / X→V
    differences gives a compact stencil that reads the T-point ghost cells
    of ψ directly:

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

    The resulting (unmasked) velocity field is discretely non-divergent:
    div(u, v) = 0 at T-points. When ``self.mask`` is set, ``u`` is
    multiplied by ``self.mask.u`` and ``v`` by ``self.mask.v``; the
    divergence-free property no longer holds in general in that case.

    Parameters
    ----------
    psi : Float[Array, "*batch Ny Nx"]
        Streamfunction at T-points.

    Returns
    -------
    tuple[Float[Array, "*batch Ny Nx"], Float[Array, "*batch Ny Nx"]]
        (u, v) — zonal velocity at U-points and meridional velocity at
        V-points.

    References
    ----------
    .. [1] louity/MQGeometry ``fd.py`` — ``grad_perp`` function.
       https://github.com/louity/MQGeometry/blob/main/fd.py
    .. [2] louity/qgsw-pytorch ``finite_diff.py`` — ``grad_perp`` function.
       https://github.com/louity/qgsw-pytorch/blob/main/src/finite_diff.py

    Examples
    --------
    >>> from finitevolx import CartesianGrid2D, Difference2D
    >>> grid = CartesianGrid2D.from_interior(8, 8, 1e3, 1e3)
    >>> diff = Difference2D(grid=grid)
    >>> psi = jnp.ones((grid.Ny, grid.Nx))
    >>> u, v = diff.grad_perp(psi)
    """
    # u[j, i+1/2] = -(ψ[j+1,i] + ψ[j+1,i+1] - ψ[j-1,i] - ψ[j-1,i+1]) / (4·dy)
    u = interior(
        -(
            psi[..., 2:, 1:-1]
            + psi[..., 2:, 2:]
            - psi[..., :-2, 1:-1]
            - psi[..., :-2, 2:]
        )
        / (4.0 * self.grid.dy),
        psi,
    )

    # v[j+1/2, i] = (ψ[j,i+1] + ψ[j+1,i+1] - ψ[j,i-1] - ψ[j+1,i-1]) / (4·dx)
    v = interior(
        (
            psi[..., 1:-1, 2:]
            + psi[..., 2:, 2:]
            - psi[..., 1:-1, :-2]
            - psi[..., 2:, :-2]
        )
        / (4.0 * self.grid.dx),
        psi,
    )

    if self.mask is not None:
        u = u * self.mask.u
        v = v * self.mask.v

    return u, v

`laplacian(h)`

Laplacian at T-points.

nabla2_h[j, i] = (h[j, i+1] - 2h[j, i] + h[j, i-1]) / dx^2 + (h[j+1, i] - 2h[j, i] + h[j-1, i]) / dy^2

Parameters:

Name	Type	Description	Default
`h`	`Float[Array, '*batch Ny Nx']`	Scalar field at T-points.	required

Returns:

Type	Description
`Float[Array, '*batch Ny Nx']`	Laplacian at T-points. When `self.mask` is set, the output is zeroed at dry T-cells via `* mask.h`.

Source code in finitevolx/_src/operators/difference.py

def laplacian(
    self, h: Float[Array, "*batch Ny Nx"]
) -> Float[Array, "*batch Ny Nx"]:
    """Laplacian at T-points.

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

    Parameters
    ----------
    h : Float[Array, "*batch Ny Nx"]
        Scalar field at T-points.

    Returns
    -------
    Float[Array, "*batch Ny Nx"]
        Laplacian at T-points.  When ``self.mask`` is set, the
        output is zeroed at dry T-cells via ``* mask.h``.
    """
    # nabla2_h[j, i] = d^2h/dx^2 + d^2h/dy^2, with each offset sliced
    # once (see `_Neighbours2D`); same arithmetic as
    #   d2x = (diff_x_fwd(h) - diff_x_bwd(h)) / dx**2
    #   d2y = (diff_y_fwd(h) - diff_y_bwd(h)) / dy**2
    nb = _neighbours(h)  # c, e, w, n, s views, each [Ny-2, Nx-2]
    d2x = ((nb.e - nb.c) - (nb.c - nb.w)) / self.grid.dx**2
    d2y = ((nb.n - nb.c) - (nb.c - nb.s)) / self.grid.dy**2
    out = interior(d2x + d2y, h)
    if self.mask is not None:
        out = out * self.mask.h
    return out

`finitevolx.Difference3D`

Bases: Module

Finite-difference operators on a 3-D Arakawa C-grid.

Operates on the horizontal (y, x) plane for each z-level. The 3-D array shape is [Nz, Ny, Nx].

Parameters:

Name	Type	Description	Default
`grid`	`CartesianGrid3D`	The underlying 3-D grid.	required
`mask`	`Mask3D or None`	Optional land/ocean mask. When provided, every method post-multiplies its output by the mask field matching the output stagger (`mask.h`, `mask.u`, or `mask.v`). `None` (default) leaves every output untouched.	required

Source code in finitevolx/_src/operators/difference.py

class Difference3D(eqx.Module):
    """Finite-difference operators on a 3-D Arakawa C-grid.

    Operates on the horizontal (y, x) plane for each z-level.
    The 3-D array shape is [Nz, Ny, Nx].

    Parameters
    ----------
    grid : CartesianGrid3D
        The underlying 3-D grid.
    mask : Mask3D or None, optional
        Optional land/ocean mask.  When provided, every method
        post-multiplies its output by the mask field matching the
        output stagger (``mask.h``, ``mask.u``, or ``mask.v``).
        ``None`` (default) leaves every output untouched.
    """

    grid: CartesianGrid3D
    mask: Mask3D | None = None

    def diff_x_T_to_U(self, h: Float[Array, "Nz Ny Nx"]) -> Float[Array, "Nz Ny Nx"]:
        """Forward x-difference over all z-levels: T -> U.

        dh_dx[k, j, i+1/2] = (h[k, j, i+1] - h[k, j, i]) / dx

        Parameters
        ----------
        h : Float[Array, "Nz Ny Nx"]
            Scalar field at T-points.

        Returns
        -------
        Float[Array, "Nz Ny Nx"]
            Forward x-difference at U-points.  When ``self.mask`` is
            set, the output is zeroed at dry U-faces via ``* mask.u``.
        """
        out = interior(diff_x_fwd_3d(h) / self.grid.dx, h)
        if self.mask is not None:
            out = out * self.mask.u
        return out

    def diff_y_T_to_V(self, h: Float[Array, "Nz Ny Nx"]) -> Float[Array, "Nz Ny Nx"]:
        """Forward y-difference over all z-levels: T -> V.

        dh_dy[k, j+1/2, i] = (h[k, j+1, i] - h[k, j, i]) / dy

        Parameters
        ----------
        h : Float[Array, "Nz Ny Nx"]
            Scalar field at T-points.

        Returns
        -------
        Float[Array, "Nz Ny Nx"]
            Forward y-difference at V-points.  When ``self.mask`` is
            set, the output is zeroed at dry V-faces via ``* mask.v``.
        """
        out = interior(diff_y_fwd_3d(h) / self.grid.dy, h)
        if self.mask is not None:
            out = out * self.mask.v
        return out

    def diff_x_U_to_T(self, u: Float[Array, "Nz Ny Nx"]) -> Float[Array, "Nz Ny Nx"]:
        """Backward x-difference over all z-levels: U -> T.

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

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

        Returns
        -------
        Float[Array, "Nz Ny Nx"]
            Backward x-difference at T-points.  When ``self.mask`` is
            set, the output is zeroed at dry T-cells via ``* mask.h``.
        """
        out = interior(diff_x_bwd_3d(u) / self.grid.dx, u)
        if self.mask is not None:
            out = out * self.mask.h
        return out

    def diff_y_V_to_T(self, v: Float[Array, "Nz Ny Nx"]) -> Float[Array, "Nz Ny Nx"]:
        """Backward y-difference over all z-levels: V -> T.

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

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

        Returns
        -------
        Float[Array, "Nz Ny Nx"]
            Backward y-difference at T-points.  When ``self.mask`` is
            set, the output is zeroed at dry T-cells via ``* mask.h``.
        """
        out = interior(diff_y_bwd_3d(v) / self.grid.dy, v)
        if self.mask is not None:
            out = out * self.mask.h
        return out

    def divergence(
        self,
        u: Float[Array, "Nz Ny Nx"],
        v: Float[Array, "Nz Ny Nx"],
    ) -> Float[Array, "Nz Ny Nx"]:
        """Horizontal divergence at T-points over all z-levels.

        delta[k, j, i] = du_dx[k, j, i] + dv_dy[k, j, i]

        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"]
            Divergence at T-points.  When ``self.mask`` is set, the
            output is zeroed at dry T-cells via ``* mask.h`` (applied
            inside the ``diff_*`` sub-calls — both sub-methods share
            the same T-point output stagger, so the mask is already
            baked into each addend).
        """
        # One fused stencil instead of two full-size outputs plus an add:
        #   self.diff_x_U_to_T(u) + self.diff_y_V_to_T(v)
        # gives the same values (the zero ghost rings add to zero and, for a
        # 0/1 mask, a*m + b*m == (a + b)*m exactly).
        # du[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   [Nz-2, Ny-2, Nx-2]
        # `interior` takes the dtype of u and the values, i.e. result_type(u, v).
        out = interior(
            diff_x_bwd_3d(u) / self.grid.dx + diff_y_bwd_3d(v) / self.grid.dy,
            u,
        )
        if self.mask is not None:
            out = out * self.mask.h
        return out

    def laplacian(self, h: Float[Array, "Nz Ny Nx"]) -> Float[Array, "Nz Ny Nx"]:
        """Horizontal Laplacian at T-points over all z-levels.

        nabla2_h[k, j, i] = (h[k, j, i+1] - 2*h[k, j, i] + h[k, j, i-1]) / dx^2
                           + (h[k, j+1, i] - 2*h[k, j, i] + h[k, j-1, i]) / dy^2

        Parameters
        ----------
        h : Float[Array, "Nz Ny Nx"]
            Scalar field at T-points.

        Returns
        -------
        Float[Array, "Nz Ny Nx"]
            Laplacian at T-points.  When ``self.mask`` is set, the
            output is zeroed at dry T-cells via ``* mask.h``.
        """
        # Each offset is sliced once (see `_Neighbours2D`); same arithmetic as
        #   d2x = (diff_x_fwd_3d(h) - diff_x_bwd_3d(h)) / dx**2
        #   d2y = (diff_y_fwd_3d(h) - diff_y_bwd_3d(h)) / dy**2
        nb = _neighbours(h[1:-1])  # z-interior levels, views [Nz-2, Ny-2, Nx-2]
        d2x = ((nb.e - nb.c) - (nb.c - nb.w)) / self.grid.dx**2
        d2y = ((nb.n - nb.c) - (nb.c - nb.s)) / self.grid.dy**2
        out = interior(d2x + d2y, h)
        if self.mask is not None:
            out = out * self.mask.h
        return out

`diff_x_T_to_U(h)`

Forward x-difference over all z-levels: T -> U.

dh_dx[k, j, i+1/2] = (h[k, j, i+1] - h[k, j, i]) / dx

Parameters:

Name	Type	Description	Default
`h`	`Float[Array, 'Nz Ny Nx']`	Scalar field at T-points.	required

Returns:

Type	Description
`Float[Array, 'Nz Ny Nx']`	Forward x-difference at U-points. When `self.mask` is set, the output is zeroed at dry U-faces via `* mask.u`.

Source code in finitevolx/_src/operators/difference.py

def diff_x_T_to_U(self, h: Float[Array, "Nz Ny Nx"]) -> Float[Array, "Nz Ny Nx"]:
    """Forward x-difference over all z-levels: T -> U.

    dh_dx[k, j, i+1/2] = (h[k, j, i+1] - h[k, j, i]) / dx

    Parameters
    ----------
    h : Float[Array, "Nz Ny Nx"]
        Scalar field at T-points.

    Returns
    -------
    Float[Array, "Nz Ny Nx"]
        Forward x-difference at U-points.  When ``self.mask`` is
        set, the output is zeroed at dry U-faces via ``* mask.u``.
    """
    out = interior(diff_x_fwd_3d(h) / self.grid.dx, h)
    if self.mask is not None:
        out = out * self.mask.u
    return out

`diff_x_U_to_T(u)`

Backward x-difference over all z-levels: U -> T.

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

Parameters:

Name	Type	Description	Default
`u`	`Float[Array, 'Nz Ny Nx']`	x-velocity at U-points.	required

Returns:

Type	Description
`Float[Array, 'Nz Ny Nx']`	Backward x-difference at T-points. When `self.mask` is set, the output is zeroed at dry T-cells via `* mask.h`.

Source code in finitevolx/_src/operators/difference.py

def diff_x_U_to_T(self, u: Float[Array, "Nz Ny Nx"]) -> Float[Array, "Nz Ny Nx"]:
    """Backward x-difference over all z-levels: U -> T.

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

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

    Returns
    -------
    Float[Array, "Nz Ny Nx"]
        Backward x-difference at T-points.  When ``self.mask`` is
        set, the output is zeroed at dry T-cells via ``* mask.h``.
    """
    out = interior(diff_x_bwd_3d(u) / self.grid.dx, u)
    if self.mask is not None:
        out = out * self.mask.h
    return out

`diff_y_T_to_V(h)`

Forward y-difference over all z-levels: T -> V.

dh_dy[k, j+1/2, i] = (h[k, j+1, i] - h[k, j, i]) / dy

Parameters:

Name	Type	Description	Default
`h`	`Float[Array, 'Nz Ny Nx']`	Scalar field at T-points.	required

Returns:

Type	Description
`Float[Array, 'Nz Ny Nx']`	Forward y-difference at V-points. When `self.mask` is set, the output is zeroed at dry V-faces via `* mask.v`.

Source code in finitevolx/_src/operators/difference.py

def diff_y_T_to_V(self, h: Float[Array, "Nz Ny Nx"]) -> Float[Array, "Nz Ny Nx"]:
    """Forward y-difference over all z-levels: T -> V.

    dh_dy[k, j+1/2, i] = (h[k, j+1, i] - h[k, j, i]) / dy

    Parameters
    ----------
    h : Float[Array, "Nz Ny Nx"]
        Scalar field at T-points.

    Returns
    -------
    Float[Array, "Nz Ny Nx"]
        Forward y-difference at V-points.  When ``self.mask`` is
        set, the output is zeroed at dry V-faces via ``* mask.v``.
    """
    out = interior(diff_y_fwd_3d(h) / self.grid.dy, h)
    if self.mask is not None:
        out = out * self.mask.v
    return out

`diff_y_V_to_T(v)`

Backward y-difference over all z-levels: V -> T.

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

Parameters:

Name	Type	Description	Default
`v`	`Float[Array, 'Nz Ny Nx']`	y-velocity at V-points.	required

Returns:

Type	Description
`Float[Array, 'Nz Ny Nx']`	Backward y-difference at T-points. When `self.mask` is set, the output is zeroed at dry T-cells via `* mask.h`.

Source code in finitevolx/_src/operators/difference.py

def diff_y_V_to_T(self, v: Float[Array, "Nz Ny Nx"]) -> Float[Array, "Nz Ny Nx"]:
    """Backward y-difference over all z-levels: V -> T.

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

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

    Returns
    -------
    Float[Array, "Nz Ny Nx"]
        Backward y-difference at T-points.  When ``self.mask`` is
        set, the output is zeroed at dry T-cells via ``* mask.h``.
    """
    out = interior(diff_y_bwd_3d(v) / self.grid.dy, v)
    if self.mask is not None:
        out = out * self.mask.h
    return out

`divergence(u, v)`

Horizontal divergence at T-points over all z-levels.

delta[k, j, i] = du_dx[k, j, i] + dv_dy[k, j, i]

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']`	Divergence at T-points. When `self.mask` is set, the output is zeroed at dry T-cells via `* mask.h` (applied inside the `diff_*` sub-calls — both sub-methods share the same T-point output stagger, so the mask is already baked into each addend).

Source code in finitevolx/_src/operators/difference.py

def divergence(
    self,
    u: Float[Array, "Nz Ny Nx"],
    v: Float[Array, "Nz Ny Nx"],
) -> Float[Array, "Nz Ny Nx"]:
    """Horizontal divergence at T-points over all z-levels.

    delta[k, j, i] = du_dx[k, j, i] + dv_dy[k, j, i]

    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"]
        Divergence at T-points.  When ``self.mask`` is set, the
        output is zeroed at dry T-cells via ``* mask.h`` (applied
        inside the ``diff_*`` sub-calls — both sub-methods share
        the same T-point output stagger, so the mask is already
        baked into each addend).
    """
    # One fused stencil instead of two full-size outputs plus an add:
    #   self.diff_x_U_to_T(u) + self.diff_y_V_to_T(v)
    # gives the same values (the zero ghost rings add to zero and, for a
    # 0/1 mask, a*m + b*m == (a + b)*m exactly).
    # du[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   [Nz-2, Ny-2, Nx-2]
    # `interior` takes the dtype of u and the values, i.e. result_type(u, v).
    out = interior(
        diff_x_bwd_3d(u) / self.grid.dx + diff_y_bwd_3d(v) / self.grid.dy,
        u,
    )
    if self.mask is not None:
        out = out * self.mask.h
    return out

`laplacian(h)`

Horizontal Laplacian at T-points over all z-levels.

nabla2_h[k, j, i] = (h[k, j, i+1] - 2h[k, j, i] + h[k, j, i-1]) / dx^2 + (h[k, j+1, i] - 2h[k, j, i] + h[k, j-1, i]) / dy^2

Parameters:

Name	Type	Description	Default
`h`	`Float[Array, 'Nz Ny Nx']`	Scalar field at T-points.	required

Returns:

Type	Description
`Float[Array, 'Nz Ny Nx']`	Laplacian at T-points. When `self.mask` is set, the output is zeroed at dry T-cells via `* mask.h`.

Source code in finitevolx/_src/operators/difference.py

def laplacian(self, h: Float[Array, "Nz Ny Nx"]) -> Float[Array, "Nz Ny Nx"]:
    """Horizontal Laplacian at T-points over all z-levels.

    nabla2_h[k, j, i] = (h[k, j, i+1] - 2*h[k, j, i] + h[k, j, i-1]) / dx^2
                       + (h[k, j+1, i] - 2*h[k, j, i] + h[k, j-1, i]) / dy^2

    Parameters
    ----------
    h : Float[Array, "Nz Ny Nx"]
        Scalar field at T-points.

    Returns
    -------
    Float[Array, "Nz Ny Nx"]
        Laplacian at T-points.  When ``self.mask`` is set, the
        output is zeroed at dry T-cells via ``* mask.h``.
    """
    # Each offset is sliced once (see `_Neighbours2D`); same arithmetic as
    #   d2x = (diff_x_fwd_3d(h) - diff_x_bwd_3d(h)) / dx**2
    #   d2y = (diff_y_fwd_3d(h) - diff_y_bwd_3d(h)) / dy**2
    nb = _neighbours(h[1:-1])  # z-interior levels, views [Nz-2, Ny-2, Nx-2]
    d2x = ((nb.e - nb.c) - (nb.c - nb.w)) / self.grid.dx**2
    d2y = ((nb.n - nb.c) - (nb.c - nb.s)) / self.grid.dy**2
    out = interior(d2x + d2y, h)
    if self.mask is not None:
        out = out * self.mask.h
    return out

Spherical

SphericalDifference2D also provides geostrophic_velocity(h, f, gravity), the public (mask-aware) entry point for the spherical geostrophic velocity.

`finitevolx.SphericalDifference2D`

Bases: Module

Spherical finite-difference operators on a 2-D Arakawa C-grid.

Parameters:

Name	Type	Description	Default
`grid`	`SphericalGrid2D`	The underlying 2-D spherical grid.	required
`mask`	`Mask2D or None`	Optional land/ocean mask. When provided, every method post-multiplies its output by the mask field matching its output stagger: T-output → `mask.h` (`diff_lon_U_to_T`, `diff_lat_V_to_T`, `diff2_lon`, `laplacian_merid`) U-output → `mask.u` (`diff_lon_T_to_U`) V-output → `mask.v` (`diff_lat_T_to_V`) X-output → `mask.xy_corner_strict` (`diff_lon_V_to_X`, `diff_lat_U_to_X`) (U, V)-output → `(mask.u, mask.v)` (`geostrophic_velocity`) Per #209 Q2, spherical 2-D operators take a Cartesian `Mask2D` rather than a dedicated `SphericalMask2D` — the mask geometry is coordinate-agnostic at the post-compute-multiply layer.	required

Source code in finitevolx/_src/operators/spherical_difference.py

class SphericalDifference2D(eqx.Module):
    """Spherical finite-difference operators on a 2-D Arakawa C-grid.

    Parameters
    ----------
    grid : SphericalGrid2D
        The underlying 2-D spherical grid.
    mask : Mask2D or None, optional
        Optional land/ocean mask.  When provided, every method
        post-multiplies its output by the mask field matching its
        output stagger:

        * T-output → ``mask.h`` (``diff_lon_U_to_T``, ``diff_lat_V_to_T``,
          ``diff2_lon``, ``laplacian_merid``)
        * U-output → ``mask.u`` (``diff_lon_T_to_U``)
        * V-output → ``mask.v`` (``diff_lat_T_to_V``)
        * X-output → ``mask.xy_corner_strict`` (``diff_lon_V_to_X``,
          ``diff_lat_U_to_X``)
        * (U, V)-output → ``(mask.u, mask.v)`` (``geostrophic_velocity``)

        Per #209 Q2, spherical 2-D operators take a Cartesian ``Mask2D``
        rather than a dedicated ``SphericalMask2D`` — the mask geometry
        is coordinate-agnostic at the post-compute-multiply layer.
    """

    grid: SphericalGrid2D
    mask: Mask2D | None = None

    # ------------------------------------------------------------------
    # Forward differences (centre/face → face/corner)
    # ------------------------------------------------------------------

    def diff_lon_T_to_U(self, h: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
        """Zonal derivative T → U on a sphere.

        dh/dx[j, i+½] = (h[j, i+1] − h[j, i]) / (R · cos(lat_U) · dlon)
        """
        raw = diff_x_fwd(h)
        cos_U = self.grid.cos_lat_U[1:-1, 1:-1]
        out = interior(_safe_div_cos(raw, cos_U, self.grid.R * self.grid.dlon), h)
        if self.mask is not None:
            out = out * self.mask.u
        return out

    def diff_lat_T_to_V(self, h: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
        """Meridional derivative T → V on a sphere.

        dh/dy[j+½, i] = (h[j+1, i] − h[j, i]) / (R · dlat)
        """
        raw = diff_y_fwd(h)
        out = interior(raw / (self.grid.R * self.grid.dlat), h)
        if self.mask is not None:
            out = out * self.mask.v
        return out

    def diff_lon_V_to_X(self, v: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
        """Zonal derivative V → X (corner) on a sphere.

        dv/dx[j+½, i+½] = (v[j+½, i+1] − v[j+½, i]) / (R · cos(lat_X) · dlon)
        """
        raw = diff_x_fwd(v)
        cos_X = self.grid.cos_lat_X[1:-1, 1:-1]
        out = interior(_safe_div_cos(raw, cos_X, self.grid.R * self.grid.dlon), v)
        if self.mask is not None:
            out = out * self.mask.xy_corner_strict
        return out

    def diff_lat_U_to_X(self, u: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
        """Meridional derivative U → X (corner) on a sphere.

        du/dy[j+½, i+½] = (u[j+1, i+½] − u[j, i+½]) / (R · dlat)
        """
        raw = diff_y_fwd(u)
        out = interior(raw / (self.grid.R * self.grid.dlat), u)
        if self.mask is not None:
            out = out * self.mask.xy_corner_strict
        return out

    # ------------------------------------------------------------------
    # Backward differences (face/corner → centre/face)
    # ------------------------------------------------------------------

    def diff_lon_U_to_T(self, u: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
        """Backward zonal derivative U → T on a sphere.

        du/dx[j, i] = (u[j, i+½] − u[j, i−½]) / (R · cos(lat_T) · dlon)
        """
        raw = diff_x_bwd(u)
        cos_T = self.grid.cos_lat_T[1:-1, 1:-1]
        out = interior(_safe_div_cos(raw, cos_T, self.grid.R * self.grid.dlon), u)
        if self.mask is not None:
            out = out * self.mask.h
        return out

    def diff_lat_V_to_T(self, v: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
        """Backward meridional derivative V → T on a sphere.

        dv/dy[j, i] = (v[j+½, i] − v[j−½, i]) / (R · dlat)
        """
        raw = diff_y_bwd(v)
        out = interior(raw / (self.grid.R * self.grid.dlat), v)
        if self.mask is not None:
            out = out * self.mask.h
        return out

    # ------------------------------------------------------------------
    # Second-order derivatives
    # ------------------------------------------------------------------

    def diff2_lon(self, h: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
        """Second zonal derivative at T-points on a sphere.

        d²h/dx² = 1/(R² cos²φ) · (h[j,i+1] − 2h[j,i] + h[j,i−1]) / dlon²
        """
        cos_T = self.grid.cos_lat_T[1:-1, 1:-1]
        # h_c sliced once (see `_Neighbours2D`); same arithmetic as
        #   (diff_x_fwd(h) - diff_x_bwd(h)) / dlon**2
        h_c, h_e, h_w = h[1:-1, 1:-1], h[1:-1, 2:], h[1:-1, :-2]
        d2h = ((h_e - h_c) - (h_c - h_w)) / self.grid.dlon**2
        out = interior(_safe_div_cos(d2h, cos_T, self.grid.R**2 * cos_T), h)
        if self.mask is not None:
            out = out * self.mask.h
        return out

    def laplacian_merid(self, h: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
        """Meridional term of the spherical Laplacian at T-points.

        1/(R² cosφ) · d/dφ(cosφ · dh/dφ)
        """
        cos_T = self.grid.cos_lat_T[1:-1, 1:-1]
        dlat = self.grid.dlat
        # dh/dlat at half-indices
        # h_c sliced once (see `_Neighbours2D`)
        h_c, h_n, h_s = h[1:-1, 1:-1], h[2:, 1:-1], h[:-2, 1:-1]
        dh_N = (h_n - h_c) / dlat  # (h[2:, 1:-1] - h[1:-1, 1:-1]) / dlat
        dh_S = (h_c - h_s) / dlat  # (h[1:-1, 1:-1] - h[:-2, 1:-1]) / dlat
        # cos(lat) at the V-faces, read from the grid (not averaged from T):
        #   cos_N = cos_lat_V[j]   (north face of cell j)
        #   cos_S = cos_lat_V[j-1] (south face of cell j)
        cos_N = self.grid.cos_lat_V[1:-1, 1:-1]
        cos_S = self.grid.cos_lat_V[:-2, 1:-1]
        d_cos_dh = (cos_N * dh_N - cos_S * dh_S) / dlat
        out = interior(_safe_div_cos(d_cos_dh, cos_T, self.grid.R**2), h)
        if self.mask is not None:
            out = out * self.mask.h
        return out

    # ------------------------------------------------------------------
    # Compound diagnostics
    # ------------------------------------------------------------------

    def geostrophic_velocity(
        self,
        h: Float[Array, "Ny Nx"],
        f: Float[Array, "Ny Nx"],
        gravity: float = 9.80665,
    ) -> tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]:
        """Geostrophic velocity from free-surface height on a sphere.

        u_g[j, i+1/2] = -g / (f_on_U[j, i+1/2] * R) * dh_dlat_U[j, i+1/2]
        v_g[j+1/2, i] =  g / (f_on_V[j+1/2, i] * R * cos_V[j+1/2, i])
                         * dh_dlon_V[j+1/2, i]

        with the compact 4-point stencils

        f_on_U[j, i+1/2]    = 1/2 * (f[j, i] + f[j, i+1])
        dh_dlat_U[j, i+1/2] = (h[j+1, i] + h[j+1, i+1]
                              - h[j-1, i] - h[j-1, i+1]) / (4 * dlat)
        f_on_V[j+1/2, i]    = 1/2 * (f[j, i] + f[j+1, i])
        cos_V[j+1/2, i]     = cos_lat_V[j, i] = cos(lat_T[j] + dlat/2)
        dh_dlon_V[j+1/2, i] = (h[j, i+1] + h[j+1, i+1]
                              - h[j, i-1] - h[j+1, i-1]) / (4 * dlon)

        Mask-aware class form of the private functional helper
        ``_geostrophic_velocity_sphere`` (not part of the public API).

        Parameters
        ----------
        h : Float[Array, "Ny Nx"]
            Free-surface height at T-points.
        f : Float[Array, "Ny Nx"]
            Coriolis parameter at T-points.
        gravity : float, optional
            Gravitational acceleration.

        Returns
        -------
        tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]
            ``(u_g, v_g)`` at U- and V-points, zero in the ghost ring and,
            when ``self.mask`` is set, at dry U- and V-faces.

        Notes
        -----
        Under a mask, the face-averaged Coriolis parameter is replaced by
        ``1`` on every dry face *before* the division
        (``f_on_U = 1`` where ``mask.u`` is False, ``f_on_V = 1`` where
        ``mask.v`` is False), so no dry face divides by zero -- whatever
        ``f`` holds on land, and even where a wet neighbour's ``f`` would
        cancel it -- keeping both the forward value and its reverse-mode
        gradient finite.  Wet faces keep their real ``f_on_face``.  The
        output is then masked with ``jnp.where``.  ``h`` is not modified:
        the 4-point stencils of wet coastal faces read ``h`` on the
        neighbouring land cells, so land ``h`` must be finite.
        """
        if self.mask is None:
            return _geostrophic_velocity_sphere(h, f, self.grid, gravity)
        u_g, v_g = _geostrophic_velocity_sphere(
            h, f, self.grid, gravity, wet_u=self.mask.u, wet_v=self.mask.v
        )
        u_g = jnp.where(self.mask.u, u_g, 0.0)
        v_g = jnp.where(self.mask.v, v_g, 0.0)
        return u_g, v_g

`diff2_lon(h)`

Second zonal derivative at T-points on a sphere.

d²h/dx² = 1/(R² cos²φ) · (h[j,i+1] − 2h[j,i] + h[j,i−1]) / dlon²

Source code in finitevolx/_src/operators/spherical_difference.py

def diff2_lon(self, h: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
    """Second zonal derivative at T-points on a sphere.

    d²h/dx² = 1/(R² cos²φ) · (h[j,i+1] − 2h[j,i] + h[j,i−1]) / dlon²
    """
    cos_T = self.grid.cos_lat_T[1:-1, 1:-1]
    # h_c sliced once (see `_Neighbours2D`); same arithmetic as
    #   (diff_x_fwd(h) - diff_x_bwd(h)) / dlon**2
    h_c, h_e, h_w = h[1:-1, 1:-1], h[1:-1, 2:], h[1:-1, :-2]
    d2h = ((h_e - h_c) - (h_c - h_w)) / self.grid.dlon**2
    out = interior(_safe_div_cos(d2h, cos_T, self.grid.R**2 * cos_T), h)
    if self.mask is not None:
        out = out * self.mask.h
    return out

`diff_lat_T_to_V(h)`

Meridional derivative T → V on a sphere.

dh/dy[j+½, i] = (h[j+1, i] − h[j, i]) / (R · dlat)

Source code in finitevolx/_src/operators/spherical_difference.py

def diff_lat_T_to_V(self, h: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
    """Meridional derivative T → V on a sphere.

    dh/dy[j+½, i] = (h[j+1, i] − h[j, i]) / (R · dlat)
    """
    raw = diff_y_fwd(h)
    out = interior(raw / (self.grid.R * self.grid.dlat), h)
    if self.mask is not None:
        out = out * self.mask.v
    return out

`diff_lat_U_to_X(u)`

Meridional derivative U → X (corner) on a sphere.

du/dy[j+½, i+½] = (u[j+1, i+½] − u[j, i+½]) / (R · dlat)

Source code in finitevolx/_src/operators/spherical_difference.py

def diff_lat_U_to_X(self, u: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
    """Meridional derivative U → X (corner) on a sphere.

    du/dy[j+½, i+½] = (u[j+1, i+½] − u[j, i+½]) / (R · dlat)
    """
    raw = diff_y_fwd(u)
    out = interior(raw / (self.grid.R * self.grid.dlat), u)
    if self.mask is not None:
        out = out * self.mask.xy_corner_strict
    return out

`diff_lat_V_to_T(v)`

Backward meridional derivative V → T on a sphere.

dv/dy[j, i] = (v[j+½, i] − v[j−½, i]) / (R · dlat)

Source code in finitevolx/_src/operators/spherical_difference.py

def diff_lat_V_to_T(self, v: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
    """Backward meridional derivative V → T on a sphere.

    dv/dy[j, i] = (v[j+½, i] − v[j−½, i]) / (R · dlat)
    """
    raw = diff_y_bwd(v)
    out = interior(raw / (self.grid.R * self.grid.dlat), v)
    if self.mask is not None:
        out = out * self.mask.h
    return out

`diff_lon_T_to_U(h)`

Zonal derivative T → U on a sphere.

dh/dx[j, i+½] = (h[j, i+1] − h[j, i]) / (R · cos(lat_U) · dlon)

Source code in finitevolx/_src/operators/spherical_difference.py

def diff_lon_T_to_U(self, h: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
    """Zonal derivative T → U on a sphere.

    dh/dx[j, i+½] = (h[j, i+1] − h[j, i]) / (R · cos(lat_U) · dlon)
    """
    raw = diff_x_fwd(h)
    cos_U = self.grid.cos_lat_U[1:-1, 1:-1]
    out = interior(_safe_div_cos(raw, cos_U, self.grid.R * self.grid.dlon), h)
    if self.mask is not None:
        out = out * self.mask.u
    return out

`diff_lon_U_to_T(u)`

Backward zonal derivative U → T on a sphere.

du/dx[j, i] = (u[j, i+½] − u[j, i−½]) / (R · cos(lat_T) · dlon)

Source code in finitevolx/_src/operators/spherical_difference.py

def diff_lon_U_to_T(self, u: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
    """Backward zonal derivative U → T on a sphere.

    du/dx[j, i] = (u[j, i+½] − u[j, i−½]) / (R · cos(lat_T) · dlon)
    """
    raw = diff_x_bwd(u)
    cos_T = self.grid.cos_lat_T[1:-1, 1:-1]
    out = interior(_safe_div_cos(raw, cos_T, self.grid.R * self.grid.dlon), u)
    if self.mask is not None:
        out = out * self.mask.h
    return out

`diff_lon_V_to_X(v)`

Zonal derivative V → X (corner) on a sphere.

dv/dx[j+½, i+½] = (v[j+½, i+1] − v[j+½, i]) / (R · cos(lat_X) · dlon)

Source code in finitevolx/_src/operators/spherical_difference.py

def diff_lon_V_to_X(self, v: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
    """Zonal derivative V → X (corner) on a sphere.

    dv/dx[j+½, i+½] = (v[j+½, i+1] − v[j+½, i]) / (R · cos(lat_X) · dlon)
    """
    raw = diff_x_fwd(v)
    cos_X = self.grid.cos_lat_X[1:-1, 1:-1]
    out = interior(_safe_div_cos(raw, cos_X, self.grid.R * self.grid.dlon), v)
    if self.mask is not None:
        out = out * self.mask.xy_corner_strict
    return out

`geostrophic_velocity(h, f, gravity=9.80665)`

Geostrophic velocity from free-surface height on a sphere.

u_g[j, i+1/2] = -g / (f_on_U[j, i+1/2] * R) * dh_dlat_U[j, i+1/2] v_g[j+1/2, i] = g / (f_on_V[j+1/2, i] * R * cos_V[j+1/2, i]) * dh_dlon_V[j+1/2, i]

with the compact 4-point stencils

f_on_U[j, i+1/2] = 1/2 * (f[j, i] + f[j, i+1]) dh_dlat_U[j, i+1/2] = (h[j+1, i] + h[j+1, i+1] - h[j-1, i] - h[j-1, i+1]) / (4 * dlat) f_on_V[j+1/2, i] = 1/2 * (f[j, i] + f[j+1, i]) cos_V[j+1/2, i] = cos_lat_V[j, i] = cos(lat_T[j] + dlat/2) dh_dlon_V[j+1/2, i] = (h[j, i+1] + h[j+1, i+1] - h[j, i-1] - h[j+1, i-1]) / (4 * dlon)

Mask-aware class form of the private functional helper _geostrophic_velocity_sphere (not part of the public API).

Parameters:

Name	Type	Description	Default
`h`	`Float[Array, 'Ny Nx']`	Free-surface height at T-points.	required
`f`	`Float[Array, 'Ny Nx']`	Coriolis parameter at T-points.	required
`gravity`	`float`	Gravitational acceleration.	`9.80665`

Returns:

Type	Description
`tuple[Float[Array, 'Ny Nx'], Float[Array, 'Ny Nx']]`	`(u_g, v_g)` at U- and V-points, zero in the ghost ring and, when `self.mask` is set, at dry U- and V-faces.

Notes

Under a mask, the face-averaged Coriolis parameter is replaced by 1 on every dry face before the division (f_on_U = 1 where mask.u is False, f_on_V = 1 where mask.v is False), so no dry face divides by zero -- whatever f holds on land, and even where a wet neighbour's f would cancel it -- keeping both the forward value and its reverse-mode gradient finite. Wet faces keep their real f_on_face. The output is then masked with jnp.where. h is not modified: the 4-point stencils of wet coastal faces read h on the neighbouring land cells, so land h must be finite.

Source code in finitevolx/_src/operators/spherical_difference.py

def geostrophic_velocity(
    self,
    h: Float[Array, "Ny Nx"],
    f: Float[Array, "Ny Nx"],
    gravity: float = 9.80665,
) -> tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]:
    """Geostrophic velocity from free-surface height on a sphere.

    u_g[j, i+1/2] = -g / (f_on_U[j, i+1/2] * R) * dh_dlat_U[j, i+1/2]
    v_g[j+1/2, i] =  g / (f_on_V[j+1/2, i] * R * cos_V[j+1/2, i])
                     * dh_dlon_V[j+1/2, i]

    with the compact 4-point stencils

    f_on_U[j, i+1/2]    = 1/2 * (f[j, i] + f[j, i+1])
    dh_dlat_U[j, i+1/2] = (h[j+1, i] + h[j+1, i+1]
                          - h[j-1, i] - h[j-1, i+1]) / (4 * dlat)
    f_on_V[j+1/2, i]    = 1/2 * (f[j, i] + f[j+1, i])
    cos_V[j+1/2, i]     = cos_lat_V[j, i] = cos(lat_T[j] + dlat/2)
    dh_dlon_V[j+1/2, i] = (h[j, i+1] + h[j+1, i+1]
                          - h[j, i-1] - h[j+1, i-1]) / (4 * dlon)

    Mask-aware class form of the private functional helper
    ``_geostrophic_velocity_sphere`` (not part of the public API).

    Parameters
    ----------
    h : Float[Array, "Ny Nx"]
        Free-surface height at T-points.
    f : Float[Array, "Ny Nx"]
        Coriolis parameter at T-points.
    gravity : float, optional
        Gravitational acceleration.

    Returns
    -------
    tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]
        ``(u_g, v_g)`` at U- and V-points, zero in the ghost ring and,
        when ``self.mask`` is set, at dry U- and V-faces.

    Notes
    -----
    Under a mask, the face-averaged Coriolis parameter is replaced by
    ``1`` on every dry face *before* the division
    (``f_on_U = 1`` where ``mask.u`` is False, ``f_on_V = 1`` where
    ``mask.v`` is False), so no dry face divides by zero -- whatever
    ``f`` holds on land, and even where a wet neighbour's ``f`` would
    cancel it -- keeping both the forward value and its reverse-mode
    gradient finite.  Wet faces keep their real ``f_on_face``.  The
    output is then masked with ``jnp.where``.  ``h`` is not modified:
    the 4-point stencils of wet coastal faces read ``h`` on the
    neighbouring land cells, so land ``h`` must be finite.
    """
    if self.mask is None:
        return _geostrophic_velocity_sphere(h, f, self.grid, gravity)
    u_g, v_g = _geostrophic_velocity_sphere(
        h, f, self.grid, gravity, wet_u=self.mask.u, wet_v=self.mask.v
    )
    u_g = jnp.where(self.mask.u, u_g, 0.0)
    v_g = jnp.where(self.mask.v, v_g, 0.0)
    return u_g, v_g

`laplacian_merid(h)`

Meridional term of the spherical Laplacian at T-points.

1/(R² cosφ) · d/dφ(cosφ · dh/dφ)

Source code in finitevolx/_src/operators/spherical_difference.py

def laplacian_merid(self, h: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
    """Meridional term of the spherical Laplacian at T-points.

    1/(R² cosφ) · d/dφ(cosφ · dh/dφ)
    """
    cos_T = self.grid.cos_lat_T[1:-1, 1:-1]
    dlat = self.grid.dlat
    # dh/dlat at half-indices
    # h_c sliced once (see `_Neighbours2D`)
    h_c, h_n, h_s = h[1:-1, 1:-1], h[2:, 1:-1], h[:-2, 1:-1]
    dh_N = (h_n - h_c) / dlat  # (h[2:, 1:-1] - h[1:-1, 1:-1]) / dlat
    dh_S = (h_c - h_s) / dlat  # (h[1:-1, 1:-1] - h[:-2, 1:-1]) / dlat
    # cos(lat) at the V-faces, read from the grid (not averaged from T):
    #   cos_N = cos_lat_V[j]   (north face of cell j)
    #   cos_S = cos_lat_V[j-1] (south face of cell j)
    cos_N = self.grid.cos_lat_V[1:-1, 1:-1]
    cos_S = self.grid.cos_lat_V[:-2, 1:-1]
    d_cos_dh = (cos_N * dh_N - cos_S * dh_S) / dlat
    out = interior(_safe_div_cos(d_cos_dh, cos_T, self.grid.R**2), h)
    if self.mask is not None:
        out = out * self.mask.h
    return out