Diffusion Operators

Harmonic (∇²) and biharmonic (∇⁴) diffusion operators for C-grid staggered models.

finitevolx.Diffusion2D

Bases: Module

Horizontal diffusion operator (flux form) on a 2-D Arakawa C-grid.

Computes ∂h/∂t = ∇·(κ ∇h) at T-points from staggered face fluxes via forward-then-backward finite differences.

Parameters:

Name	Type	Description	Default
grid	CartesianGrid2D	The underlying 2-D grid.	required
mask	Mask2D or None	Optional land/ocean mask. When provided, the three-step intermediate-flux masking pattern is applied inside both :meth:__call__ and :meth:fluxes: flux_x *= mask.u at the U-face stage, flux_y *= mask.v at the V-face stage, tendency *= mask.h on the final output (__call__ only). Unlike most other operators in this package, diffusion cannot use the simpler post-compute pattern because the divergence at wet T-cells adjacent to land would otherwise be contaminated by polluted dry-side face fluxes. None (default) matches the pre-existing unmasked behaviour bit for bit.	required

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import CartesianGrid2D, Diffusion2D
>>> grid = CartesianGrid2D.from_interior(8, 8, 1.0, 1.0)
>>> diff_op = Diffusion2D(grid=grid)
>>> h = jnp.ones((grid.Ny, grid.Nx))
>>> tendency = diff_op(h, kappa=1e-3)  # zero for constant tracer
>>> flux_x, flux_y = diff_op.fluxes(h, kappa=1e-3)

Source code in finitevolx/_src/diffusion/diffusion.py

class Diffusion2D(eqx.Module):
    """Horizontal diffusion operator (flux form) on a 2-D Arakawa C-grid.

    Computes ∂h/∂t = ∇·(κ ∇h) at T-points from staggered face fluxes via
    forward-then-backward finite differences.

    Parameters
    ----------
    grid : CartesianGrid2D
        The underlying 2-D grid.
    mask : Mask2D or None, optional
        Optional land/ocean mask.  When provided, the three-step
        intermediate-flux masking pattern is applied inside both
        :meth:`__call__` and :meth:`fluxes`:

        * ``flux_x *= mask.u`` at the U-face stage,
        * ``flux_y *= mask.v`` at the V-face stage,
        * tendency ``*= mask.h`` on the final output (``__call__`` only).

        Unlike most other operators in this package, diffusion cannot
        use the simpler post-compute pattern because the divergence at
        wet T-cells adjacent to land would otherwise be contaminated
        by polluted dry-side face fluxes.  ``None`` (default) matches
        the pre-existing unmasked behaviour bit for bit.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import CartesianGrid2D, Diffusion2D
    >>> grid = CartesianGrid2D.from_interior(8, 8, 1.0, 1.0)
    >>> diff_op = Diffusion2D(grid=grid)
    >>> h = jnp.ones((grid.Ny, grid.Nx))
    >>> tendency = diff_op(h, kappa=1e-3)  # zero for constant tracer
    >>> flux_x, flux_y = diff_op.fluxes(h, kappa=1e-3)
    """

    grid: CartesianGrid2D
    mask: Mask2D | None = None

    def __call__(
        self,
        h: Float[Array, "Ny Nx"],
        kappa: float | Float[Array, "Ny Nx"],
        wall: str = "closed",
    ) -> Float[Array, "Ny Nx"]:
        """Diffusion tendency ∂h/∂t = ∇·(κ ∇h) at T-points.

        dh[j, i] = (flux_x[j, i+1/2] - flux_x[j, i-1/2]) / dx
                 + (flux_y[j+1/2, i] - flux_y[j-1/2, i]) / dy

        where:
            flux_x[j, i+1/2] = κ[j, i+1/2] * (h[j, i+1] - h[j, i]) / dx
            flux_y[j+1/2, i] = κ[j+1/2, i] * (h[j+1, i] - h[j, i]) / dy

        with the face diffusivity ``κ[j, i+1/2] = (κ[j, i] + κ[j, i+1]) / 2``
        for a T-point field (a scalar ``κ`` is used as is).

        Parameters
        ----------
        h : Float[Array, "Ny Nx"]
            Tracer field at T-points.
        kappa : float or Float[Array, "Ny Nx"]
            Diffusion coefficient (scalar or T-point field).  A field is
            averaged to the faces (arithmetic mean of the two T cells).
            With ``wall="open"`` the wall faces read its ghost ring, which
            the caller fills consistently with the tracer BC (see the
            module docstring).
        wall : {"closed", "open"}, default "closed"
            Domain-wall face treatment.  ``"closed"`` (default, unchanged)
            leaves the four wall faces at zero, giving no-flux (closed-wall)
            BCs — the ghost ring has no effect on the horizontal diffusion
            term.  ``"open"`` also assembles the wall-face fluxes
            ``κ·(h_edge − h_ghost)/dx`` from the caller-filled ghost ring, so
            lateral Dirichlet / periodic boundaries produce a diffusive flux
            through the wall.  The flux field stays single-valued, so the
            tendency is exactly flux-conservative.  ``"open"`` is not
            supported together with a ``mask``.

        Returns
        -------
        Float[Array, "Ny Nx"]
            Diffusion tendency at T-points.  When ``self.mask`` is set,
            the intermediate-flux masking pattern described in the
            class docstring is applied.
        """
        _check_diffusion_wall(wall, self.mask)
        if self.mask is None:
            return _diffusion_2d_impl(
                h,
                kappa,
                self.grid.dx,
                self.grid.dy,
                mh=None,
                mu=None,
                mv=None,
                wall=wall,
            )
        return _diffusion_2d_impl(
            h,
            kappa,
            self.grid.dx,
            self.grid.dy,
            mh=self.mask.h,
            mu=self.mask.u,
            mv=self.mask.v,
            wall=wall,
        )

    def fluxes(
        self,
        h: Float[Array, "Ny Nx"],
        kappa: float | Float[Array, "Ny Nx"],
        wall: str = "closed",
    ) -> tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]:
        """Diagnostic diffusive face fluxes at U- and V-points.

        Returns the east-face and north-face diffusive fluxes before the
        divergence step, useful for flux-conservative time-stepping and
        diagnostics.

            flux_x[j, i+1/2] = κ[j, i+1/2] * (h[j, i+1] - h[j, i]) / dx  (U-points)
            flux_y[j+1/2, i] = κ[j+1/2, i] * (h[j+1, i] - h[j, i]) / dy  (V-points)

        with ``κ`` averaged to the faces as in :meth:`__call__`.

        Parameters
        ----------
        h : Float[Array, "Ny Nx"]
            Tracer field at T-points.
        kappa : float or Float[Array, "Ny Nx"]
            Diffusion coefficient (scalar or T-point field).
        wall : {"closed", "open"}, default "closed"
            Domain-wall face treatment (see :meth:`__call__`).  ``"closed"``
            leaves wall faces at zero; ``"open"`` also fills them from the
            ghost ring.  ``"open"`` is not supported together with a ``mask``.

        Returns
        -------
        tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]
            ``(flux_x, flux_y)`` — east-face fluxes at U-points and
            north-face fluxes at V-points.  When ``self.mask`` is set,
            ``flux_x`` is multiplied by ``mask.u`` and ``flux_y`` by
            ``mask.v``.
        """
        _check_diffusion_wall(wall, self.mask)
        if self.mask is None:
            return _diffusion_2d_fluxes_impl(
                h, kappa, self.grid.dx, self.grid.dy, mu=None, mv=None, wall=wall
            )
        return _diffusion_2d_fluxes_impl(
            h,
            kappa,
            self.grid.dx,
            self.grid.dy,
            mu=self.mask.u,
            mv=self.mask.v,
            wall=wall,
        )

