Advection Operators

Flux-form advection operators for scalar transport on staggered C-grids.

`finitevolx.Advection1D`

Bases: Module

1-D advection operator.

Parameters:

Name	Type	Description	Default
`grid`	`CartesianGrid1D`	The underlying 1-D grid.	required
`mask`	`Mask1D or None`	Optional land/ocean mask. When provided and `__call__` is invoked with a mask-dispatchable method (WENO3/5, any TVD limiter), `upwind_flux` is used with the pre-built adaptive stencil hierarchy, falling back to lower-order stencils near coastlines. For non-dispatchable methods (`naive`, `upwind1/2/3`, `weno7`, `weno9`) the unmasked code path runs and the final tendency is post-multiplied by `mask.h` so dry T-cells stay exactly zero. `None` (default) matches the pre-existing unmasked behaviour bit for bit.	`None`

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

class Advection1D(eqx.Module):
    """1-D advection operator.

    Parameters
    ----------
    grid : CartesianGrid1D
        The underlying 1-D grid.
    mask : Mask1D or None, optional
        Optional land/ocean mask.  When provided and ``__call__`` is
        invoked with a mask-dispatchable method (WENO3/5, any TVD
        limiter), ``upwind_flux`` is used with the pre-built adaptive
        stencil hierarchy, falling back to lower-order stencils near
        coastlines.  For non-dispatchable methods (``naive``,
        ``upwind1/2/3``, ``weno7``, ``weno9``) the unmasked code path
        runs and the final tendency is post-multiplied by ``mask.h``
        so dry T-cells stay exactly zero.  ``None`` (default) matches
        the pre-existing unmasked behaviour bit for bit.
    """

    grid: CartesianGrid1D
    mask: Mask1D | None
    recon: Reconstruction1D
    _mask_hierarchy_x: dict[int, Bool[Array, "Nx"]] | None

    def __init__(
        self,
        grid: CartesianGrid1D,
        mask: Mask1D | None = None,
    ) -> None:
        self.grid = grid
        self.mask = mask
        self.recon = Reconstruction1D(grid=grid)
        if mask is not None:
            self._mask_hierarchy_x = mask.get_adaptive_masks(
                stencil_sizes=_HIERARCHY_SIZES
            )
        else:
            self._mask_hierarchy_x = None

    def __call__(
        self,
        h: Float[Array, "Nx"],
        u: Float[Array, "Nx"],
        method: str = "upwind1",
        wall: str = "closed",
    ) -> Float[Array, "Nx"]:
        """Advective tendency -d(h*u)/dx at T-points.

        dh[i] = -(fe[i+1/2] - fe[i-1/2]) / dx

        Parameters
        ----------
        h : Float[Array, "Nx"]
            Scalar at T-points.
        u : Float[Array, "Nx"]
            Velocity at U-points.
        method : str
            Reconstruction method: ``'naive'``, ``'upwind1'``, ``'upwind2'``,
            ``'upwind3'``, ``'weno3'``, ``'weno5'``, ``'weno7'``, ``'weno9'``,
            or a flux-limiter TVD scheme: ``'minmod'``, ``'van_leer'``,
            ``'superbee'``, ``'mc'``.
        wall : {"closed", "open"}, default "closed"
            Lateral-boundary treatment.  ``"closed"`` (default, unchanged)
            writes the tendency only on the strict interior ``[2:-2]`` and
            treats the domain walls as no-flux — the wall-adjacent interior
            cells hold stale values.  ``"open"`` writes the full interior
            ``[1:-1]`` and computes the two domain-wall face fluxes with a
            first-order upwind reconstruction that reads the ghost ring, so
            Dirichlet / outflow / periodic boundaries (applied by the caller
            to the ghost cells beforehand) drive the advective transport.
            The interior faces keep the requested high-order ``method``; the
            flux field stays single-valued, so mass is conserved exactly.
            ``"open"`` is not supported together with a ``mask``.

        Returns
        -------
        Float[Array, "Nx"]
            Advective tendency at T-points.
        """
        _validate_wall(wall, self.mask)
        # ── masked path: route through upwind_flux with adaptive stencils ─
        if self._mask_hierarchy_x is not None and method in _MASK_DISPATCHABLE:
            rfx, sizes = _rec_funcs_for_method_1d(self.recon, method)
            mh = narrow_mask_hierarchy(self._mask_hierarchy_x, sizes)
            fe = upwind_flux(h, u, dim=0, rec_funcs=rfx, mask_hierarchy=mh)
        # ── unmasked path: existing per-method dispatch ─────────────────
        elif method == "naive":
            fe = self.recon.naive_x(h, u)
        elif method == "upwind1":
            fe = self.recon.upwind1_x(h, u)
        elif method == "upwind2":
            fe = self.recon.upwind2_x(h, u)
        elif method == "upwind3":
            fe = self.recon.upwind3_x(h, u)
        elif method == "weno3":
            fe = self.recon.weno3_x(h, u)
        elif method == "weno5":
            fe = self.recon.weno5_x(h, u)
        elif method == "weno7":
            fe = self.recon.weno7_x(h, u)
        elif method == "weno9":
            fe = self.recon.weno9_x(h, u)
        elif method in _TVD_LIMITERS:
            fe = self.recon.tvd_x(h, u, limiter=method)
        else:
            raise ValueError(f"Unknown method: {method!r}")

        out = jnp.zeros_like(h)
        # dh[i] = -(fe[i+1/2] - fe[i-1/2]) / dx
        # fe[i] represents the flux at the east face of cell i (at i+1/2)
        # For cell i, we need fe[i] (east) and fe[i-1] (west)
        if wall == "open":
            # Overwrite the two domain-wall faces (west face 1/2 == fe[0],
            # east face Nx-3/2 == fe[-2]) with a first-order upwind flux that
            # reads the ghost ring, then write the full interior [1:-1].
            fe = fe.at[0].set(_wall_upwind_flux(u[0], h[0], h[1]))
            fe = fe.at[-2].set(_wall_upwind_flux(u[-2], h[-2], h[-1]))
            out = out.at[1:-1].set(-(fe[1:-1] - fe[0:-2]) / self.grid.dx)
        else:
            # Only use face fluxes that are defined by the reconstruction
            # scheme, avoiding the ghost-ring entries fe[0] and fe[-1].
            out = out.at[2:-2].set(-(fe[2:-2] - fe[1:-3]) / self.grid.dx)
        if self.mask is not None:
            out = out * self.mask.h
        return out

