Iterative Solvers & Preconditioners

Preconditioned Conjugate Gradient solver and preconditioner factories for masked-domain elliptic problems.

Conjugate Gradient

finitevolx.solve_cg(matvec, rhs, x0=None, preconditioner=None, rtol=1e-06, atol=None, max_steps=500, *, throw=False)

Preconditioned Conjugate Gradient solver for symmetric linear operators.

Solves A(x) = rhs where matvec implements the symmetric linear operator A. The operator need not be explicitly formed; only matrix-vector products are required.

The solve is delegated to :class:lineax.CG, which provides a numerically stabilised PCG implementation with periodic residual recomputation to guard against floating-point drift.

The operator A is assumed to be negative definite, which is the case for the Helmholtz operator (nabla^2 - lambda) when lambda > 0. Note that lambda < 0 yields an indefinite operator and can make CG stall (info.converged is then False); use lambda = 0 only if the operator is strictly negative semidefinite via the mask (e.g. the masked Laplacian with homogeneous Dirichlet BCs on a non-trivial domain). The preconditioner, if supplied, should approximate A^{-1} (the negative-definite inverse). The required sign-flip for lineax.CG is handled internally.

Parameters:

Name	Type	Description	Default
matvec	callable	Function implementing the symmetric linear operator A. Signature: matvec(x: Array) -> Array.	required
rhs	Float[Array, '...']	Right-hand side of the linear system.	required
x0	Float[Array, '...'] or None	Initial guess. Defaults to zeros.	None
preconditioner	callable or None	Optional approximate inverse of A. Signature: preconditioner(r: Array) -> Array. When None, no preconditioning is applied (identity).	None
rtol	float	Relative convergence tolerance. Default: 1e-6.	1e-06
atol	float or None	Absolute convergence tolerance. None (the default) resolves it from the dtype: 1e-6 in float64; below float64, the solve runs on the rescaled system A y = rhs / s with s = max|rhs_i| and atol = √ε, and returns s y. See Notes.	None
max_steps	int or None	Maximum CG iterations. None means unlimited. Default: 500.	500
throw	bool	Raise when CG does not converge (as info.converged defines it), through :func:equinox.error_if, so it also works under jit. Default False: the iterate is returned and info.converged reports the failure, so check it.	False

Notes

Convergence. max_steps exhaustion never raises unless throw=True. info.converged is checked against the explicit residual of the un-preconditioned system, r = rhs - A(x): max|r_i| <= atol + rtol max|rhs_i| with the atol / rtol the solve used (s · √ε for the rescaled float32 solve). lineax alone reports success when CG stalls, e.g. after 0–1 iterations on an indefinite operator (lambda_ < 0) or with a preconditioner of the wrong sign; this check reports converged=False there.

Masks. The returned x can be non-zero outside the mask: a preconditioner that ignores the mask (the spectral ones) leaks into the dry points, which the masked operator then never sees. Multiply by the mask before using x; the wrappers (:func:streamfunction_from_vorticity, :func:pressure_from_divergence, :func:pv_inversion) do.

Gradients. The solve is differentiable with respect to anything matvec and rhs close over (lambda_, dx, …), through lineax's implicit derivative. x0 and the preconditioner do not change the converged solution, so their gradient is stopped; that is what lets a preconditioner built from the same lambda_ / dx stay in a differentiated solve.

Tolerances. lineax's CG stops when every entry satisfies |r_i| <= atol + rtol |rhs_i|. Where rhs_i = 0 (land, the ghost ring, quiet regions) that is |r_i| <= atol with an absolute atol. A fixed 1e-6 sits at float32's rounding floor once |rhs| ~ 1 and below it for an elliptic right-hand side of order 1/dx², so float32 CG never meets it: it runs to max_steps and can diverge with a spectral preconditioner. Rescaling by s = max|rhs_i| makes the effective atol = √ε · max|rhs_i| relative to the right-hand side (and unit-agnostic). s carries no gradient: it only changes the iteration count. This is the policy gaussx's CG strategies use (gaussx gh-639).

In float64 the default stays the historical absolute atol = 1e-6, which is not unit-agnostic: a right-hand side far below 1e-6 in magnitude (e.g. SI vorticity ~1e-9) is "converged" at the zero iterate. Pass an atol scaled to the right-hand side in that case.

Returns:

Name	Type	Description
x	Float[Array, '...']	Approximate solution, same shape as rhs.
info	CGInfo	Named tuple with fields iterations, residual_norm, and converged.

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

