Convenience Wrappers

High-level functions for common elliptic PDE inversions on C-grids. Each wrapper dispatches to spectral, CG, or capacitance solvers depending on the method argument.

Streamfunction from Vorticity

finitevolx.streamfunction_from_vorticity(zeta, dx, dy, bc='dst', lambda_=0.0, method='spectral', mask=None, capacitance_solver=None, preconditioner=None, known_values=None, known_mask=None)

Invert the vorticity–streamfunction relation ∇²ψ − λψ = ζ.

Solves the Poisson (λ = 0) or Helmholtz (λ ≠ 0) equation to recover the streamfunction from relative vorticity.

Three solver methods are available:

• "spectral" — Direct spectral solver (DST/DCT/FFT) for rectangular domains. Selected by bc. Default. ψ lives at X points: the unknowns are sliced out of the full array so that the wall is the C-grid wall ("dst": DST-I on [1:-2, 1:-2], ψ = 0 on the wall nodes 0 and N-2; "dct": DCT-I on [0:-1, 0:-1]; "fft": [1:-1, 1:-1]), the same wall method="cg" imposes with a Mask2D. The other entries of the result are zero.
• "cg" — Preconditioned Conjugate Gradient for masked / irregular domains. Requires mask. Uses a spectral preconditioner by default, or a custom one via preconditioner.
• "capacitance" — Capacitance matrix method for masked domains. Requires a pre-built :class:CapacitanceSolver via capacitance_solver.

Parameters:

Name	Type	Description	Default
zeta	Float[Array, 'Ny Nx']	Relative vorticity (right-hand side).	required
dx	float	Grid spacings.	required
dy	float	Grid spacings.	required
bc	('dst', 'dct', 'fft')	Boundary condition on the wall / coast. "dst" (Dirichlet, ψ = 0 on the wall nodes) is the standard choice for streamfunction inversion and the only one the mask-based methods accept for ψ; "dct" and "fft" are for method="spectral" only (the others raise rather than ignore them). With "dct" (∂ψ/∂n = 0 on the walls) the wall nodes themselves are unknowns, so zeta is read there too (X rows / columns 0 and N-2): an operator that writes only [1:-1, 1:-1] leaves row / column 0 at zero, which costs the second-order accuracy; fill the wall values of zeta first. A :class:BoundaryConditionSet is also accepted. Its face types select the transform (all Dirichlet1D -> "dst", all Neumann1D with zero value -> "dct", all Periodic1D -> "fft"; mixing raises), its mask is used by the mask-based methods (method="spectral" with a bc.mask raises), and non-zero Dirichlet1D values become known_values on each face's wall nodes -- the coast cells just outside the wet cells, X rows / columns 0 and N-2 of a rectangle (explicit known_values take precedence). An all-zero Dirichlet set is the ordinary homogeneous "dst" solve.	"dst"
lambda_	float	Helmholtz parameter. Use 0.0 for the pure Poisson problem (streamfunction from vorticity). Non-zero values arise in QG PV inversion: (∇² − λ)ψ = q.	0.0
method	('spectral', 'cg', 'capacitance')	Solver method. Default: "spectral".	"spectral"
mask	Float[Array, 'Ny Nx'] or Mask2D or None	Domain mask. Required for method="cg". When an :class:Mask2D is passed the psi staggering mask is extracted automatically.	None
capacitance_solver	CapacitanceSolver or None	Pre-built capacitance solver. Required for method="capacitance".	None
preconditioner	callable or None	Custom preconditioner for method="cg". Signature: preconditioner(r: Array) -> Array. When None, a spectral preconditioner (FFT-based) is used automatically.	None
known_values	Float[Array, 'Ny Nx'] or None	Prescribed (inhomogeneous Dirichlet) values of the solution, used on the coast ring -- the dry cells adjacent to a wet cell, where method="cg" puts its ψ = 0 wall (for a Mask2D, the C-grid wall nodes) -- and at any known_mask cells; values elsewhere are ignored. The solution equals known_values exactly at those cells and every other wet cell is solved for. None (default) keeps the homogeneous solve unchanged. Mask-based methods require mask. For method="capacitance" build the solver on the same wet mask: its unknowns are then exactly the solve domain. method="spectral" requires bc="dst" and no mask: it solves the ψ basin of the rectangle (wet X nodes [1:-2, 1:-2], values on the wall nodes 0 and N-2), matching method="cg" with the rectangle's Mask2D.	None
known_mask	Bool[Array, 'Ny Nx'] or None	Extra interior wet cells with known values (e.g. sparse observations), pinned alongside the boundary ring. method="cg" only.	None

Returns:

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

Source code in finitevolx/_src/solvers/elliptic.py