`call(h, u, method='upwind1', wall='closed')`

Advective tendency -d(h*u)/dx at T-points.

dh[i] = -(fe[i+1/2] - fe[i-1/2]) / dx

Parameters:

Name	Type	Description	Default
`h`	`Float[Array, 'Nx']`	Scalar at T-points.	required
`u`	`Float[Array, 'Nx']`	Velocity at U-points.	required
`method`	`str`	Reconstruction method: `'naive'`, `'upwind1'`, `'upwind2'`, `'upwind3'`, `'weno3'`, `'weno5'`, `'weno7'`, `'weno9'`, or a flux-limiter TVD scheme: `'minmod'`, `'van_leer'`, `'superbee'`, `'mc'`.	`'upwind1'`
`wall`	`('closed', 'open')`	Lateral-boundary treatment. `"closed"` (default, unchanged) writes the tendency only on the strict interior `[2:-2]` and treats the domain walls as no-flux — the wall-adjacent interior cells hold stale values. `"open"` writes the full interior `[1:-1]` and computes the two domain-wall face fluxes with a first-order upwind reconstruction that reads the ghost ring, so Dirichlet / outflow / periodic boundaries (applied by the caller to the ghost cells beforehand) drive the advective transport. The interior faces keep the requested high-order `method`; the flux field stays single-valued, so mass is conserved exactly. `"open"` is not supported together with a `mask`.	`"closed"`

Returns:

Type	Description
`Float[Array, 'Nx']`	Advective tendency at T-points.

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

def __call__(
    self,
    h: Float[Array, "Nx"],
    u: Float[Array, "Nx"],
    method: str = "upwind1",
    wall: str = "closed",
) -> Float[Array, "Nx"]:
    """Advective tendency -d(h*u)/dx at T-points.

    dh[i] = -(fe[i+1/2] - fe[i-1/2]) / dx

    Parameters
    ----------
    h : Float[Array, "Nx"]
        Scalar at T-points.
    u : Float[Array, "Nx"]
        Velocity at U-points.
    method : str
        Reconstruction method: ``'naive'``, ``'upwind1'``, ``'upwind2'``,
        ``'upwind3'``, ``'weno3'``, ``'weno5'``, ``'weno7'``, ``'weno9'``,
        or a flux-limiter TVD scheme: ``'minmod'``, ``'van_leer'``,
        ``'superbee'``, ``'mc'``.
    wall : {"closed", "open"}, default "closed"
        Lateral-boundary treatment.  ``"closed"`` (default, unchanged)
        writes the tendency only on the strict interior ``[2:-2]`` and
        treats the domain walls as no-flux — the wall-adjacent interior
        cells hold stale values.  ``"open"`` writes the full interior
        ``[1:-1]`` and computes the two domain-wall face fluxes with a
        first-order upwind reconstruction that reads the ghost ring, so
        Dirichlet / outflow / periodic boundaries (applied by the caller
        to the ghost cells beforehand) drive the advective transport.
        The interior faces keep the requested high-order ``method``; the
        flux field stays single-valued, so mass is conserved exactly.
        ``"open"`` is not supported together with a ``mask``.

    Returns
    -------
    Float[Array, "Nx"]
        Advective tendency at T-points.
    """
    _validate_wall(wall, self.mask)
    # ── masked path: route through upwind_flux with adaptive stencils ─
    if self._mask_hierarchy_x is not None and method in _MASK_DISPATCHABLE:
        rfx, sizes = _rec_funcs_for_method_1d(self.recon, method)
        mh = narrow_mask_hierarchy(self._mask_hierarchy_x, sizes)
        fe = upwind_flux(h, u, dim=0, rec_funcs=rfx, mask_hierarchy=mh)
    # ── unmasked path: existing per-method dispatch ─────────────────
    elif method == "naive":
        fe = self.recon.naive_x(h, u)
    elif method == "upwind1":
        fe = self.recon.upwind1_x(h, u)
    elif method == "upwind2":
        fe = self.recon.upwind2_x(h, u)
    elif method == "upwind3":
        fe = self.recon.upwind3_x(h, u)
    elif method == "weno3":
        fe = self.recon.weno3_x(h, u)
    elif method == "weno5":
        fe = self.recon.weno5_x(h, u)
    elif method == "weno7":
        fe = self.recon.weno7_x(h, u)
    elif method == "weno9":
        fe = self.recon.weno9_x(h, u)
    elif method in _TVD_LIMITERS:
        fe = self.recon.tvd_x(h, u, limiter=method)
    else:
        raise ValueError(f"Unknown method: {method!r}")

    out = jnp.zeros_like(h)
    # dh[i] = -(fe[i+1/2] - fe[i-1/2]) / dx
    # fe[i] represents the flux at the east face of cell i (at i+1/2)
    # For cell i, we need fe[i] (east) and fe[i-1] (west)
    if wall == "open":
        # Overwrite the two domain-wall faces (west face 1/2 == fe[0],
        # east face Nx-3/2 == fe[-2]) with a first-order upwind flux that
        # reads the ghost ring, then write the full interior [1:-1].
        fe = fe.at[0].set(_wall_upwind_flux(u[0], h[0], h[1]))
        fe = fe.at[-2].set(_wall_upwind_flux(u[-2], h[-2], h[-1]))
        out = out.at[1:-1].set(-(fe[1:-1] - fe[0:-2]) / self.grid.dx)
    else:
        # Only use face fluxes that are defined by the reconstruction
        # scheme, avoiding the ghost-ring entries fe[0] and fe[-1].
        out = out.at[2:-2].set(-(fe[2:-2] - fe[1:-3]) / self.grid.dx)
    if self.mask is not None:
        out = out * self.mask.h
    return out

