Multigrid Helmholtz Solver

Geometric multigrid solver for variable-coefficient Helmholtz equations on masked Arakawa C-grids.

Solver

finitevolx.MultigridSolver

Bases: Module

Geometric multigrid V-cycle solver for the variable-coefficient Helmholtz equation::

div(c(x,y) grad u) - lambda u = rhs

Supports spatially varying coefficients c(x, y) on staggered faces and masked (irregular) domains.

V-Cycle Algorithm

The V-cycle is a recursive algorithm that visits coarser grids to correct low-frequency error that the smoother cannot resolve::

Level 0 (fine)    *---smooth---*-----------*---smooth---*
                       | restrict    prolong |
Level 1           .....*---smooth---*---smooth---*.....
                            | restrict  prolong |
Level 2 (coarse)  ..........*--exact solve--*..........

At each level:

1. Pre-smooth (nu_1 weighted Jacobi iterations): damp high-frequency error.
2. Compute residual: r = rhs - A u.
3. Restrict the residual to the coarse grid (full weighting, R = Pᵀ / 4).
4. Recurse: solve for the error on the coarse grid, whose operator is the Galerkin product A_c = R A P.
5. Prolongate the coarse correction back to the fine grid (bilinear interpolation P).
6. Post-smooth (nu_2 weighted Jacobi iterations): damp any high-frequency error introduced by the prolongation.

The coarsest level is solved exactly (an inverse precomputed at build time). With n_pre == n_post one V-cycle from a zero guess, r -> v_cycle(0, r), is a symmetric linear operator, so it is a valid CG preconditioner.

The recursion is statically unrolled at JAX trace time because each level has a different array shape. All integer parameters (n_levels, n_pre, etc.) are eqx.field(static=True), so the unrolled structure is visible to the XLA compiler.

Differentiation Modes

Three solve modes trade off backward-pass cost vs gradient accuracy:

• __call__ — Implicit differentiation via jax.lax.custom_linear_solve(symmetric=True). The backward pass solves the adjoint system A^T v = dL/du with one multigrid call. Since A is symmetric, this costs the same as the forward pass. O(1) memory, exact gradients for the linear system being solved (gradient accuracy is limited by how well the V-cycles approximate A^{-1}, i.e. depends on n_cycles and smoother settings).

• solve_onestep — One-step differentiation (Bolte, Pauwels & Vaiter, NeurIPS 2023). Runs K V-cycles, applies stop_gradient after K-1, then autodiffs through the last cycle only. O(1 V-cycle) memory, approximate gradients with error O(rho).

• solve_unrolled — Unrolled differentiation via jax.lax.fori_loop. Backward replays all K iterations. O(K) memory, exact through-iteration gradients (reproduces the forward computation exactly, so gradient accuracy matches the forward solve accuracy).

Attributes:

Name	Type	Description
levels	tuple of MultigridLevel	Precomputed level data, finest (index 0) to coarsest (index L-1).
n_levels	int	Number of multigrid levels.
bottom_inverse	Float[Array, 'n_wet n_wet']	Inverse of the coarsest operator on its wet points (the pseudo-inverse, with a relative eigenvalue cut-off, for the singular Neumann lambda = 0 problem).
bottom_index	Int[Array, ' n_wet']	Flat (row-major) indices of the coarsest level's wet points.
n_pre, n_post	int	Pre- and post-smoothing iterations (weighted Jacobi). Keep them equal for a symmetric V-cycle.
omega	float	Jacobi relaxation weight.
n_cycles	int	Number of V-cycles per solve.
bc	{'dirichlet', 'neumann'}	Wall condition on the wet–dry faces the hierarchy was built with (see :func:build_multigrid_solver).
null_labels, null_counts	Array or None	Connected-component labels and wet counts of the finest mask, set only for bc="neumann" with lambda_ = 0, where the operator is singular; the solve then returns the zero-mean solution on each component (and projects rhs onto the solvable, zero-mean part).

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

