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.

Uses random probing vectors to approximate the action of A^{-1} via a low-rank Nyström factorisation. The resulting preconditioner is a callable that can be passed to :func:~finitevolx._src.solvers.iterative.solve_cg.

Algorithm

Draw a Gaussian random matrix Ω ∈ ℝ^{n × k} (k = rank).
QR-factorize: Ω = Q R to get orthonormal Q ∈ ℝ^{n × k}.
Compute Y = A Q (k matvec applications).
Form the small matrix B = Q^T Y ∈ ℝ^{k × k}.
Compute B = U S U^T (eigendecomposition of the small matrix).
The preconditioner applies M^{-1} x ≈ (Q U S^{-1} U^T Q^T) x.

The operator A is assumed to be symmetric negative definite. The eigenvalues of B will be negative; their absolute values are used internally for inversion.

Parameters:

Name	Type	Description	Default
`matvec`	`callable`	Function implementing the symmetric 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). Higher values give a better preconditioner but cost more setup time. 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.

    Uses random probing vectors to approximate the action of ``A^{-1}`` via
    a low-rank Nyström factorisation.  The resulting preconditioner is a
    callable that can be passed to :func:`~finitevolx._src.solvers.iterative.solve_cg`.

    Algorithm
    ---------
    1. Draw a Gaussian random matrix ``Ω ∈ ℝ^{n × k}`` (``k = rank``).
    2. QR-factorize: ``Ω = Q R`` to get orthonormal ``Q ∈ ℝ^{n × k}``.
    3. Compute ``Y = A Q`` (``k`` matvec applications).
    4. Form the small matrix ``B = Q^T Y ∈ ℝ^{k × k}``.
    5. Compute ``B = U S U^T`` (eigendecomposition of the small matrix).
    6. The preconditioner applies ``M^{-1} x ≈ (Q U S^{-1} U^T Q^T) x``.

    The operator ``A`` is assumed to be symmetric **negative definite**.
    The eigenvalues of ``B`` will be negative; their absolute values are
    used internally for inversion.

    Parameters
    ----------
    matvec : callable
        Function implementing the symmetric 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).  Higher values
        give a better preconditioner but cost more setup time.
        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

    # Clamp rank to problem size
    k = min(rank, n)

    # Step 1: Random probing matrix Omega in R^{n x k}, then QR for stability
    omega = jax.random.normal(key, (n, k))
    Q, _ = jnp.linalg.qr(omega)  # Q in R^{n x k}, orthonormal columns

    # Step 2: Y = A Q  (k matvec applications)
    def _apply_col(col: Float[Array, " n"]) -> Float[Array, " n"]:
        return matvec(col.reshape(Ny, Nx)).ravel()

    Y = eqx.filter_vmap(_apply_col, in_axes=1, out_axes=1)(Q)  # [n, k]

    # Step 3: Small matrix B = Q^T Y = Q^T A Q in R^{k x k}
    B = Q.T @ Y  # [k, k]

    # Step 4: Eigendecomposition of B (symmetric)
    eigvals, U = jnp.linalg.eigh(B)  # eigvals sorted ascending

    # For a negative-definite A, eigvals should be negative.
    # Invert using absolute values, with a floor for numerical safety.
    abs_eigvals = jnp.abs(eigvals)
    eps = jnp.finfo(abs_eigvals.dtype).eps * n
    s_inv = jnp.where(abs_eigvals > eps, 1.0 / abs_eigvals, 0.0)
    # Preserve the sign: A^{-1} is also negative definite
    s_inv = -s_inv  # negate so preconditioner ≈ A^{-1} (negative)

    # Basis vectors: W = Q U has orthonormal columns (product of orthonormal
    # matrices), so M^{-1} = W diag(s_inv) W^T is a proper spectral
    # decomposition of the rank-k approximate inverse.
    W = Q @ U  # [n, k], orthonormal columns

    # Full-rank correction: for directions orthogonal to the captured subspace,
    # apply a scalar scaling instead of mapping to zero.  Without this, CG
    # falsely "converges" in the preconditioned norm while the true residual
    # remains ~1.0 (see #187).
    #
    # Random probing preferentially captures the largest-magnitude eigenvalues
    # of A; the uncaptured directions tend to have smaller-magnitude eigenvalues
    # and therefore *larger*-magnitude inverses.  Using s_inv[-1] (the largest-
    # magnitude captured inverse, from the smallest captured |eigenvalue|) as
    # the fallback keeps the preconditioned spectrum of uncaptured directions
    # close to 1, which is what CG needs for fast convergence.
    #
    # Rewrite:  M^{-1} x = a*x + W*((s_inv - a)*(W^T x))
    # which avoids explicit projection and is efficient.
    alpha_shift = s_inv[-1]  # best estimate for uncaptured (small |eig|) directions
    s_eff = s_inv - alpha_shift  # extra scaling for the captured subspace

    def _preconditioner(r: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
        r_flat = r.ravel()  # [n]
        coeffs = W.T @ r_flat  # [k]
        result = alpha_shift * r_flat + W @ (s_eff * coeffs)  # [n]
        return result.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}"
    )

Iterative Solvers & Preconditioners

Conjugate Gradient

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

finitevolx.CGInfo

converged instance-attribute

iterations instance-attribute

residual_norm instance-attribute

Preconditioners

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

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

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

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

`finitevolx.CGInfo`

`converged` `instance-attribute`

`iterations` `instance-attribute`

`residual_norm` `instance-attribute`

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

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

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