`finitevolx.Advection2D`

Bases: Module

2-D advection operator.

Parameters:

Name	Type	Description	Default
`grid`	`CartesianGrid2D`	The underlying 2-D grid.	required
`mask`	`Mask2D or None`	Optional land/ocean mask. When provided, the `(2, 4, 6)` adaptive stencil hierarchies for both directions are pre-built once in `__init__` and reused on every call — this is the main reason the API moved from a per-call keyword to a class field. For mask-dispatchable methods (WENO3/5, WENOz5, any TVD limiter), `upwind_flux` is used with the stored hierarchy narrowed to the method's stencil sizes. For non-dispatchable methods the unmasked code path runs and the final tendency is post-multiplied by `mask.h`. `None` (default) matches the pre-existing unmasked behaviour bit for bit.	`None`

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

class Advection2D(eqx.Module):
    """2-D advection operator.

    Parameters
    ----------
    grid : CartesianGrid2D
        The underlying 2-D grid.
    mask : Mask2D or None, optional
        Optional land/ocean mask.  When provided, the ``(2, 4, 6)``
        adaptive stencil hierarchies for both directions are
        pre-built once in ``__init__`` and reused on every call —
        this is the main reason the API moved from a per-call
        keyword to a class field.  For mask-dispatchable methods
        (WENO3/5, WENOz5, any TVD limiter), ``upwind_flux`` is used
        with the stored hierarchy narrowed to the method's stencil
        sizes.  For non-dispatchable methods the unmasked code path
        runs and the final tendency is post-multiplied by ``mask.h``.
        ``None`` (default) matches the pre-existing unmasked
        behaviour bit for bit.
    """

    grid: CartesianGrid2D
    mask: Mask2D | None
    recon: Reconstruction2D
    _mask_hierarchy_x: dict[int, Bool[Array, "Ny Nx"]] | None
    _mask_hierarchy_y: dict[int, Bool[Array, "Ny Nx"]] | None

    def __init__(
        self,
        grid: CartesianGrid2D,
        mask: Mask2D | None = None,
    ) -> None:
        self.grid = grid
        self.mask = mask
        self.recon = Reconstruction2D(grid=grid)
        if mask is not None:
            self._mask_hierarchy_x = mask.get_adaptive_masks(
                direction="x", stencil_sizes=_HIERARCHY_SIZES
            )
            self._mask_hierarchy_y = mask.get_adaptive_masks(
                direction="y", stencil_sizes=_HIERARCHY_SIZES
            )
        else:
            self._mask_hierarchy_x = None
            self._mask_hierarchy_y = None

    def __call__(
        self,
        h: Float[Array, "Ny Nx"],
        u: Float[Array, "Ny Nx"],
        v: Float[Array, "Ny Nx"],
        method: str = "upwind1",
        wall: str = "closed",
    ) -> Float[Array, "Ny Nx"]:
        """Advective tendency -div(h * u_vec) at T-points.

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

        Parameters
        ----------
        h : Float[Array, "Ny Nx"]
            Scalar at T-points.
        u : Float[Array, "Ny Nx"]
            x-velocity at U-points.
        v : Float[Array, "Ny Nx"]
            y-velocity at V-points.
        method : str
            Reconstruction method: ``'naive'``, ``'upwind1'``, ``'upwind2'``,
            ``'upwind3'``, ``'weno3'``, ``'weno5'``, ``'weno7'``, ``'weno9'``,
            ``'wenoz5'``, or a flux-limiter TVD scheme: ``'minmod'``,
            ``'van_leer'``, ``'superbee'``, ``'mc'``.
        wall : {"closed", "open"}, default "closed"
            Lateral-boundary treatment.  ``"closed"`` (default, unchanged)
            writes the tendency only on the strict interior ``[2:-2, 2:-2]``
            and treats all four domain walls as no-flux.  ``"open"`` writes
            the full interior ``[1:-1, 1:-1]`` and computes the four
            domain-wall face fluxes with a first-order upwind reconstruction
            that reads the ghost ring, so Dirichlet / outflow / periodic
            lateral boundaries (applied by the caller to the ghost cells
            beforehand) drive horizontal transport.  Interior faces keep the
            requested high-order ``method`` and the flux field stays
            single-valued, so mass is conserved exactly.  ``"open"`` is not
            supported together with a ``mask``.

        Returns
        -------
        Float[Array, "Ny Nx"]
            Advective tendency at T-points.
        """
        _validate_wall(wall, self.mask)
        # ── masked path: route through upwind_flux with pre-built hier. ──
        mh_x = self._mask_hierarchy_x
        mh_y = self._mask_hierarchy_y
        if mh_x is not None and mh_y is not None and method in _MASK_DISPATCHABLE:
            rfx, rfy, sizes = _rec_funcs_for_method_2d(self.recon, method)
            mask_x = narrow_mask_hierarchy(mh_x, sizes)
            mask_y = narrow_mask_hierarchy(mh_y, sizes)
            fe = upwind_flux(h, u, dim=1, rec_funcs=rfx, mask_hierarchy=mask_x)
            fn = upwind_flux(h, v, dim=0, rec_funcs=rfy, mask_hierarchy=mask_y)
        # ── unmasked path: existing dispatch ────────────────────────────
        elif method == "naive":
            fe = self.recon.naive_x(h, u)
            fn = self.recon.naive_y(h, v)
        elif method == "upwind1":
            fe = self.recon.upwind1_x(h, u)
            fn = self.recon.upwind1_y(h, v)
        elif method == "upwind2":
            fe = self.recon.upwind2_x(h, u)
            fn = self.recon.upwind2_y(h, v)
        elif method == "upwind3":
            fe = self.recon.upwind3_x(h, u)
            fn = self.recon.upwind3_y(h, v)
        elif method == "weno3":
            fe = self.recon.weno3_x(h, u)
            fn = self.recon.weno3_y(h, v)
        elif method == "weno5":
            fe = self.recon.weno5_x(h, u)
            fn = self.recon.weno5_y(h, v)
        elif method == "wenoz5":
            fe = self.recon.wenoz5_x(h, u)
            fn = self.recon.wenoz5_y(h, v)
        elif method == "weno7":
            fe = self.recon.weno7_x(h, u)
            fn = self.recon.weno7_y(h, v)
        elif method == "weno9":
            fe = self.recon.weno9_x(h, u)
            fn = self.recon.weno9_y(h, v)
        elif method in _TVD_LIMITERS:
            fe = self.recon.tvd_x(h, u, limiter=method)
            fn = self.recon.tvd_y(h, v, limiter=method)
        else:
            raise ValueError(f"Unknown method: {method!r}")

        # dh[j, i] = -( (fe[j, i+1/2] - fe[j, i-1/2])/dx
        #             + (fn[j+1/2, i] - fn[j-1/2, i])/dy )
        # fe[j,i] is flux at east face of cell [j,i], fn[j,i] is flux at north face
        # For cell [j,i], we need fe[j,i] (east) and fe[j,i-1] (west),
        #                      and fn[j,i] (north) and fn[j-1,i] (south)
        if wall == "open":
            # Overwrite the west/east (fe cols 0, -2) and south/north
            # (fn rows 0, -2) domain-wall faces with a first-order upwind
            # flux that reads the ghost ring, then write the full interior.
            fe = fe.at[:, 0].set(_wall_upwind_flux(u[:, 0], h[:, 0], h[:, 1]))
            fe = fe.at[:, -2].set(_wall_upwind_flux(u[:, -2], h[:, -2], h[:, -1]))
            fn = fn.at[0, :].set(_wall_upwind_flux(v[0, :], h[0, :], h[1, :]))
            fn = fn.at[-2, :].set(_wall_upwind_flux(v[-2, :], h[-2, :], h[-1, :]))
            out = jnp.zeros_like(h)
            out = out.at[1:-1, 1:-1].set(
                -(
                    (fe[1:-1, 1:-1] - fe[1:-1, 0:-2]) / self.grid.dx
                    + (fn[1:-1, 1:-1] - fn[0:-2, 1:-1]) / self.grid.dy
                )
            )
        else:
            # Only use face fluxes defined by the reconstruction scheme,
            # avoiding ghost-ring flux entries.
            out = interior(
                -(
                    (fe[2:-2, 2:-2] - fe[2:-2, 1:-3]) / self.grid.dx
                    + (fn[2:-2, 2:-2] - fn[1:-3, 2:-2]) / self.grid.dy
                ),
                h,
                ghost=2,
            )
        if self.mask is not None:
            out = out * self.mask.h
        return out

