← Back to Mesh-based solvers Module · scimba_jax

Multigrid

Without a preconditioner, a Krylov solve gets slower as the mesh is refined, for two reasons: there are more unknowns, and each iteration reduces the error less. A multigrid preconditioner addresses the second issue. It solves a coarser version of the same problem to correct the smooth part of the error at a low cost, and uses a simple local sweep (a smoother) to reduce the remaining error.

Hierarchy, smoother and Krylov solve

\(u \leftarrow u + \omega\, D^{-1}(f - Au)\)

The update above is a damped Jacobi sweep, with omega=2/3. This smoother quickly reduces the oscillatory part of the error, and leaves the smooth part, which is cheaper to correct on the coarser level. build_hierarchy_structured builds a sequence of the same discretization at coarser and coarser resolutions, together with the operators that transfer a solution or a residual between two consecutive levels. MG combines this hierarchy with a smoother (here DampedJacobiSmoother) into a single object, which is given to LinearKrylov like any other preconditioner. 

from scimba_jax.linear_approximation.transfer.hierarchy import (
    build_hierarchy_structured,
)
from scimba_jax.linear_approximation.solvers.multigrid import MG
from scimba_jax.linear_approximation.solvers.smoothers import DampedJacobiSmoother
from scimba_jax.linear_approximation.solvers.newton import LinearKrylov
def make_scheme(n_cells):
    # FEM scheme from the FEM page, at the resolution requested by the
    # hierarchy (one call per level, coarsest first).
    ...
    return EllipticFEscheme(model, variables)
# 4 levels, 64x64 down to 8x8, same discretization at each resolution.
hierarchy = build_hierarchy_structured(make_scheme, (64, 64), n_levels=4)
mg = MG(hierarchy, DampedJacobiSmoother(omega=2.0 / 3.0), nu_pre=2, nu_post=2)
# The multigrid is used as a preconditioner of the usual Krylov solver.
scheme = hierarchy.finest.scheme
solver = LinearKrylov(tol=1e-10, cg_solver="lax", preconditioner=mg)
scheme, report = EllipticFEscheme.solve(scheme, solver=solver, return_report=True)
# 10 CG iterations to 1e-10, independently of the mesh size (an
# unpreconditioned CG needs O(1/h) iterations).

Shifted Laplacian preconditioner for the Helmholtz equation

\(M = -\Delta - (k^2 + i\,\varepsilon),\quad \varepsilon = 2k\)

For an indefinite operator, multigrid cannot efficiently correct the smooth part of the error. The cycle is therefore applied to a shifted operator, damped by i ε, and is used to precondition BiCGSTAB on the original problem. shifted_laplacian_solver builds this preconditioner from a single factory of the problem, with an additive shift of order k, two levels and a damped Jacobi smoother. The classical shift βk² stalls at k = 100. At k = 100, three levels amplify the error instead of reducing it, since the coarsest grid no longer resolves the wave. The coarse level is then a quarter of the fine one, and its dense LU factorization dominates the cost: 9 s to factorise, then 100 ms per cycle. This cost is amortised over a batch of right-hand sides: four sources under jax.vmap cost 1.15 times as much as one. 

from scimba_jax.linear_approximation.solvers import shifted_laplacian_solver
from scimba_jax.physical_models.classical_weakform.helmholtz_weak_form import (
    HelmholtzWeakForm,
)
from scimba_jax.physical_models.weak_boundary_conditions import Sommerfeld
def scheme_factory(n_cells, shift=0.0):
    # A single factory: the problem for shift=0, the preconditioner operator
    # -Delta - (k^2 + i shift) for shift=eps (kappa and shift may be x -> float).
    model = AbstractPhysicalWeakModel(dim=2)
    model.add_weak_form("main", HelmholtzWeakForm(2, K**2, f, kappa_im=shift))
    for side in ("west", "east", "south", "north"):
        model.add_boundary_condition(side, Sommerfeld(K))  # d_n u - i k u = 0
    return EllipticFEscheme(model, make_variables(n_cells))
# Two levels, Jacobi omega=0.5, V(1,1), BiCGSTAB and shift 2k, as in the paper.
solver = shifted_laplacian_solver(scheme_factory, (160, 160), shift=2 * K)
scheme, report = EllipticFEscheme.solve(
    scheme_factory((160, 160)), solver=solver, return_report=True
)
# k = 100, 51 842 dofs: 35 iterations to 1e-8 (plain BiCGSTAB stalls).
# For a batch of sources, use solve_pure under jax.vmap with a shared multigrid.

The multigrid presented here is linear: it preconditions a fixed operator. When the operator changes at every iteration, two other methods build on it. The Nonlinear solvers page presents Newton preconditioned by a multigrid that follows the Jacobian, and FAS, a nonlinear multigrid that does not require the Jacobian.