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=1e-06, max_steps=500)

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 will cause CG to fail; 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	Absolute convergence tolerance. Default: 1e-6.	1e-06
max_steps	int or None	Maximum CG iterations. None means unlimited. Default: 500.	500

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 = 1e-6,
    max_steps: int | None = 500,
) -> 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 will cause CG to fail; 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
        Absolute convergence tolerance.  Default: 1e-6.
    max_steps : int or None
        Maximum CG iterations.  ``None`` means unlimited.  Default: 500.

    Returns
    -------
    x : Float[Array, "..."]
        Approximate solution, same shape as *rhs*.
    info : CGInfo
        Named tuple with fields ``iterations``, ``residual_norm``, and
        ``converged``.
    """
    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 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]
        )

    sol = lx.linear_solve(operator, rhs, solver=solver, options=options)

    x_out = sol.value
    res_norm = jnp.linalg.norm(matvec(x_out) - rhs)
    info = CGInfo(
        iterations=sol.stats["num_steps"],
        residual_norm=res_norm,
        converged=sol.result == lx.RESULTS.successful,
    )
    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 final residual ``A(x) - rhs``."""
    converged: bool
    """True if the solver converged within the requested tolerances."""

converged instance-attribute

True if the solver converged within the requested tolerances.

iterations instance-attribute

Number of PCG iterations performed.

residual_norm instance-attribute

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

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 = 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 :class:gaussx.NystromPreconditioner, 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 probes the PSD operator -A (obtaining (-A)^{-1}) and negates the result to recover A^{-1}. The resulting callable can be passed straight to :func:~finitevolx._src.solvers.iterative.solve_cg.

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
    :class:`gaussx.NystromPreconditioner`, 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 probes the PSD
    operator ``-A`` (obtaining ``(-A)^{-1}``) and negates the result to recover
    ``A^{-1}``. The resulting callable can be passed straight to
    :func:`~finitevolx._src.solvers.iterative.solve_cg`.

    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
    )
    preconditioner = gaussx.NystromPreconditioner.from_operator(
        operator, rank=rank, key=key
    )
    minv = preconditioner.as_operator()  # applies (-A)^{-1} on flat vectors

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

    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}"
    )