A hierarchy, a smoother, and a Krylov solve
\(u \leftarrow u + \omega\, D^{-1}(f - Au)\)
A damped Jacobi sweep — the smoother above, with omega=2/3
— is a fixed point of that update; it kills the OSCILLATORY part of
the error fast and leaves the smooth part, which the coarser level
is cheap at fixing instead.
build_hierarchy_structured builds a stack of the SAME
discretization at coarser and coarser resolutions and works out how
to move a solution — or a residual — between two consecutive
levels. MG combines that hierarchy with a smoother
(here DampedJacobiSmoother) into ONE object that plugs
into LinearKrylov exactly where any other
preconditioner would.
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):
# The FEM scheme from the FEM page, at whatever resolution the
# hierarchy asks for -- 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 a preconditioner: same Krylov solver as anywhere else.
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 -- flat in the mesh size, where an
# unpreconditioned CG needs O(1/h) of them.
Helmholtz: the shifted-Laplacian preconditioner
\(M = -\Delta - (k^2 + i\,\varepsilon),\quad \varepsilon = 2k\)
An indefinite operator has no smooth error for a multigrid to fix,
so the cycle runs on a SHIFTED one instead — damped by
i ε — and preconditions BiCGSTAB on the
real problem. shifted_laplacian_solver builds it from
one factory of the problem: an additive shift of order
k (the classical βk² stalls at
k = 100), two levels, damped Jacobi. Measured at
k = 100: three levels amplify instead of contracting
(the coarsest grid no longer resolves the wave), so the coarse
level is a quarter of the fine one and its dense LU is the whole
cost — 9 s to factorise, 100 ms per cycle. A batch of
right-hand sides amortises it: four sources under
jax.vmap cost 1.15× 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):
# ONE factory: the problem at shift=0, the preconditioner's operator
# -Delta - (k^2 + i shift) at 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: the paper's recipe, shift 2k.
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, where plain BiCGSTAB stalls.
# A batch of sources: solve_pure under jax.vmap, the multigrid shared.
A multigrid built here is LINEAR: it preconditions one fixed operator. Two more pieces build on it once the operator itself changes at every iteration — see Nonlinear solvers for Newton preconditioned by a multigrid that follows the Jacobian, and for FAS, the nonlinear multigrid that needs no Jacobian at all.