__call__(h, kappa, wall='closed')

Diffusion tendency ∂h/∂t = ∇·(κ ∇h) at T-points.

dh[j, i] = (flux_x[j, i+1/2] - flux_x[j, i-1/2]) / dx + (flux_y[j+1/2, i] - flux_y[j-1/2, i]) / dy

where: flux_x[j, i+1/2] = κ[j, i+1/2] * (h[j, i+1] - h[j, i]) / dx flux_y[j+1/2, i] = κ[j+1/2, i] * (h[j+1, i] - h[j, i]) / dy

with the face diffusivity κ[j, i+1/2] = (κ[j, i] + κ[j, i+1]) / 2 for a T-point field (a scalar κ is used as is).

Parameters:

Name	Type	Description	Default
h	Float[Array, 'Ny Nx']	Tracer field at T-points.	required
kappa	float or Float[Array, 'Ny Nx']	Diffusion coefficient (scalar or T-point field). A field is averaged to the faces (arithmetic mean of the two T cells). With wall="open" the wall faces read its ghost ring, which the caller fills consistently with the tracer BC (see the module docstring).	required
wall	('closed', 'open')	Domain-wall face treatment. "closed" (default, unchanged) leaves the four wall faces at zero, giving no-flux (closed-wall) BCs — the ghost ring has no effect on the horizontal diffusion term. "open" also assembles the wall-face fluxes κ·(h_edge − h_ghost)/dx from the caller-filled ghost ring, so lateral Dirichlet / periodic boundaries produce a diffusive flux through the wall. The flux field stays single-valued, so the tendency is exactly flux-conservative. "open" is not supported together with a mask.	"closed"

Returns:

Type	Description
Float[Array, 'Ny Nx']	Diffusion tendency at T-points. When self.mask is set, the intermediate-flux masking pattern described in the class docstring is applied.

Source code in finitevolx/_src/diffusion/diffusion.py

def __call__(
    self,
    h: Float[Array, "Ny Nx"],
    kappa: float | Float[Array, "Ny Nx"],
    wall: str = "closed",
) -> Float[Array, "Ny Nx"]:
    """Diffusion tendency ∂h/∂t = ∇·(κ ∇h) at T-points.

    dh[j, i] = (flux_x[j, i+1/2] - flux_x[j, i-1/2]) / dx
             + (flux_y[j+1/2, i] - flux_y[j-1/2, i]) / dy

    where:
        flux_x[j, i+1/2] = κ[j, i+1/2] * (h[j, i+1] - h[j, i]) / dx
        flux_y[j+1/2, i] = κ[j+1/2, i] * (h[j+1, i] - h[j, i]) / dy

    with the face diffusivity ``κ[j, i+1/2] = (κ[j, i] + κ[j, i+1]) / 2``
    for a T-point field (a scalar ``κ`` is used as is).

    Parameters
    ----------
    h : Float[Array, "Ny Nx"]
        Tracer field at T-points.
    kappa : float or Float[Array, "Ny Nx"]
        Diffusion coefficient (scalar or T-point field).  A field is
        averaged to the faces (arithmetic mean of the two T cells).
        With ``wall="open"`` the wall faces read its ghost ring, which
        the caller fills consistently with the tracer BC (see the
        module docstring).
    wall : {"closed", "open"}, default "closed"
        Domain-wall face treatment.  ``"closed"`` (default, unchanged)
        leaves the four wall faces at zero, giving no-flux (closed-wall)
        BCs — the ghost ring has no effect on the horizontal diffusion
        term.  ``"open"`` also assembles the wall-face fluxes
        ``κ·(h_edge − h_ghost)/dx`` from the caller-filled ghost ring, so
        lateral Dirichlet / periodic boundaries produce a diffusive flux
        through the wall.  The flux field stays single-valued, so the
        tendency is exactly flux-conservative.  ``"open"`` is not
        supported together with a ``mask``.

    Returns
    -------
    Float[Array, "Ny Nx"]
        Diffusion tendency at T-points.  When ``self.mask`` is set,
        the intermediate-flux masking pattern described in the
        class docstring is applied.
    """
    _check_diffusion_wall(wall, self.mask)
    if self.mask is None:
        return _diffusion_2d_impl(
            h,
            kappa,
            self.grid.dx,
            self.grid.dy,
            mh=None,
            mu=None,
            mv=None,
            wall=wall,
        )
    return _diffusion_2d_impl(
        h,
        kappa,
        self.grid.dx,
        self.grid.dy,
        mh=self.mask.h,
        mu=self.mask.u,
        mv=self.mask.v,
        wall=wall,
    )

fluxes(h, kappa, wall='closed')

Diagnostic diffusive face fluxes at U- and V-points.

Returns the east-face and north-face diffusive fluxes before the divergence step, useful for flux-conservative time-stepping and diagnostics.