def solve_cg(
    matvec: Callable[[Float[Array, ...]], Float[Array, ...]],
    rhs: Float[Array, ...],
    x0: Float[Array, ...] | None = None,
    preconditioner: Callable[[Float[Array, ...]], Float[Array, ...]] | None = None,
    rtol: float = 1e-6,
    atol: float | None = None,
    max_steps: int | None = 500,
    *,
    throw: bool = False,
) -> tuple[Float[Array, ...], CGInfo]:
    """Preconditioned Conjugate Gradient solver for symmetric linear operators.

    Solves ``A(x) = rhs`` where ``matvec`` implements the symmetric linear
    operator ``A``.  The operator need not be explicitly formed; only
    matrix-vector products are required.

    The solve is delegated to :class:`lineax.CG`, which provides a
    numerically stabilised PCG implementation with periodic residual
    recomputation to guard against floating-point drift.

    The operator ``A`` is assumed to be *negative definite*, which is
    the case for the Helmholtz operator ``(nabla^2 - lambda)`` when
    ``lambda > 0``.  Note that ``lambda < 0`` yields an indefinite operator
    and can make CG stall (``info.converged`` is then ``False``); use
    ``lambda = 0`` only if the operator is strictly negative semidefinite
    via the mask (e.g. the masked Laplacian with homogeneous Dirichlet BCs
    on a non-trivial domain).
    The preconditioner, if supplied, should approximate ``A^{-1}``
    (the negative-definite inverse).  The required sign-flip for
    ``lineax.CG`` is handled internally.

    Parameters
    ----------
    matvec : callable
        Function implementing the symmetric linear operator ``A``.
        Signature: ``matvec(x: Array) -> Array``.
    rhs : Float[Array, "..."]
        Right-hand side of the linear system.
    x0 : Float[Array, "..."] or None
        Initial guess.  Defaults to zeros.
    preconditioner : callable or None
        Optional approximate inverse of ``A``.
        Signature: ``preconditioner(r: Array) -> Array``.
        When ``None``, no preconditioning is applied (identity).
    rtol : float
        Relative convergence tolerance.  Default: 1e-6.
    atol : float or None
        Absolute convergence tolerance.  ``None`` (the default) resolves it
        from the dtype: ``1e-6`` in float64; below float64, the solve runs on
        the rescaled system ``A y = rhs / s`` with ``s = max|rhs_i|`` and
        ``atol = √ε``, and returns ``s y``.  See Notes.
    max_steps : int or None
        Maximum CG iterations.  ``None`` means unlimited.  Default: 500.
    throw : bool
        Raise when CG does not converge (as ``info.converged`` defines it),
        through :func:`equinox.error_if`, so it also works under ``jit``.
        Default ``False``: the iterate is returned and ``info.converged``
        reports the failure, so check it.

    Notes
    -----
    **Convergence.**  ``max_steps`` exhaustion never raises unless
    ``throw=True``.  ``info.converged`` is checked against the explicit
    residual of the *un-preconditioned* system, ``r = rhs - A(x)``:
    ``max|r_i| <= atol + rtol max|rhs_i|`` with the ``atol`` / ``rtol`` the
    solve used (``s · √ε`` for the rescaled float32 solve).  lineax alone
    reports success when CG stalls, e.g. after 0–1 iterations on an
    indefinite operator (``lambda_ < 0``) or with a preconditioner of the
    wrong sign; this check reports ``converged=False`` there.

    **Masks.**  The returned ``x`` can be non-zero outside the mask: a
    preconditioner that ignores the mask (the spectral ones) leaks into the
    dry points, which the masked operator then never sees.  Multiply by the
    mask before using ``x``; the wrappers (:func:`streamfunction_from_vorticity`,
    :func:`pressure_from_divergence`, :func:`pv_inversion`) do.

    **Gradients.**  The solve is differentiable with respect to anything
    ``matvec`` and ``rhs`` close over (``lambda_``, ``dx``, …), through
    lineax's implicit derivative.  ``x0`` and the preconditioner do not
    change the converged solution, so their gradient is stopped; that is
    what lets a preconditioner built from the same ``lambda_`` / ``dx``
    stay in a differentiated solve.

    **Tolerances.**  lineax's CG stops when every entry satisfies
    ``|r_i| <= atol + rtol |rhs_i|``.  Where ``rhs_i = 0`` (land, the ghost
    ring, quiet regions) that is ``|r_i| <= atol`` with an *absolute*
    ``atol``.  A fixed ``1e-6`` sits at float32's rounding floor once
    ``|rhs| ~ 1`` and below it for an elliptic right-hand side of order
    ``1/dx²``, so float32 CG never meets it: it runs to ``max_steps`` and can
    diverge with a spectral preconditioner.  Rescaling by ``s = max|rhs_i|``
    makes the effective ``atol = √ε · max|rhs_i|`` relative to the
    right-hand side (and unit-agnostic).  ``s`` carries no gradient: it only
    changes the iteration count.  This is the policy gaussx's CG strategies
    use (gaussx gh-639).

    In float64 the default stays the historical *absolute* ``atol = 1e-6``,
    which is not unit-agnostic: a right-hand side far below ``1e-6`` in
    magnitude (e.g. SI vorticity ``~1e-9``) is "converged" at the zero
    iterate.  Pass an ``atol`` scaled to the right-hand side in that case.

    Returns
    -------
    x : Float[Array, "..."]
        Approximate solution, same shape as *rhs*.
    info : CGInfo
        Named tuple with fields ``iterations``, ``residual_norm``, and
        ``converged``.
    """
    atol, scale = _cg_atol_and_scale(atol, rhs)
    rhs_solve = rhs if scale is None else rhs / scale

    zero = jnp.zeros_like(rhs)
    # Tag as NSD; caller is responsible for ensuring A is actually negative (semi)definite.
    operator = lx.FunctionLinearOperator(
        matvec, zero, tags=[lx.negative_semidefinite_tag]
    )
    solver = lx.CG(rtol=rtol, atol=atol, max_steps=max_steps)

    options: dict = {}
    if x0 is not None:
        options["y0"] = x0 if scale is None else x0 / scale
    if preconditioner is not None:
        # lineax CG internally negates the NSD operator to make it positive, so it
        # expects a preconditioner that approximates (-A)^{-1} = -A^{-1} (positive).
        # The caller's `preconditioner` approximates A^{-1} (negative), so we negate
        # it here; this is transparent to the caller.
        def _lx_precond(r: Array) -> Array:
            return -preconditioner(r)  # negate: maps caller's A^{-1} -> (-A)^{-1}

        options["preconditioner"] = lx.FunctionLinearOperator(
            _lx_precond, zero, tags=[lx.positive_semidefinite_tag]
        )
    # lineax refuses tangents through ``options``.  Neither the initial guess
    # nor the preconditioner changes the converged solution x = A⁻¹ b, so its
    # derivative (implicit function theorem, through ``operator``) needs no
    # tangent from them: stop the gradient of their closed-over arrays
    # (e.g. ``lambda_`` / ``dx`` captured by a spectral preconditioner).
    options = jax.tree_util.tree_map(
        lambda leaf: jax.lax.stop_gradient(leaf) if eqx.is_array(leaf) else leaf,
        options,
    )

    # throw=False: max_steps / breakdown are reported through `converged`.
    sol = lx.linear_solve(
        operator, rhs_solve, solver=solver, options=options, throw=False
    )

    x_out = sol.value if scale is None else sol.value * scale
    # Explicit residual of the un-preconditioned system, r = b - A x  [Ny, Nx].
    # lineax's own success flag is not enough: it also exits "successfully"
    # on a stalled residual (indefinite A, wrong-sign preconditioner).
    residual = rhs - matvec(x_out)
    # ‖r‖∞ ≤ atol' + rtol ‖b‖∞, implied by lineax's entrywise stopping rule
    # |r_i| ≤ atol + rtol |b_i| (atol' = s · atol on the rescaled system).
    atol_eff = atol if scale is None else atol * scale
    tol = atol_eff + rtol * jnp.max(jnp.abs(rhs), initial=0.0)
    converged = (sol.result == lx.RESULTS.successful) & (
        jnp.max(jnp.abs(residual), initial=0.0) <= tol
    )
    if throw:
        x_out = eqx.error_if(
            x_out,
            jnp.logical_not(converged),
            "solve_cg did not converge: max_steps reached, or the residual "
            "|b - A x| exceeds atol + rtol |b| (an indefinite operator, e.g. "
            "lambda_ < 0, or a preconditioner of the wrong sign).",
        )
    info = CGInfo(
        iterations=sol.stats["num_steps"],
        residual_norm=jnp.linalg.norm(residual),
        converged=converged,
    )
    return x_out, info

