Separate reconstructed and in-place PETSc checkpoint reloads

src/underworld3/checkpoint/disk_snapshot.py

@@ -485,6 +485,15 @@ def read_snapshot(model, path: str) -> None:
     import h5py
 
     md = read_snapshot_metadata(path)
+    write_size = int(md.get("mpi_ranks_at_write", 1))
+    if write_size != int(uw.mpi.size):
+        raise ValueError(
+            f"snapshot at {path} was written on {write_size} MPI rank(s); "
+            f"this run uses {uw.mpi.size}. Exact disk restart requires the "
+            "same rank count; use mesh.write_timestep/read_timestep for "
+            "coordinate-remapped field transfer."
+        )
+
     bulk_dir = _bulk_dir_for(path)
     if not os.path.isdir(bulk_dir):
         raise FileNotFoundError(
@@ -526,6 +535,7 @@ def read_snapshot(model, path: str) -> None:
                     var.read_checkpoint(
                         short_external_file,
                         data_name=var_name,
+                        same_layout=True,
                     )
 
         # Phase 3a: restore state-bearer dataclasses.

src/underworld3/discretisation/discretisation_mesh.py

@@ -4836,9 +4836,10 @@ def write_timestep(
         - ``create_xdmf=True`` writes ParaView/XDMF output. Variable files also
           receive ``/vertex_fields`` or ``/cell_fields`` compatibility groups,
           and rank 0 writes the companion ``.xdmf`` file.
-        - ``petsc_reload=True`` writes PETSc DMPlex section/vector metadata into
-          the same per-variable HDF5 files. These files can then be loaded with
-          ``MeshVariable.read_checkpoint()`` for PETSc-native same-mesh reload.
+        - ``petsc_reload=True`` writes PETSc DMPlex section/local-vector
+          metadata and an in-place global-vector payload into the same
+          per-variable HDF5 files. These files can then be loaded with
+          ``MeshVariable.read_checkpoint()`` for exact restart.
 
         Common choices are:
 
@@ -5015,7 +5016,7 @@ def _write_petsc_reload_variable(self, viewer, var):
             subdm.destroy()
 
     def _write_petsc_reload_file(self, checkpoint_file, variables, mode="w"):
-        """Write PETSc DMPlex section/vector reload metadata."""
+        """Write DMPlex reload metadata and in-place vector payloads."""
 
         old_dm_name = self.dm.getName()
         self.dm.setName("uw_mesh")
@@ -5037,6 +5038,23 @@ def _write_petsc_reload_file(self, checkpoint_file, variables, mode="w"):
             if old_dm_name is not None:
                 self.dm.setName(old_dm_name)
 
+        viewer = PETSc.ViewerHDF5().create(
+            checkpoint_file, "a", comm=PETSc.COMM_WORLD
+        )
+        try:
+            viewer.pushGroup("/uw_checkpoint")
+            for var in variables:
+                var._sync_lvec_to_gvec()
+                checkpoint_vec = PETSc.Vec().createWithArray(
+                    var._gvec.array_r, comm=PETSc.COMM_WORLD
+                )
+                checkpoint_vec.setName(var.clean_name)
+                viewer(checkpoint_vec)
+                checkpoint_vec.destroy()
+            viewer.popGroup()
+        finally:
+            viewer.destroy()
+
     @timing.routine_timer_decorator
     def write_checkpoint(
         self,

src/underworld3/discretisation/discretisation_mesh_variables.py

@@ -1449,14 +1449,21 @@ def read_checkpoint(
         self,
         filename: str,
         data_name: Optional[str] = None,
+        same_layout: bool = False,
     ):
         """Load this mesh variable from PETSc reload output.
 
-        This is an exact PETSc DMPlex section/vector reload path. It does not
-        use the coordinate/KDTree remapping used by ``read_timestep()``. New
-        output should be written with ``Mesh.write_timestep(...,
-        petsc_reload=True)``; legacy ``Mesh.write_checkpoint()`` files are also
-        supported.
+        The default path restores DMPlex section/local-vector data through the
+        topology migration SF, so a mesh reconstructed from its checkpoint may
+        have a different parallel DOF ordering. Set ``same_layout=True`` only
+        for an in-place restore onto the exact mesh object that wrote the file;
+        that path reloads the saved global vector directly.
+
+        This method does not use the coordinate/KDTree remapping provided by
+        ``read_timestep()``. New output should be written with
+        ``Mesh.write_timestep(..., petsc_reload=True)``; legacy
+        ``Mesh.write_checkpoint()`` files are also supported by the default
+        DMPlex path.
         """
 
         if data_name is None:
@@ -1465,6 +1472,27 @@ def read_checkpoint(
         if self._lvec is None:
             self._set_vec(available=True)
 
+        if same_layout:
+            import h5py
+
+            if uw.mpi.rank == 0:
+                with h5py.File(filename, "r") as checkpoint_h5:
+                    has_direct_vector = (
+                        "uw_checkpoint" in checkpoint_h5
+                        and data_name in checkpoint_h5["uw_checkpoint"]
+                    )
+            else:
+                has_direct_vector = None
+            has_direct_vector = uw.mpi.comm.bcast(
+                has_direct_vector,
+                root=0,
+            )
+            if not has_direct_vector:
+                raise RuntimeError(
+                    f"{filename} has no in-place checkpoint vector for "
+                    f"{data_name!r}. Reload it with same_layout=False."
+                )
+
         indexset, subdm = self.mesh.dm.createSubDM(self.field_id)
         sectiondm = self.mesh.dm.clone()
         viewer = PETSc.ViewerHDF5().create(filename, "r", comm=PETSc.COMM_WORLD)
@@ -1481,40 +1509,53 @@ def read_checkpoint(
             self._lvec.setName(data_name)
             self._gvec.setName(data_name)
 
-            from underworld3.cython.petsc_discretisation import (
-                petsc_dmplex_load_local_vector,
-            )
+            if same_layout:
+                checkpoint_vec = PETSc.Vec().createMPI(
+                    (self._gvec.getLocalSize(), self._gvec.getSize()),
+                    comm=PETSc.COMM_WORLD,
+                )
+                checkpoint_vec.setName(data_name)
+                viewer.pushGroup("/uw_checkpoint")
+                checkpoint_vec.load(viewer)
+                viewer.popGroup()
+                self._gvec.array[...] = checkpoint_vec.array_r
+                checkpoint_vec.destroy()
+                subdm.globalToLocal(self._gvec, self._lvec, addv=False)
+            else:
+                from underworld3.cython.petsc_discretisation import (
+                    petsc_dmplex_load_local_vector,
+                )
 
-            loaded_lvec = petsc_dmplex_load_local_vector(
-                self.mesh.dm, viewer, sectiondm, self.mesh.sf, data_name
-            )
+                loaded_lvec = petsc_dmplex_load_local_vector(
+                    self.mesh.dm, viewer, sectiondm, self.mesh.sf, data_name
+                )
 
-            source_section = sectiondm.getSection()
-            target_section = subdm.getSection()
-            source_array = loaded_lvec.array_r
-            target_array = self._lvec.array
-            p_start, p_end = target_section.getChart()
-
-            for point in range(p_start, p_end):
-                target_dof = target_section.getDof(point)
-                if target_dof == 0:
-                    continue
-
-                source_dof = source_section.getDof(point)
-                if source_dof < target_dof:
-                    raise RuntimeError(
-                        f"Checkpoint section has {source_dof} dofs for point {point}, "
-                        f"but target variable requires {target_dof}."
+                source_section = sectiondm.getSection()
+                target_section = subdm.getSection()
+                source_array = loaded_lvec.array_r
+                target_array = self._lvec.array
+                p_start, p_end = target_section.getChart()
+
+                for point in range(p_start, p_end):
+                    target_dof = target_section.getDof(point)
+                    if target_dof == 0:
+                        continue
+
+                    source_dof = source_section.getDof(point)
+                    if source_dof < target_dof:
+                        raise RuntimeError(
+                            f"Checkpoint section has {source_dof} dofs for point "
+                            f"{point}, but target variable requires {target_dof}."
+                        )
+
+                    source_offset = source_section.getOffset(point)
+                    target_offset = target_section.getOffset(point)
+                    target_array[target_offset : target_offset + target_dof] = (
+                        source_array[source_offset : source_offset + target_dof]
                     )
 
-                source_offset = source_section.getOffset(point)
-                target_offset = target_section.getOffset(point)
-                target_array[target_offset : target_offset + target_dof] = (
-                    source_array[source_offset : source_offset + target_dof]
-                )
-
-            loaded_lvec.destroy()
-            self._sync_lvec_to_gvec()
+                loaded_lvec.destroy()
+                self._sync_lvec_to_gvec()
         finally:
             self._lvec.setName(old_lvec_name)
             self._gvec.setName(old_vec_name)
@@ -1526,6 +1567,11 @@ def read_checkpoint(
             indexset.destroy()
             subdm.destroy()
 
+        # The mesh-wide auxiliary vector packs every registered field and may
+        # still contain values from before this reload. Force the next residual
+        # assembly to rebuild it from the restored per-variable vectors.
+        self.mesh._stale_lvec = True
+
         return
 
     @property

tests/parallel/ptest_0010_snapshot_disk.py

@@ -1,10 +1,7 @@
 """Parallel (MPI) test of the on-disk snapshot path (v1.1).
 
-Phase 6 of the snapshot toolkit: per-rank swarm sidecars. The mesh
-+ mesh-variable disk path is already parallel-correct via #146's
-PETSc-collective HDF5 viewer; the swarm sidecar layer needs its
-own per-rank file per swarm. This ptest exercises both together at
-multi-rank.
+Phase 6 of the snapshot toolkit: exact same-rank PETSc mesh-variable reload and
+per-rank swarm sidecars. This ptest exercises both layers together under MPI.
 
 Run (4 ranks exercises cross-rank distribution of swarm particles):
 
@@ -19,7 +16,10 @@
      state (verified by per-rank attrs on the sidecar).
   3. Round-trip is exact: scribble all variables + swarm coords +
      swarm-var data, model.load_state(file=...), gathered (gid, x, y,
-     material) tables sorted by gid are np.array_equal.
+     material) tables sorted by gid are np.array_equal, and every restored
+     mesh-variable dof matches its coordinate-defined analytic value.
+  4. Loading the PETSc field onto a newly reconstructed checkpoint mesh is
+     also exact, so reload does not depend on the original global DOF order.
 """
 
 import os
@@ -66,8 +66,7 @@ def build():
 
 
 def global_sorted_state(T, swarm, gid, material):
-    """Gather (gid, x, y, material, T-value-by-coord-bin) across ranks
-    + sort by gid → order/rank-independent canonical view."""
+    """Return rank-independent swarm state and mesh-field error."""
     g = gid.data[:, 0].copy()
     coords = swarm._particle_coordinates.data.copy()
     m = material.data[:, 0].copy()
@@ -79,19 +78,12 @@ def global_sorted_state(T, swarm, gid, material):
     order = np.argsort(full[:, 0], kind="stable")
     swarm_state = full[order]
 
-    # T round-trip check: gather partition-invariant scalars
-    # (max, sum) rather than the full (coord, value) table — DOFs at
-    # partition boundaries are visible to multiple ranks and would
-    # appear duplicated/reordered in a gathered table, even though
-    # the underlying data is bit-exact.
     t_arr = np.asarray(T.array[...]).reshape(-1)
-    t_max = comm.allreduce(float(t_arr.max()) if t_arr.size else -np.inf,
-                           op=MPI.MAX)
-    t_min = comm.allreduce(float(t_arr.min()) if t_arr.size else np.inf,
-                           op=MPI.MIN)
-    # bit-exact float sum across ranks is non-deterministic in general
-    # (non-associative); use min/max as bit-exact invariants instead.
-    return swarm_state, (t_max, t_min)
+    t_coords = np.asarray(T.coords)
+    expected = t_coords[:, 0] - t_coords[:, 1]
+    local_error = float(np.max(np.abs(t_arr - expected))) if t_arr.size else 0.0
+    global_error = comm.allreduce(local_error, op=MPI.MAX)
+    return swarm_state, global_error
 
 
 def main():
@@ -112,10 +104,11 @@ def main():
     model.save_state(file=wrapper)
     comm.Barrier()
 
+    bulk = os.path.join(tmp, "parrun.snap.bulk")
+
     # Check files on rank 0
     files_ok = True
     if rank == 0:
-        bulk = os.path.join(tmp, "parrun.snap.bulk")
         files = sorted(os.listdir(bulk))
         per_rank = [f for f in files if ".swarm.rank" in f]
         # Expect one swarm sidecar per rank
@@ -146,13 +139,38 @@ def main():
     post_swarm, post_T = global_sorted_state(T, swarm, gid, material)
 
     swarm_ok = np.array_equal(pre_swarm, post_swarm)
-    # T is checked via partition-invariant min/max scalars (see note
-    # in global_sorted_state — gathered DOFs include partition-
-    # boundary duplicates that resist a global-table comparison).
-    T_ok = (pre_T == post_T)
+    T_ok = pre_T == 0.0 and post_T == 0.0
     count_ok = pre_count == post_count
     tracker_ok = (model.tracker.time == 1.5 and model.tracker.step == 7)
 
+    bulk_files = sorted(os.listdir(bulk))
+    mesh_file = os.path.join(
+        bulk,
+        next(name for name in bulk_files if name.endswith(".mesh.00000.h5")),
+    )
+    temperature_file = os.path.join(
+        bulk,
+        next(name for name in bulk_files if name.endswith(".T.00000.h5")),
+    )
+    reloaded_mesh = uw.discretisation.Mesh(mesh_file)
+    reloaded_temperature = uw.discretisation.MeshVariable(
+        "T_reloaded",
+        reloaded_mesh,
+        1,
+        degree=1,
+    )
+    reloaded_temperature.read_checkpoint(temperature_file, data_name="T")
+    reloaded_values = np.asarray(reloaded_temperature.array[...]).reshape(-1)
+    reloaded_coords = np.asarray(reloaded_temperature.coords)
+    reloaded_expected = reloaded_coords[:, 0] - reloaded_coords[:, 1]
+    local_reloaded_error = (
+        float(np.max(np.abs(reloaded_values - reloaded_expected)))
+        if reloaded_values.size
+        else 0.0
+    )
+    reloaded_error = comm.allreduce(local_reloaded_error, op=MPI.MAX)
+    reconstructed_mesh_ok = reloaded_error == 0.0
+
     if rank == 0:
         print(f"[ranks={size}] particles total = {pre_count}", flush=True)
         print(f"  P1 disk wrapper + per-rank sidecars present: {files_ok}",
@@ -161,16 +179,21 @@ def main():
               flush=True)
         print(f"  P3 swarm (coords + gid + material) exact:    {swarm_ok}",
               flush=True)
-        print(f"  P4 T (mesh-variable DOFs) exact:             {T_ok}",
+        print(f"  P4 T (mesh-variable DOFs) exact:             {T_ok} "
+              f"(max error={post_T:.3e})",
               flush=True)
         print(f"  P5 tracker state restored:                   {tracker_ok}",
               flush=True)
+        print(f"  P6 reconstructed-mesh field exact:           "
+              f"{reconstructed_mesh_ok} (max error={reloaded_error:.3e})",
+              flush=True)
 
         assert files_ok
         assert count_ok
         assert swarm_ok
         assert T_ok
         assert tracker_ok
+        assert reconstructed_mesh_ok
         print(f"[ranks={size}] PASS", flush=True)