Capacitance Matrix Method

Extends fast spectral solvers to irregular / masked domains using the Sherman-Morrison capacitance correction (Buzbee, Golub & Nielson 1970). The unknowns are every wet cell of the mask, with ψ = 0 on the coast (the dry cells next to water): the same problem masked_laplacian and method="cg" define.

Solver

Note

CapacitanceSolver is re-exported from spectraldiffx, so its docstring describes spectraldiffx's own builder (unknowns: the wet cells off the inner ring, ψ = 0 on that ring). finitevolX's build_capacitance_solver fills the same class with every wet cell as an unknown and ψ = 0 on the coast.

finitevolx.CapacitanceSolver

Bases: Module

Spectral Poisson/Helmholtz solver for masked irregular domains.

Solves (∇² − λ) ψ = f on the wet cells of mask with ψ = 0 on the inner boundary (wet cells adjacent to a dry cell) and outside the mask, using the five-point finite-difference Laplacian.

Method (Buzbee, Golub & Nielson 1970, via gaussx)

Let B be the rectangular operator (periodic / Dirichlet / Neumann by base_bc), I the interior cells (wet, not on the inner boundary) and C the inner-boundary cells, which are exactly the cells outside I that the stencil of I reaches. The masked problem is B[I][:, I] ψ_I = f_I. The capacitance method solves it with two fast rectangular solves per call: find y with y[C] = 0 and (B y)_k = f_k for every k ∉ C (point sources on C absorb the constraints, via an |C| × |C| capacitance matrix factorised once at construction). For a singular base (fft/dct with λ = 0) the constant null vector is included in the capacitance system, so the PDE holds exactly in the interior (gh-87).

Construct with :func:build_capacitance_solver.

Attributes:

Name	Type	Description
operator	MaskedOperator	B[I][:, I] with its precomputed capacitance factorisation.
interior_indices	Int[Array, 'Ni']	Flat (row-major) indices of the interior cells I.
shape	tuple[int, int]	Grid shape (Ny, Nx).
dx, dy	float	Grid spacings (static; baked into operator).
lambda_	float	Helmholtz parameter (static).
base_bc	str	Rectangular base: "fft", "dst" or "dct" (static).

Source code in .venv/lib/python3.12/site-packages/spectraldiffx/_src/fourier/capacitance.py