`call(h, u, v, method='upwind1', wall='closed')`

Advective tendency -div(h * u_vec) at T-points.

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

Parameters:

Name	Type	Description	Default
`h`	`Float[Array, 'Ny Nx']`	Scalar at T-points.	required
`u`	`Float[Array, 'Ny Nx']`	x-velocity at U-points.	required
`v`	`Float[Array, 'Ny Nx']`	y-velocity at V-points.	required
`method`	`str`	Reconstruction method: `'naive'`, `'upwind1'`, `'upwind2'`, `'upwind3'`, `'weno3'`, `'weno5'`, `'weno7'`, `'weno9'`, `'wenoz5'`, or a flux-limiter TVD scheme: `'minmod'`, `'van_leer'`, `'superbee'`, `'mc'`.	`'upwind1'`
`wall`	`('closed', 'open')`	Lateral-boundary treatment. `"closed"` (default, unchanged) writes the tendency only on the strict interior `[2:-2, 2:-2]` and treats all four domain walls as no-flux. `"open"` writes the full interior `[1:-1, 1:-1]` and computes the four domain-wall face fluxes with a first-order upwind reconstruction that reads the ghost ring, so Dirichlet / outflow / periodic lateral boundaries (applied by the caller to the ghost cells beforehand) drive horizontal transport. Interior faces keep the requested high-order `method` and the flux field stays single-valued, so mass is conserved exactly. `"open"` is not supported together with a `mask`.	`"closed"`

Returns:

Type	Description
`Float[Array, 'Ny Nx']`	Advective tendency at T-points.

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

