Geometric multigrid preconditioner (MPI, 3D coarsening, operator re-discretization)

[spossann]

Replaces the previous multigrid attempt (#60, and the earlier state of this PR) with a new geometric multigrid that runs under MPI. Its V-cycle is used as a preconditioner for CG.

Paired feectools PR: struphy-hub/feectools#91 (into devel-tiny). The submodule points to its branch ddm-coarsen; the feectools pin in pyproject.toml can be raised once 0.3.0 is on PyPI.

What it does

MultiGridSolver(A, derham, domain, MultiGridOptions(...)) solves A x = b for a symmetric positive (semi-)definite A on one Derham space. Typical example: sigma * M0 + grad.T @ M1 @ grad.

• Grid hierarchy (multigrid/hierarchy.py): halves the number of elements in every direction where possible, and stops in a direction that gets too small (semi-coarsening). The MPI decomposition of each coarse level is aligned with the fine one, via DomainDecomposition.coarsen and the new Derham(..., domain_decomposition=...) argument.

• Grid transfer (multigrid/transfer.py):

  • SplineProlongation is the exact embedding of the nested spline spaces, for any degree, periodic or clamped, B- and D-splines, with homogeneous Dirichlet BCs.
  • It is applied one direction at a time on the local ghosted arrays.
  • The restriction is its transpose.

• Coarse operators (multigrid/coarsen.py): the fine operator is walked as an expression tree and rebuilt on each coarse grid. No string recipes are needed.

  • Sums, compositions, scalings and block operators are rebuilt from their children; derivative, boundary and identity operators are recreated on the coarse spaces.
  • Mass and basis projection operators are recreated from their new to_dict()/from_dict().
  • Coarse pieces are cached, so changing a scalar re-assembles nothing.

• Smoothers (multigrid/smoothers.py), selected via MultiGridOptions:

  • Chebyshev (default) with an approximate mass inverse, the exact Jacobi diagonal, or the identity inside;
  • damped Jacobi;
  • a fixed number of PCG steps.

  The exact diagonals of composite operators are computed by colored probing.

• V-cycle (multigrid/preconditioner.py): symmetric pre- and post-smoothing. The coarsest level is solved directly on every rank (or by CG). There is an optional nullspace="constants" for singular problems.

• Propagators: ImplicitDiffusion and PoissonSolve get precond="MultiGrid" and a multigrid: MultiGridOptions field. The preconditioner is only updated when sigma_1 (e.g. sigma_1/dt) changes. The other precond values keep their previous behavior (pc=None).

• Removed: the old multigrid_solver.py (eval-string operators, serial-only transfer, 2D coarsening) and the 3,246-line research file feec/tests/test_multigrid.py.

Results

Poisson, CG iterations to a relative tolerance of 1e-8, default options:

Case	Multigrid + CG	Plain CG
2D, 16² → 128², p = 2, 3, Dirichlet or periodic, 1 and 4 ranks	7 at every size	21 → 149
3D, 8³ → 32³, p = 2, 3, 4 ranks	6 – 9	28 – 135

Notes and known limitations

• Which smoother. The default mass-inverse smoother is robust in the spline degree but degrades on strongly curved mappings (Colella α = 0.1: 35 iterations). The Jacobi-based smoother is robust to the mapping (10 iterations there) but slower at degree 3 in 3D (24–31 iterations). It also costs a one-time setup of up to (2p+1)³ operator applications per level.
• 1-forms and 2-forms. Problems such as curl-curl will need dedicated smoothers (Hiptmair / Arnold–Falk–Winther type). The coarsening and grid transfer already support them.
• Polar splines are not supported (clear error). The transfer operators are NumPy-only.
• Coarsest level. The coarse matrix is assembled by applying the operator to every unit vector, which is fine for small coarse grids.
• Performance with Dirichlet BCs. With several ranks, the existing MassMatrixPreconditioner is the main cost: it calls SuperLU on many right-hand sides. A banded solver in feectools would help.

Tests

• New in linear_algebra/tests/, run serially and on 4 ranks:
  • test_multigrid_transfer.py: exactness of the transfer, R = Pᵀ, R M_h P = M_H.
  • test_multigrid_coarsen.py: to_dict round trips, R A P = A_H for GᵀM₁G and CᵀM₂C.
  • test_multigrid_solver.py: smoother symmetry, V-cycle symmetric positive definite and contracting, iteration counts that don't grow with the grid, update.
• New in propagators/tests/test_poisson.py: PoissonSolve with multigrid on the Colella mapping (periodic, Dirichlet, Neumann), and ImplicitDiffusion with a changing dt.
• No new test folders, so no CI shard changes.
• Existing tests re-run: mass / basis-op tests (89 passed), Poisson / gyrokinetic Poisson (92 passed).

🤖 Generated with Claude Code

[spossann]

Hi @mateomarin97 - we have done some house cleaning recently on Github and had to move to a new base branch devel-clean. From now on, please continue working on the new branch new-multigrid-clean from this PR. It has exactly the same code as origin/new-multigrid (I merged with devel and pushed).

In case you still have local code on the old branch that you do not want to lose, you can

git fetch
git checkout new-multigrid-clean
git checkout new-multigrid -- .
git commit -m 'add all local changes'