def streamfunction_from_vorticity(
    zeta: Float[Array, "Ny Nx"],
    dx: float,
    dy: float,
    bc: _BCLike = "dst",
    lambda_: float = 0.0,
    method: str = "spectral",
    mask: _MaskLike | None = None,
    capacitance_solver: CapacitanceSolver | None = None,
    preconditioner: _PrecondLike | None = None,
    known_values: Float[Array, "Ny Nx"] | None = None,
    known_mask: Array | None = None,
) -> Float[Array, "Ny Nx"]:
    r"""Invert the vorticity–streamfunction relation ∇²ψ − λψ = ζ.

    Solves the Poisson (λ = 0) or Helmholtz (λ ≠ 0) equation to recover the
    streamfunction from relative vorticity.

    Three solver methods are available:

    * ``"spectral"`` — Direct spectral solver (DST/DCT/FFT) for rectangular
      domains.  Selected by *bc*.  Default.  ψ lives at X points: the
      unknowns are sliced out of the full array so that the wall is the
      C-grid wall (``"dst"``: DST-I on ``[1:-2, 1:-2]``, ψ = 0 on the wall
      nodes ``0`` and ``N-2``; ``"dct"``: DCT-I on ``[0:-1, 0:-1]``;
      ``"fft"``: ``[1:-1, 1:-1]``), the same wall ``method="cg"`` imposes
      with a ``Mask2D``.  The other entries of the result are zero.
    * ``"cg"`` — Preconditioned Conjugate Gradient for masked / irregular
      domains.  Requires *mask*.  Uses a spectral preconditioner by default,
      or a custom one via *preconditioner*.
    * ``"capacitance"`` — Capacitance matrix method for masked domains.
      Requires a pre-built :class:`CapacitanceSolver` via
      *capacitance_solver*.

    Parameters
    ----------
    zeta : Float[Array, "Ny Nx"]
        Relative vorticity (right-hand side).
    dx, dy : float
        Grid spacings.
    bc : {"dst", "dct", "fft"} or BoundaryConditionSet
        Boundary condition on the wall / coast.  ``"dst"`` (Dirichlet,
        ψ = 0 on the wall nodes) is the standard choice for streamfunction
        inversion and the only one the mask-based methods accept for ψ;
        ``"dct"`` and ``"fft"`` are for ``method="spectral"`` only (the
        others raise rather than ignore them).  With ``"dct"`` (∂ψ/∂n = 0 on
        the walls) the wall nodes themselves are unknowns, so *zeta* is read
        there too (X rows / columns 0 and N-2): an operator that writes only
        ``[1:-1, 1:-1]`` leaves row / column 0 at zero, which costs the
        second-order accuracy; fill the wall values of *zeta* first.
        A :class:`BoundaryConditionSet` is also accepted.  Its face types
        select the transform (all ``Dirichlet1D`` -> ``"dst"``, all
        ``Neumann1D`` with zero value -> ``"dct"``, all ``Periodic1D`` ->
        ``"fft"``; mixing raises), its ``mask`` is used by the mask-based
        methods (``method="spectral"`` with a ``bc.mask`` raises), and
        non-zero ``Dirichlet1D`` values become ``known_values`` on each
        face's wall nodes -- the coast cells just outside the wet cells, X
        rows / columns ``0`` and ``N-2`` of a rectangle (explicit
        ``known_values`` take precedence).  An all-zero Dirichlet set is the
        ordinary homogeneous ``"dst"`` solve.
    lambda_ : float
        Helmholtz parameter.  Use 0.0 for the pure Poisson problem
        (streamfunction from vorticity).  Non-zero values arise in QG PV
        inversion: (∇² − λ)ψ = q.
    method : {"spectral", "cg", "capacitance"}
        Solver method.  Default: ``"spectral"``.
    mask : Float[Array, "Ny Nx"] or Mask2D or None
        Domain mask.  Required for ``method="cg"``.  When an
        :class:`Mask2D` is passed the ``psi`` staggering mask is
        extracted automatically.
    capacitance_solver : CapacitanceSolver or None
        Pre-built capacitance solver.  Required for
        ``method="capacitance"``.
    preconditioner : callable or None
        Custom preconditioner for ``method="cg"``.  Signature:
        ``preconditioner(r: Array) -> Array``.  When ``None``, a spectral
        preconditioner (FFT-based) is used automatically.
    known_values : Float[Array, "Ny Nx"] or None
        Prescribed (inhomogeneous Dirichlet) values of the solution, used
        on the coast ring -- the dry cells adjacent to a wet cell, where
        ``method="cg"`` puts its ψ = 0 wall (for a ``Mask2D``, the C-grid
        wall nodes) -- and at any ``known_mask`` cells; values elsewhere are
        ignored.  The solution equals ``known_values`` exactly at those
        cells and every other wet cell is solved for.  ``None`` (default)
        keeps the homogeneous solve unchanged.  Mask-based methods require
        *mask*.  For ``method="capacitance"`` build the solver on the same
        wet *mask*: its unknowns are then exactly the solve domain.
        ``method="spectral"`` requires ``bc="dst"`` and no *mask*: it solves
        the ψ basin of the rectangle (wet X nodes ``[1:-2, 1:-2]``, values
        on the wall nodes ``0`` and ``N-2``), matching ``method="cg"`` with
        the rectangle's ``Mask2D``.
    known_mask : Bool[Array, "Ny Nx"] or None
        Extra interior wet cells with known values (e.g. sparse
        observations), pinned alongside the boundary ring.
        ``method="cg"`` only.

    Returns
    -------
    Float[Array, "Ny Nx"]
        Streamfunction ψ.
    """
    bc, mask, known_values = _resolve_bc(
        bc, method, mask, known_values, zeta.shape, zeta.dtype, "X"
    )
    if known_values is not None:
        return _solve_with_known_values(
            zeta,
            known_values,
            dx,
            dy,
            lambda_,
            bc,
            method,
            mask,
            known_mask,
            capacitance_solver,
            preconditioner,
            "X",
        )
    return _solve_dispatch(
        zeta,
        dx,
        dy,
        lambda_,
        bc,
        method,
        mask,
        capacitance_solver,
        preconditioner,
        "X",
    )

Pressure from Divergence