def __call__(
    self,
    h: Float[Array, "Ny Nx"],
    u: Float[Array, "Ny Nx"],
    v: Float[Array, "Ny Nx"],
    method: str = "upwind1",
    wall: str = "closed",
) -> Float[Array, "Ny Nx"]:
    """Advective tendency -div(h * u_vec) at T-points.

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

    Parameters
    ----------
    h : Float[Array, "Ny Nx"]
        Scalar at T-points.
    u : Float[Array, "Ny Nx"]
        x-velocity at U-points.
    v : Float[Array, "Ny Nx"]
        y-velocity at V-points.
    method : str
        Reconstruction method: ``'naive'``, ``'upwind1'``, ``'upwind2'``,
        ``'upwind3'``, ``'weno3'``, ``'weno5'``, ``'weno7'``, ``'weno9'``,
        ``'wenoz5'``, or a flux-limiter TVD scheme: ``'minmod'``,
        ``'van_leer'``, ``'superbee'``, ``'mc'``.
    wall : {"closed", "open"}, default "closed"
        Lateral-boundary treatment.  ``"closed"`` (default, unchanged)
        writes the tendency only on the strict interior ``[2:-2, 2:-2]``
        and treats all four domain walls as no-flux.  ``"open"`` writes
        the full interior ``[1:-1, 1:-1]`` and computes the four
        domain-wall face fluxes with a first-order upwind reconstruction
        that reads the ghost ring, so Dirichlet / outflow / periodic
        lateral boundaries (applied by the caller to the ghost cells
        beforehand) drive horizontal transport.  Interior faces keep the
        requested high-order ``method`` and the flux field stays
        single-valued, so mass is conserved exactly.  ``"open"`` is not
        supported together with a ``mask``.

    Returns
    -------
    Float[Array, "Ny Nx"]
        Advective tendency at T-points.
    """
    _validate_wall(wall, self.mask)
    # ── masked path: route through upwind_flux with pre-built hier. ──
    mh_x = self._mask_hierarchy_x
    mh_y = self._mask_hierarchy_y
    if mh_x is not None and mh_y is not None and method in _MASK_DISPATCHABLE:
        rfx, rfy, sizes = _rec_funcs_for_method_2d(self.recon, method)
        mask_x = narrow_mask_hierarchy(mh_x, sizes)
        mask_y = narrow_mask_hierarchy(mh_y, sizes)
        fe = upwind_flux(h, u, dim=1, rec_funcs=rfx, mask_hierarchy=mask_x)
        fn = upwind_flux(h, v, dim=0, rec_funcs=rfy, mask_hierarchy=mask_y)
    # ── unmasked path: existing dispatch ────────────────────────────
    elif method == "naive":
        fe = self.recon.naive_x(h, u)
        fn = self.recon.naive_y(h, v)
    elif method == "upwind1":
        fe = self.recon.upwind1_x(h, u)
        fn = self.recon.upwind1_y(h, v)
    elif method == "upwind2":
        fe = self.recon.upwind2_x(h, u)
        fn = self.recon.upwind2_y(h, v)
    elif method == "upwind3":
        fe = self.recon.upwind3_x(h, u)
        fn = self.recon.upwind3_y(h, v)
    elif method == "weno3":
        fe = self.recon.weno3_x(h, u)
        fn = self.recon.weno3_y(h, v)
    elif method == "weno5":
        fe = self.recon.weno5_x(h, u)
        fn = self.recon.weno5_y(h, v)
    elif method == "wenoz5":
        fe = self.recon.wenoz5_x(h, u)
        fn = self.recon.wenoz5_y(h, v)
    elif method == "weno7":
        fe = self.recon.weno7_x(h, u)
        fn = self.recon.weno7_y(h, v)
    elif method == "weno9":
        fe = self.recon.weno9_x(h, u)
        fn = self.recon.weno9_y(h, v)
    elif method in _TVD_LIMITERS:
        fe = self.recon.tvd_x(h, u, limiter=method)
        fn = self.recon.tvd_y(h, v, limiter=method)
    else:
        raise ValueError(f"Unknown method: {method!r}")

    # dh[j, i] = -( (fe[j, i+1/2] - fe[j, i-1/2])/dx
    #             + (fn[j+1/2, i] - fn[j-1/2, i])/dy )
    # fe[j,i] is flux at east face of cell [j,i], fn[j,i] is flux at north face
    # For cell [j,i], we need fe[j,i] (east) and fe[j,i-1] (west),
    #                      and fn[j,i] (north) and fn[j-1,i] (south)
    if wall == "open":
        # Overwrite the west/east (fe cols 0, -2) and south/north
        # (fn rows 0, -2) domain-wall faces with a first-order upwind
        # flux that reads the ghost ring, then write the full interior.
        fe = fe.at[:, 0].set(_wall_upwind_flux(u[:, 0], h[:, 0], h[:, 1]))
        fe = fe.at[:, -2].set(_wall_upwind_flux(u[:, -2], h[:, -2], h[:, -1]))
        fn = fn.at[0, :].set(_wall_upwind_flux(v[0, :], h[0, :], h[1, :]))
        fn = fn.at[-2, :].set(_wall_upwind_flux(v[-2, :], h[-2, :], h[-1, :]))
        out = jnp.zeros_like(h)
        out = out.at[1:-1, 1:-1].set(
            -(
                (fe[1:-1, 1:-1] - fe[1:-1, 0:-2]) / self.grid.dx
                + (fn[1:-1, 1:-1] - fn[0:-2, 1:-1]) / self.grid.dy
            )
        )
    else:
        # Only use face fluxes defined by the reconstruction scheme,
        # avoiding ghost-ring flux entries.
        out = interior(
            -(
                (fe[2:-2, 2:-2] - fe[2:-2, 1:-3]) / self.grid.dx
                + (fn[2:-2, 2:-2] - fn[1:-3, 2:-2]) / self.grid.dy
            ),
            h,
            ghost=2,
        )
    if self.mask is not None:
        out = out * self.mask.h
    return out

`finitevolx.Advection3D`

Bases: Module

3-D advection operator (horizontal plane per z-level).

Parameters:

Name	Type	Description	Default
`grid`	`CartesianGrid3D`	The underlying 3-D grid.	required
`mask`	`Mask3D or None`	Optional 3-D land/ocean mask. Pre-builds native 3-D adaptive stencil hierarchies `(2, 4, 6)` in both horizontal directions in `__init__`. For mask-dispatchable methods (WENO3/5, any TVD limiter), `upwind_flux` is used with the stored 3-D hierarchies narrowed to the method's stencil sizes — the z axis is treated as a batch dimension throughout. For non-dispatchable methods the unmasked code path runs and the final tendency is post-multiplied by `mask.h`. Pivoted from `Mask2D` to `Mask3D` per issue #209 Q4 for type-uniformity with the rest of the 3-D operator suite. A vertically-homogeneous 3-D mask (the 2-D mask broadcast over z) reproduces the old 2-D-mask-broadcast behaviour. A genuinely vertically-varying mask is now natively supported.	`None`

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

class Advection3D(eqx.Module):
    """3-D advection operator (horizontal plane per z-level).

    Parameters
    ----------
    grid : CartesianGrid3D
        The underlying 3-D grid.
    mask : Mask3D or None, optional
        Optional 3-D land/ocean mask.  Pre-builds native 3-D adaptive
        stencil hierarchies ``(2, 4, 6)`` in both horizontal directions
        in ``__init__``.  For mask-dispatchable methods (WENO3/5, any
        TVD limiter), ``upwind_flux`` is used with the stored 3-D
        hierarchies narrowed to the method's stencil sizes — the z
        axis is treated as a batch dimension throughout.  For
        non-dispatchable methods the unmasked code path runs and the
        final tendency is post-multiplied by ``mask.h``.

        Pivoted from ``Mask2D`` to ``Mask3D`` per issue #209 Q4 for
        type-uniformity with the rest of the 3-D operator suite.  A
        vertically-homogeneous 3-D mask (the 2-D mask broadcast over
        z) reproduces the old 2-D-mask-broadcast behaviour.  A
        genuinely vertically-varying mask is now natively supported.
    """

    grid: CartesianGrid3D
    mask: Mask3D | None
    recon: Reconstruction3D
    _mask_hierarchy_x: dict[int, Bool[Array, "Nz Ny Nx"]] | None
    _mask_hierarchy_y: dict[int, Bool[Array, "Nz Ny Nx"]] | None

    def __init__(
        self,
        grid: CartesianGrid3D,
        mask: Mask3D | None = None,
    ) -> None:
        self.grid = grid
        self.mask = mask
        self.recon = Reconstruction3D(grid=grid)
        if mask is not None:
            self._mask_hierarchy_x = mask.get_adaptive_masks(
                direction="x", stencil_sizes=_HIERARCHY_SIZES
            )
            self._mask_hierarchy_y = mask.get_adaptive_masks(
                direction="y", stencil_sizes=_HIERARCHY_SIZES
            )
        else:
            self._mask_hierarchy_x = None
            self._mask_hierarchy_y = None

    def __call__(
        self,
        h: Float[Array, "Nz Ny Nx"],
        u: Float[Array, "Nz Ny Nx"],
        v: Float[Array, "Nz Ny Nx"],
        method: str = "upwind1",
        wall: str = "closed",
    ) -> Float[Array, "Nz Ny Nx"]:
        """Advective tendency -div(h * u_vec) at T-points over all z-levels.

        Parameters
        ----------
        h : Float[Array, "Nz Ny Nx"]
            Scalar at T-points.
        u : Float[Array, "Nz Ny Nx"]
            x-velocity at U-points.
        v : Float[Array, "Nz Ny Nx"]
            y-velocity at V-points.
        method : str
            Reconstruction method: ``'naive'``, ``'upwind1'``, ``'weno3'``,
            ``'weno5'``, ``'weno7'``, ``'weno9'``, or a flux-limiter TVD
            scheme: ``'minmod'``, ``'van_leer'``, ``'superbee'``, ``'mc'``.
        wall : {"closed", "open"}, default "closed"
            Lateral-boundary treatment for the horizontal plane.
            ``"closed"`` (default, unchanged) writes the horizontal tendency
            only on the strict interior ``[1:-1, 2:-2, 2:-2]`` and treats the
            four lateral walls as no-flux — the wall-adjacent interior ring
            holds stale values.  ``"open"`` writes the full interior
            ``[1:-1, 1:-1, 1:-1]`` and computes the lateral domain-wall face
            fluxes with a first-order upwind reconstruction that reads the
            ghost ring, so Dirichlet / outflow / periodic lateral boundaries
            (applied by the caller to the ghost cells beforehand) drive
            horizontal transport.  Interior faces keep the requested
            high-order ``method`` and the flux field stays single-valued, so
            mass is conserved exactly.  The z axis is a batch dimension and is
            unaffected.  ``"open"`` is not supported together with a ``mask``.

        Returns
        -------
        Float[Array, "Nz Ny Nx"]
            Advective tendency at T-points.
        """
        _validate_wall(wall, self.mask)
        # ── masked path: route through upwind_flux with pre-built 3-D hier. ──
        mh_x = self._mask_hierarchy_x
        mh_y = self._mask_hierarchy_y
        if mh_x is not None and mh_y is not None and method in _MASK_DISPATCHABLE:
            rfx, rfy, sizes = _rec_funcs_for_method_3d(self.recon, method)
            mask_x = narrow_mask_hierarchy(mh_x, sizes)
            mask_y = narrow_mask_hierarchy(mh_y, sizes)
            # 3-D arrays: x-flux axis is 2, y-flux axis is 1.
            fe = upwind_flux(h, u, dim=2, rec_funcs=rfx, mask_hierarchy=mask_x)
            fn = upwind_flux(h, v, dim=1, rec_funcs=rfy, mask_hierarchy=mask_y)
        # ── unmasked path: existing dispatch ────────────────────────────
        elif method == "naive":
            fe = self.recon.naive_x(h, u)
            fn = self.recon.naive_y(h, v)
        elif method == "upwind1":
            fe = self.recon.upwind1_x(h, u)
            fn = self.recon.upwind1_y(h, v)
        elif method == "weno3":
            fe = self.recon.weno3_x(h, u)
            fn = self.recon.weno3_y(h, v)
        elif method == "weno5":
            fe = self.recon.weno5_x(h, u)
            fn = self.recon.weno5_y(h, v)
        elif method == "weno7":
            fe = self.recon.weno7_x(h, u)
            fn = self.recon.weno7_y(h, v)
        elif method == "weno9":
            fe = self.recon.weno9_x(h, u)
            fn = self.recon.weno9_y(h, v)
        elif method in _TVD_LIMITERS:
            fe = self.recon.tvd_x(h, u, limiter=method)
            fn = self.recon.tvd_y(h, v, limiter=method)
        else:
            raise ValueError(f"Unknown method: {method!r}")

        out = jnp.zeros_like(h)
        # dh[k, j, i] = -( (fe[k,j,i+1/2] - fe[k,j,i-1/2])/dx
        #                 + (fn[k,j+1/2,i] - fn[k,j-1/2,i])/dy )
        if wall == "open":
            # Overwrite the west/east (fe cols 0, -2) and south/north
            # (fn rows 0, -2) lateral wall faces with a first-order upwind
            # flux that reads the ghost ring, then write the full horizontal
            # interior.  z is a batch dimension and keeps the full [1:-1].
            fe = fe.at[:, :, 0].set(
                _wall_upwind_flux(u[:, :, 0], h[:, :, 0], h[:, :, 1])
            )
            fe = fe.at[:, :, -2].set(
                _wall_upwind_flux(u[:, :, -2], h[:, :, -2], h[:, :, -1])
            )
            fn = fn.at[:, 0, :].set(
                _wall_upwind_flux(v[:, 0, :], h[:, 0, :], h[:, 1, :])
            )
            fn = fn.at[:, -2, :].set(
                _wall_upwind_flux(v[:, -2, :], h[:, -2, :], h[:, -1, :])
            )
            out = out.at[1:-1, 1:-1, 1:-1].set(
                -(
                    (fe[1:-1, 1:-1, 1:-1] - fe[1:-1, 1:-1, 0:-2]) / self.grid.dx
                    + (fn[1:-1, 1:-1, 1:-1] - fn[1:-1, 0:-2, 1:-1]) / self.grid.dy
                )
            )
        else:
            # Reconstruction writes to [1:-1, 1:-1, 1:-1]; the west flux at i=0
            # and south flux at j=0 are ghost cells (value 0, not filled).
            # Consistent with 1D/2D operators, skip the ghost-adjacent interior
            # ring in the horizontal plane so we never read ghost flux cells.
            # All z-levels are independent, so z uses the full interior [1:-1].
            out = out.at[1:-1, 2:-2, 2:-2].set(
                -(
                    (fe[1:-1, 2:-2, 2:-2] - fe[1:-1, 2:-2, 1:-3]) / self.grid.dx
                    + (fn[1:-1, 2:-2, 2:-2] - fn[1:-1, 1:-3, 2:-2]) / self.grid.dy
                )
            )
        if self.mask is not None:
            out = out * self.mask.h
        return out

