Restore FMG mesh hierarchy transparently on checkpoint reload

docs/developer/design/fmg-checkpoint-hierarchy.md

@@ -0,0 +1,117 @@
+---
+title: "Persisting the FMG mesh hierarchy across checkpoints"
+---
+
+# Persisting the geometric-multigrid hierarchy across checkpoints
+
+## Problem
+
+A mesh built with `refinement=N` carries a geometric refinement hierarchy in
+`mesh.dm_hierarchy`, which the Stokes/scalar/vector solvers use for geometric
+**Full Multigrid (FMG)** — the anisotropy-robust preconditioner of choice on
+adapted meshes. But the checkpoint load path (`_from_plexh5`) reconstructs only
+the single saved DMPlex, so a **reloaded mesh has `dm_hierarchy = [dm]`** (one
+level) and FMG silently falls back to GAMG after a restart. This note records the
+design that restores the hierarchy on reload, and the experiments that shaped it.
+
+## Why store the coarse levels (not reconstruct, not refine-from-label)
+
+Several approaches were prototyped and rejected:
+
+- **Reconstruct the coarse mesh from a per-node "level" label.** The coarse
+  *topology* is recoverable bit-exact from labelled nodes (the all-midpoint
+  "central" fine cells map 1:1 to coarse cells). **But** a coarse DMPlex built
+  from scratch (`createFromCellList`) does not reproduce the original's internal
+  cone-orientation / edge ordering, and PETSc's nested multigrid interpolator
+  needs the **canonical `refine()` numbering**. Splicing a reconstructed coarse
+  in throws `PETSc has generated inconsistent data` (err 77).
+- **Store only the coarsest level + `refine()` back up on reload.** Works for the
+  hierarchy, but then the rebuilt fine has to be reconciled (numbering + field
+  data) with the saved deformed mesh — fragile.
+
+The winning insight: a **loaded** coarse DM preserves the canonical numbering
+(`topologyLoad` is faithful), so it can be `setCoarseDM`-linked directly under the
+working fine and FMG just works — **no refine, no node-moving, no reconstruction**.
+This is exactly the live `clone_dm_hierarchy` pattern with *load* swapped for
+*clone*. Validated: `refine(stored L0) == saved fine` bit-exact in 2D and 3D, and a
+reloaded hierarchy drives `pc_type=mg` to convergence.
+
+## On-disk format: a single coarsest sidecar
+
+PETSc's `HDF5_PETSC` `DMView` writes to fixed top-level groups (`/topology`,
+`/geometry`, `/labels`) — it is **not namespaced by DM name**. Writing a second
+DMPlex into the same file (PETSc viewer append, *or* an h5py-injected subgroup
+that the PETSc reader then ignores) corrupts the file (a reload BUS-errors). So
+the hierarchy is stored in **one extra single-DM file** beside the main
+checkpoint, holding only the **coarsest** level:
+
+```
+mymesh.h5                  # the working/fine mesh (unchanged, fully compatible)
+mymesh.hierarchy.L0.h5     # coarsest level only
+```
+
+The intermediate coarse levels are not stored — on reload they are rebuilt by
+`refine()`-ing the coarsest `N-1` times (they come back canonically numbered,
+which is all the co-located nested interpolation needs). The main file's
+`metadata` group gains `hierarchy_coarse_levels = N-1` (the refinement depth).
+Old checkpoints (attribute absent) and plain meshes (no hierarchy) write no
+sidecar and reload exactly as before.
+
+## Reload and the link-free working `dm`
+
+On reload the coarse levels are loaded and spliced:
+`dm_hierarchy = [L0, …, L_{N-2}, fine]`, linked with `setCoarseDM`. One subtlety:
+the mesh's **working `self.dm` must be a link-free clone** of the finest level
+(mirroring the `refinement` construction branch). If `self.dm` itself carries a
+coarse-DM link, `mesh.update_lvec()`'s `createFieldDecomposition` recurses into
+the 0-field coarse levels and fails (`requested fields 1 > DM fields 0`). The
+linked hierarchy lives in `dm_hierarchy`; the solver clones it
+(`clone_dm_hierarchy`) for its own multigrid setup.
+
+## Parallel: co-location via the Simple partitioner
+
+Works in serial **and** parallel through the same reload path. The hazard in
+parallel is that the coarse sidecars and the fine reload on **independent
+partitions**; linking incompatibly-partitioned levels sends the interpolator into
+a cross-rank point-location spin (observed before the fix: np=2, rank 0 at 99% CPU
+indefinitely, rank 1 idle).
+
+The fix needs no custom partition math. The fine carries the **canonical
+refinement numbering** — coarse cell `c`'s children are fine cells
+`c·numSubcells + r`, laid out contiguously right after `c`. So if the fine *and*
+every coarse level are distributed with PETSc's **Simple** partitioner (equal
+contiguous splits of `[0, Ncells)`), the fine split at `k·Nf/p` lines up with the
+coarse split at `k·Nc/p` (since `Nf = numSubcells·Nc`): **each rank's coarse cells
+and their fine children land on the same rank.** The multigrid interpolation is
+then rank-local — no cross-partition communication, no hang — and the levels are a
+genuine per-rank refinement, so the exact **nested** interpolator applies (the fine
+levels are flagged via `DMPlexSetRegularRefinement`).
+
+Trade-off: hierarchy meshes reload with a Simple (contiguous) partition rather than
+the default graph partition. For refinement meshes the canonical ordering is
+reasonably coherent, and field reload is coordinate-matched (partition-agnostic),
+so correctness is unaffected; partition-quality tuning can come later. Plain
+(non-hierarchy) meshes are untouched — they keep the default partitioner.
+
+## Implementation
+
+All in `src/underworld3/discretisation/discretisation_mesh.py`:
+
+- `_hierarchy_sidecar_name()` — sidecar path convention.
+- `Mesh.write()` — writes `metadata/hierarchy_coarse_levels` and one sidecar
+  holding the coarsest level (collective).
+- `Mesh.__init__` `.h5` branch — loads the coarsest sidecar and rebuilds the
+  intermediate coarse levels by `refine()` (serial and parallel), stashing the
+  list on `self._sidecar_coarse_levels`.
+- `Mesh.__init__` hierarchy section — distributes fine + coarse with the Simple
+  partitioner (co-location), splices them under the working dm, flags the fine
+  levels as regular refinements, re-establishes the link-free clone.
+- `petsc_dm_{set,get}_regular_refinement` in `cython/petsc_discretisation.pyx` —
+  wraps `DMPlexSetRegularRefinement` (not exposed by petsc4py) so reloaded levels
+  take the exact nested interpolation path.
+
+No new user-facing surface: the same `Mesh(file)` reload transparently restores the
+hierarchy when the checkpoint has one, and behaves exactly as before when it does
+not.
+
+Tests: `tests/test_0004_checkpoint_fmg_hierarchy.py`.