finitevolx.pressure_from_divergence(div_u, dx, dy, bc='dct', method='spectral', mask=None, capacitance_solver=None, preconditioner=None, known_values=None, known_mask=None)

Solve ∇²p = ∇·u for the pressure correction.

Used in pressure-projection methods (Chorin splitting) where the divergence of the provisional velocity field must be removed.

p lives at T points, and every method puts the wall on the cell faces: the spectral path solves on the T interior [1:-1, 1:-1] ("dct": DCT-II, ∂p/∂n = 0; "dst": DST-II, p = 0; "fft": periodic) and returns a zero ghost ring; method="cg" solves on a mask (a Mask2D gives its h mask) with the same face conditions on every wet–dry face. method="capacitance" and known values pin values at cell centres, not on faces, so they raise here.

Parameters:

Name	Type	Description	Default
div_u	Float[Array, 'Ny Nx']	Divergence of the velocity field (right-hand side).	required
dx	float	Grid spacings.	required
dy	float	Grid spacings.	required
bc	('dct', 'dst', 'fft')	Boundary condition on the wall / coast, honoured by every method or rejected. "dct" (Neumann, ∂p/∂n = 0 on the cell faces) is the standard choice for pressure with solid walls; with method="cg" it is zero flux through every wet–dry face of the mask, and the solution is the zero-mean one on each connected wet region (the problem fixes p only up to a constant per region). "dst" is p = 0 on the faces, for both methods. "fft" (periodic) is method="spectral" only. A :class:BoundaryConditionSet is accepted as in :func:streamfunction_from_vorticity, except that non-zero Dirichlet1D faces raise (see known_values).	"dct"
method	('spectral', 'cg')	Solver method. Default: "spectral". "capacitance" raises: the capacitance solver holds the unknown at zero on the dry cell centres (the ψ wall), not on the faces.	"spectral"
mask	Float[Array, 'Ny Nx'] or Mask2D or None	Domain mask. Required for method="cg". A :class:Mask2D gives its T-point mask h.	None
capacitance_solver	None	Not supported (kept for signature symmetry with the other wrappers).	None
preconditioner	callable or None	Custom preconditioner for method="cg".	None
known_values	None	Not supported for pressure: a T-point Dirichlet value belongs on the cell face, while known values are pinned at cell centres (the ghost cells, half a cell outside the face), which would be a different, first-order wall. Passing either raises.	None
known_mask	None	Not supported for pressure: a T-point Dirichlet value belongs on the cell face, while known values are pinned at cell centres (the ghost cells, half a cell outside the face), which would be a different, first-order wall. Passing either raises.	None

Returns:

Type	Description
Float[Array, 'Ny Nx']	Pressure field p.

Source code in finitevolx/_src/solvers/elliptic.py

def pressure_from_divergence(
    div_u: Float[Array, "Ny Nx"],
    dx: float,
    dy: float,
    bc: _BCLike = "dct",
    method: str = "spectral",
    mask: _MaskLike | None = None,
    capacitance_solver: CapacitanceSolver | None = None,
    preconditioner: _PrecondLike | None = None,
    known_values: Float[Array, "Ny Nx"] | None = None,
    known_mask: Array | None = None,
) -> Float[Array, "Ny Nx"]:
    r"""Solve ∇²p = ∇·u for the pressure correction.

    Used in pressure-projection methods (Chorin splitting) where the
    divergence of the provisional velocity field must be removed.

    p lives at T points, and every method puts the wall on the cell faces:
    the spectral path solves on the T interior ``[1:-1, 1:-1]``
    (``"dct"``: DCT-II, ∂p/∂n = 0; ``"dst"``: DST-II, p = 0; ``"fft"``:
    periodic) and returns a zero ghost ring; ``method="cg"`` solves on a
    mask (a ``Mask2D`` gives its ``h`` mask) with the same face conditions
    on every wet–dry face.  ``method="capacitance"`` and known values pin
    values at cell centres, not on faces, so they raise here.

    Parameters
    ----------
    div_u : Float[Array, "Ny Nx"]
        Divergence of the velocity field (right-hand side).
    dx, dy : float
        Grid spacings.
    bc : {"dct", "dst", "fft"} or BoundaryConditionSet
        Boundary condition on the wall / coast, honoured by every method
        or rejected.  ``"dct"`` (Neumann, ∂p/∂n = 0 on the cell faces) is
        the standard choice for pressure with solid walls; with
        ``method="cg"`` it is zero flux through every wet–dry face of the
        mask, and the solution is the zero-mean one on each connected wet
        region (the problem fixes p only up to a constant per region).
        ``"dst"`` is p = 0 on the faces, for both methods.  ``"fft"``
        (periodic) is ``method="spectral"`` only.  A
        :class:`BoundaryConditionSet` is accepted as in
        :func:`streamfunction_from_vorticity`, except that non-zero
        ``Dirichlet1D`` faces raise (see *known_values*).
    method : {"spectral", "cg"}
        Solver method.  Default: ``"spectral"``.  ``"capacitance"`` raises:
        the capacitance solver holds the unknown at zero on the dry cell
        centres (the ψ wall), not on the faces.
    mask : Float[Array, "Ny Nx"] or Mask2D or None
        Domain mask.  Required for ``method="cg"``.  A :class:`Mask2D` gives
        its T-point mask ``h``.
    capacitance_solver : None
        Not supported (kept for signature symmetry with the other wrappers).
    preconditioner : callable or None
        Custom preconditioner for ``method="cg"``.
    known_values, known_mask : None
        Not supported for pressure: a T-point Dirichlet value belongs on the
        cell face, while known values are pinned at cell centres (the ghost
        cells, half a cell outside the face), which would be a different,
        first-order wall.  Passing either raises.

    Returns
    -------
    Float[Array, "Ny Nx"]
        Pressure field p.
    """
    if method == "capacitance":
        raise ValueError(
            "method='capacitance' is not available for the T-point pressure "
            "solve: the capacitance solver holds the unknown at zero on the dry "
            "cell centres next to the wet ones (the psi wall), not on the cell "
            "faces where the pressure wall is.  Use method='cg' with a mask."
        )
    if known_mask is not None and known_values is None:
        raise ValueError(
            "known_mask needs known_values (and pressure accepts neither)."
        )
    bc, mask, known_values = _resolve_bc(
        bc, method, mask, known_values, div_u.shape, div_u.dtype, "T"
    )
    if known_values is not None:
        return _solve_with_known_values(
            div_u,
            known_values,
            dx,
            dy,
            0.0,
            bc,
            method,
            mask,
            known_mask,
            capacitance_solver,
            preconditioner,
            "T",
        )
    return _solve_dispatch(
        div_u,
        dx,
        dy,
        0.0,
        bc,
        method,
        mask,
        capacitance_solver,
        preconditioner,
        "T",
    )