`call(h, u, v, method='upwind1', wall='closed')`

Advective tendency -div(h * u_vec) at T-points over all z-levels.

Parameters:

Name	Type	Description	Default
`h`	`Float[Array, 'Nz Ny Nx']`	Scalar at T-points.	required
`u`	`Float[Array, 'Nz Ny Nx']`	x-velocity at U-points.	required
`v`	`Float[Array, 'Nz Ny Nx']`	y-velocity at V-points.	required
`method`	`str`	Reconstruction method: `'naive'`, `'upwind1'`, `'weno3'`, `'weno5'`, `'weno7'`, `'weno9'`, or a flux-limiter TVD scheme: `'minmod'`, `'van_leer'`, `'superbee'`, `'mc'`.	`'upwind1'`
`wall`	`('closed', 'open')`	Lateral-boundary treatment for the horizontal plane. `"closed"` (default, unchanged) writes the horizontal tendency only on the strict interior `[1:-1, 2:-2, 2:-2]` and treats the four lateral walls as no-flux — the wall-adjacent interior ring holds stale values. `"open"` writes the full interior `[1:-1, 1:-1, 1:-1]` and computes the lateral domain-wall face fluxes with a first-order upwind reconstruction that reads the ghost ring, so Dirichlet / outflow / periodic lateral boundaries (applied by the caller to the ghost cells beforehand) drive horizontal transport. Interior faces keep the requested high-order `method` and the flux field stays single-valued, so mass is conserved exactly. The z axis is a batch dimension and is unaffected. `"open"` is not supported together with a `mask`.	`"closed"`

Returns:

Type	Description
`Float[Array, 'Nz Ny Nx']`	Advective tendency at T-points.

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