src/underworld3/cython/petsc_discretisation.pyx

@@ -163,6 +163,24 @@ def petsc_dm_filter_by_label(incoming_dm, label_name, label_value):
         return subdm
 
 
+def petsc_dm_set_regular_refinement(dm, regular=True):
+        """Flag a DMPlex as having been produced by uniform (regular) refinement
+        of its coarse DM.
+
+        This is the flag PETSc's geometric multigrid checks
+        (``DMCreateInterpolation_Plex``) to take the exact *nested* interpolation
+        path (``DMPlexComputeInterpolatorNested``, which maps coarse cell ``c`` to
+        fine children ``c*numSubcells + r``) instead of the point-location
+        *general* path. Only valid when the dm's point numbering follows the
+        canonical ``refine()`` convention and its coarse DM is set
+        (``setCoarseDM``). Not exposed by petsc4py, hence this wrapper.
+        """
+        cdef DM c_dm = dm
+        cdef PetscBool flag = PETSC_TRUE if regular else PETSC_FALSE
+        CHKERRQ( DMPlexSetRegularRefinement(c_dm.dm, flag) )
+        return
+
+
 
 # This is not cython, does it need to be here or in discretisation.py ?
 

src/underworld3/cython/petsc_extras.pxi

@@ -55,6 +55,7 @@ cdef extern from "petsc.h" nogil:
     # PetscErrorCode DMPlexSetSNESLocalFEM( PetscDM, void *, void *, void *)
     # PetscErrorCode DMPlexSetSNESLocalFEM( PetscDM, PetscBool, void *)
     PetscErrorCode DMPlexComputeGeometryFVM( PetscDM dm, PetscVec *cellgeom, PetscVec *facegeom)