PV Inversion

finitevolx.pv_inversion(pv, dx, dy, lambda_, bc='dst', method='spectral', mask=None, capacitance_solver=None, preconditioner=None, known_values=None, known_mask=None)

QG potential-vorticity inversion: solve (∇² − λ)ψ = q.

Supports batched / multi-layer PV fields. When lambda_ is a 1-D array of shape (nl,), each layer is solved with its own Helmholtz parameter (e.g. 1/Rd² per vertical mode from :func:~finitevolx.decompose_vertical_modes).

Solver selection follows the same three-method dispatch as :func:streamfunction_from_vorticity.

Parameters:

Name	Type	Description	Default
pv	Float[Array, '... Ny Nx']	Potential-vorticity field. Leading dimensions are batched.	required
dx	float	Grid spacings.	required
dy	float	Grid spacings.	required
lambda_	float or Float[Array, ' nl']	Helmholtz parameter(s). Scalar for a single layer; array of shape (nl,) for multi-layer inversion.	required
bc	('dst', 'dct', 'fft')	Boundary condition on the wall / coast, as in :func:streamfunction_from_vorticity ("dst" for the mask-based methods). A :class:BoundaryConditionSet is accepted as there.	"dst"
method	('spectral', 'cg', 'capacitance')	Solver method. Default: "spectral".	"spectral"
mask	Float[Array, 'Ny Nx'] or Mask2D or None	Domain mask. Required for method="cg".	None
capacitance_solver	CapacitanceSolver or None	Pre-built capacitance solver. Required for method="capacitance".	None
preconditioner	callable or None	Custom preconditioner for method="cg".	None
known_values	Float[Array, '... Ny Nx'] or None	Prescribed Dirichlet values (see :func:streamfunction_from_vorticity), broadcast against pv: a (Ny, Nx) field applies to every layer, (nl, Ny, Nx) gives per-layer values. Each layer's lift uses its own lambda_.	None
known_mask	Bool[Array, 'Ny Nx'] or None	Extra interior wet cells with known values, shared by all layers. method="cg" only.	None

Returns:

Type	Description
Float[Array, '... Ny Nx']	Streamfunction ψ, same shape as pv.

Source code in finitevolx/_src/solvers/elliptic.py