flux_x[j, i+1/2] = κ[j, i+1/2] * (h[j, i+1] - h[j, i]) / dx  (U-points)
flux_y[j+1/2, i] = κ[j+1/2, i] * (h[j+1, i] - h[j, i]) / dy  (V-points)

with κ averaged to the faces as in :meth:__call__.

Parameters:

Name	Type	Description	Default
h	Float[Array, 'Ny Nx']	Tracer field at T-points.	required
kappa	float or Float[Array, 'Ny Nx']	Diffusion coefficient (scalar or T-point field).	required
wall	('closed', 'open')	Domain-wall face treatment (see :meth:__call__). "closed" leaves wall faces at zero; "open" also fills them from the ghost ring. "open" is not supported together with a mask.	"closed"

Returns:

Type	Description
tuple[Float[Array, 'Ny Nx'], Float[Array, 'Ny Nx']]	(flux_x, flux_y) — east-face fluxes at U-points and north-face fluxes at V-points. When self.mask is set, flux_x is multiplied by mask.u and flux_y by mask.v.

Source code in finitevolx/_src/diffusion/diffusion.py

def fluxes(
    self,
    h: Float[Array, "Ny Nx"],
    kappa: float | Float[Array, "Ny Nx"],
    wall: str = "closed",
) -> tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]:
    """Diagnostic diffusive face fluxes at U- and V-points.

    Returns the east-face and north-face diffusive fluxes before the
    divergence step, useful for flux-conservative time-stepping and
    diagnostics.

        flux_x[j, i+1/2] = κ[j, i+1/2] * (h[j, i+1] - h[j, i]) / dx  (U-points)
        flux_y[j+1/2, i] = κ[j+1/2, i] * (h[j+1, i] - h[j, i]) / dy  (V-points)

    with ``κ`` averaged to the faces as in :meth:`__call__`.

    Parameters
    ----------
    h : Float[Array, "Ny Nx"]
        Tracer field at T-points.
    kappa : float or Float[Array, "Ny Nx"]
        Diffusion coefficient (scalar or T-point field).
    wall : {"closed", "open"}, default "closed"
        Domain-wall face treatment (see :meth:`__call__`).  ``"closed"``
        leaves wall faces at zero; ``"open"`` also fills them from the
        ghost ring.  ``"open"`` is not supported together with a ``mask``.

    Returns
    -------
    tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]
        ``(flux_x, flux_y)`` — east-face fluxes at U-points and
        north-face fluxes at V-points.  When ``self.mask`` is set,
        ``flux_x`` is multiplied by ``mask.u`` and ``flux_y`` by
        ``mask.v``.
    """
    _check_diffusion_wall(wall, self.mask)
    if self.mask is None:
        return _diffusion_2d_fluxes_impl(
            h, kappa, self.grid.dx, self.grid.dy, mu=None, mv=None, wall=wall
        )
    return _diffusion_2d_fluxes_impl(
        h,
        kappa,
        self.grid.dx,
        self.grid.dy,
        mu=self.mask.u,
        mv=self.mask.v,
        wall=wall,
    )

finitevolx.Diffusion3D

Bases: Module

Horizontal diffusion operator (flux form) on a 3-D Arakawa C-grid.