class CapacitanceSolver(eqx.Module):
    """Spectral Poisson/Helmholtz solver for masked irregular domains.

    Solves ``(∇² − λ) ψ = f`` on the wet cells of ``mask`` with ``ψ = 0`` on
    the *inner boundary* (wet cells adjacent to a dry cell) and outside the
    mask, using the five-point finite-difference Laplacian.

    Method (Buzbee, Golub & Nielson 1970, via gaussx)
    --------------------------------------------------
    Let ``B`` be the rectangular operator (periodic / Dirichlet / Neumann by
    ``base_bc``), ``I`` the interior cells (wet, not on the inner boundary)
    and ``C`` the inner-boundary cells, which are exactly the cells outside
    ``I`` that the stencil of ``I`` reaches. The masked problem is
    ``B[I][:, I] ψ_I = f_I``. The capacitance method solves it with two fast
    rectangular solves per call: find ``y`` with ``y[C] = 0`` and
    ``(B y)_k = f_k`` for every ``k ∉ C`` (point sources on ``C`` absorb the
    constraints, via an ``|C| × |C|`` capacitance matrix factorised once at
    construction). For a singular base (``fft``/``dct`` with λ = 0) the
    constant null vector is included in the capacitance system, so the PDE
    holds exactly in the interior (gh-87).

    Construct with :func:`build_capacitance_solver`.

    Attributes
    ----------
    operator : gaussx.MaskedOperator
        ``B[I][:, I]`` with its precomputed capacitance factorisation.
    interior_indices : Int[Array, "Ni"]
        Flat (row-major) indices of the interior cells ``I``.
    shape : tuple[int, int]
        Grid shape ``(Ny, Nx)``.
    dx, dy : float
        Grid spacings (static; baked into ``operator``).
    lambda_ : float
        Helmholtz parameter (static).
    base_bc : str
        Rectangular base: ``"fft"``, ``"dst"`` or ``"dct"`` (static).
    """

    operator: gaussx.MaskedOperator
    interior_indices: Int[Array, " Ni"]
    shape: tuple[int, int] = eqx.field(static=True)
    dx: float = eqx.field(static=True)
    dy: float = eqx.field(static=True)
    lambda_: float = eqx.field(static=True)
    base_bc: str = eqx.field(static=True)

    def __call__(self, rhs: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
        """Solve (∇² − λ)ψ = rhs on the masked domain.

        Parameters
        ----------
        rhs : Float[Array, "Ny Nx"]
            Right-hand side on the full rectangular grid. Only the interior
            cells are used; values on the inner boundary and outside the
            mask are ignored.

        Returns
        -------
        Float[Array, "Ny Nx"]
            Solution ψ on the full grid: the masked solve on the interior,
            exactly zero on the inner boundary and outside the mask.
        """
        ny, nx = self.shape
        f_interior = rhs.reshape(ny * nx)[self.interior_indices]
        psi_interior = gaussx.solve(self.operator, f_interior)
        psi = jnp.zeros(ny * nx, dtype=psi_interior.dtype)
        return psi.at[self.interior_indices].set(psi_interior).reshape(ny, nx)

__call__(rhs)

Solve (∇² − λ)ψ = rhs on the masked domain.

Parameters:

Name	Type	Description	Default
rhs	Float[Array, 'Ny Nx']	Right-hand side on the full rectangular grid. Only the interior cells are used; values on the inner boundary and outside the mask are ignored.	required

Returns:

Type	Description
Float[Array, 'Ny Nx']	Solution ψ on the full grid: the masked solve on the interior, exactly zero on the inner boundary and outside the mask.

Source code in .venv/lib/python3.12/site-packages/spectraldiffx/_src/fourier/capacitance.py

def __call__(self, rhs: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
    """Solve (∇² − λ)ψ = rhs on the masked domain.

    Parameters
    ----------
    rhs : Float[Array, "Ny Nx"]
        Right-hand side on the full rectangular grid. Only the interior
        cells are used; values on the inner boundary and outside the
        mask are ignored.

    Returns
    -------
    Float[Array, "Ny Nx"]
        Solution ψ on the full grid: the masked solve on the interior,
        exactly zero on the inner boundary and outside the mask.
    """
    ny, nx = self.shape
    f_interior = rhs.reshape(ny * nx)[self.interior_indices]
    psi_interior = gaussx.solve(self.operator, f_interior)
    psi = jnp.zeros(ny * nx, dtype=psi_interior.dtype)
    return psi.at[self.interior_indices].set(psi_interior).reshape(ny, nx)

Factory

finitevolx.build_capacitance_solver(mask, dx, dy, lambda_=0.0, base_bc='dst')

Pre-compute the capacitance matrix and return a ready-to-use solver.

The solver's unknowns are every wet cell of mask, with ψ = 0 on the coast ring (the dry cells 4-adjacent to a wet cell) and outside: the same problem method="cg" solves through :func:masked_laplacian on that mask (for ψ with a :class:Mask2D, the wall on the C-grid wall nodes).

The masked operator is the wet sub-block B[W][:, W] of the rectangular five-point operator B = (∇² − λ) (:func:~finitevolx.masked_helmholtz_operator with bc=base_bc), a gaussx.MaskedOperator whose coupling set is the coast ring. This is an offline function: it factorises the N_c × N_c capacitance matrix (N_c = number of coast cells) with N_c rectangular solves. The returned :class:CapacitanceSolver is then two rectangular solves plus an O(N_c²) back-substitution per right-hand side and is JIT-compilable.

Parameters:

Name	Type	Description	Default
mask	np.ndarray of bool shape (Ny, Nx), or Mask2D	Physical domain mask. True = wet (ocean/fluid), False = dry (land/walls). A :class:Mask2D gives its ψ (X-point) mask xy_corner_strict.	required
dx	float	Grid spacing in x.	required
dy	float	Grid spacing in y.	required
lambda_	float	Helmholtz parameter λ ≥ 0. Use 0.0 for pure Poisson.	0.0
base_bc	('dst', 'fft', 'dct')	Rectangular base operator: Dirichlet (DST-I, the default: non-singular for every λ ≥ 0), periodic (FFT) or Neumann (DCT-II). The base only enters through cells beyond the array edge, so every base gives the masked_laplacian solution as long as no wet cell touches the array edge (true for any mask with a dry ghost ring); "fft" and "dct" raise when one does. A singular base ("fft" / "dct" at λ = 0) is handled exactly: gaussx borders the capacitance system with the base's null vector.	"dst"

Returns:

Type	Description
CapacitanceSolver	A callable equinox Module with all precomputed arrays baked in; it returns ψ on the full grid, zero on every dry cell.

Raises:

Type	Description
ValueError	For an unknown base_bc, a mask with no wet cells, a mask with no dry cells (nothing to correct: use the rectangular solver), or a "fft" / "dct" base with a wet cell on the array edge.

Examples:

>>> mask = np.zeros((34, 34), dtype=bool)
>>> mask[1:-1, 1:-1] = True
>>> solver = build_capacitance_solver(Mask2D.from_mask(mask), dx, dy, lambda_=4.0)
>>> psi = streamfunction_from_vorticity(
...     zeta, dx, dy, lambda_=4.0, method="capacitance", capacitance_solver=solver
... )

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

def build_capacitance_solver(
    mask: np.ndarray | Mask2D,
    dx: float,
    dy: float,
    lambda_: float = 0.0,
    base_bc: str = "dst",
) -> CapacitanceSolver:
    """Pre-compute the capacitance matrix and return a ready-to-use solver.

    The solver's unknowns are **every wet cell** of *mask*, with ψ = 0 on the
    coast ring (the dry cells 4-adjacent to a wet cell) and outside: the same
    problem ``method="cg"`` solves through :func:`masked_laplacian` on that
    mask (for ψ with a :class:`Mask2D`, the wall on the C-grid wall nodes).

    The masked operator is the wet sub-block ``B[W][:, W]`` of the
    rectangular five-point operator ``B = (∇² − λ)``
    (:func:`~finitevolx.masked_helmholtz_operator` with ``bc=base_bc``), a
    ``gaussx.MaskedOperator`` whose coupling set is the coast ring.  This is
    an **offline** function: it factorises the ``N_c × N_c`` capacitance
    matrix (``N_c`` = number of coast cells) with ``N_c`` rectangular solves.
    The returned :class:`CapacitanceSolver` is then two rectangular solves
    plus an ``O(N_c²)`` back-substitution per right-hand side and is
    JIT-compilable.

    Parameters
    ----------
    mask : np.ndarray of bool shape (Ny, Nx), or Mask2D
        Physical domain mask.  ``True`` = wet (ocean/fluid), ``False`` = dry
        (land/walls).  A :class:`Mask2D` gives its ψ (X-point) mask
        ``xy_corner_strict``.
    dx : float
        Grid spacing in x.
    dy : float
        Grid spacing in y.
    lambda_ : float
        Helmholtz parameter λ ≥ 0.  Use ``0.0`` for pure Poisson.
    base_bc : {"dst", "fft", "dct"}
        Rectangular base operator: Dirichlet (DST-I, the default:
        non-singular for every λ ≥ 0), periodic (FFT) or Neumann (DCT-II).
        The base only enters through cells beyond the array edge, so every
        base gives the ``masked_laplacian`` solution as long as no wet cell
        touches the array edge (true for any mask with a dry ghost ring);
        ``"fft"`` and ``"dct"`` raise when one does.  A singular base
        (``"fft"`` / ``"dct"`` at λ = 0) is handled exactly: gaussx borders
        the capacitance system with the base's null vector.

    Returns
    -------
    CapacitanceSolver
        A callable equinox Module with all precomputed arrays baked in; it
        returns ψ on the full grid, zero on every dry cell.

    Raises
    ------
    ValueError
        For an unknown ``base_bc``, a mask with no wet cells, a mask with
        no dry cells (nothing to correct: use the rectangular solver), or a
        ``"fft"`` / ``"dct"`` base with a wet cell on the array edge.

    Examples
    --------
    >>> mask = np.zeros((34, 34), dtype=bool)
    >>> mask[1:-1, 1:-1] = True
    >>> solver = build_capacitance_solver(Mask2D.from_mask(mask), dx, dy, lambda_=4.0)
    >>> psi = streamfunction_from_vorticity(
    ...     zeta, dx, dy, lambda_=4.0, method="capacitance", capacitance_solver=solver
    ... )
    """
    if base_bc not in ("fft", "dst", "dct"):
        raise ValueError(f"base_bc must be 'fft', 'dst', or 'dct'; got {base_bc!r}")
    if isinstance(mask, Mask2D):
        mask = np.asarray(mask.xy_corner_strict)
    wet = np.asarray(mask) > 0.5  # [Ny, Nx], concrete
    if not wet.any():
        raise ValueError("The mask has no wet (True) cells.")
    if wet.all():
        raise ValueError(
            "The mask has no dry (False) cells, so there is no coast to "
            f"correct for; use solve_helmholtz_{base_bc} directly."
        )
    edge = wet[0].any() | wet[-1].any() | wet[:, 0].any() | wet[:, -1].any()
    if base_bc != "dst" and edge:
        raise ValueError(
            f"base_bc={base_bc!r} reaches beyond the array edge, where "
            "masked_laplacian (method='cg') has psi = 0, but the mask has wet "
            "cells on the edge: use base_bc='dst' or a mask with a dry ring."
        )
    # The wet block of B = (∇² - λ) on the rectangle (the same eigenvalues as
    # the spectral solvers), coupled to the dry cells through the coast ring.
    operator = masked_helmholtz_operator(
        jnp.asarray(wet), dx, dy, lambda_=lambda_, bc=base_bc
    )
    return CapacitanceSolver(
        operator=operator,
        interior_indices=jnp.asarray(np.flatnonzero(wet.ravel())),
        shape=wet.shape,
        dx=float(dx),
        dy=float(dy),
        lambda_=float(lambda_),
        base_bc=base_bc,
    )

Masked Laplacian

finitevolx.masked_laplacian(psi, mask, dx, dy, lambda_=0.0)

Apply the masked discrete Helmholtz operator (∇² − λ)·ψ.

Enforces homogeneous Dirichlet conditions at the mask boundary by zeroing psi outside the mask before applying the 5-point stencil. The output is also zeroed outside the mask.

Neighbours beyond the array edge are zero (dry), so the array edge is a homogeneous Dirichlet boundary one node outside the array, wherever the mask is wet there; nothing wraps around. On an all-ones mask this is the dense Dirichlet 5-point operator that the DST-I solvers invert, so make_spectral_preconditioner(..., bc="dst") matches it exactly on a rectangle; the FFT preconditioner (periodic) is an approximation.

Parameters:

Name	Type	Description	Default
psi	Float[Array, 'Ny Nx']	Field to which the operator is applied.	required
mask	Float[Array, 'Ny Nx'] or Mask2D	Binary mask: 1 inside the physical domain, 0 outside (land/exterior). When an :class:Mask2D is passed, the psi staggering mask is used (all four surrounding h-cells must be wet).	required
dx	float	Grid spacing in x.	required
dy	float	Grid spacing in y.	required
lambda_	float	Helmholtz parameter. Default: 0.0 (pure Laplacian).	0.0

Returns:

Type	Description
Float[Array, 'Ny Nx']	Result of (∇² − λ)·(ψ·mask), zeroed outside the mask.

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

def masked_laplacian(
    psi: Float[Array, "Ny Nx"],
    mask: Float[Array, "Ny Nx"] | Mask2D,
    dx: float,
    dy: float,
    lambda_: float | Array = 0.0,
) -> Float[Array, "Ny Nx"]:
    """Apply the masked discrete Helmholtz operator (∇² − λ)·ψ.

    Enforces homogeneous Dirichlet conditions at the mask boundary by zeroing
    *psi* outside the mask before applying the 5-point stencil.  The output
    is also zeroed outside the mask.

    Neighbours beyond the array edge are zero (dry), so the array edge is a
    homogeneous Dirichlet boundary one node outside the array, wherever the
    mask is wet there; nothing wraps around.  On an all-ones mask this is
    the dense Dirichlet 5-point operator that the DST-I solvers invert, so
    ``make_spectral_preconditioner(..., bc="dst")`` matches it exactly on a
    rectangle; the FFT preconditioner (periodic) is an approximation.

    Parameters
    ----------
    psi : Float[Array, "Ny Nx"]
        Field to which the operator is applied.
    mask : Float[Array, "Ny Nx"] or Mask2D
        Binary mask: 1 inside the physical domain, 0 outside (land/exterior).
        When an :class:`Mask2D` is passed, the ``psi`` staggering
        mask is used (all four surrounding h-cells must be wet).
    dx : float
        Grid spacing in x.
    dy : float
        Grid spacing in y.
    lambda_ : float
        Helmholtz parameter.  Default: 0.0 (pure Laplacian).

    Returns
    -------
    Float[Array, "Ny Nx"]
        Result of (∇² − λ)·(ψ·mask), zeroed outside the mask.
    """
    if isinstance(mask, Mask2D):
        mask_arr = mask.xy_corner_strict.astype(psi.dtype)
    else:
        mask_arr = mask
    psi_m = psi * mask_arr  # enforce zero outside domain  [Ny, Nx]
    # Zero ring outside the array (Dirichlet, not periodic); pad once and
    # bind each neighbour view once, so the stencil lowers to one fused pass.
    p = jnp.pad(psi_m, 1)  # [Ny+2, Nx+2]
    east = p[1:-1, 2:]  # ψ[j, i+1]  [Ny, Nx]
    west = p[1:-1, :-2]  # ψ[j, i-1]  [Ny, Nx]
    north = p[2:, 1:-1]  # ψ[j+1, i]  [Ny, Nx]
    south = p[:-2, 1:-1]  # ψ[j-1, i]  [Ny, Nx]
    # ∇²ψ[j, i] = (ψ[j, i+1] + ψ[j, i-1] - 2ψ[j, i]) / dx²
    #           + (ψ[j+1, i] + ψ[j-1, i] - 2ψ[j, i]) / dy²
    lap = (east + west - 2.0 * psi_m) / dx**2 + (north + south - 2.0 * psi_m) / dy**2
    return (lap - lambda_ * psi_m) * mask_arr  # [Ny, Nx]