def pv_inversion(
    pv: Float[Array, "... Ny Nx"],
    dx: float,
    dy: float,
    lambda_: float | Float[Array, " nl"],
    bc: _BCLike = "dst",
    method: str = "spectral",
    mask: _MaskLike | None = None,
    capacitance_solver: CapacitanceSolver | None = None,
    preconditioner: _PrecondLike | None = None,
    known_values: Float[Array, "... Ny Nx"] | None = None,
    known_mask: Array | None = None,
) -> Float[Array, "... Ny Nx"]:
    r"""QG potential-vorticity inversion: solve (∇² − λ)ψ = q.

    Supports batched / multi-layer PV fields.  When *lambda_* is a 1-D
    array of shape ``(nl,)``, each layer is solved with its own Helmholtz
    parameter (e.g. 1/Rd² per vertical mode from
    :func:`~finitevolx.decompose_vertical_modes`).

    Solver selection follows the same three-method dispatch as
    :func:`streamfunction_from_vorticity`.

    Parameters
    ----------
    pv : Float[Array, "... Ny Nx"]
        Potential-vorticity field.  Leading dimensions are batched.
    dx, dy : float
        Grid spacings.
    lambda_ : float or Float[Array, " nl"]
        Helmholtz parameter(s).  Scalar for a single layer; array of
        shape ``(nl,)`` for multi-layer inversion.
    bc : {"dst", "dct", "fft"} or BoundaryConditionSet
        Boundary condition on the wall / coast, as in
        :func:`streamfunction_from_vorticity` (``"dst"`` for the mask-based
        methods).  A :class:`BoundaryConditionSet` is accepted as there.
    method : {"spectral", "cg", "capacitance"}
        Solver method.  Default: ``"spectral"``.
    mask : Float[Array, "Ny Nx"] or Mask2D or None
        Domain mask.  Required for ``method="cg"``.
    capacitance_solver : CapacitanceSolver or None
        Pre-built capacitance solver.  Required for
        ``method="capacitance"``.
    preconditioner : callable or None
        Custom preconditioner for ``method="cg"``.
    known_values : Float[Array, "... Ny Nx"] or None
        Prescribed Dirichlet values (see
        :func:`streamfunction_from_vorticity`), broadcast against *pv*: a
        ``(Ny, Nx)`` field applies to every layer, ``(nl, Ny, Nx)`` gives
        per-layer values.  Each layer's lift uses its own *lambda_*.
    known_mask : Bool[Array, "Ny Nx"] or None
        Extra interior wet cells with known values, shared by all layers.
        ``method="cg"`` only.

    Returns
    -------
    Float[Array, "... Ny Nx"]
        Streamfunction ψ, same shape as *pv*.
    """
    bc, mask, known_values = _resolve_bc(
        bc, method, mask, known_values, pv.shape, pv.dtype, "X"
    )
    lam = jnp.asarray(lambda_)
    # A Python-number lambda stays concrete (the singular-Neumann projection
    # keys on a concrete zero); a traced scalar is passed through as is, so
    # jit / vmap over a scalar lambda_ work.
    lam_scalar = lambda_ if isinstance(lambda_, numbers.Real) else lam
    kv = None if known_values is None else jnp.broadcast_to(known_values, pv.shape)

    if lam.ndim == 0 and kv is not None:

        def _solve_one_known(
            rhs: Float[Array, "Ny Nx"], kv_i: Float[Array, "Ny Nx"]
        ) -> Float[Array, "Ny Nx"]:
            return _solve_with_known_values(
                rhs,
                kv_i,
                dx,
                dy,
                lam_scalar,
                bc,
                method,
                mask,
                known_mask,
                capacitance_solver,
                preconditioner,
                "X",
            )

        if pv.ndim == 2:
            return _solve_one_known(pv, kv)
        shape = pv.shape
        flat = pv.reshape(-1, shape[-2], shape[-1])
        kv_flat = kv.reshape(-1, shape[-2], shape[-1])
        return eqx.filter_vmap(_solve_one_known)(flat, kv_flat).reshape(shape)

    if lam.ndim == 0:
        # Scalar lambda: vmap over all leading dims if present
        if pv.ndim == 2:
            return _solve_dispatch(
                pv,
                dx,
                dy,
                lam_scalar,
                bc,
                method,
                mask,
                capacitance_solver,
                preconditioner,
                "X",
            )
        # Flatten leading dims, solve each, reshape
        shape = pv.shape
        flat = pv.reshape(-1, shape[-2], shape[-1])

        def _solve_one(rhs: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
            return _solve_dispatch(
                rhs,
                dx,
                dy,
                lam_scalar,
                bc,
                method,
                mask,
                capacitance_solver,
                preconditioner,
                "X",
            )

        out = eqx.filter_vmap(_solve_one)(flat)
        return out.reshape(shape)

    # Array lambda: leading dim must match lam.shape[0]
    if pv.ndim < 3:
        raise ValueError(
            f"pv must have at least 3 dims when lambda_ is an array, "
            f"got shape {pv.shape}"
        )
    nl = lam.shape[0]
    if pv.shape[-3] != nl:
        raise ValueError(
            f"pv.shape[-3]={pv.shape[-3]} does not match lambda_ length {nl}"
        )

    # Solve each layer with its own lambda.  lam_i is a JAX tracer inside
    # vmap, so the per-layer solves below call the tracer-safe building
    # blocks (masked_laplacian + solve_cg, or _solve_spectral, which has no
    # Python branch on lambda) with the same X staggering as _solve_dispatch.
    if method == "capacitance":
        raise ValueError(
            "method='capacitance' does not support array-valued lambda_; "
            "solve each layer separately or use method='spectral' or 'cg' "
            "for multi-layer problems."
        )

    # Known values: shared solve domain; each layer lifts with its own lambda.
    domain = (
        None
        if kv is None
        else _known_value_domain(pv.shape, pv.dtype, method, bc, mask, known_mask, "X")
    )

    if method == "cg":
        _check_cg_bc(bc, "X")
        mask_arr = _resolve_mask_arr(mask, pv.dtype, "X")
        if mask_arr is None:
            raise ValueError("method='cg' requires a mask")
        # With known values the solve runs on the effective domain.
        solve_mask = (
            mask_arr if domain is None else domain.effective_mask.astype(pv.dtype)
        )

        def _solve_layer(
            rhs: Float[Array, "Ny Nx"], lam_i: float
        ) -> Float[Array, "Ny Nx"]:
            return _cg_masked_solve(
                rhs, solve_mask, dx, dy, lam_i, bc, preconditioner, "X"
            )

    elif method == "spectral":
        # With known values the ψ basin's solve domain is the DST-I rectangle.
        layer_bc = bc if domain is None else "dst"
        if ("X", layer_bc) not in _SPECTRAL_STAGGERED:
            raise ValueError(f"bc must be 'fft', 'dst', or 'dct'; got {bc!r}")

        def _solve_layer(
            rhs: Float[Array, "Ny Nx"], lam_i: float
        ) -> Float[Array, "Ny Nx"]:
            return _solve_spectral(rhs, dx, dy, lam_i, layer_bc, "X")

    else:
        raise ValueError(
            f"method must be 'spectral', 'cg', or 'capacitance'; got {method!r}"
        )

    # Flatten any leading batch dims: (..., nl, Ny, Nx) -> (batch, nl, Ny, Nx)
    shape = pv.shape
    ny, nx = shape[-2], shape[-1]
    pv_4d = pv.reshape(-1, nl, ny, nx)

    if domain is None or kv is None:
        # vmap over layer axis (pairing each layer with its lambda)
        _solve_layers = eqx.filter_vmap(_solve_layer, in_axes=(0, 0))

        # vmap over the (flattened) batch axis
        out_4d = eqx.filter_vmap(lambda batch: _solve_layers(batch, lam))(pv_4d)
        return out_4d.reshape(shape)

    known_domain = domain

    def _solve_layer_known(
        rhs: Float[Array, "Ny Nx"], kv_i: Float[Array, "Ny Nx"], lam_i: float
    ) -> Float[Array, "Ny Nx"]:
        # psi = lift + psi_hom, psi_hom solving the lambda_i-corrected RHS
        rhs_c, value_lift = lift_rhs(known_domain, rhs, kv_i, dx, dy, lam_i)
        return reconstruct_from_lift(
            known_domain, _solve_layer(rhs_c, lam_i), value_lift
        )

    kv_4d = jnp.reshape(kv, (-1, nl, ny, nx))
    _solve_layers_known = eqx.filter_vmap(_solve_layer_known, in_axes=(0, 0, 0))
    out_4d = eqx.filter_vmap(lambda batch, k: _solve_layers_known(batch, k, lam))(
        pv_4d, kv_4d
    )
    return out_4d.reshape(shape)

Known Values (Inhomogeneous Dirichlet)

Building blocks behind the wrappers' known_values / known_mask arguments, for use with your own solver. The known values sit on the coast ring (the dry cells adjacent to the wet mask: the C-grid wall nodes for ψ).

finitevolx.coast_ring(mask)

Dry cells adjacent to at least one wet cell (the coast ring).

This is where :func:~finitevolx.masked_laplacian imposes its homogeneous Dirichlet wall (a dry neighbour reads as zero), so it is where prescribed (known) values sit. For ψ on a from_interior rectangle (Mask2D.xy_corner_strict wet on [1:-2, 1:-2]) it is the C-grid wall: X rows / columns 0 and N-2 (corners excluded).

Uses pad-then-slice (cells outside the array count as dry), so it never wraps around.

Parameters:

Name	Type	Description	Default
mask	Float[Array, 'Ny Nx'] or Bool[Array, 'Ny Nx']	Binary wet/dry mask (1 / True = wet, 0 / False = dry).	required

Returns:

Type	Description
Bool[Array, 'Ny Nx']	True at dry cells that neighbour at least one wet cell.

Examples:

>>> import jax.numpy as jnp
>>> mask = jnp.zeros((5, 5), dtype=bool).at[1:-2, 1:-2].set(True)
>>> coast_ring(mask).astype(int)[:, :4]
Array([[0, 1, 1, 0],
       [1, 0, 0, 1],
       [1, 0, 0, 1],
       [0, 1, 1, 0],
       [0, 0, 0, 0]], dtype=int32)

Source code in finitevolx/_src/solvers/inhomogeneous.py

def coast_ring(mask: Float[Array, "Ny Nx"]) -> Bool[Array, "Ny Nx"]:
    """Dry cells adjacent to at least one wet cell (the coast ring).

    This is where :func:`~finitevolx.masked_laplacian` imposes its
    homogeneous Dirichlet wall (a dry neighbour reads as zero), so it is
    where prescribed (known) values sit.  For ψ on a ``from_interior``
    rectangle (``Mask2D.xy_corner_strict`` wet on ``[1:-2, 1:-2]``) it is the
    C-grid wall: X rows / columns ``0`` and ``N-2`` (corners excluded).

    Uses pad-then-slice (cells outside the array count as dry), so it never
    wraps around.

    Parameters
    ----------
    mask : Float[Array, "Ny Nx"] or Bool[Array, "Ny Nx"]
        Binary wet/dry mask (1 / True = wet, 0 / False = dry).

    Returns
    -------
    Bool[Array, "Ny Nx"]
        True at dry cells that neighbour at least one wet cell.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> mask = jnp.zeros((5, 5), dtype=bool).at[1:-2, 1:-2].set(True)
    >>> coast_ring(mask).astype(int)[:, :4]
    Array([[0, 1, 1, 0],
           [1, 0, 0, 1],
           [1, 0, 0, 1],
           [0, 1, 1, 0],
           [0, 0, 0, 0]], dtype=int32)
    """
    wet = jnp.asarray(mask) > 0.5
    # coast = the inner ring of the dry set: dry cells 4-adjacent to a wet one
    # (boundary_ring pads with "not dry", i.e. cells outside the array are dry).
    return boundary_ring(~wet)

finitevolx.boundary_ring(mask)

Wet cells adjacent to at least one dry cell (the inner ring).

Uses pad-then-slice to avoid jnp.roll wrap-around artifacts. JIT-compatible but intended for offline (one-shot) use — the mask is typically fixed at setup time.

Parameters:

Name	Type	Description	Default
mask	Float[Array, 'Ny Nx'] or Bool[Array, 'Ny Nx']	Binary wet/dry mask (1 / True = wet, 0 / False = dry).	required

Returns:

Type	Description
Bool[Array, 'Ny Nx']	True at wet cells that neighbor at least one dry cell.

Source code in finitevolx/_src/solvers/inhomogeneous.py

def boundary_ring(mask: Float[Array, "Ny Nx"]) -> Bool[Array, "Ny Nx"]:
    """Wet cells adjacent to at least one dry cell (the inner ring).

    Uses pad-then-slice to avoid ``jnp.roll`` wrap-around artifacts.
    JIT-compatible but intended for offline (one-shot) use — the mask
    is typically fixed at setup time.

    Parameters
    ----------
    mask : Float[Array, "Ny Nx"] or Bool[Array, "Ny Nx"]
        Binary wet/dry mask (1 / True = wet, 0 / False = dry).

    Returns
    -------
    Bool[Array, "Ny Nx"]
        True at wet cells that neighbor at least one dry cell.
    """
    wet = mask > 0.5
    dry = ~wet
    # padded_dry[j+1, i+1] = dry[j, i]; cells outside the array count as wet.
    padded_dry = jnp.pad(dry, pad_width=1, mode="constant", constant_values=False)
    # adjacent_to_dry[j, i] = dry[j+1, i] | dry[j-1, i] | dry[j, i+1] | dry[j, i-1]
    adjacent_to_dry = (
        padded_dry[2:, 1:-1]  # dry[j+1, i]  (north neighbour)
        | padded_dry[:-2, 1:-1]  # dry[j-1, i]  (south neighbour)
        | padded_dry[1:-1, 2:]  # dry[j, i+1]  (east neighbour)
        | padded_dry[1:-1, :-2]  # dry[j, i-1]  (west neighbour)
    )
    return wet & adjacent_to_dry

finitevolx.SolveDomain

Bases: Module

Derived masks for the lifting trick.

The known values sit on the coast ring -- the dry cells adjacent to the wet mask, where masked_laplacian / method="cg" put the homogeneous wall -- plus any wet known_mask cells. The effective solve domain is every other wet cell::

all_known = coast_ring | (known_mask & wet)
effective_mask = wet & ~known_mask

Parameters:

Name	Type	Description	Default
mask	Float[Array, 'Ny Nx'] or Mask2D	Binary wet/dry mask (1 = wet, 0 = dry). A :class:Mask2D gives its ψ (X-point) mask xy_corner_strict.	required
known_mask	Bool[Array, 'Ny Nx'] or None	Locations of additional known values (e.g., sparse interior observations). Merged with the coast ring.	None

Attributes:

Name	Type	Description
wet_mask	Bool[Array, 'Ny Nx']	The wet cells.
coast_ring	Bool[Array, 'Ny Nx']	Dry cells adjacent to a wet cell, where the boundary values sit.
boundary_ring	Bool[Array, 'Ny Nx']	Wet cells adjacent to a dry cell (informational; they are unknowns).
all_known	Bool[Array, 'Ny Nx']	Cells whose value is prescribed.
effective_mask	Bool[Array, 'Ny Nx']	Wet cells that are solved for.

Source code in finitevolx/_src/solvers/inhomogeneous.py

class SolveDomain(eqx.Module):
    """Derived masks for the lifting trick.

    The known values sit on the coast ring -- the dry cells adjacent to the
    wet mask, where ``masked_laplacian`` / ``method="cg"`` put the
    homogeneous wall -- plus any wet ``known_mask`` cells.  The effective
    solve domain is every other wet cell::

        all_known = coast_ring | (known_mask & wet)
        effective_mask = wet & ~known_mask

    Parameters
    ----------
    mask : Float[Array, "Ny Nx"] or Mask2D
        Binary wet/dry mask (1 = wet, 0 = dry).  A :class:`Mask2D` gives its
        ψ (X-point) mask ``xy_corner_strict``.
    known_mask : Bool[Array, "Ny Nx"] or None
        Locations of additional known values (e.g., sparse interior
        observations).  Merged with the coast ring.

    Attributes
    ----------
    wet_mask : Bool[Array, "Ny Nx"]
        The wet cells.
    coast_ring : Bool[Array, "Ny Nx"]
        Dry cells adjacent to a wet cell, where the boundary values sit.
    boundary_ring : Bool[Array, "Ny Nx"]
        Wet cells adjacent to a dry cell (informational; they are unknowns).
    all_known : Bool[Array, "Ny Nx"]
        Cells whose value is prescribed.
    effective_mask : Bool[Array, "Ny Nx"]
        Wet cells that are solved for.
    """

    wet_mask: Bool[Array, "Ny Nx"]
    coast_ring: Bool[Array, "Ny Nx"]
    boundary_ring: Bool[Array, "Ny Nx"]
    all_known: Bool[Array, "Ny Nx"]
    effective_mask: Bool[Array, "Ny Nx"]

    def __init__(
        self,
        mask: Float[Array, "Ny Nx"] | Mask2D,
        known_mask: Bool[Array, "Ny Nx"] | None = None,
    ):
        # Keep the masks boolean: a float32 copy would downcast a weakly
        # typed float64 field it multiplies (the lifting uses jnp.where).
        if isinstance(mask, Mask2D):
            mask_arr = jnp.asarray(mask.xy_corner_strict, dtype=bool)
        else:
            mask_arr = jnp.asarray(mask) > 0.5

        self.wet_mask = mask_arr
        self.coast_ring = coast_ring(mask_arr)
        self.boundary_ring = boundary_ring(mask_arr)

        if known_mask is not None:
            known_wet = jnp.asarray(known_mask, dtype=bool) & self.wet_mask
            self.all_known = self.coast_ring | known_wet
            self.effective_mask = self.wet_mask & ~known_wet
        else:
            self.all_known = self.coast_ring
            self.effective_mask = self.wet_mask

finitevolx.KnownValueLifting

Bases: Module

Pre/post-processing wrapper for the lifting trick.

Does not own or call a solver. The user builds their solver separately and calls it themselves. This operator only sandwiches the solve with pre- and post-processing.

The solver's unknowns must be exactly domain.effective_mask, with zero on every other cell (masked_laplacian on domain.effective_mask: a CG or multigrid operator built on it, or, without a known_mask, a capacitance solver built on domain.wet_mask). The known values sit on the coast ring (the dry cells adjacent to the wet mask) and on the wet known_mask cells.

Parameters:

Name	Type	Description	Default
domain	SolveDomain	Precomputed domain masks.	required
dx	float	Grid spacing in x.	required
dy	float	Grid spacing in y.	required
lambda_	float	Helmholtz parameter.	required

Source code in finitevolx/_src/solvers/inhomogeneous.py

class KnownValueLifting(eqx.Module):
    """Pre/post-processing wrapper for the lifting trick.

    Does **not** own or call a solver.  The user builds their solver
    separately and calls it themselves.  This operator only sandwiches the
    solve with pre- and post-processing.

    The solver's unknowns must be exactly ``domain.effective_mask``, with
    zero on every other cell (``masked_laplacian`` on
    ``domain.effective_mask``: a CG or multigrid operator built on it, or,
    without a ``known_mask``, a capacitance solver built on
    ``domain.wet_mask``).
    The known values sit on the coast ring (the dry cells adjacent to the
    wet mask) and on the wet ``known_mask`` cells.

    Parameters
    ----------
    domain : SolveDomain
        Precomputed domain masks.
    dx : float
        Grid spacing in x.
    dy : float
        Grid spacing in y.
    lambda_ : float
        Helmholtz parameter.
    """

    domain: SolveDomain
    dx: float = eqx.field(static=True)
    dy: float = eqx.field(static=True)
    lambda_: float = eqx.field(static=True)

    def preprocess(
        self,
        rhs: Float[Array, "Ny Nx"],
        known_values: Float[Array, "Ny Nx"],
    ) -> tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]:
        """Correct RHS for known values.

        Returns ``(rhs_corrected, value_lift)``.  Feed ``rhs_corrected``
        to your solver, then pass both to :meth:`postprocess`.

        Parameters
        ----------
        rhs : Float[Array, "Ny Nx"]
            Original right-hand side.
        known_values : Float[Array, "Ny Nx"]
            Prescribed values at known cells.  Values at non-known cells
            are ignored.

        Returns
        -------
        rhs_corrected : Float[Array, "Ny Nx"]
            Corrected RHS restricted to the effective solve domain.
        value_lift : Float[Array, "Ny Nx"]
            Lifting function — add to homogeneous solution to get the
            full solution.
        """
        return lift_rhs(self.domain, rhs, known_values, self.dx, self.dy, self.lambda_)

    def postprocess(
        self,
        psi_hom: Float[Array, "Ny Nx"],
        value_lift: Float[Array, "Ny Nx"],
    ) -> Float[Array, "Ny Nx"]:
        """Reconstruct full solution from homogeneous solve + lift."""
        return reconstruct_from_lift(self.domain, psi_hom, value_lift)