+    PetscErrorCode DMPlexSetRegularRefinement( PetscDM dm, PetscBool regular)
     PetscErrorCode MatInterpolate(PetscMat A, PetscVec x, PetscVec y)
     PetscErrorCode DMSetLocalSection(PetscDM, PetscSection)
 

src/underworld3/discretisation/discretisation_mesh.py

@@ -166,6 +166,20 @@ def _from_plexh5(
         return sf0, h5plex
 
 
+def _hierarchy_sidecar_name(mesh_filename, level):
+    """Filename for a coarse hierarchy sidecar of a mesh checkpoint.
+
+    The geometric-multigrid (FMG) hierarchy is persisted as a single extra
+    single-DM HDF5 file holding the **coarsest** level (``level=0``) beside the
+    main mesh checkpoint — PETSc's ``HDF5_PETSC`` format does not support several
+    DMPlex objects in one file. The intermediate coarse levels are rebuilt by
+    refinement on reload, so only ``mymesh.hierarchy.L0.h5`` is written. The
+    ``level`` argument is kept for forward-compatibility.
+    """
+    base, ext = os.path.splitext(mesh_filename)
+    return f"{base}.hierarchy.L{level}{ext}"
+
+
 class Mesh(Stateful, uw_object):
     r"""
     Unstructured mesh with PETSc DMPlex backend.
@@ -408,6 +422,35 @@ def __init__(
 
                 f.close()
 
+                # Restore the geometric-multigrid (FMG) coarse hierarchy from
+                # sidecar files, if present. Works in serial and parallel: the
+                # coarse levels and the fine are co-located on reload by the
+                # Simple partitioner (see the hierarchy-construction branch), so
+                # the multigrid interpolation is rank-local and does not hang.
+                self._sidecar_coarse_levels = None
+                try:
+                    with h5py.File(plex_or_meshfile, "r") as fh:
+                        n_coarse = int(
+                            fh["metadata"].attrs.get("hierarchy_coarse_levels", 0)
+                        )
+                except (KeyError, OSError):
+                    n_coarse = 0
+                if n_coarse > 0:
+                    sidecar = _hierarchy_sidecar_name(plex_or_meshfile, 0)
+                    if os.path.isfile(sidecar):
+                        # Load the coarsest level and rebuild the intermediate
+                        # coarse levels by refining it (n_coarse - 1 times). They
+                        # come back canonically numbered — all the co-located
+                        # nested interpolation needs — and the hierarchy branch
+                        # then Simple-distributes each level. (refine() propagates
+                        # the boundary labels from the coarsest.)
+                        coarsest = _from_plexh5(sidecar, PETSc.COMM_WORLD)
+                        levels = [coarsest]
+                        for _ in range(n_coarse - 1):
+                            levels[-1].setRefinementUniform()
+                            levels.append(levels[-1].refine())
+                        self._sidecar_coarse_levels = levels
+
                 # Do not call setFromOptions() here. DMPlexTopologyLoad()
                 # returns the topology SF needed to reload checkpoint fields.
                 # setFromOptions() can repartition/reorder the DM before UW
@@ -529,7 +572,55 @@ class replacement_boundaries(Enum):
 
         uw.mpi.barrier()
 
-        if not refinement is None and refinement > 0:
+        if getattr(self, "_sidecar_coarse_levels", None):
+
+            # Reloaded mesh with a persisted FMG hierarchy: splice the loaded
+            # coarse levels under the working (fine) dm and link them so PETSc's
+            # geometric multigrid sees a refinement hierarchy. The loaded coarse
+            # DMs keep the canonical refine() numbering the nested interpolator
+            # needs (a *reconstructed* coarse would not — see the
+            # checkpoint-hierarchy design note).
+            #
+            # Co-location (serial and parallel): distribute the fine AND every
+            # coarse level with the Simple (contiguous-by-canonical-index)
+            # partitioner. Because the fine carries the canonical refinement
+            # numbering (coarse cell c -> fine cells c*numSubcells+r, contiguous),
+            # equal contiguous splits put each rank's coarse cells and their fine
+            # children on the same rank. The multigrid interpolation is then
+            # rank-local — this is what avoids the cross-partition hang that a
+            # default (graph) partitioner produces on independently-loaded levels.
+            from underworld3.cython.petsc_discretisation import (
+                petsc_dm_set_regular_refinement,
+            )
+
+            for _hdm in [self.dm] + list(self._sidecar_coarse_levels):
+                if not _hdm.isDistributed():
+                    _hdm.getPartitioner().setType(PETSc.Partitioner.Type.SIMPLE)
+            if not self.dm.isDistributed():
+                self.sf1 = self.dm.distribute()
+            for _cdm in self._sidecar_coarse_levels:
+                if not _cdm.isDistributed():
+                    _cdm.distribute()
+
+            self.dm_hierarchy = list(self._sidecar_coarse_levels) + [self.dm]
+            for i in range(len(self.dm_hierarchy) - 1):
+                self.dm_hierarchy[i + 1].setCoarseDM(self.dm_hierarchy[i])
+                # The fine level IS a uniform refinement of the coarse (canonical
+                # numbering, co-located) — flag it so PETSc builds the exact
+                # nested interpolator rather than falling back to point location.
+                petsc_dm_set_regular_refinement(self.dm_hierarchy[i + 1], True)
+
+            self.dm_h = self.dm_hierarchy[-1]
+            self.dm_h.setName("uw_hierarchical_dm")
+
+            # Working dm is a link-free clone of the finest level (mirrors the
+            # refinement branch). It must NOT carry a coarse-DM link or
+            # mesh.update_lvec()'s createFieldDecomposition recurses into the
+            # 0-field coarse levels and fails.
+            self.dm = self.dm_h.clone()
+            self._sidecar_coarse_levels = None
+
+        elif not refinement is None and refinement > 0:
 
             self.dm.setRefinementUniform()
 
@@ -3347,8 +3438,34 @@ def write(self, filename: str, index: Optional[int] = None):
                     }
                     g.attrs["coordinate_units"] = json.dumps(coord_units_dict)
 
+                # Number of coarse multigrid levels in the hierarchy (= number
+                # of refinements from the stored coarsest level up to the fine
+                # mesh). Used on reload to rebuild the intermediate levels.
+                g.attrs["hierarchy_coarse_levels"] = len(self.dm_hierarchy) - 1
+
                 f.close()
 
+        # Persist the geometric-multigrid (FMG) hierarchy as a SINGLE sidecar
+        # holding the coarsest level only. On reload the intermediate coarse
+        # levels are rebuilt by refining it (they come back canonically numbered,
+        # which is all the co-located nested interpolation needs). Without this
+        # file a reloaded mesh has a single level and falls back to GAMG. One
+        # single-DM HDF5 file (PETSc's HDF5_PETSC format holds one DMPlex per
+        # file). Collective write. See _hierarchy_sidecar_name and the .h5 reload.
+        if len(self.dm_hierarchy) > 1:
+            coarse_dm = self.dm_hierarchy[0]
+            sidecar = _hierarchy_sidecar_name(filename, 0)
+            cviewer = PETSc.ViewerHDF5().create(sidecar, "w", comm=PETSc.COMM_WORLD)
+            cviewer.pushFormat(PETSc.Viewer.Format.HDF5_PETSC)
+            saved_name = coarse_dm.getName()
+            coarse_dm.setName("uw_mesh")  # _from_plexh5 loads the DM named "uw_mesh"
+            try:
+                cviewer(coarse_dm)
+            finally:
+                coarse_dm.setName(saved_name)
+                cviewer.popFormat()
+                cviewer.destroy()
+
     def vtk(self, filename: str):
         """
         Save mesh to the specified file

tests/parallel/mpi_runner.sh

@@ -38,3 +38,8 @@ echo "ptest 0010 snapshot on-disk -np 3 (uneven)"
 mpirun -np 3 $PYTHON ./ptest_0010_snapshot_disk.py
 echo "ptest 0010 snapshot on-disk -np 4"
 mpirun -np 4 $PYTHON ./ptest_0010_snapshot_disk.py
+
+echo "ptest 0004 checkpoint FMG hierarchy -np 2"
+mpirun -np 2 $PYTHON ./ptest_0004_checkpoint_fmg_hierarchy.py
+echo "ptest 0004 checkpoint FMG hierarchy -np 3 (uneven partition)"
+mpirun -np 3 $PYTHON ./ptest_0004_checkpoint_fmg_hierarchy.py

tests/parallel/ptest_0004_checkpoint_fmg_hierarchy.py

@@ -0,0 +1,77 @@
+"""Parallel (MPI) test: FMG mesh hierarchy survives a checkpoint round-trip.
+
+A mesh built with ``refinement`` carries a geometric-multigrid hierarchy. On
+reload from a checkpoint, ``Mesh(file)`` must transparently restore that
+hierarchy in parallel too — without the cross-partition interpolation hang that
+a naive (graph-partitioned) reload of independently-loaded levels produces.
+
+The fix co-locates the levels by distributing the fine and every coarse level
+with PETSc's Simple partitioner: because the fine carries canonical refinement
+numbering, equal contiguous splits put each rank's coarse cells and their fine
+children on the same rank, so the multigrid interpolation is rank-local.
+
+Run:
+
+    cd tests/parallel
+    mpirun -np 2 python ./ptest_0004_checkpoint_fmg_hierarchy.py
+
+Asserts (checked on rank 0):
+  1. The reloaded mesh has its hierarchy restored (levels > 1).
+  2. Geometric FMG converges on the reloaded mesh — proving the levels are
+     co-located and linked (a hang would block here, caught by the CI timeout).
+"""
+
+import os
+import glob
+
+import underworld3 as uw
+from petsc4py import PETSc
+
+rank = uw.mpi.rank
+size = uw.mpi.size
+
+fn = "/tmp/_ptest_0004_fmg_ckpt.h5"
+if rank == 0:
+    for p in glob.glob("/tmp/_ptest_0004_fmg_ckpt*"):
+        os.remove(p)
+uw.mpi.barrier()
+
+# Build a refinement mesh (3-level hierarchy) and checkpoint it.
+mesh = uw.meshing.UnstructuredSimplexBox(
+    minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0),
+    cellSize=0.2, refinement=2, qdegree=2,
+)
+mesh.write(fn)
+uw.mpi.barrier()
+
+# Reload in parallel — same call as a non-hierarchical mesh.
+mesh2 = uw.discretisation.Mesh(fn)
+assert len(mesh2.dm_hierarchy) == len(mesh.dm_hierarchy) == 3, (
+    f"hierarchy not restored: got {len(mesh2.dm_hierarchy)} levels"
+)
+
+# Geometric FMG must converge on the reloaded mesh (no cross-partition hang).
+poisson = uw.systems.Poisson(mesh2)
+poisson.constitutive_model = uw.constitutive_models.DiffusionModel
+poisson.constitutive_model.Parameters.diffusivity = 1
+poisson.f = 0.0
+poisson.add_dirichlet_bc(0.0, "Bottom")
+poisson.add_dirichlet_bc(1.0, "Top")
+for k, v in {
+    "pc_type": "mg", "pc_mg_type": "full", "pc_mg_galerkin": "both",
+    "mg_levels_ksp_type": "chebyshev", "mg_levels_pc_type": "sor",
+    "mg_coarse_pc_type": "redundant", "mg_coarse_redundant_pc_type": "lu",
+}.items():
+    poisson.petsc_options[k] = v
+poisson.solve()
+
+assert poisson.petsc_options.getString("pc_type") == "mg"
+assert poisson.snes.getConvergedReason() > 0
+
+if rank == 0:
+    print(
+        f"ptest_0004 OK (np={size}): hierarchy restored to "
+        f"{len(mesh2.dm_hierarchy)} levels, FMG converged in "
+        f"{poisson.snes.getKSP().getIterationNumber()} iters",
+        flush=True,
+    )