def __call__(
    self,
    h: Float[Array, "Nz Ny Nx"],
    u: Float[Array, "Nz Ny Nx"],
    v: Float[Array, "Nz Ny Nx"],
    method: str = "upwind1",
    wall: str = "closed",
) -> Float[Array, "Nz Ny Nx"]:
    """Advective tendency -div(h * u_vec) at T-points over all z-levels.

    Parameters
    ----------
    h : Float[Array, "Nz Ny Nx"]
        Scalar at T-points.
    u : Float[Array, "Nz Ny Nx"]
        x-velocity at U-points.
    v : Float[Array, "Nz Ny Nx"]
        y-velocity at V-points.
    method : str
        Reconstruction method: ``'naive'``, ``'upwind1'``, ``'weno3'``,
        ``'weno5'``, ``'weno7'``, ``'weno9'``, or a flux-limiter TVD
        scheme: ``'minmod'``, ``'van_leer'``, ``'superbee'``, ``'mc'``.
    wall : {"closed", "open"}, default "closed"
        Lateral-boundary treatment for the horizontal plane.
        ``"closed"`` (default, unchanged) writes the horizontal tendency
        only on the strict interior ``[1:-1, 2:-2, 2:-2]`` and treats the
        four lateral walls as no-flux — the wall-adjacent interior ring
        holds stale values.  ``"open"`` writes the full interior
        ``[1:-1, 1:-1, 1:-1]`` and computes the lateral domain-wall face
        fluxes with a first-order upwind reconstruction that reads the
        ghost ring, so Dirichlet / outflow / periodic lateral boundaries
        (applied by the caller to the ghost cells beforehand) drive
        horizontal transport.  Interior faces keep the requested
        high-order ``method`` and the flux field stays single-valued, so
        mass is conserved exactly.  The z axis is a batch dimension and is
        unaffected.  ``"open"`` is not supported together with a ``mask``.

    Returns
    -------
    Float[Array, "Nz Ny Nx"]
        Advective tendency at T-points.
    """
    _validate_wall(wall, self.mask)
    # ── masked path: route through upwind_flux with pre-built 3-D hier. ──
    mh_x = self._mask_hierarchy_x
    mh_y = self._mask_hierarchy_y
    if mh_x is not None and mh_y is not None and method in _MASK_DISPATCHABLE:
        rfx, rfy, sizes = _rec_funcs_for_method_3d(self.recon, method)
        mask_x = narrow_mask_hierarchy(mh_x, sizes)
        mask_y = narrow_mask_hierarchy(mh_y, sizes)
        # 3-D arrays: x-flux axis is 2, y-flux axis is 1.
        fe = upwind_flux(h, u, dim=2, rec_funcs=rfx, mask_hierarchy=mask_x)
        fn = upwind_flux(h, v, dim=1, rec_funcs=rfy, mask_hierarchy=mask_y)
    # ── unmasked path: existing dispatch ────────────────────────────
    elif method == "naive":
        fe = self.recon.naive_x(h, u)
        fn = self.recon.naive_y(h, v)
    elif method == "upwind1":
        fe = self.recon.upwind1_x(h, u)
        fn = self.recon.upwind1_y(h, v)
    elif method == "weno3":
        fe = self.recon.weno3_x(h, u)
        fn = self.recon.weno3_y(h, v)
    elif method == "weno5":
        fe = self.recon.weno5_x(h, u)
        fn = self.recon.weno5_y(h, v)
    elif method == "weno7":
        fe = self.recon.weno7_x(h, u)
        fn = self.recon.weno7_y(h, v)
    elif method == "weno9":
        fe = self.recon.weno9_x(h, u)
        fn = self.recon.weno9_y(h, v)
    elif method in _TVD_LIMITERS:
        fe = self.recon.tvd_x(h, u, limiter=method)
        fn = self.recon.tvd_y(h, v, limiter=method)
    else:
        raise ValueError(f"Unknown method: {method!r}")

    out = jnp.zeros_like(h)
    # dh[k, j, i] = -( (fe[k,j,i+1/2] - fe[k,j,i-1/2])/dx
    #                 + (fn[k,j+1/2,i] - fn[k,j-1/2,i])/dy )
    if wall == "open":
        # Overwrite the west/east (fe cols 0, -2) and south/north
        # (fn rows 0, -2) lateral wall faces with a first-order upwind
        # flux that reads the ghost ring, then write the full horizontal
        # interior.  z is a batch dimension and keeps the full [1:-1].
        fe = fe.at[:, :, 0].set(
            _wall_upwind_flux(u[:, :, 0], h[:, :, 0], h[:, :, 1])
        )
        fe = fe.at[:, :, -2].set(
            _wall_upwind_flux(u[:, :, -2], h[:, :, -2], h[:, :, -1])
        )
        fn = fn.at[:, 0, :].set(
            _wall_upwind_flux(v[:, 0, :], h[:, 0, :], h[:, 1, :])
        )
        fn = fn.at[:, -2, :].set(
            _wall_upwind_flux(v[:, -2, :], h[:, -2, :], h[:, -1, :])
        )
        out = out.at[1:-1, 1:-1, 1:-1].set(
            -(
                (fe[1:-1, 1:-1, 1:-1] - fe[1:-1, 1:-1, 0:-2]) / self.grid.dx
                + (fn[1:-1, 1:-1, 1:-1] - fn[1:-1, 0:-2, 1:-1]) / self.grid.dy
            )
        )
    else:
        # Reconstruction writes to [1:-1, 1:-1, 1:-1]; the west flux at i=0
        # and south flux at j=0 are ghost cells (value 0, not filled).
        # Consistent with 1D/2D operators, skip the ghost-adjacent interior
        # ring in the horizontal plane so we never read ghost flux cells.
        # All z-levels are independent, so z uses the full interior [1:-1].
        out = out.at[1:-1, 2:-2, 2:-2].set(
            -(
                (fe[1:-1, 2:-2, 2:-2] - fe[1:-1, 2:-2, 1:-3]) / self.grid.dx
                + (fn[1:-1, 2:-2, 2:-2] - fn[1:-1, 1:-3, 2:-2]) / self.grid.dy
            )
        )
    if self.mask is not None:
        out = out * self.mask.h
    return out

Advection Operators

finitevolx.Advection1D

__call__(h, u, method='upwind1', wall='closed')

finitevolx.Advection2D

__call__(h, u, v, method='upwind1', wall='closed')

finitevolx.Advection3D

__call__(h, u, v, method='upwind1', wall='closed')

`finitevolx.Advection1D`

`call(h, u, method='upwind1', wall='closed')`

`finitevolx.Advection2D`

`call(h, u, v, method='upwind1', wall='closed')`

`finitevolx.Advection3D`

`call(h, u, v, method='upwind1', wall='closed')`