finitevolx.CGInfo

Bases: NamedTuple

Convergence diagnostics returned by :func:solve_cg.

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

class CGInfo(NamedTuple):
    """Convergence diagnostics returned by :func:`solve_cg`."""

    iterations: int
    """Number of PCG iterations performed."""
    residual_norm: float
    """L2 norm of the explicit final residual ``rhs - A(x)``."""
    converged: bool
    """True if lineax reported success *and* the explicit residual satisfies
    ``max|rhs - A(x)| <= atol + rtol max|rhs|``."""

converged instance-attribute

True if lineax reported success and the explicit residual satisfies max|rhs - A(x)| <= atol + rtol max|rhs|.

iterations instance-attribute

Number of PCG iterations performed.

residual_norm instance-attribute

L2 norm of the explicit final residual rhs - A(x).

Preconditioners

finitevolx.make_spectral_preconditioner(dx, dy, lambda_=0.0, bc='fft')

Return a spectral-solve preconditioner for PCG.

The preconditioner applies the rectangular spectral solver as an approximate inverse of the Helmholtz operator (∇² − λ). It is particularly effective when the physical domain is a subset of the rectangle (masked-domain problems).

Parameters:

Name	Type	Description	Default
dx	float	Grid spacing in x.	required
dy	float	Grid spacing in y.	required
lambda_	float	Helmholtz parameter. Must match the parameter used in the main problem. Default: 0.0 (Poisson).	0.0
bc	('fft', 'dst', 'dct')	Spectral solver type to use as the preconditioner. Default: "fft" (periodic).	"fft"

Returns:

Type	Description
Callable	A function M_inv(r: Array) -> Array that applies the approximate inverse.

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