postprocess(psi_hom, value_lift)

Reconstruct full solution from homogeneous solve + lift.

Source code in finitevolx/_src/solvers/inhomogeneous.py

def postprocess(
    self,
    psi_hom: Float[Array, "Ny Nx"],
    value_lift: Float[Array, "Ny Nx"],
) -> Float[Array, "Ny Nx"]:
    """Reconstruct full solution from homogeneous solve + lift."""
    return reconstruct_from_lift(self.domain, psi_hom, value_lift)

preprocess(rhs, known_values)

Correct RHS for known values.

Returns (rhs_corrected, value_lift). Feed rhs_corrected to your solver, then pass both to :meth:postprocess.

Parameters:

Name	Type	Description	Default
rhs	Float[Array, 'Ny Nx']	Original right-hand side.	required
known_values	Float[Array, 'Ny Nx']	Prescribed values at known cells. Values at non-known cells are ignored.	required

Returns:

Name	Type	Description
rhs_corrected	Float[Array, 'Ny Nx']	Corrected RHS restricted to the effective solve domain.
value_lift	Float[Array, 'Ny Nx']	Lifting function — add to homogeneous solution to get the full solution.

Source code in finitevolx/_src/solvers/inhomogeneous.py

def preprocess(
    self,
    rhs: Float[Array, "Ny Nx"],
    known_values: Float[Array, "Ny Nx"],
) -> tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]:
    """Correct RHS for known values.

    Returns ``(rhs_corrected, value_lift)``.  Feed ``rhs_corrected``
    to your solver, then pass both to :meth:`postprocess`.

    Parameters
    ----------
    rhs : Float[Array, "Ny Nx"]
        Original right-hand side.
    known_values : Float[Array, "Ny Nx"]
        Prescribed values at known cells.  Values at non-known cells
        are ignored.

    Returns
    -------
    rhs_corrected : Float[Array, "Ny Nx"]
        Corrected RHS restricted to the effective solve domain.
    value_lift : Float[Array, "Ny Nx"]
        Lifting function — add to homogeneous solution to get the
        full solution.
    """
    return lift_rhs(self.domain, rhs, known_values, self.dx, self.dy, self.lambda_)