Applies ∂h/∂t = ∇·(κ ∇h) independently at 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 3-D land/ocean mask. When provided, the intermediate flux masking pattern from :class:Diffusion2D is applied at every z-level via vmap with per-z slices of mask.h, mask.u, mask.v (Pattern B per issue #209 — the only way to get correct divergence at wet T-cells adjacent to coastlines). None (default) matches the pre-existing unmasked behaviour bit for bit.	required

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import CartesianGrid3D, Diffusion3D
>>> grid = CartesianGrid3D.from_interior(4, 8, 8, 1.0, 1.0, 1.0)
>>> diff_op = Diffusion3D(grid=grid)
>>> h = jnp.ones((grid.Nz, grid.Ny, grid.Nx))
>>> tendency = diff_op(h, kappa=1e-3)  # zero for constant tracer

Source code in finitevolx/_src/diffusion/diffusion.py

class Diffusion3D(eqx.Module):
    """Horizontal diffusion operator (flux form) on a 3-D Arakawa C-grid.

    Applies ∂h/∂t = ∇·(κ ∇h) independently at 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 3-D land/ocean mask.  When provided, the intermediate
        flux masking pattern from :class:`Diffusion2D` is applied at
        every z-level via vmap with per-z slices of ``mask.h``,
        ``mask.u``, ``mask.v`` (Pattern B per issue #209 — the only
        way to get correct divergence at wet T-cells adjacent to
        coastlines).  ``None`` (default) matches the pre-existing
        unmasked behaviour bit for bit.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import CartesianGrid3D, Diffusion3D
    >>> grid = CartesianGrid3D.from_interior(4, 8, 8, 1.0, 1.0, 1.0)
    >>> diff_op = Diffusion3D(grid=grid)
    >>> h = jnp.ones((grid.Nz, grid.Ny, grid.Nx))
    >>> tendency = diff_op(h, kappa=1e-3)  # zero for constant tracer
    """

    grid: CartesianGrid3D
    mask: Mask3D | None = None

    def __call__(
        self,
        h: Float[Array, "Nz Ny Nx"],
        kappa: float | Float[Array, "Nz Ny Nx"],
        wall: str = "closed",
    ) -> Float[Array, "Nz Ny Nx"]:
        """Diffusion tendency ∂h/∂t = ∇·(κ ∇h) at T-points over all z-levels.

        Applies the horizontal diffusion stencil independently at each
        z-level.  Only interior cells ``[1:-1, 1:-1, 1:-1]`` are written;
        the ghost ring is left as zero.

        Parameters
        ----------
        h : Float[Array, "Nz Ny Nx"]
            Tracer field at T-points.
        kappa : float or Float[Array, "Nz Ny Nx"]
            Diffusion coefficient (scalar or T-point field).
        wall : {"closed", "open"}, default "closed"
            Lateral domain-wall face treatment, applied identically at every
            z-level (see :class:`Diffusion2D`).  ``"closed"`` (default,
            unchanged) gives no-flux lateral walls; ``"open"`` assembles the
            lateral wall-face fluxes from the caller-filled ghost ring so
            Dirichlet / periodic lateral boundaries drive horizontal
            diffusion.  ``"open"`` is not supported together with a ``mask``.

        Returns
        -------
        Float[Array, "Nz Ny Nx"]
            Diffusion tendency at T-points.  When ``self.mask`` is set,
            the intermediate-flux masking pattern is applied per-z-slice.
        """
        _check_diffusion_wall(wall, self.mask)
        dx, dy = self.grid.dx, self.grid.dy
        kappa_arr = jnp.asarray(kappa)
        kappa_ax = 0 if kappa_arr.ndim >= 3 else None

        if self.mask is None:
            # Unmasked path: vmap the free function over z-levels.
            def _apply_unmasked(h_k, kap_k):
                return _diffusion_2d_impl(
                    h_k, kap_k, dx, dy, mh=None, mu=None, mv=None, wall=wall
                )

            out = eqx.filter_vmap(_apply_unmasked, in_axes=(0, kappa_ax))(h, kappa_arr)
            return zero_z_ghosts(out)

        # Masked path: vmap with per-z slices of mask.h / mask.u / mask.v.
        mh = self.mask.h  # (Nz, Ny, Nx)
        mu = self.mask.u
        mv = self.mask.v

        def _apply_masked(h_k, kap_k, mh_k, mu_k, mv_k):
            return _diffusion_2d_impl(h_k, kap_k, dx, dy, mh=mh_k, mu=mu_k, mv=mv_k)

        out = eqx.filter_vmap(_apply_masked, in_axes=(0, kappa_ax, 0, 0, 0))(
            h, kappa_arr, mh, mu, mv
        )
        return zero_z_ghosts(out)

    def fluxes(
        self,
        h: Float[Array, "Nz Ny Nx"],
        kappa: float | Float[Array, "Nz Ny Nx"],
        wall: str = "closed",
    ) -> tuple[Float[Array, "Nz Ny Nx"], Float[Array, "Nz Ny Nx"]]:
        """Diagnostic diffusive face fluxes at U- and V-points, all z-levels.

        Parameters
        ----------
        h : Float[Array, "Nz Ny Nx"]
            Tracer field at T-points.
        kappa : float or Float[Array, "Nz Ny Nx"]
            Diffusion coefficient (scalar or T-point field).
        wall : {"closed", "open"}, default "closed"
            Lateral domain-wall face treatment (see :class:`Diffusion2D`).
            ``"open"`` is not supported together with a ``mask``.

        Returns
        -------
        tuple[Float[Array, "Nz Ny Nx"], Float[Array, "Nz Ny Nx"]]
            ``(flux_x, flux_y)`` — east-face fluxes at U-points and
            north-face fluxes at V-points.  When ``self.mask`` is set,
            ``flux_x`` is multiplied by ``mask.u`` and ``flux_y`` by
            ``mask.v`` at each z-level.
        """
        _check_diffusion_wall(wall, self.mask)
        dx, dy = self.grid.dx, self.grid.dy
        kappa_arr = jnp.asarray(kappa)
        kappa_ax = 0 if kappa_arr.ndim >= 3 else None

        if self.mask is None:

            def _apply_unmasked(h_k, kap_k):
                return _diffusion_2d_fluxes_impl(
                    h_k, kap_k, dx, dy, mu=None, mv=None, wall=wall
                )

            fx, fy = eqx.filter_vmap(_apply_unmasked, in_axes=(0, kappa_ax))(
                h, kappa_arr
            )
            return zero_z_ghosts(fx), zero_z_ghosts(fy)

        mu = self.mask.u
        mv = self.mask.v

        def _apply_masked(h_k, kap_k, mu_k, mv_k):
            return _diffusion_2d_fluxes_impl(h_k, kap_k, dx, dy, mu=mu_k, mv=mv_k)

        fx, fy = eqx.filter_vmap(_apply_masked, in_axes=(0, kappa_ax, 0, 0))(
            h, kappa_arr, mu, mv
        )
        return zero_z_ghosts(fx), zero_z_ghosts(fy)

__call__(h, kappa, wall='closed')

Diffusion tendency ∂h/∂t = ∇·(κ ∇h) at T-points over all z-levels.

Applies the horizontal diffusion stencil independently at each z-level. Only interior cells [1:-1, 1:-1, 1:-1] are written; the ghost ring is left as zero.

Parameters:

Name	Type	Description	Default
h	Float[Array, 'Nz Ny Nx']	Tracer field at T-points.	required
kappa	float or Float[Array, 'Nz Ny Nx']	Diffusion coefficient (scalar or T-point field).	required
wall	('closed', 'open')	Lateral domain-wall face treatment, applied identically at every z-level (see :class:Diffusion2D). "closed" (default, unchanged) gives no-flux lateral walls; "open" assembles the lateral wall-face fluxes from the caller-filled ghost ring so Dirichlet / periodic lateral boundaries drive horizontal diffusion. "open" is not supported together with a mask.	"closed"

Returns:

Type	Description
Float[Array, 'Nz Ny Nx']	Diffusion tendency at T-points. When self.mask is set, the intermediate-flux masking pattern is applied per-z-slice.

Source code in finitevolx/_src/diffusion/diffusion.py

def __call__(
    self,
    h: Float[Array, "Nz Ny Nx"],
    kappa: float | Float[Array, "Nz Ny Nx"],
    wall: str = "closed",
) -> Float[Array, "Nz Ny Nx"]:
    """Diffusion tendency ∂h/∂t = ∇·(κ ∇h) at T-points over all z-levels.

    Applies the horizontal diffusion stencil independently at each
    z-level.  Only interior cells ``[1:-1, 1:-1, 1:-1]`` are written;
    the ghost ring is left as zero.

    Parameters
    ----------
    h : Float[Array, "Nz Ny Nx"]
        Tracer field at T-points.
    kappa : float or Float[Array, "Nz Ny Nx"]
        Diffusion coefficient (scalar or T-point field).
    wall : {"closed", "open"}, default "closed"
        Lateral domain-wall face treatment, applied identically at every
        z-level (see :class:`Diffusion2D`).  ``"closed"`` (default,
        unchanged) gives no-flux lateral walls; ``"open"`` assembles the
        lateral wall-face fluxes from the caller-filled ghost ring so
        Dirichlet / periodic lateral boundaries drive horizontal
        diffusion.  ``"open"`` is not supported together with a ``mask``.

    Returns
    -------
    Float[Array, "Nz Ny Nx"]
        Diffusion tendency at T-points.  When ``self.mask`` is set,
        the intermediate-flux masking pattern is applied per-z-slice.
    """
    _check_diffusion_wall(wall, self.mask)
    dx, dy = self.grid.dx, self.grid.dy
    kappa_arr = jnp.asarray(kappa)
    kappa_ax = 0 if kappa_arr.ndim >= 3 else None

    if self.mask is None:
        # Unmasked path: vmap the free function over z-levels.
        def _apply_unmasked(h_k, kap_k):
            return _diffusion_2d_impl(
                h_k, kap_k, dx, dy, mh=None, mu=None, mv=None, wall=wall
            )

        out = eqx.filter_vmap(_apply_unmasked, in_axes=(0, kappa_ax))(h, kappa_arr)
        return zero_z_ghosts(out)

    # Masked path: vmap with per-z slices of mask.h / mask.u / mask.v.
    mh = self.mask.h  # (Nz, Ny, Nx)
    mu = self.mask.u
    mv = self.mask.v

    def _apply_masked(h_k, kap_k, mh_k, mu_k, mv_k):
        return _diffusion_2d_impl(h_k, kap_k, dx, dy, mh=mh_k, mu=mu_k, mv=mv_k)

    out = eqx.filter_vmap(_apply_masked, in_axes=(0, kappa_ax, 0, 0, 0))(
        h, kappa_arr, mh, mu, mv
    )
    return zero_z_ghosts(out)

fluxes(h, kappa, wall='closed')

Diagnostic diffusive face fluxes at U- and V-points, all z-levels.

Parameters:

Name	Type	Description	Default
h	Float[Array, 'Nz Ny Nx']	Tracer field at T-points.	required
kappa	float or Float[Array, 'Nz Ny Nx']	Diffusion coefficient (scalar or T-point field).	required
wall	('closed', 'open')	Lateral domain-wall face treatment (see :class:Diffusion2D). "open" is not supported together with a mask.	"closed"

Returns:

Type	Description
tuple[Float[Array, 'Nz Ny Nx'], Float[Array, 'Nz Ny Nx']]	(flux_x, flux_y) — east-face fluxes at U-points and north-face fluxes at V-points. When self.mask is set, flux_x is multiplied by mask.u and flux_y by mask.v at each z-level.

Source code in finitevolx/_src/diffusion/diffusion.py

def fluxes(
    self,
    h: Float[Array, "Nz Ny Nx"],
    kappa: float | Float[Array, "Nz Ny Nx"],
    wall: str = "closed",
) -> tuple[Float[Array, "Nz Ny Nx"], Float[Array, "Nz Ny Nx"]]:
    """Diagnostic diffusive face fluxes at U- and V-points, all z-levels.

    Parameters
    ----------
    h : Float[Array, "Nz Ny Nx"]
        Tracer field at T-points.
    kappa : float or Float[Array, "Nz Ny Nx"]
        Diffusion coefficient (scalar or T-point field).
    wall : {"closed", "open"}, default "closed"
        Lateral domain-wall face treatment (see :class:`Diffusion2D`).
        ``"open"`` is not supported together with a ``mask``.

    Returns
    -------
    tuple[Float[Array, "Nz Ny Nx"], Float[Array, "Nz Ny Nx"]]
        ``(flux_x, flux_y)`` — east-face fluxes at U-points and
        north-face fluxes at V-points.  When ``self.mask`` is set,
        ``flux_x`` is multiplied by ``mask.u`` and ``flux_y`` by
        ``mask.v`` at each z-level.
    """
    _check_diffusion_wall(wall, self.mask)
    dx, dy = self.grid.dx, self.grid.dy
    kappa_arr = jnp.asarray(kappa)
    kappa_ax = 0 if kappa_arr.ndim >= 3 else None

    if self.mask is None:

        def _apply_unmasked(h_k, kap_k):
            return _diffusion_2d_fluxes_impl(
                h_k, kap_k, dx, dy, mu=None, mv=None, wall=wall
            )

        fx, fy = eqx.filter_vmap(_apply_unmasked, in_axes=(0, kappa_ax))(
            h, kappa_arr
        )
        return zero_z_ghosts(fx), zero_z_ghosts(fy)

    mu = self.mask.u
    mv = self.mask.v

    def _apply_masked(h_k, kap_k, mu_k, mv_k):
        return _diffusion_2d_fluxes_impl(h_k, kap_k, dx, dy, mu=mu_k, mv=mv_k)

    fx, fy = eqx.filter_vmap(_apply_masked, in_axes=(0, kappa_ax, 0, 0))(
        h, kappa_arr, mu, mv
    )
    return zero_z_ghosts(fx), zero_z_ghosts(fy)

finitevolx.BiharmonicDiffusion2D

Bases: Module

Biharmonic (∇⁴) diffusion operator on a 2-D Arakawa C-grid.

Computes the local biharmonic diffusion tendency:

∂h/∂t|_diff = −κ · ∇⁴h

where ∇⁴h = ∇²(∇²h) is implemented as two successive flux-form Laplacians via :class:Diffusion2D. The negative sign ensures that a positive κ provides dissipation (the operator damps high-wavenumber modes).

Scale-selective property: for a Fourier mode with wavenumber k, the harmonic tendency scales as −κ_h · k² while the biharmonic tendency scales as −κ_bi · k⁴. Biharmonic diffusion therefore damps small scales much more strongly than large scales.

Only interior cells [1:-1, 1:-1] are written; the ghost ring is zero. The caller is responsible for boundary conditions.

Notes

Both Laplacian passes use the closed-wall (wall="closed") :class:Diffusion2D, whose wall-face fluxes are zero: the four domain walls are no-flux for h and for ∇²h, and neither the ghost ring of h nor that of the intermediate ∇²h (zero) is ever read. The operator is therefore conservative (Σ_interior ∂h/∂t = 0) and symmetric negative semi-definite on the interior.

Parameters:

Name	Type	Description	Default
grid	CartesianGrid2D	The underlying 2-D grid.	required
mask	Mask2D or None	Optional land/ocean mask. When provided, the tendency is −κ · Dm(Dm(h)) with Dm the masked no-flux Laplacian of :class:Diffusion2D (fluxes zeroed on faces with mask.u / mask.v dry, output zeroed where mask.h is dry): the coast is a no-flux wall for h and for ∇²h, exactly like the domain walls. Dry inputs are zeroed with jnp.where first, so land may hold any value (NaN, 1e10) without entering the wet tendency or its gradient. The masked operator is conservative and symmetric on the wet cells, and has no land→wet coupling. None (default) gives the unmasked operator.	None

References

.. [1] MITgcm Biharmonic Mixing: https://mitgcm.readthedocs.io/en/latest/optionals/packages/mixing.html#biharmonic-mixing .. [2] Leith, C. E. (1968). Diffusion approximation for two-dimensional turbulence. Physics of Fluids, 11(3), 671-673.

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import CartesianGrid2D, BiharmonicDiffusion2D
>>> grid = CartesianGrid2D.from_interior(8, 8, 1.0, 1.0)
>>> op = BiharmonicDiffusion2D(grid=grid)
>>> h = jnp.ones((grid.Ny, grid.Nx))
>>> tend = op(h, kappa=1e-6)  # zero for constant field
>>> tend.shape
(10, 10)

Source code in finitevolx/_src/diffusion/diffusion.py

class BiharmonicDiffusion2D(eqx.Module):
    """Biharmonic (∇⁴) diffusion operator on a 2-D Arakawa C-grid.

    Computes the local biharmonic diffusion tendency:

        ∂h/∂t|_diff = −κ · ∇⁴h

    where ∇⁴h = ∇²(∇²h) is implemented as two successive flux-form
    Laplacians via :class:`Diffusion2D`.  The negative sign ensures that a
    positive κ provides dissipation (the operator damps high-wavenumber
    modes).

    Scale-selective property: for a Fourier mode with wavenumber **k**, the
    harmonic tendency scales as ``−κ_h · k²`` while the biharmonic tendency
    scales as ``−κ_bi · k⁴``.  Biharmonic diffusion therefore damps small
    scales much more strongly than large scales.

    Only interior cells ``[1:-1, 1:-1]`` are written; the ghost ring is
    zero.  The caller is responsible for boundary conditions.

    Notes
    -----
    Both Laplacian passes use the closed-wall (``wall="closed"``)
    :class:`Diffusion2D`, whose wall-face fluxes are zero: the four domain
    walls are no-flux for ``h`` **and** for ``∇²h``, and neither the ghost
    ring of ``h`` nor that of the intermediate ``∇²h`` (zero) is ever read.
    The operator is therefore conservative (``Σ_interior ∂h/∂t = 0``) and
    symmetric negative semi-definite on the interior.

    Parameters
    ----------
    grid : CartesianGrid2D
        The underlying 2-D grid.
    mask : Mask2D or None, optional
        Optional land/ocean mask.  When provided, the tendency is
        ``−κ · Dm(Dm(h))`` with ``Dm`` the masked no-flux Laplacian of
        :class:`Diffusion2D` (fluxes zeroed on faces with ``mask.u`` /
        ``mask.v`` dry, output zeroed where ``mask.h`` is dry): the coast
        is a no-flux wall for ``h`` and for ``∇²h``, exactly like the
        domain walls.  Dry inputs are zeroed with ``jnp.where`` first, so
        land may hold any value (``NaN``, ``1e10``) without entering the
        wet tendency or its gradient.  The masked operator is
        conservative and symmetric on the wet cells, and has no land→wet
        coupling.  ``None`` (default) gives the unmasked operator.

    References
    ----------
    .. [1] MITgcm Biharmonic Mixing:
           https://mitgcm.readthedocs.io/en/latest/optionals/packages/mixing.html#biharmonic-mixing
    .. [2] Leith, C. E. (1968). Diffusion approximation for two-dimensional
           turbulence. *Physics of Fluids*, 11(3), 671-673.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import CartesianGrid2D, BiharmonicDiffusion2D
    >>> grid = CartesianGrid2D.from_interior(8, 8, 1.0, 1.0)
    >>> op = BiharmonicDiffusion2D(grid=grid)
    >>> h = jnp.ones((grid.Ny, grid.Nx))
    >>> tend = op(h, kappa=1e-6)  # zero for constant field
    >>> tend.shape
    (10, 10)
    """

    grid: CartesianGrid2D
    mask: Mask2D | None
    _harm: Diffusion2D

    def __init__(
        self,
        grid: CartesianGrid2D,
        mask: Mask2D | None = None,
    ) -> None:
        self.grid = grid
        self.mask = mask
        # Both passes are the masked no-flux Laplacian Dm (#292).
        self._harm = Diffusion2D(grid=grid, mask=mask)

    def __call__(
        self,
        h: Float[Array, "Ny Nx"],
        kappa: float,
    ) -> Float[Array, "Ny Nx"]:
        """Apply biharmonic diffusion and return the tendency -kappa * nabla^4 h.

        Parameters
        ----------
        h : Float[Array, "Ny Nx"]
            Scalar tracer field at T-points, shape ``[Ny, Nx]``.
        kappa : float
            Biharmonic diffusion coefficient (kappa >= 0 gives dissipation).

        Returns
        -------
        Float[Array, "Ny Nx"]
            Tendency -kappa * nabla^4 h at T-points, same shape as ``h``.
            Ghost cells are zero.  When ``self.mask`` is set, it is
            ``-kappa * Dm(Dm(h))`` (see the class docstring) and zero on
            dry cells.
        """
        if self.mask is not None:
            # land may hold NaN / 1e10: zero it before any stencil reads it
            # (NaN * 0 = NaN in the flux masking, and in its gradient)
            h = jnp.where(self.mask.h, h, 0.0)
        # First Laplacian pass: kappa=1.0 gives pure nabla^2 h (Dm h when
        # masked: zero on dry cells, no flux through the coast).
        # _materialize: the second pass reads lap1 at four offsets; keep it as
        # one buffer rather than letting XLA recompute it per offset.
        lap1 = _materialize(self._harm(h, kappa=1.0))
        # Second Laplacian pass: nabla^2(nabla^2 h) = nabla^4 h; the output
        # is already zero on dry cells when masked.
        lap2 = self._harm(lap1, kappa=1.0)
        return -kappa * lap2

__call__(h, kappa)

Apply biharmonic diffusion and return the tendency -kappa * nabla^4 h.

Parameters:

Name	Type	Description	Default
h	Float[Array, 'Ny Nx']	Scalar tracer field at T-points, shape [Ny, Nx].	required
kappa	float	Biharmonic diffusion coefficient (kappa >= 0 gives dissipation).	required

Returns:

Type	Description
Float[Array, 'Ny Nx']	Tendency -kappa * nabla^4 h at T-points, same shape as h. Ghost cells are zero. When self.mask is set, it is -kappa * Dm(Dm(h)) (see the class docstring) and zero on dry cells.

Source code in finitevolx/_src/diffusion/diffusion.py

def __call__(
    self,
    h: Float[Array, "Ny Nx"],
    kappa: float,
) -> Float[Array, "Ny Nx"]:
    """Apply biharmonic diffusion and return the tendency -kappa * nabla^4 h.

    Parameters
    ----------
    h : Float[Array, "Ny Nx"]
        Scalar tracer field at T-points, shape ``[Ny, Nx]``.
    kappa : float
        Biharmonic diffusion coefficient (kappa >= 0 gives dissipation).

    Returns
    -------
    Float[Array, "Ny Nx"]
        Tendency -kappa * nabla^4 h at T-points, same shape as ``h``.
        Ghost cells are zero.  When ``self.mask`` is set, it is
        ``-kappa * Dm(Dm(h))`` (see the class docstring) and zero on
        dry cells.
    """
    if self.mask is not None:
        # land may hold NaN / 1e10: zero it before any stencil reads it
        # (NaN * 0 = NaN in the flux masking, and in its gradient)
        h = jnp.where(self.mask.h, h, 0.0)
    # First Laplacian pass: kappa=1.0 gives pure nabla^2 h (Dm h when
    # masked: zero on dry cells, no flux through the coast).
    # _materialize: the second pass reads lap1 at four offsets; keep it as
    # one buffer rather than letting XLA recompute it per offset.
    lap1 = _materialize(self._harm(h, kappa=1.0))
    # Second Laplacian pass: nabla^2(nabla^2 h) = nabla^4 h; the output
    # is already zero on dry cells when masked.
    lap2 = self._harm(lap1, kappa=1.0)
    return -kappa * lap2

finitevolx.BiharmonicDiffusion3D

Bases: Module

Biharmonic (nabla^4) diffusion operator on a 3-D Arakawa C-grid.

Applies the horizontal biharmonic Laplacian independently at each z-level:

dh/dt|_diff = -kappa * nabla^4_h h

where nabla^4_h = nabla^2_h(nabla^2_h) denotes the horizontal biharmonic operator, implemented as two successive :class:Diffusion3D passes.

Only interior cells [1:-1, 1:-1, 1:-1] are written; the ghost ring is zero. The caller is responsible for boundary conditions.

Notes

Both passes are closed-wall: the lateral walls are no-flux for h and ∇²h and no ghost cell is read. See :class:BiharmonicDiffusion2D notes for details.

Parameters:

Name	Type	Description	Default
grid	CartesianGrid3D	The underlying 3-D grid.	required
mask	Mask3D or None	Optional 3-D land/ocean mask. When provided, the tendency is −κ · Dm(Dm(h)) at every z-level with Dm the masked no-flux :class:Diffusion3D, and dry inputs are zeroed first — same semantics as :class:BiharmonicDiffusion2D.	None

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import CartesianGrid3D, BiharmonicDiffusion3D
>>> grid = CartesianGrid3D.from_interior(4, 8, 8, 1.0, 1.0, 1.0)
>>> op = BiharmonicDiffusion3D(grid=grid)
>>> h = jnp.ones((grid.Nz, grid.Ny, grid.Nx))
>>> tend = op(h, kappa=1e-6)
>>> tend.shape
(6, 10, 10)

Source code in finitevolx/_src/diffusion/diffusion.py

class BiharmonicDiffusion3D(eqx.Module):
    """Biharmonic (nabla^4) diffusion operator on a 3-D Arakawa C-grid.

    Applies the horizontal biharmonic Laplacian independently at each
    z-level:

        dh/dt|_diff = -kappa * nabla^4_h h

    where nabla^4_h = nabla^2_h(nabla^2_h) denotes the horizontal biharmonic
    operator, implemented as two successive :class:`Diffusion3D` passes.

    Only interior cells ``[1:-1, 1:-1, 1:-1]`` are written; the ghost ring
    is zero.  The caller is responsible for boundary conditions.

    Notes
    -----
    Both passes are closed-wall: the lateral walls are no-flux for ``h``
    and ``∇²h`` and no ghost cell is read.  See
    :class:`BiharmonicDiffusion2D` notes for details.

    Parameters
    ----------
    grid : CartesianGrid3D
        The underlying 3-D grid.
    mask : Mask3D or None, optional
        Optional 3-D land/ocean mask.  When provided, the tendency is
        ``−κ · Dm(Dm(h))`` at every z-level with ``Dm`` the masked no-flux
        :class:`Diffusion3D`, and dry inputs are zeroed first — same
        semantics as :class:`BiharmonicDiffusion2D`.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import CartesianGrid3D, BiharmonicDiffusion3D
    >>> grid = CartesianGrid3D.from_interior(4, 8, 8, 1.0, 1.0, 1.0)
    >>> op = BiharmonicDiffusion3D(grid=grid)
    >>> h = jnp.ones((grid.Nz, grid.Ny, grid.Nx))
    >>> tend = op(h, kappa=1e-6)
    >>> tend.shape
    (6, 10, 10)
    """

    grid: CartesianGrid3D
    mask: Mask3D | None
    _harm: Diffusion3D

    def __init__(
        self,
        grid: CartesianGrid3D,
        mask: Mask3D | None = None,
    ) -> None:
        self.grid = grid
        self.mask = mask
        # Both passes are the masked no-flux Laplacian Dm (#292).
        self._harm = Diffusion3D(grid=grid, mask=mask)

    def __call__(
        self,
        h: Float[Array, "Nz Ny Nx"],
        kappa: float,
    ) -> Float[Array, "Nz Ny Nx"]:
        """Apply horizontal biharmonic diffusion: -kappa * nabla^4_h h.

        Parameters
        ----------
        h : Float[Array, "Nz Ny Nx"]
            Scalar tracer field at T-points, shape ``[Nz, Ny, Nx]``.
        kappa : float
            Biharmonic diffusion coefficient (kappa >= 0 gives dissipation).

        Returns
        -------
        Float[Array, "Nz Ny Nx"]
            Tendency -kappa * nabla^4_h h at T-points, same shape as ``h``.
            When ``self.mask`` is set, it is ``-kappa * Dm(Dm(h))`` and zero
            on dry cells.
        """
        if self.mask is not None:
            # zero land (NaN / 1e10) before any stencil reads it
            h = jnp.where(self.mask.h, h, 0.0)
        # _materialize: the second pass reads lap1 at four offsets; keep it as
        # one buffer rather than letting XLA recompute it per offset.
        lap1 = _materialize(self._harm(h, kappa=1.0))
        lap2 = self._harm(lap1, kappa=1.0)
        return -kappa * lap2

__call__(h, kappa)

Apply horizontal biharmonic diffusion: -kappa * nabla^4_h h.

Parameters:

Name	Type	Description	Default
h	Float[Array, 'Nz Ny Nx']	Scalar tracer field at T-points, shape [Nz, Ny, Nx].	required
kappa	float	Biharmonic diffusion coefficient (kappa >= 0 gives dissipation).	required

Returns:

Type	Description
Float[Array, 'Nz Ny Nx']	Tendency -kappa * nabla^4_h h at T-points, same shape as h. When self.mask is set, it is -kappa * Dm(Dm(h)) and zero on dry cells.

Source code in finitevolx/_src/diffusion/diffusion.py

def __call__(
    self,
    h: Float[Array, "Nz Ny Nx"],
    kappa: float,
) -> Float[Array, "Nz Ny Nx"]:
    """Apply horizontal biharmonic diffusion: -kappa * nabla^4_h h.

    Parameters
    ----------
    h : Float[Array, "Nz Ny Nx"]
        Scalar tracer field at T-points, shape ``[Nz, Ny, Nx]``.
    kappa : float
        Biharmonic diffusion coefficient (kappa >= 0 gives dissipation).

    Returns
    -------
    Float[Array, "Nz Ny Nx"]
        Tendency -kappa * nabla^4_h h at T-points, same shape as ``h``.
        When ``self.mask`` is set, it is ``-kappa * Dm(Dm(h))`` and zero
        on dry cells.
    """
    if self.mask is not None:
        # zero land (NaN / 1e10) before any stencil reads it
        h = jnp.where(self.mask.h, h, 0.0)
    # _materialize: the second pass reads lap1 at four offsets; keep it as
    # one buffer rather than letting XLA recompute it per offset.
    lap1 = _materialize(self._harm(h, kappa=1.0))
    lap2 = self._harm(lap1, kappa=1.0)
    return -kappa * lap2

finitevolx.diffusion_2d(h, kappa, dx, dy)

Horizontal tracer diffusion tendency at T-points (flux form).

Computes ∂h/∂t = ∇·(κ ∇h) = ∂/∂x(κ ∂h/∂x) + ∂/∂y(κ ∂h/∂y) at interior T-points using forward-then-backward finite differences.

Only interior cells [1:-1, 1:-1] are written; the ghost ring is left as zero. East and north boundary faces are not computed, giving no-flux (closed-wall) BCs at all four domain walls by default.

This is the mask-free functional form. For masked diffusion, use :class:Diffusion2D with a mask= class attribute; it applies the mask via the intermediate-flux pattern described in the module docstring.

Parameters:

Name	Type	Description	Default
h	Float[Array, 'Ny Nx']	Tracer field at T-points.	required
kappa	float or Float[Array, 'Ny Nx']	Diffusion coefficient. May be a scalar or an array of the same shape as h (spatially varying diffusivity at T-points). When kappa is an array, each face flux uses the arithmetic mean of the two T-cells the face separates, κ[j, i+1/2] = (κ[j, i] + κ[j, i+1]) / 2 (second order); no ghost value of kappa is read.	required
dx	float	Grid spacing in x.	required
dy	float	Grid spacing in y.	required

Returns:

Type	Description
Float[Array, 'Ny Nx']	Diffusion tendency ∂h/∂t at T-points, same shape as h.

Examples:

>>> import jax.numpy as jnp
>>> h = jnp.zeros((10, 10))
>>> tendency = diffusion_2d(h, kappa=1.0, dx=0.1, dy=0.1)
>>> tendency.shape
(10, 10)

Source code in finitevolx/_src/diffusion/diffusion.py

def diffusion_2d(
    h: Float[Array, "Ny Nx"],
    kappa: float | Float[Array, "Ny Nx"],
    dx: float,
    dy: float,
) -> Float[Array, "Ny Nx"]:
    """Horizontal tracer diffusion tendency at T-points (flux form).

    Computes ∂h/∂t = ∇·(κ ∇h) = ∂/∂x(κ ∂h/∂x) + ∂/∂y(κ ∂h/∂y)
    at interior T-points using forward-then-backward finite differences.

    Only interior cells ``[1:-1, 1:-1]`` are written; the ghost ring is left
    as zero.  East and north boundary faces are not computed, giving no-flux
    (closed-wall) BCs at all four domain walls by default.

    This is the mask-free functional form.  For masked diffusion, use
    :class:`Diffusion2D` with a ``mask=`` class attribute; it applies the
    mask via the intermediate-flux pattern described in the module
    docstring.

    Parameters
    ----------
    h : Float[Array, "Ny Nx"]
        Tracer field at T-points.
    kappa : float or Float[Array, "Ny Nx"]
        Diffusion coefficient.  May be a scalar or an array of the same shape
        as ``h`` (spatially varying diffusivity at T-points).  When ``kappa``
        is an array, each face flux uses the arithmetic mean of the two
        T-cells the face separates, ``κ[j, i+1/2] = (κ[j, i] + κ[j, i+1]) / 2``
        (second order); no ghost value of ``kappa`` is read.
    dx : float
        Grid spacing in x.
    dy : float
        Grid spacing in y.

    Returns
    -------
    Float[Array, "Ny Nx"]
        Diffusion tendency ∂h/∂t at T-points, same shape as ``h``.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> h = jnp.zeros((10, 10))
    >>> tendency = diffusion_2d(h, kappa=1.0, dx=0.1, dy=0.1)
    >>> tendency.shape
    (10, 10)
    """
    return _diffusion_2d_impl(h, kappa, dx, dy, mh=None, mu=None, mv=None)