def make_spectral_preconditioner(
    dx: float,
    dy: float,
    lambda_: float | Array = 0.0,
    bc: str = "fft",
) -> Callable[[Float[Array, "Ny Nx"]], Float[Array, "Ny Nx"]]:
    """Return a spectral-solve preconditioner for PCG.

    The preconditioner applies the rectangular spectral solver as an
    approximate inverse of the Helmholtz operator ``(∇² − λ)``.  It is
    particularly effective when the physical domain is a subset of the
    rectangle (masked-domain problems).

    Parameters
    ----------
    dx : float
        Grid spacing in x.
    dy : float
        Grid spacing in y.
    lambda_ : float
        Helmholtz parameter.  Must match the parameter used in the main
        problem.  Default: 0.0 (Poisson).
    bc : {"fft", "dst", "dct"}
        Spectral solver type to use as the preconditioner.
        Default: ``"fft"`` (periodic).

    Returns
    -------
    Callable
        A function ``M_inv(r: Array) -> Array`` that applies the
        approximate inverse.
    """
    if bc not in {"fft", "dst", "dct"}:
        raise ValueError(f"bc must be 'fft', 'dst', or 'dct'; got {bc!r}")

    from finitevolx._src.solvers.spectral import _spectral_solve

    def _preconditioner(r: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
        return _spectral_solve(r, dx, dy, lambda_, bc)

    return _preconditioner

finitevolx.make_nystrom_preconditioner(matvec, shape, rank=50, key=None)

Build a randomized Nyström preconditioner for a symmetric operator.

Delegates the low-rank Nyström construction to :func:gaussx.randomized_nystrom and :class:gaussx.NystromPreconditioner (Frangella, Tropp & Udell, 2023), then adapts it to the finitevolX convention: the returned callable operates on 2-D fields and approximates A^{-1} for the symmetric negative-definite operator A.

gaussx works with positive-semidefinite operators, so this sketches the PSD operator B = -A as B ≈ U Λ̂ Uᵀ (rank l) and applies

M⁻¹ r = -[ U (Λ̂ + μ)⁻¹ Uᵀ r + (r - U Uᵀ r) / (λ̂_l + μ) ],

i.e. the exact inverse on the captured top of the spectrum and a scalar 1 / (λ̂_l + μ) on its complement. μ = √ε · λ̂_1 is a relative regulariser (ε the machine epsilon of the dtype), so the preconditioner rescales with the operator (contract 6) and stays well-defined for a singular A (e.g. a Neumann Poisson problem).

The Nyström sketch captures the largest eigenvalues of -A. For a Laplacian these are the grid-scale modes, so it deflates the top of the spectrum but leaves the smooth modes (the small eigenvalues that set the condition number) to CG: use make_spectral_preconditioner or make_multigrid_preconditioner when iteration count matters.

Parameters:

Name	Type	Description	Default
matvec	callable	Function implementing the symmetric (negative-definite) linear operator A. Signature: matvec(x: Array) -> Array.	required
shape	(int, int)	Spatial shape (Ny, Nx) of the 2-D fields.	required
rank	int	Number of probing vectors (approximation rank). Default: 50.	50
key	Array or None	PRNG key for the random probing matrix. If None, uses jax.random.PRNGKey(0).	None

Returns:

Type	Description
Callable	A function M_inv(r: Array) -> Array that applies the approximate inverse.

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

def make_nystrom_preconditioner(
    matvec: Callable[[Float[Array, "Ny Nx"]], Float[Array, "Ny Nx"]],
    shape: tuple[int, int],
    rank: int = 50,
    key: jax.Array | None = None,
) -> Callable[[Float[Array, "Ny Nx"]], Float[Array, "Ny Nx"]]:
    r"""Build a randomized Nyström preconditioner for a symmetric operator.

    Delegates the low-rank Nyström construction to
    :func:`gaussx.randomized_nystrom` and :class:`gaussx.NystromPreconditioner`
    (Frangella, Tropp & Udell, 2023), then adapts it to the finitevolX
    convention: the returned callable operates on 2-D fields and approximates
    ``A^{-1}`` for the symmetric **negative-definite** operator ``A``.

    gaussx works with positive-semidefinite operators, so this sketches the
    PSD operator ``B = -A`` as ``B ≈ U Λ̂ Uᵀ`` (rank ``l``) and applies

    ``M⁻¹ r = -[ U (Λ̂ + μ)⁻¹ Uᵀ r + (r - U Uᵀ r) / (λ̂_l + μ) ]``,

    i.e. the exact inverse on the captured top of the spectrum and a scalar
    ``1 / (λ̂_l + μ)`` on its complement. ``μ = √ε · λ̂_1`` is a relative
    regulariser (``ε`` the machine epsilon of the dtype), so the
    preconditioner rescales with the operator (contract 6) and stays
    well-defined for a singular ``A`` (e.g. a Neumann Poisson problem).

    The Nyström sketch captures the *largest* eigenvalues of ``-A``. For a
    Laplacian these are the grid-scale modes, so it deflates the top of the
    spectrum but leaves the smooth modes (the small eigenvalues that set the
    condition number) to CG: use ``make_spectral_preconditioner`` or
    ``make_multigrid_preconditioner`` when iteration count matters.

    Parameters
    ----------
    matvec : callable
        Function implementing the symmetric (negative-definite) linear operator
        ``A``.  Signature: ``matvec(x: Array) -> Array``.
    shape : (int, int)
        Spatial shape ``(Ny, Nx)`` of the 2-D fields.
    rank : int
        Number of probing vectors (approximation rank).  Default: 50.
    key : jax.Array or None
        PRNG key for the random probing matrix.  If ``None``, uses
        ``jax.random.PRNGKey(0)``.

    Returns
    -------
    Callable
        A function ``M_inv(r: Array) -> Array`` that applies the
        approximate inverse.
    """
    if key is None:
        key = jax.random.PRNGKey(0)

    Ny, Nx = shape
    n = Ny * Nx

    # gaussx Nyström expects a PSD operator; probe -A (which is PSD) on flat
    # vectors, then negate the approximate inverse to recover A^{-1}.
    def _neg_flat_matvec(v: Float[Array, " n"]) -> Float[Array, " n"]:
        return -matvec(v.reshape(Ny, Nx)).ravel()

    operator = gaussx.as_linear_operator(
        _neg_flat_matvec, shape=(n, n), positive_semidefinite=True
    )
    # B = -A ≈ U Λ̂ Uᵀ, U: [n, l], Λ̂: [l] descending
    approx = gaussx.randomized_nystrom(operator, min(rank, n), key=key)
    eigenvalues = approx.d
    # μ = √ε · λ̂_1: relative, so it rescales with the operator (no absolute ε)
    shift = jnp.sqrt(jnp.finfo(eigenvalues.dtype).eps) * eigenvalues[0]
    preconditioner = gaussx.NystromPreconditioner(
        basis=approx.U, eigenvalues=eigenvalues, shift=shift
    )
    # P⁻¹ = (λ̂_l + μ) [U (Λ̂ + μ)⁻¹ Uᵀ + (I - U Uᵀ) / (λ̂_l + μ)]
    minv = preconditioner.as_operator()
    scale = eigenvalues[-1] + shift  # undo the (λ̂_l + μ) factor → ≈ (-A)⁻¹

    def _preconditioner(r: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
        return -minv.mv(r.reshape(n)).reshape(Ny, Nx) / scale

    return _preconditioner

finitevolx.make_preconditioner(kind, *, dx=None, dy=None, lambda_=0.0, bc='fft', matvec=None, shape=None, rank=50, key=None, mg_solver=None)

Convenience factory that builds a preconditioner by name.

Dispatches to :func:make_spectral_preconditioner, :func:make_nystrom_preconditioner, or :func:make_multigrid_preconditioner based on kind.

Parameters:

Name	Type	Description	Default
kind	('spectral', 'nystrom', 'multigrid')	Which preconditioner to build.	"spectral"
dx	float or None	Grid spacings. Required for kind="spectral".	None
dy	float or None	Grid spacings. Required for kind="spectral".	None
lambda_	float	Helmholtz parameter. Used by kind="spectral". Default: 0.0.	0.0
bc	('fft', 'dst', 'dct')	Spectral solver type. Used by kind="spectral". Default: "fft".	"fft"
matvec	callable or None	Operator A. Required for kind="nystrom".	None
shape	(int, int) or None	Spatial shape (Ny, Nx). Required for kind="nystrom".	None
rank	int	Approximation rank. Used by kind="nystrom". Default: 50.	50
key	Array or None	PRNG key. Used by kind="nystrom".	None
mg_solver	MultigridSolver or None	Pre-built multigrid solver. Required for kind="multigrid".	None

Returns:

Type	Description
callable	A function M_inv(r: Array) -> Array that applies the approximate inverse, compatible with :func:solve_cg.

Raises:

Type	Description
ValueError	If kind is unknown or required arguments are missing.

Examples:

>>> pc = make_preconditioner("spectral", dx=dx, dy=dy, lambda_=10.0)
>>> pc = make_preconditioner("nystrom", matvec=A, shape=(64, 64), rank=30)
>>> pc = make_preconditioner("multigrid", mg_solver=mg)

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

def make_preconditioner(
    kind: str,
    *,
    dx: float | None = None,
    dy: float | None = None,
    lambda_: float = 0.0,
    bc: str = "fft",
    matvec: Callable[[Float[Array, "Ny Nx"]], Float[Array, "Ny Nx"]] | None = None,
    shape: tuple[int, int] | None = None,
    rank: int = 50,
    key: jax.Array | None = None,
    mg_solver: MultigridSolver | None = None,
) -> Callable[[Float[Array, "Ny Nx"]], Float[Array, "Ny Nx"]]:
    """Convenience factory that builds a preconditioner by name.

    Dispatches to :func:`make_spectral_preconditioner`,
    :func:`make_nystrom_preconditioner`, or
    :func:`make_multigrid_preconditioner` based on *kind*.

    Parameters
    ----------
    kind : {"spectral", "nystrom", "multigrid"}
        Which preconditioner to build.

    dx, dy : float or None
        Grid spacings.  Required for ``kind="spectral"``.
    lambda_ : float
        Helmholtz parameter.  Used by ``kind="spectral"``.  Default: 0.0.
    bc : {"fft", "dst", "dct"}
        Spectral solver type.  Used by ``kind="spectral"``.  Default: ``"fft"``.

    matvec : callable or None
        Operator ``A``.  Required for ``kind="nystrom"``.
    shape : (int, int) or None
        Spatial shape ``(Ny, Nx)``.  Required for ``kind="nystrom"``.
    rank : int
        Approximation rank.  Used by ``kind="nystrom"``.  Default: 50.
    key : jax.Array or None
        PRNG key.  Used by ``kind="nystrom"``.

    mg_solver : MultigridSolver or None
        Pre-built multigrid solver.  Required for ``kind="multigrid"``.

    Returns
    -------
    callable
        A function ``M_inv(r: Array) -> Array`` that applies the
        approximate inverse, compatible with :func:`solve_cg`.

    Raises
    ------
    ValueError
        If *kind* is unknown or required arguments are missing.

    Examples
    --------
    >>> pc = make_preconditioner("spectral", dx=dx, dy=dy, lambda_=10.0)
    >>> pc = make_preconditioner("nystrom", matvec=A, shape=(64, 64), rank=30)
    >>> pc = make_preconditioner("multigrid", mg_solver=mg)
    """
    if kind == "spectral":
        if dx is None or dy is None:
            raise ValueError("kind='spectral' requires dx and dy")
        return make_spectral_preconditioner(dx, dy, lambda_=lambda_, bc=bc)

    if kind == "nystrom":
        if matvec is None or shape is None:
            raise ValueError("kind='nystrom' requires matvec and shape")
        return make_nystrom_preconditioner(matvec, shape, rank=rank, key=key)

    if kind == "multigrid":
        if mg_solver is None:
            raise ValueError("kind='multigrid' requires mg_solver")
        return make_multigrid_preconditioner(mg_solver)

    raise ValueError(
        f"kind must be 'spectral', 'nystrom', or 'multigrid'; got {kind!r}"
    )

Helmholtz operator as a gaussx operator

The rectangular and masked Helmholtz operators as structured gaussx linear operators, for gaussx.solve, gaussx.logdet, gaussx.sqrt and SPDE (Matérn) priors on the model grid.

finitevolx.helmholtz_operator(shape, dx, dy, lambda_=0.0, bc='fft')

The rectangular Helmholtz operator (∇² − λ) as a gaussx operator.

The 5-point FD2 Laplacian on a rectangle with periodic ("fft"), homogeneous Dirichlet ("dst", DST-I) or homogeneous Neumann ("dct", DCT-II) boundaries is diagonalised by that transform:

A = V⁻¹ Λ V, Λ[j, i] = Λʸ_j + Λˣ_i − λ

with the same eigenvalues (fft_eigenvalues, dst1_eigenvalues, dct2_eigenvalues) as the spectral solvers, so gaussx.solve(A, f) equals solve_helmholtz_{bc}(f, dx, dy, λ). gaussx.logdet, gaussx.sqrt, gaussx.inv and gaussx.trace are elementwise in Λ (gaussx.diag_inv is not: it densifies), and the algebra stays closed (-A, A - c I, c A), so a PSD κ² − ∇² is -helmholtz_operator(..., lambda_=κ²).

Parameters:

Name	Type	Description	Default
shape	(int, int)	Field shape (Ny, Nx); the operator acts on flat vectors of length Ny·Nx (row-major), the layout the spectral solvers use.	required
dx	float	Grid spacings.	required
dy	float	Grid spacings.	required
lambda_	float	Helmholtz parameter λ ≥ 0 (contract 5). Default 0 (Poisson). A concrete λ ≥ 0 tags the operator negative semidefinite.	0.0
bc	('fft', 'dst', 'dct')	Rectangle boundary condition, as for the spectral solvers.	"fft"

Notes

The Poisson null mode. With λ = 0 and bc="fft" or "dct" the constant is in the null space: Λ[0, 0] = 0. gaussx treats an exactly-zero eigenvalue as the pseudo-inverse (a zero coefficient), so gaussx.solve and gaussx.inv return the zero-mean solution — the same gauge as solve_helmholtz_fft / _dct — and stay finite, while gaussx.logdet is -inf (the slogdet convention for a singular operator). Use λ > 0 (or bc="dst") for a log-determinant. A right-hand side with a non-zero mean has no solution; the pseudo-inverse silently drops that mean, exactly as the spectral solvers do.

Returns:

Type	Description
DiagonalizedOperator	Symmetric, (Ny·Nx, Ny·Nx), never materialised.

Examples:

>>> A = helmholtz_operator((Ny, Nx), dx, dy, lambda_=1.0, bc="dst")
>>> psi = gaussx.solve(A, f.ravel()).reshape(Ny, Nx)  # = solve_helmholtz_dst
>>> logdet = gaussx.logdet(-A)  # log det(1 - ∇²)

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

def helmholtz_operator(
    shape: tuple[int, int],
    dx: float,
    dy: float,
    lambda_: float = 0.0,
    bc: str = "fft",
) -> gaussx.DiagonalizedOperator:
    r"""The rectangular Helmholtz operator ``(∇² − λ)`` as a gaussx operator.

    The 5-point FD2 Laplacian on a rectangle with periodic (``"fft"``),
    homogeneous Dirichlet (``"dst"``, DST-I) or homogeneous Neumann
    (``"dct"``, DCT-II) boundaries is diagonalised by that transform:

    ``A = V⁻¹ Λ V``,  ``Λ[j, i] = Λʸ_j + Λˣ_i − λ``

    with the same eigenvalues (``fft_eigenvalues``, ``dst1_eigenvalues``,
    ``dct2_eigenvalues``) as the spectral solvers, so ``gaussx.solve(A, f)``
    equals ``solve_helmholtz_{bc}(f, dx, dy, λ)``. ``gaussx.logdet``,
    ``gaussx.sqrt``, ``gaussx.inv`` and ``gaussx.trace`` are elementwise in
    ``Λ`` (``gaussx.diag_inv`` is not: it densifies), and the algebra stays
    closed (``-A``, ``A - c I``, ``c A``), so a
    PSD ``κ² − ∇²`` is ``-helmholtz_operator(..., lambda_=κ²)``.

    Parameters
    ----------
    shape : (int, int)
        Field shape ``(Ny, Nx)``; the operator acts on flat vectors of length
        ``Ny·Nx`` (row-major), the layout the spectral solvers use.
    dx, dy : float
        Grid spacings.
    lambda_ : float
        Helmholtz parameter ``λ ≥ 0`` (contract 5). Default ``0`` (Poisson).
        A concrete ``λ ≥ 0`` tags the operator negative semidefinite.
    bc : {"fft", "dst", "dct"}
        Rectangle boundary condition, as for the spectral solvers.

    Notes
    -----
    **The Poisson null mode.** With ``λ = 0`` and ``bc="fft"`` or ``"dct"``
    the constant is in the null space: ``Λ[0, 0] = 0``. gaussx treats an
    exactly-zero eigenvalue as the pseudo-inverse (a zero coefficient), so
    ``gaussx.solve`` and ``gaussx.inv`` return the zero-mean solution — the
    same gauge as ``solve_helmholtz_fft`` / ``_dct`` — and stay finite, while
    ``gaussx.logdet`` is ``-inf`` (the ``slogdet`` convention for a singular
    operator). Use ``λ > 0`` (or ``bc="dst"``) for a log-determinant. A
    right-hand side with a non-zero mean has no solution; the pseudo-inverse
    silently drops that mean, exactly as the spectral solvers do.

    Returns
    -------
    gaussx.DiagonalizedOperator
        Symmetric, ``(Ny·Nx, Ny·Nx)``, never materialised.

    Examples
    --------
    >>> A = helmholtz_operator((Ny, Nx), dx, dy, lambda_=1.0, bc="dst")
    >>> psi = gaussx.solve(A, f.ravel()).reshape(Ny, Nx)  # = solve_helmholtz_dst
    >>> logdet = gaussx.logdet(-A)  # log det(1 - ∇²)
    """
    if bc not in _EIGENVALUES:
        raise ValueError(f"bc must be 'fft', 'dst', or 'dct'; got {bc!r}")
    ny, nx = (int(n) for n in shape)
    eig_1d = _EIGENVALUES[bc]
    # Λ[j, i] = Λʸ_j + Λˣ_i - λ, [Ny, Nx]
    eigenvalues = eig_1d(ny, dy)[:, None] + eig_1d(nx, dx)[None, :] - lambda_
    if bc == "fft":
        # A x = ifft2(Λ ⊙ fft2(x)); gaussx checks Λ is conjugate-even (real here)
        return gaussx.circulant_from_symbol(eigenvalues, tags=_sign_tags(lambda_))
    forward, inverse = _ORTHO_PAIRS[bc]
    return gaussx.DiagonalizedOperator(
        eigenvalues,
        forward,
        inverse,
        (ny, nx),
        normal=True,
        tags=_sign_tags(lambda_),
    )

finitevolx.masked_helmholtz_operator(mask, dx, dy, lambda_=0.0, bc='dst')

The masked Helmholtz operator of :func:masked_laplacian as a gaussx operator.

masked_laplacian applies the 5-point (∇² − λ) to ψ zeroed outside the mask and outside the array (homogeneous Dirichlet on the coast and one node beyond the array edge) and keeps the wet rows. As a matrix that is the wet sub-block A[m][:, m] of the Dirichlet rectangular operator A (:func:helmholtz_operator with bc="dst"). This returns that sub-block as a gaussx.MaskedOperator acting on the wet values only, with the coupling ring (the dry cells the stencil reaches) precomputed, so gaussx.solve is the capacitance method: exact, two DST solves per right-hand side.

The mask is concrete and the construction runs on the host (it builds the |C| × |C| capacitance factorisation); build it once per geometry and reuse it, like build_capacitance_solver.

Parameters:

Name	Type	Description	Default
mask	array of shape (Ny, Nx), or Mask2D	Wet cells (True / 1). A :class:Mask2D contributes its xy_corner_strict (ψ-point) mask, as in masked_laplacian.	required
dx	float	Grid spacings.	required
dy	float	Grid spacings.	required
lambda_	float	Helmholtz parameter λ ≥ 0. Default 0 (Poisson); the Dirichlet base operator is non-singular at λ = 0.	0.0
bc	('dst', 'fft', 'dct')	Rectangular base operator whose solves the capacitance method uses. The base enters the masked block only through cells beyond the array edge: with "dst" (default) the block is exactly masked_laplacian's; "fft" (periodic wrap) and "dct" give the same block only when no wet cell touches the array edge. A singular base ("fft" / "dct" at λ = 0) is handled exactly by gaussx (it borders the capacitance system with the base's null vector).	"dst"

Returns:

Type	Description
MaskedOperator	(n_wet, n_wet) with n_wet = mask.sum(); vectors are psi.ravel()[mask.ravel()].

Examples:

>>> A_m = masked_helmholtz_operator(mask, dx, dy, lambda_=1.0)
>>> wet = np.asarray(mask, bool).ravel()
>>> psi_wet = gaussx.solve(A_m, f.ravel()[wet])
>>> psi = jnp.zeros(mask.size).at[wet].set(psi_wet).reshape(mask.shape)

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

def masked_helmholtz_operator(
    mask: Float[Array, "Ny Nx"] | Bool[Array, "Ny Nx"] | Mask2D,
    dx: float,
    dy: float,
    lambda_: float = 0.0,
    bc: str = "dst",
) -> gaussx.MaskedOperator:
    r"""The masked Helmholtz operator of :func:`masked_laplacian` as a gaussx operator.

    ``masked_laplacian`` applies the 5-point ``(∇² − λ)`` to ``ψ`` zeroed
    outside the mask and outside the array (homogeneous Dirichlet on the
    coast and one node beyond the array edge) and keeps the wet rows. As a
    matrix that is the wet sub-block ``A[m][:, m]`` of the Dirichlet
    rectangular operator ``A`` (:func:`helmholtz_operator` with
    ``bc="dst"``). This returns that sub-block as a ``gaussx.MaskedOperator``
    acting on the wet values only, with the coupling ring (the dry cells the
    stencil reaches) precomputed, so ``gaussx.solve`` is the capacitance
    method: exact, two DST solves per right-hand side.

    The mask is concrete and the construction runs on the host (it builds
    the ``|C| × |C|`` capacitance factorisation); build it once per geometry
    and reuse it, like ``build_capacitance_solver``.

    Parameters
    ----------
    mask : array of shape (Ny, Nx), or Mask2D
        Wet cells (``True`` / ``1``). A :class:`Mask2D` contributes its
        ``xy_corner_strict`` (ψ-point) mask, as in ``masked_laplacian``.
    dx, dy : float
        Grid spacings.
    lambda_ : float
        Helmholtz parameter ``λ ≥ 0``. Default ``0`` (Poisson); the
        Dirichlet base operator is non-singular at ``λ = 0``.
    bc : {"dst", "fft", "dct"}
        Rectangular base operator whose solves the capacitance method uses.
        The base enters the masked block only through cells beyond the
        array edge: with ``"dst"`` (default) the block is exactly
        ``masked_laplacian``'s; ``"fft"`` (periodic wrap) and ``"dct"``
        give the same block only when no wet cell touches the array edge.
        A singular base (``"fft"`` / ``"dct"`` at ``λ = 0``) is handled
        exactly by gaussx (it borders the capacitance system with the
        base's null vector).

    Returns
    -------
    gaussx.MaskedOperator
        ``(n_wet, n_wet)`` with ``n_wet = mask.sum()``; vectors are
        ``psi.ravel()[mask.ravel()]``.

    Examples
    --------
    >>> A_m = masked_helmholtz_operator(mask, dx, dy, lambda_=1.0)
    >>> wet = np.asarray(mask, bool).ravel()
    >>> psi_wet = gaussx.solve(A_m, f.ravel()[wet])
    >>> psi = jnp.zeros(mask.size).at[wet].set(psi_wet).reshape(mask.shape)
    """
    if isinstance(mask, Mask2D):
        mask = mask.xy_corner_strict
    wet = np.asarray(mask).astype(bool)  # [Ny, Nx], concrete
    # B = (∇² - λ) on the rectangle; "dst" matches masked_laplacian's zero pad
    base = helmholtz_operator(wet.shape, dx, dy, lambda_=lambda_, bc=bc)
    flat = jnp.asarray(wet.ravel())  # [Ny*Nx], the unknowns
    # Coupling set: the dry cells the stencil of a wet cell reaches (the coast
    # ring), wrapping around the array only for the periodic base.
    coupling = gaussx.grid_coupling_indices(jnp.asarray(wet), periodic=bc == "fft")
    return gaussx.MaskedOperator(
        base,
        flat,
        flat,
        coupling_indices=coupling,
        tags=_sign_tags(lambda_) | {lx.symmetric_tag},
    )