Linear approximation: FEM, DG and multigrid

scimba_jax.linear_approximation is the JAX discretisation layer. It combines meshes, finite-element or discontinuous-Galerkin spaces, weak forms, nonlinear solvers and preconditioners. The important separation is:

  1. a mesh describes geometry and quadrature;

  2. a space describes degrees of freedom and basis functions;

  3. a Galerkin scheme assembles a residual and its linearisation;

  4. a solver consumes that residual, optionally with a preconditioner.

The code keeps geometric/topological preparation outside JAX traces whenever possible. This matters particularly for curved, block-structured and multigrid problems: a repeated setup or a new traced shape costs more than a few arithmetic operations in a Krylov iteration.

Meshes and patch domains

Mesh is a structured mesh on a reference hypercube, mapped into physical coordinates. Use it for Cartesian grids and smooth mapped geometries. UnstructuredMesh stores tensor-order quadrilateral/hexahedral cells and their geometric nodes; it is the route for Gmsh meshes and curved boundaries.

BlockStructuredMesh combines one structured Mesh per macro patch. It contains the face pairing between patches, including orientation and supported non-conforming interfaces. A patch is geometrically independent, but its interface is part of the discrete problem:

  • BlockStructuredDGscheme applies the same numerical flux as on an interior DG face; splitting a conforming DG mesh into patches should therefore retain the same discrete solution.

  • BlockStructuredFEscheme glues continuous spaces with a Nitsche/SIPG interface term. It needs an explicit SIPGFlux; there is intentionally no penalty default.

For a mapped mesh, use local cell sizes in penalty terms. A global h is not reliable when cells have different physical sizes after mapping.

CG-FEM and DG

EllipticFEscheme is the continuous Galerkin scheme: degrees of freedom are shared by neighbouring cells and Dirichlet data are imposed strongly through a lifted residual. EllipticDGscheme is the discontinuous counterpart: each cell owns its DOFs and a numerical flux imposes coupling and boundary data.

DG fluxes are selected explicitly. SIPGFlux is the symmetric interior penalty choice for elliptic diffusion; NIPGFlux and BabuSkaZlamalFlux provide nonsymmetric variants; UpwindFlux and RiemannFlux address transport/hyperbolic terms. On an interface, both sides must use compatible flux data and the physical cell length scale.

Transfer and hierarchy

Multigrid is built around a Hierarchy, ordered coarsest to finest. A Level contains both the scheme and its flattened residual/JVP. Each consecutive pair has an AbstractTransfer:

  • StructuredTransfer for nodal tensor-product structured grids;

  • NestedTransfer for nodal meshes with an exact child table;

  • CellwiseTransfer for DG, whose DOFs belong to cells;

  • PatchTransfer for block containers, obtained by applying the rule patch by patch.

restrict_dual is always the transpose of prolongation. It is the restriction for residuals and preserves the Galerkin relation A_H = P^T A_h P. restrict_primal is separate because nonlinear FAS must restrict a solution, not a residual.

from scimba_jax.linear_approximation.transfer.hierarchy import (
    build_hierarchy_structured,
)

# The factory constructs the same physical problem at every resolution.
hierarchy = build_hierarchy_structured(
    scheme_factory=make_scheme,
    n_cells_finest=(64, 64),
    n_levels=4,
)

scheme_factory is called once per level. Its physical model, boundary conditions and basis family must describe the same PDE at each resolution. For a block mesh it may return a block scheme; the builder coarsens each patch and selects PatchTransfer automatically.

For unstructured grids, use build_hierarchy_nested only with a MeshHierarchy generated by nested refinement. Independently generated Gmsh meshes at h and h / 2 are not nested and do not define a valid prolongation table.

Multigrid

MG is both a solver and an AbstractPreconditioner. Its expensive work is done once at construction/setup: smoother state for each non-coarse level, and a dense LU factorisation of the assembled coarse operator. A cycle thereafter only smooths, transfers and solves that already factorised coarse system.

from scimba_jax.linear_approximation.solvers.multigrid import MG
from scimba_jax.linear_approximation.solvers.smoothers import DampedJacobiSmoother

mg = MG(
    hierarchy,
    smoother=DampedJacobiSmoother(),
    nu_pre=2,
    nu_post=2,
    cycle="V",
)
preconditioner = mg.build()

Use MG.solve(b) for standalone linear multigrid, or pass mg to the Krylov path as a preconditioner. coarse="rediscretised" is the default: every level uses its own scheme and it is the required route for future nonlinear FAS. coarse="galerkin" instead assembles P^T A P from the fine operator; use it only when a coarse rediscretisation fails to represent strongly varying or learned coefficients.

The hierarchy and coarse factors belong to one operator. Rebuild MG when the operator changes materially. In a nonlinear sequence, refresh it based on a measured degradation in Krylov iterations, not on a guessed iteration count.

Smoothers and setup cost

DampedJacobiSmoother is the baseline and constructs an exact diagonal once per level. ChebyshevSmoother estimates a spectral interval once per level and is usually more robust for stretched spectra such as SIPG penalties or Nitsche interfaces. SchwarzSmoother uses overlapping local patches when pointwise smoothing is inadequate.

Do not replace the exact diagonal by a global mesh-size heuristic by default: on a mapped mesh, the relevant local penalty scales with the physical cell size. A cheaper setup that worsens every subsequent Krylov solve is rarely a win unless there is exactly one solve per operator.