class MultigridSolver(eqx.Module):
    r"""Geometric multigrid V-cycle solver for the variable-coefficient
    Helmholtz equation::

        div(c(x,y) grad u) - lambda u = rhs

    Supports spatially varying coefficients ``c(x, y)`` on staggered faces
    and masked (irregular) domains.

    V-Cycle Algorithm
    -----------------
    The V-cycle is a recursive algorithm that visits coarser grids to
    correct low-frequency error that the smoother cannot resolve::

        Level 0 (fine)    *---smooth---*-----------*---smooth---*
                               | restrict    prolong |
        Level 1           .....*---smooth---*---smooth---*.....
                                    | restrict  prolong |
        Level 2 (coarse)  ..........*--exact solve--*..........

    At each level:

    1. **Pre-smooth** (nu_1 weighted Jacobi iterations): damp
       high-frequency error.
    2. **Compute residual**: ``r = rhs - A u``.
    3. **Restrict** the residual to the coarse grid (full weighting,
       ``R = Pᵀ / 4``).
    4. **Recurse**: solve for the error on the coarse grid, whose operator
       is the Galerkin product ``A_c = R A P``.
    5. **Prolongate** the coarse correction back to the fine grid
       (bilinear interpolation ``P``).
    6. **Post-smooth** (nu_2 weighted Jacobi iterations): damp any
       high-frequency error introduced by the prolongation.

    The coarsest level is solved exactly (an inverse precomputed at build
    time).  With ``n_pre == n_post`` one V-cycle from a zero guess,
    ``r -> v_cycle(0, r)``, is a symmetric linear operator, so it is a valid
    CG preconditioner.

    The recursion is **statically unrolled** at JAX trace time because
    each level has a different array shape.  All integer parameters
    (``n_levels``, ``n_pre``, etc.) are ``eqx.field(static=True)``,
    so the unrolled structure is visible to the XLA compiler.

    Differentiation Modes
    ---------------------
    Three solve modes trade off backward-pass cost vs gradient accuracy:

    * ``__call__`` — **Implicit differentiation** via
      ``jax.lax.custom_linear_solve(symmetric=True)``.  The backward pass
      solves the adjoint system ``A^T v = dL/du`` with one multigrid call.
      Since ``A`` is symmetric, this costs the same as the forward pass.
      O(1) memory, exact gradients for the linear system being solved
      (gradient accuracy is limited by how well the V-cycles approximate
      ``A^{-1}``, i.e. depends on ``n_cycles`` and smoother settings).

    * ``solve_onestep`` — **One-step differentiation** (Bolte, Pauwels &
      Vaiter, NeurIPS 2023).  Runs K V-cycles, applies ``stop_gradient``
      after K-1, then autodiffs through the last cycle only.  O(1 V-cycle)
      memory, approximate gradients with error O(rho).

    * ``solve_unrolled`` — **Unrolled differentiation** via
      ``jax.lax.fori_loop``.  Backward replays all K iterations.
      O(K) memory, exact through-iteration gradients (reproduces the
      forward computation exactly, so gradient accuracy matches the
      forward solve accuracy).

    Attributes
    ----------
    levels : tuple of MultigridLevel
        Precomputed level data, finest (index 0) to coarsest (index L-1).
    n_levels : int
        Number of multigrid levels.
    bottom_inverse : Float[Array, "n_wet n_wet"]
        Inverse of the coarsest operator on its wet points (the
        pseudo-inverse, with a relative eigenvalue cut-off, for the singular
        Neumann ``lambda = 0`` problem).
    bottom_index : Int[Array, " n_wet"]
        Flat (row-major) indices of the coarsest level's wet points.
    n_pre, n_post : int
        Pre- and post-smoothing iterations (weighted Jacobi).  Keep them
        equal for a symmetric V-cycle.
    omega : float
        Jacobi relaxation weight.
    n_cycles : int
        Number of V-cycles per solve.
    bc : {"dirichlet", "neumann"}
        Wall condition on the wet–dry faces the hierarchy was built with
        (see :func:`build_multigrid_solver`).
    null_labels, null_counts : Array or None
        Connected-component labels and wet counts of the finest mask, set
        only for ``bc="neumann"`` with ``lambda_ = 0``, where the operator
        is singular; the solve then returns the zero-mean solution on each
        component (and projects ``rhs`` onto the solvable, zero-mean part).
    """

    levels: tuple[MultigridLevel, ...]
    n_levels: int = eqx.field(static=True)
    bottom_inverse: Float[Array, "n_wet n_wet"]
    bottom_index: Int[Array, " n_wet"]
    n_pre: int = eqx.field(static=True, default=3)
    n_post: int = eqx.field(static=True, default=3)
    omega: float = eqx.field(static=True, default=0.8)
    n_cycles: int = eqx.field(static=True, default=5)
    bc: str = eqx.field(static=True, default="dirichlet")
    null_labels: Array | None = None
    null_counts: Array | None = None

    # -- Public API ----------------------------------------------------------

    def __call__(self, rhs: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
        r"""Solve ``A u = rhs`` with implicit differentiation.

        The forward pass runs K V-cycles (identical to :meth:`solve_unrolled`).

        The backward pass uses ``jax.lax.custom_linear_solve`` with
        ``symmetric=True`` to compute gradients via the implicit function
        theorem (IFT) rather than unrolling through V-cycle iterations.

        For a scalar loss ``L(u)``, the gradient w.r.t. the RHS is::

            dL/d(rhs) = A^{-T} dL/du = A^{-1} dL/du   (since A = A^T)

        This adjoint solve is just another multigrid call — so the
        backward pass costs the same as the forward pass, with O(1) extra
        memory (no iteration history stored).

        Parameters
        ----------
        rhs : Float[Array, "Ny Nx"]
            Right-hand side of the linear system.

        Returns
        -------
        Float[Array, "Ny Nx"]
            Approximate solution ``u``.
        """
        level = self.levels[0]

        # _matvec defines the linear operator A for custom_linear_solve
        def _matvec(u: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
            return _apply_operator(u, level)

        # _solve provides the forward solve; custom_linear_solve will also
        # call it for the backward (adjoint) pass since symmetric=True
        def _solve(
            _matvec_fn: Callable,
            b: Float[Array, "Ny Nx"],
        ) -> Float[Array, "Ny Nx"]:
            return self._run_vcycles(b)

        return jax.lax.custom_linear_solve(_matvec, rhs, solve=_solve, symmetric=True)

    def solve_onestep(self, rhs: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
        r"""Solve with one-step differentiation (Bolte et al., NeurIPS 2023).

        Runs K V-cycles to convergence, but autodiff only sees the **last**
        cycle.  The first K-1 cycles are wrapped in ``jax.lax.stop_gradient``
        so they contribute no backward-pass cost.

        The forward result is identical to :meth:`solve_unrolled`.  The
        gradient approximation error is O(rho) where rho is the per-cycle
        convergence rate (about 0.07 per cycle on the default settings,
        measured on 34² to 258² basins; the one-step gradient is then within
        5 % of the implicit one).

        Gradient structure::

            u_0 = 0
            u_1 = V(u_0, rhs)
            ...
            u_{K-1} = V(u_{K-2}, rhs)      <-- stop_gradient here
            u_K     = V(u_{K-1}, rhs)       <-- autodiff traces only this

        Parameters
        ----------
        rhs : Float[Array, "Ny Nx"]
            Right-hand side.

        Returns
        -------
        Float[Array, "Ny Nx"]
            Approximate solution.

        References
        ----------
        Bolte, Pauwels & Vaiter (NeurIPS 2023). "One-step differentiation
        of iterative algorithms." https://arxiv.org/abs/2305.13768
        """
        u = jnp.zeros_like(rhs)

        # Run K-1 cycles with stop_gradient: the forward pass converges
        # normally, but no gradient graph is built for these iterations
        if self.n_cycles > 1:

            def _body(_: int, u: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
                return self.v_cycle(u, rhs)

            u = jax.lax.fori_loop(0, self.n_cycles - 1, _body, u)
            u = jax.lax.stop_gradient(u)

        # Final V-cycle: JAX autodiff traces through this one only.
        # The gradient cost is O(1 V-cycle) regardless of n_cycles.
        return self.v_cycle(u, rhs)

    def solve_unrolled(self, rhs: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
        """Solve by unrolling all V-cycles through ``lax.fori_loop``.

        The backward pass differentiates through every iteration, storing
        intermediate states for replay.  This costs O(n_cycles) memory.

        Use this mode when you specifically need gradients through the
        iteration dynamics itself.  For most applications, prefer
        ``__call__`` (implicit differentiation) which gives exact gradients
        at O(1) memory cost.

        Parameters
        ----------
        rhs : Float[Array, "Ny Nx"]
            Right-hand side.

        Returns
        -------
        Float[Array, "Ny Nx"]
            Approximate solution.
        """
        return self._run_vcycles(rhs)

    def v_cycle(
        self,
        u: Float[Array, "Ny Nx"],
        rhs: Float[Array, "Ny Nx"],
        level_idx: int = 0,
    ) -> Float[Array, "Ny Nx"]:
        """Execute a single multigrid V-cycle starting at *level_idx*.

        Algorithm::

            if coarsest level:
                return A_c^{-1} rhs                # exact bottom solve

            u = jacobi(u, rhs, n_pre)             # 1. pre-smooth
            r = rhs - A(u)                        # 2. compute residual
            r_c = restrict(r)                     # 3. R = P^T / 4
            e_c = v_cycle(0, r_c, level+1)        # 4. recurse on A_c = R A P
            u = u + prolongate(e_c)               # 5. correct with coarse error
            u = jacobi(u, rhs, n_post)            # 6. post-smooth

        The recursion unrolls statically at JAX trace time because
        ``level_idx`` and ``n_levels`` are Python ints (static fields).

        Parameters
        ----------
        u : Float[Array, "Ny Nx"]
            Initial guess (typically zeros for the error equation on
            coarse grids, or the current iterate on the fine grid).
        rhs : Float[Array, "Ny Nx"]
            Right-hand side (original RHS on the fine grid, or the
            restricted residual on coarser grids).
        level_idx : int
            Current level (0 = finest, n_levels-1 = coarsest).

        Returns
        -------
        Float[Array, "Ny Nx"]
            Improved solution after one V-cycle.
        """
        # Singular Neumann problem (lambda = 0): solve in the zero-mean gauge.
        # Projecting the input and output of the finest cycle keeps the
        # V-cycle linear and symmetric (P V P with P an orthogonal projection).
        if level_idx == 0 and self.null_labels is not None:
            rhs = self._project(rhs)
            return self._project(self._v_cycle_from(u, rhs, level_idx))
        return self._v_cycle_from(u, rhs, level_idx)

    def _project(self, u: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
        """Zero-mean projection on each wet component (Neumann, lambda = 0)."""
        assert self.null_labels is not None and self.null_counts is not None
        return _project_zero_mean(
            u, self.levels[0].mask, self.null_labels, self.null_counts
        )

    def _bottom_solve(self, rhs: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
        """Exact solve on the coarsest level: ``u = A_c^{-1} rhs`` on wet points."""
        flat = rhs.ravel()  # [Ny_c*Nx_c]
        idx = self.bottom_index  # [n_wet], distinct flat indices
        u_wet = self.bottom_inverse.astype(rhs.dtype) @ flat[idx]  # [n_wet]
        # unique_indices: the scatter is then transposable (jax.linear_transpose)
        u = jnp.zeros_like(flat).at[idx].set(u_wet, unique_indices=True)
        return u.reshape(rhs.shape)

    def _v_cycle_from(
        self,
        u: Float[Array, "Ny Nx"],
        rhs: Float[Array, "Ny Nx"],
        level_idx: int,
    ) -> Float[Array, "Ny Nx"]:
        """One V-cycle at *level_idx* without the zero-mean projection."""
        level = self.levels[level_idx]

        # --- Coarsest level: exact solve (the guess u is not needed) ---
        if level_idx == self.n_levels - 1:
            return self._bottom_solve(rhs)

        # 1. Pre-smooth: damp high-frequency error on the current grid
        u = _weighted_jacobi(u, rhs, level, self.omega, self.n_pre)

        # 2. Compute residual: r = f - A(u)
        #    After smoothing, r contains mostly low-frequency error that
        #    the smoother cannot resolve at this grid resolution.
        mask = level.mask.astype(rhs.dtype)  # the hierarchy is float64
        r = (rhs - _apply_operator(u, level)) * mask

        # 3. Restrict the residual to the coarse grid: r_c = P^T r / 4
        coarse_level = self.levels[level_idx + 1]
        assert level.prolong_scale is not None
        r_coarse = _restrict(r, level.prolong_scale, coarse_level.mask)

        # 4. Recurse: solve A_c e_c = r_c (A_c = R A P) for the error,
        #    starting from zero.
        e_coarse = self.v_cycle(jnp.zeros_like(r_coarse), r_coarse, level_idx + 1)

        # 5. Prolongate the coarse correction and add it: u <- u + P e_c
        e_fine = _prolongate(e_coarse, coarse_level.mask, level.prolong_scale)
        u = (u + e_fine) * mask

        # 6. Post-smooth: clean up high-frequency error introduced by
        #    the coarse-to-fine interpolation
        return _weighted_jacobi(u, rhs, level, self.omega, self.n_post)

    # -- Internal ------------------------------------------------------------

    def _run_vcycles(self, rhs: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
        """Run *n_cycles* V-cycles from a zero initial guess.

        Uses ``jax.lax.fori_loop`` so that the iteration count does not
        increase the traced program size (unlike a Python for-loop, which
        would unroll each cycle into separate XLA operations).
        """

        def _body(_: int, u: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
            return self.v_cycle(u, rhs)

        u0 = jnp.zeros_like(rhs)
        return jax.lax.fori_loop(0, self.n_cycles, _body, u0)

__call__(rhs)

Solve A u = rhs with implicit differentiation.

The forward pass runs K V-cycles (identical to :meth:solve_unrolled).

The backward pass uses jax.lax.custom_linear_solve with symmetric=True to compute gradients via the implicit function theorem (IFT) rather than unrolling through V-cycle iterations.

For a scalar loss L(u), the gradient w.r.t. the RHS is::

dL/d(rhs) = A^{-T} dL/du = A^{-1} dL/du   (since A = A^T)

This adjoint solve is just another multigrid call — so the backward pass costs the same as the forward pass, with O(1) extra memory (no iteration history stored).

Parameters:

Name	Type	Description	Default
rhs	Float[Array, 'Ny Nx']	Right-hand side of the linear system.	required

Returns:

Type	Description
Float[Array, 'Ny Nx']	Approximate solution u.

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

def __call__(self, rhs: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
    r"""Solve ``A u = rhs`` with implicit differentiation.

    The forward pass runs K V-cycles (identical to :meth:`solve_unrolled`).

    The backward pass uses ``jax.lax.custom_linear_solve`` with
    ``symmetric=True`` to compute gradients via the implicit function
    theorem (IFT) rather than unrolling through V-cycle iterations.

    For a scalar loss ``L(u)``, the gradient w.r.t. the RHS is::

        dL/d(rhs) = A^{-T} dL/du = A^{-1} dL/du   (since A = A^T)

    This adjoint solve is just another multigrid call — so the
    backward pass costs the same as the forward pass, with O(1) extra
    memory (no iteration history stored).

    Parameters
    ----------
    rhs : Float[Array, "Ny Nx"]
        Right-hand side of the linear system.

    Returns
    -------
    Float[Array, "Ny Nx"]
        Approximate solution ``u``.
    """
    level = self.levels[0]

    # _matvec defines the linear operator A for custom_linear_solve
    def _matvec(u: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
        return _apply_operator(u, level)

    # _solve provides the forward solve; custom_linear_solve will also
    # call it for the backward (adjoint) pass since symmetric=True
    def _solve(
        _matvec_fn: Callable,
        b: Float[Array, "Ny Nx"],
    ) -> Float[Array, "Ny Nx"]:
        return self._run_vcycles(b)

    return jax.lax.custom_linear_solve(_matvec, rhs, solve=_solve, symmetric=True)

solve_onestep(rhs)

Solve with one-step differentiation (Bolte et al., NeurIPS 2023).

Runs K V-cycles to convergence, but autodiff only sees the last cycle. The first K-1 cycles are wrapped in jax.lax.stop_gradient so they contribute no backward-pass cost.

The forward result is identical to :meth:solve_unrolled. The gradient approximation error is O(rho) where rho is the per-cycle convergence rate (about 0.07 per cycle on the default settings, measured on 34² to 258² basins; the one-step gradient is then within 5 % of the implicit one).

Gradient structure::

u_0 = 0
u_1 = V(u_0, rhs)
...
u_{K-1} = V(u_{K-2}, rhs)      <-- stop_gradient here
u_K     = V(u_{K-1}, rhs)       <-- autodiff traces only this

Parameters:

Name	Type	Description	Default
rhs	Float[Array, 'Ny Nx']	Right-hand side.	required

Returns:

Type	Description
Float[Array, 'Ny Nx']	Approximate solution.

References

Bolte, Pauwels & Vaiter (NeurIPS 2023). "One-step differentiation of iterative algorithms." https://arxiv.org/abs/2305.13768

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

def solve_onestep(self, rhs: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
    r"""Solve with one-step differentiation (Bolte et al., NeurIPS 2023).

    Runs K V-cycles to convergence, but autodiff only sees the **last**
    cycle.  The first K-1 cycles are wrapped in ``jax.lax.stop_gradient``
    so they contribute no backward-pass cost.

    The forward result is identical to :meth:`solve_unrolled`.  The
    gradient approximation error is O(rho) where rho is the per-cycle
    convergence rate (about 0.07 per cycle on the default settings,
    measured on 34² to 258² basins; the one-step gradient is then within
    5 % of the implicit one).

    Gradient structure::

        u_0 = 0
        u_1 = V(u_0, rhs)
        ...
        u_{K-1} = V(u_{K-2}, rhs)      <-- stop_gradient here
        u_K     = V(u_{K-1}, rhs)       <-- autodiff traces only this

    Parameters
    ----------
    rhs : Float[Array, "Ny Nx"]
        Right-hand side.

    Returns
    -------
    Float[Array, "Ny Nx"]
        Approximate solution.

    References
    ----------
    Bolte, Pauwels & Vaiter (NeurIPS 2023). "One-step differentiation
    of iterative algorithms." https://arxiv.org/abs/2305.13768
    """
    u = jnp.zeros_like(rhs)

    # Run K-1 cycles with stop_gradient: the forward pass converges
    # normally, but no gradient graph is built for these iterations
    if self.n_cycles > 1:

        def _body(_: int, u: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
            return self.v_cycle(u, rhs)

        u = jax.lax.fori_loop(0, self.n_cycles - 1, _body, u)
        u = jax.lax.stop_gradient(u)

    # Final V-cycle: JAX autodiff traces through this one only.
    # The gradient cost is O(1 V-cycle) regardless of n_cycles.
    return self.v_cycle(u, rhs)

solve_unrolled(rhs)

Solve by unrolling all V-cycles through lax.fori_loop.

The backward pass differentiates through every iteration, storing intermediate states for replay. This costs O(n_cycles) memory.

Use this mode when you specifically need gradients through the iteration dynamics itself. For most applications, prefer __call__ (implicit differentiation) which gives exact gradients at O(1) memory cost.

Parameters:

Name	Type	Description	Default
rhs	Float[Array, 'Ny Nx']	Right-hand side.	required

Returns:

Type	Description
Float[Array, 'Ny Nx']	Approximate solution.

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

def solve_unrolled(self, rhs: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
    """Solve by unrolling all V-cycles through ``lax.fori_loop``.

    The backward pass differentiates through every iteration, storing
    intermediate states for replay.  This costs O(n_cycles) memory.

    Use this mode when you specifically need gradients through the
    iteration dynamics itself.  For most applications, prefer
    ``__call__`` (implicit differentiation) which gives exact gradients
    at O(1) memory cost.

    Parameters
    ----------
    rhs : Float[Array, "Ny Nx"]
        Right-hand side.

    Returns
    -------
    Float[Array, "Ny Nx"]
        Approximate solution.
    """
    return self._run_vcycles(rhs)

v_cycle(u, rhs, level_idx=0)

Execute a single multigrid V-cycle starting at level_idx.

Algorithm::

if coarsest level:
    return A_c^{-1} rhs                # exact bottom solve

u = jacobi(u, rhs, n_pre)             # 1. pre-smooth
r = rhs - A(u)                        # 2. compute residual
r_c = restrict(r)                     # 3. R = P^T / 4
e_c = v_cycle(0, r_c, level+1)        # 4. recurse on A_c = R A P
u = u + prolongate(e_c)               # 5. correct with coarse error
u = jacobi(u, rhs, n_post)            # 6. post-smooth

The recursion unrolls statically at JAX trace time because level_idx and n_levels are Python ints (static fields).

Parameters:

Name	Type	Description	Default
u	Float[Array, 'Ny Nx']	Initial guess (typically zeros for the error equation on coarse grids, or the current iterate on the fine grid).	required
rhs	Float[Array, 'Ny Nx']	Right-hand side (original RHS on the fine grid, or the restricted residual on coarser grids).	required
level_idx	int	Current level (0 = finest, n_levels-1 = coarsest).	0

Returns:

Type	Description
Float[Array, 'Ny Nx']	Improved solution after one V-cycle.

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

def v_cycle(
    self,
    u: Float[Array, "Ny Nx"],
    rhs: Float[Array, "Ny Nx"],
    level_idx: int = 0,
) -> Float[Array, "Ny Nx"]:
    """Execute a single multigrid V-cycle starting at *level_idx*.

    Algorithm::

        if coarsest level:
            return A_c^{-1} rhs                # exact bottom solve

        u = jacobi(u, rhs, n_pre)             # 1. pre-smooth
        r = rhs - A(u)                        # 2. compute residual
        r_c = restrict(r)                     # 3. R = P^T / 4
        e_c = v_cycle(0, r_c, level+1)        # 4. recurse on A_c = R A P
        u = u + prolongate(e_c)               # 5. correct with coarse error
        u = jacobi(u, rhs, n_post)            # 6. post-smooth

    The recursion unrolls statically at JAX trace time because
    ``level_idx`` and ``n_levels`` are Python ints (static fields).

    Parameters
    ----------
    u : Float[Array, "Ny Nx"]
        Initial guess (typically zeros for the error equation on
        coarse grids, or the current iterate on the fine grid).
    rhs : Float[Array, "Ny Nx"]
        Right-hand side (original RHS on the fine grid, or the
        restricted residual on coarser grids).
    level_idx : int
        Current level (0 = finest, n_levels-1 = coarsest).

    Returns
    -------
    Float[Array, "Ny Nx"]
        Improved solution after one V-cycle.
    """
    # Singular Neumann problem (lambda = 0): solve in the zero-mean gauge.
    # Projecting the input and output of the finest cycle keeps the
    # V-cycle linear and symmetric (P V P with P an orthogonal projection).
    if level_idx == 0 and self.null_labels is not None:
        rhs = self._project(rhs)
        return self._project(self._v_cycle_from(u, rhs, level_idx))
    return self._v_cycle_from(u, rhs, level_idx)

Factory

finitevolx.build_multigrid_solver(mask, dx, dy, lambda_=0.0, coeff=None, n_levels=None, n_pre=3, n_post=3, omega=0.8, n_cycles=5, bc='dirichlet')

Build a multigrid solver with precomputed level hierarchies.

This is an offline function (runs once on CPU with NumPy / SciPy) that constructs the entire multigrid hierarchy:

1. Finest operator: the 5-point stencil of div(c grad u) - lambda u, face coefficients the mean of the two cell values, wet–dry faces set by bc (below).
2. Coarsening: each dimension of n points becomes (n + 1) // 2, coarse point J on fine point 2J — odd and even sizes alike, so a standard N + 2 grid is as deep as N. A coarse point is wet if its coincident fine point is wet; a point whose interpolation support (fine points 2J-1 .. 2J+1 in each direction) only touches wet fine points is added when its prolongation column is linearly independent of the others, so narrow channels keep a coarse unknown and P keeps full column rank (:func:_coarse_mask).
3. Transfers: bilinear prolongation P between the nested points (zero at dry coarse points for Dirichlet; renormalised over the wet ones for Neumann) and full-weighting restriction R = Pᵀ / 4.
4. Galerkin coarse operators A_c = R A P (9-point stencils), which follow the fine geometry and coefficient exactly.
5. Exact bottom solve: the inverse of the coarsest operator on its wet points (at most a few dozen unknowns for square grids).

Grid hierarchy example (34x34, a 32-cell interior plus ghost ring)::

Level 0:  34 x 34   (dx,    dy)     <- finest (solve here)
Level 1:  17 x 17   (2*dx,  2*dy)
Level 2:   9 x  9   (4*dx,  4*dy)
Level 3:   5 x  5   (8*dx,  8*dy)   <- coarsest (exact solve)

The returned MultigridSolver is an immutable equinox.Module with frozen JAX arrays. All subsequent calls (forward solves, gradients, JIT compilation) use the precomputed hierarchy.

Boundary condition ~~~~~~~~~~~~~~~~~~ bc sets what happens on every face between a wet cell and a dry one (land, or outside the array):

• "dirichlet" (default): u = 0 on the dry neighbour, the face keeps its coefficient. With coeff=None the finest-level operator is exactly :func:~finitevolx.masked_laplacian with dry neighbours zero — the streamfunction problem, psi = 0 on the wall / coast X nodes of Mask2D.xy_corner_strict. Non-singular for every lambda_ >= 0.
• "neumann": zero flux through wet–dry faces — the pressure / free-surface problem at T points (pass mask.h). At lambda_ = 0 the operator is singular (a constant per connected wet component); the solver then projects rhs and the solution onto zero mean on each component.

Parameters:

Name	Type	Description	Default
mask	array, shape (Ny, Nx), or Mask2D	Domain mask (1 = fluid, 0 = land). None is not accepted; pass np.ones((Ny, Nx)) for a rectangular domain. When a :class:Mask2D is passed, the ψ mask xy_corner_strict is extracted (pair it with the default bc="dirichlet"); for a T-point problem pass mask.h.	required
dx	float	Fine-grid spacings (metres or non-dimensional).	required
dy	float	Fine-grid spacings (metres or non-dimensional).	required
lambda_	float	Helmholtz parameter (>= 0). Use 0.0 for pure Poisson (Laplacian only). For QG PV inversion (bc="dirichlet"), lambda_ = 1 / Rd**2.	0.0
coeff	array, shape (Ny, Nx), or None	Spatially varying coefficient c(x, y) at cell centres. None -> constant coefficient = 1 everywhere (reduces to the standard constant-coefficient Helmholtz operator).	None
n_levels	int or None	Number of multigrid levels. None -> auto-detect by coarsening while both dimensions stay >= 4 (see :func:_compute_n_levels). Any shape works with any n_levels >= 1; 1 is a direct solve.	None
n_pre	int	Number of pre- and post-smoothing Jacobi iterations per V-cycle. Keep them equal: the V-cycle is then a symmetric operator (needed when it preconditions CG). Default: 3 each.	3
n_post	int	Number of pre- and post-smoothing Jacobi iterations per V-cycle. Keep them equal: the V-cycle is then a symmetric operator (needed when it preconditions CG). Default: 3 each.	3
omega	float	Jacobi relaxation weight (0 < omega < 1). Default 0.8, the weight that damps the high-frequency half of the 5-point spectrum best (by at least 3/5 per sweep).	0.8
n_cycles	int	Number of V-cycles applied per solve. Each cycle reduces the residual by 0.06-0.07 on the default settings (measured on 34² to 258² circular basins, independent of the grid size), so 5 cycles reduce it by about 3e-7. An island is slower: 0.14 / 0.17 per cycle for bc="dirichlet" in a circular basin with an island at 130² / 258² (5 cycles: 7e-5 / 4e-4). One-cell channels and islands are slower still for bc="neumann" (0.1-0.5 per cycle, measured on 34² masks); use the V-cycle as a CG preconditioner there. Default: 5.	5
bc	('dirichlet', 'neumann')	Condition on wet–dry faces (see above). Default: "dirichlet".	"dirichlet"

Returns:

Type	Description
MultigridSolver	Ready-to-use solver (JIT-compilable equinox.Module).

Raises:

Type	Description
ValueError	If n_levels < 1 or bc is not "dirichlet" or "neumann".

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

def build_multigrid_solver(
    mask: np.ndarray | Float[Array, "Ny Nx"] | Mask2D,
    dx: float,
    dy: float,
    lambda_: float = 0.0,
    coeff: np.ndarray | Float[Array, "Ny Nx"] | None = None,
    n_levels: int | None = None,
    n_pre: int = 3,
    n_post: int = 3,
    omega: float = 0.8,
    n_cycles: int = 5,
    bc: Literal["dirichlet", "neumann"] = "dirichlet",
) -> MultigridSolver:
    r"""Build a multigrid solver with precomputed level hierarchies.

    This is an **offline** function (runs once on CPU with NumPy / SciPy)
    that constructs the entire multigrid hierarchy:

    1. **Finest operator**: the 5-point stencil of ``div(c grad u) -
       lambda u``, face coefficients the mean of the two cell values, wet–dry
       faces set by ``bc`` (below).
    2. **Coarsening**: each dimension of ``n`` points becomes
       ``(n + 1) // 2``, coarse point ``J`` on fine point ``2J`` — odd and
       even sizes alike, so a standard ``N + 2`` grid is as deep as ``N``.
       A coarse point is wet if its coincident fine point is wet; a point
       whose interpolation support (fine points ``2J-1 .. 2J+1`` in each
       direction) only touches wet fine points is added when its
       prolongation column is linearly independent of the others, so
       narrow channels keep a coarse unknown and ``P`` keeps full column
       rank (:func:`_coarse_mask`).
    3. **Transfers**: bilinear prolongation ``P`` between the nested points
       (zero at dry coarse points for Dirichlet; renormalised over the wet
       ones for Neumann) and full-weighting restriction ``R = Pᵀ / 4``.
    4. **Galerkin coarse operators** ``A_c = R A P`` (9-point stencils),
       which follow the fine geometry and coefficient exactly.
    5. **Exact bottom solve**: the inverse of the coarsest operator on its
       wet points (at most a few dozen unknowns for square grids).

    Grid hierarchy example (34x34, a 32-cell interior plus ghost ring)::

        Level 0:  34 x 34   (dx,    dy)     <- finest (solve here)
        Level 1:  17 x 17   (2*dx,  2*dy)
        Level 2:   9 x  9   (4*dx,  4*dy)
        Level 3:   5 x  5   (8*dx,  8*dy)   <- coarsest (exact solve)

    The returned ``MultigridSolver`` is an immutable ``equinox.Module``
    with frozen JAX arrays.  All subsequent calls (forward solves,
    gradients, JIT compilation) use the precomputed hierarchy.

    Boundary condition
    ~~~~~~~~~~~~~~~~~~
    ``bc`` sets what happens on every face between a wet cell and a dry
    one (land, or outside the array):

    * ``"dirichlet"`` (default): ``u = 0`` on the dry neighbour, the face
      keeps its coefficient.  With ``coeff=None`` the finest-level operator
      is exactly :func:`~finitevolx.masked_laplacian` with dry neighbours
      zero — the streamfunction problem, ``psi = 0`` on the wall / coast
      X nodes of ``Mask2D.xy_corner_strict``.  Non-singular for every
      ``lambda_ >= 0``.
    * ``"neumann"``: zero flux through wet–dry faces — the pressure /
      free-surface problem at T points (pass ``mask.h``).  At
      ``lambda_ = 0`` the operator is singular (a constant per connected
      wet component); the solver then projects ``rhs`` and the solution
      onto zero mean on each component.

    Parameters
    ----------
    mask : array, shape (Ny, Nx), or Mask2D
        Domain mask (1 = fluid, 0 = land).  ``None`` is *not* accepted;
        pass ``np.ones((Ny, Nx))`` for a rectangular domain.
        When a :class:`Mask2D` is passed, the ψ mask
        ``xy_corner_strict`` is extracted (pair it with the default
        ``bc="dirichlet"``); for a T-point problem pass ``mask.h``.
    dx, dy : float
        Fine-grid spacings (metres or non-dimensional).
    lambda_ : float
        Helmholtz parameter (>= 0).  Use 0.0 for pure Poisson (Laplacian
        only).  For QG PV inversion (``bc="dirichlet"``),
        ``lambda_ = 1 / Rd**2``.
    coeff : array, shape (Ny, Nx), or None
        Spatially varying coefficient ``c(x, y)`` at cell centres.
        ``None`` -> constant coefficient = 1 everywhere (reduces to the
        standard constant-coefficient Helmholtz operator).
    n_levels : int or None
        Number of multigrid levels.  ``None`` -> auto-detect by coarsening
        while both dimensions stay >= 4 (see :func:`_compute_n_levels`).
        Any shape works with any ``n_levels >= 1``; ``1`` is a direct
        solve.
    n_pre, n_post : int
        Number of pre- and post-smoothing Jacobi iterations per V-cycle.
        Keep them equal: the V-cycle is then a symmetric operator (needed
        when it preconditions CG).  Default: 3 each.
    omega : float
        Jacobi relaxation weight (0 < omega < 1).  Default 0.8, the
        weight that damps the high-frequency half of the 5-point spectrum
        best (by at least 3/5 per sweep).
    n_cycles : int
        Number of V-cycles applied per solve.  Each cycle reduces the
        residual by 0.06-0.07 on the default settings (measured on 34² to
        258² circular basins, independent of the grid size), so 5 cycles
        reduce it by about 3e-7.  An island is slower: 0.14 / 0.17 per
        cycle for ``bc="dirichlet"`` in a circular basin with an island at
        130² / 258² (5 cycles: 7e-5 / 4e-4).  One-cell channels and islands
        are slower still for ``bc="neumann"`` (0.1-0.5 per cycle, measured on
        34² masks); use the V-cycle as a CG preconditioner there.
        Default: 5.
    bc : {"dirichlet", "neumann"}
        Condition on wet–dry faces (see above).  Default: ``"dirichlet"``.

    Returns
    -------
    MultigridSolver
        Ready-to-use solver (JIT-compilable ``equinox.Module``).

    Raises
    ------
    ValueError
        If ``n_levels < 1`` or ``bc`` is not ``"dirichlet"`` or
        ``"neumann"``.
    """
    if bc not in ("dirichlet", "neumann"):
        raise ValueError(f"bc must be 'dirichlet' or 'neumann'; got {bc!r}")

    # --- Extract mask as a NumPy float64 array ---
    if isinstance(mask, Mask2D):
        mask_np = np.asarray(mask.xy_corner_strict, dtype=np.float64)
    else:
        mask_np = np.asarray(mask, dtype=np.float64)

    ny, nx = mask_np.shape

    # --- Auto-detect or validate number of levels ---
    if n_levels is None:
        n_levels = _compute_n_levels(ny, nx)
    if n_levels < 1:
        raise ValueError(f"n_levels must be >= 1; got {n_levels}")

    # --- Default coefficient: c(x,y) = 1 everywhere ---
    if coeff is None:
        coeff_np = np.ones_like(mask_np)
    else:
        coeff_np = np.asarray(coeff, dtype=np.float64)
    coeff_np = coeff_np * mask_np  # zero outside domain

    # --- Finest operator, then Galerkin products down the hierarchy ---
    cur_mask = mask_np
    cur_dx, cur_dy = float(dx), float(dy)
    stencil = _fine_stencil(coeff_np, mask_np, cur_dx, cur_dy, float(lambda_), bc)
    offsets = _OFFSETS_5
    a = _stencil_to_sparse(stencil, offsets)  # [Ny*Nx, Ny*Nx]

    levels: list[MultigridLevel] = []
    for lev in range(n_levels):
        scale = None
        next_a = next_mask = None
        if lev < n_levels - 1:
            # P = diag(s) P0 diag(m_c);  A_c = R A P with R = P^T / 4
            next_mask, scale, p = _transfer_operators(cur_mask, bc)
            next_a = (0.25 * (p.T @ a @ p)).tocsr()

        levels.append(
            MultigridLevel(
                mask=jnp.asarray(cur_mask),
                stencil=jnp.asarray(stencil),
                inv_diagonal=jnp.asarray(_inverse_diagonal(stencil[0], cur_mask)),
                prolong_scale=None if scale is None else jnp.asarray(scale),
                dx=cur_dx,
                dy=cur_dy,
                lambda_=float(lambda_),
                offsets=offsets,
            )
        )

        if next_a is not None and next_mask is not None:
            cur_mask, a = next_mask, next_a
            offsets = _OFFSETS_9
            stencil = _sparse_to_stencil(a, cur_mask.shape)
            cur_dx *= 2.0  # grid spacing doubles at each coarser level
            cur_dy *= 2.0

    # --- Exact bottom solve: dense inverse on the coarsest wet points ---
    wet = np.flatnonzero(cur_mask.ravel() > 0)
    if wet.size > _MAX_BOTTOM:
        raise ValueError(
            f"the coarsest level ({cur_mask.shape[0]} x {cur_mask.shape[1]}) has "
            f"{wet.size} wet points, more than the {_MAX_BOTTOM} the dense bottom "
            "solve allows; use more levels (n_levels=None picks them)"
        )
    a_wet = a.toarray()[np.ix_(wet, wet)]
    bottom_inverse = _symmetric_inverse(a_wet)

    # --- Neumann at lambda = 0: constants per wet component are the null space
    null_labels = null_counts = None
    if bc == "neumann" and lambda_ == 0.0:
        null_labels, null_counts = _wet_components(mask_np)

    return MultigridSolver(
        levels=tuple(levels),
        n_levels=n_levels,
        bottom_inverse=jnp.asarray(bottom_inverse),
        bottom_index=jnp.asarray(wet),
        n_pre=n_pre,
        n_post=n_post,
        omega=omega,
        n_cycles=n_cycles,
        bc=bc,
        null_labels=null_labels,
        null_counts=null_counts,
    )

Preconditioner

finitevolx.make_multigrid_preconditioner(mg_solver)

Return a preconditioner closure that applies a single multigrid V-cycle.

The returned callable approximates A^{-1} r by running one V-cycle from a zero initial guess, which is sufficient as a preconditioner (it doesn't need to converge — it just needs to be a good approximation).

This is compatible with :func:~finitevolx._src.solvers.iterative.solve_cg: pass the returned closure as the preconditioner argument. The V-cycle is a symmetric operator (restriction is the scaled transpose of prolongation, the coarse operators are Galerkin products and the bottom solve is exact), as CG requires, provided the solver smooths as much before as after (n_pre == n_post). CG then converges in a few iterations independent of the grid size: 6 at 128² and 256² on a circular basin (rtol = 1e-8), against 33 and 47 with the FFT preconditioner.

Parameters:

Name	Type	Description	Default
mg_solver	MultigridSolver	A pre-built multigrid solver (from :func:~finitevolx._src.solvers.multigrid.build_multigrid_solver).	required

Returns:

Type	Description
callable	preconditioner(r) -> approx_solution, where r has shape (Ny, Nx) and the output has the same shape. For gaussx's solvers, use :func:make_gaussx_multigrid_preconditioner.

Raises:

Type	Description
ValueError	If mg_solver.n_pre != mg_solver.n_post (the V-cycle would not be symmetric).

Examples:

>>> mg = build_multigrid_solver(mask, dx, dy, lambda_=10.0)
>>> precond = make_multigrid_preconditioner(mg)
>>> u, info = solve_cg(A, rhs, preconditioner=precond)

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

def make_multigrid_preconditioner(
    mg_solver: MultigridSolver,
) -> Callable[[Float[Array, "Ny Nx"]], Float[Array, "Ny Nx"]]:
    """Return a preconditioner closure that applies a single multigrid V-cycle.

    The returned callable approximates ``A^{-1} r`` by running one V-cycle
    from a zero initial guess, which is sufficient as a preconditioner
    (it doesn't need to converge — it just needs to be a good approximation).

    This is compatible with :func:`~finitevolx._src.solvers.iterative.solve_cg`:
    pass the returned closure as the ``preconditioner`` argument.  The
    V-cycle is a symmetric operator (restriction is the scaled transpose of
    prolongation, the coarse operators are Galerkin products and the bottom
    solve is exact), as CG requires, provided the solver smooths as much
    before as after (``n_pre == n_post``).  CG then converges in a few
    iterations independent of the grid size: 6 at 128² and 256² on a
    circular basin (``rtol = 1e-8``), against 33 and 47 with the FFT
    preconditioner.

    Parameters
    ----------
    mg_solver : MultigridSolver
        A pre-built multigrid solver
        (from :func:`~finitevolx._src.solvers.multigrid.build_multigrid_solver`).

    Returns
    -------
    callable
        ``preconditioner(r) -> approx_solution``, where ``r`` has shape
        ``(Ny, Nx)`` and the output has the same shape.  For gaussx's
        solvers, use :func:`make_gaussx_multigrid_preconditioner`.

    Raises
    ------
    ValueError
        If ``mg_solver.n_pre != mg_solver.n_post`` (the V-cycle would not be
        symmetric).

    Examples
    --------
    >>> mg = build_multigrid_solver(mask, dx, dy, lambda_=10.0)
    >>> precond = make_multigrid_preconditioner(mg)
    >>> u, info = solve_cg(A, rhs, preconditioner=precond)
    """

    if mg_solver.n_pre != mg_solver.n_post:
        raise ValueError(
            "a CG preconditioner must be symmetric: build the multigrid solver "
            f"with n_pre == n_post (got {mg_solver.n_pre} and {mg_solver.n_post})"
        )

    def _preconditioner(r: Float[Array, "Ny Nx"]) -> Float[Array, "Ny Nx"]:
        return mg_solver.v_cycle(jnp.zeros_like(r), r)

    return _preconditioner

finitevolx.make_gaussx_multigrid_preconditioner(mg_solver)

The multigrid V-cycle as a symmetric gaussx preconditioner.

Wraps a pre-built :class:~finitevolx.MultigridSolver in gaussx.OperatorPreconditioner, so the geometric V-cycle plugs into gaussx's strategies (gaussx.CGSolver(preconditioner=...), gaussx.PreconditionedCGSolver, gaussx.linear_solve).

gaussx works with positive semidefinite systems, so it solves (−A) x = −b for finitevolX's negative (semi)definite Helmholtz operator A = ∇·(c∇) − λ. Let V(v) = −Vcycle(0, v) ≈ (−A)⁻¹ v. With restriction the scaled transpose of prolongation and equal pre- and post-smoothing (the defaults), V is already symmetric; this applies the two-step iteration (V, then its adjoint), with K = −A,

M⁻¹ v = V v + Vᵀ v − Vᵀ K V v,

with Vᵀ the exact transpose of the (linear) V-cycle from jax.linear_transpose: two V-cycles and one operator application per use. M⁻¹ is symmetric by construction, and since I − M⁻¹K = (I − VᵀK)(I − VK) it is positive definite exactly when one V-cycle reduces the error in the energy norm, ‖I − VK‖_K < 1 — i.e. when the multigrid itself converges. The default build_multigrid_solver settings do (about 0.07 residual reduction per cycle); a divergent configuration (one smoothing sweep with omega near 1, high coefficient contrast) makes no symmetric form of the V-cycle positive definite, so check the solver converges on its own first. It costs two V-cycles per application and takes fewer CG steps than the raw V-cycle, which :func:make_multigrid_preconditioner uses.

The V-cycle runs in the hierarchy's dtype (float64); each application returns the dtype of the vector it was given. The structure comes from the system operator, so pass it when calling the preconditioner directly: pc(v, operator).

For finitevolX's own :func:~finitevolx.solve_cg, which takes a callable approximating A⁻¹ on 2-D fields, use :func:make_multigrid_preconditioner on the same solver.

Parameters:

Name	Type	Description	Default
mg_solver	MultigridSolver	A pre-built solver from :func:~finitevolx.build_multigrid_solver.	required

Returns:

Type	Description
OperatorPreconditioner	Applies M⁻¹ on flat vectors of length Ny·Nx (row-major).

Examples:

>>> mg = build_multigrid_solver(mask, dx, dy, lambda_=10.0)
>>> neg_A = gaussx.as_linear_operator(
...     lambda v: -A(v.reshape(Ny, Nx)).ravel(),
...     shape=(Ny * Nx, Ny * Nx),
...     positive_semidefinite=True,
... )
>>> pc = make_gaussx_multigrid_preconditioner(mg)
>>> x = gaussx.CGSolver(preconditioner=pc).solve(neg_A, -rhs.ravel())

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

def make_gaussx_multigrid_preconditioner(
    mg_solver: MultigridSolver,
) -> gaussx.OperatorPreconditioner:
    r"""The multigrid V-cycle as a symmetric gaussx preconditioner.

    Wraps a pre-built :class:`~finitevolx.MultigridSolver` in
    ``gaussx.OperatorPreconditioner``, so the geometric V-cycle plugs into
    gaussx's strategies (``gaussx.CGSolver(preconditioner=...)``,
    ``gaussx.PreconditionedCGSolver``, ``gaussx.linear_solve``).

    gaussx works with positive semidefinite systems, so it solves
    ``(−A) x = −b`` for finitevolX's negative (semi)definite Helmholtz
    operator ``A = ∇·(c∇) − λ``. Let ``V(v) = −Vcycle(0, v) ≈ (−A)⁻¹ v``.
    With restriction the scaled transpose of prolongation and equal pre- and
    post-smoothing (the defaults), ``V`` is already symmetric; this applies
    the two-step iteration (V, then its adjoint), with ``K = −A``,

    ``M⁻¹ v = V v + Vᵀ v − Vᵀ K V v``,

    with ``Vᵀ`` the exact transpose of the (linear) V-cycle from
    ``jax.linear_transpose``: two V-cycles and one operator application per
    use. ``M⁻¹`` is symmetric by construction, and since
    ``I − M⁻¹K = (I − VᵀK)(I − VK)`` it is positive definite exactly when
    one V-cycle reduces the error in the energy norm, ``‖I − VK‖_K < 1`` —
    i.e. when the multigrid itself converges. The default
    ``build_multigrid_solver`` settings do (about 0.07 residual reduction per
    cycle); a divergent configuration (one smoothing sweep with ``omega``
    near 1, high coefficient contrast) makes no symmetric form of the
    V-cycle positive definite, so check the solver converges on its own
    first. It costs two V-cycles per application and takes fewer CG steps
    than the raw V-cycle, which :func:`make_multigrid_preconditioner` uses.

    The V-cycle runs in the hierarchy's dtype (float64); each application
    returns the dtype of the vector it was given. The structure comes from
    the system operator, so pass it when calling the preconditioner directly:
    ``pc(v, operator)``.

    For finitevolX's own :func:`~finitevolx.solve_cg`, which takes a callable
    approximating ``A⁻¹`` on 2-D fields, use
    :func:`make_multigrid_preconditioner` on the same solver.

    Parameters
    ----------
    mg_solver : MultigridSolver
        A pre-built solver from :func:`~finitevolx.build_multigrid_solver`.

    Returns
    -------
    gaussx.OperatorPreconditioner
        Applies ``M⁻¹`` on flat vectors of length ``Ny·Nx`` (row-major).

    Examples
    --------
    >>> mg = build_multigrid_solver(mask, dx, dy, lambda_=10.0)
    >>> neg_A = gaussx.as_linear_operator(
    ...     lambda v: -A(v.reshape(Ny, Nx)).ravel(),
    ...     shape=(Ny * Nx, Ny * Nx),
    ...     positive_semidefinite=True,
    ... )
    >>> pc = make_gaussx_multigrid_preconditioner(mg)
    >>> x = gaussx.CGSolver(preconditioner=pc).solve(neg_A, -rhs.ravel())
    """
    from finitevolx._src.solvers.multigrid import _apply_operator

    level0 = mg_solver.levels[0]
    mask = level0.mask  # [Ny, Nx], the hierarchy's dtype
    shape = mask.shape

    def _v_cycle(v: Float[Array, " n"]) -> Float[Array, " n"]:
        # V v = -Vcycle(0, v) ~ K^-1 v with K = -A, [Ny*Nx]; linear in v
        r = v.reshape(shape)  # [Ny, Nx]
        return -mg_solver.v_cycle(jnp.zeros_like(r), r).ravel()

    def _k(v: Float[Array, " n"]) -> Float[Array, " n"]:
        # K v = -A v = -(div(c grad) - lambda) v on the finest level, [Ny*Nx]
        return -_apply_operator(v.reshape(shape), level0).ravel()

    _v_cycle_t = jax.linear_transpose(
        _v_cycle, jax.ShapeDtypeStruct((mask.size,), mask.dtype)
    )

    def _approx_inverse(v: Float[Array, " n"]) -> Float[Array, " n"]:
        w = v.astype(mask.dtype)  # [Ny*Nx], hierarchy dtype
        vw = _v_cycle(w)  # V w
        # M^-1 w = V w + V^T w - V^T K V w = V w + V^T (w - K V w)
        (vt,) = _v_cycle_t(w - _k(vw))
        return (vw + vt).astype(v.dtype)

    return gaussx.OperatorPreconditioner(approx_inverse=_approx_inverse)