From 1eedb339ef2721eba07e4a5e831d64ec4ee44718 Mon Sep 17 00:00:00 2001 From: Tyagi Date: Thu, 3 Sep 2026 11:17:35 +1000 Subject: [PATCH] Separate reconstructed and in-place PETSc field reloads Restore reconstructed checkpoint meshes through DMPlex section/local-vector metadata and the topology migration SF, because raw global-vector ordering is not stable across mesh reconstruction. Preserve an explicit same-layout vector path for in-place disk-snapshot restoration and reject snapshot rank-count changes. Add serial format coverage and a four-rank regression that verifies both in-place snapshot restore and fresh reconstructed-mesh reload exactly reproduce coordinate-defined fields. --- src/underworld3/checkpoint/disk_snapshot.py | 10 ++ .../discretisation/discretisation_mesh.py | 26 +++- .../discretisation_mesh_variables.py | 116 ++++++++++++------ tests/parallel/ptest_0010_snapshot_disk.py | 75 +++++++---- tests/test_0010_snapshot_disk_format.py | 73 ++++++++++- 5 files changed, 233 insertions(+), 67 deletions(-) diff --git a/src/underworld3/checkpoint/disk_snapshot.py b/src/underworld3/checkpoint/disk_snapshot.py index c0cdac286..9cbf7c22a 100644 --- a/src/underworld3/checkpoint/disk_snapshot.py +++ b/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. diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index d3c67611b..d22d8b190 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/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, diff --git a/src/underworld3/discretisation/discretisation_mesh_variables.py b/src/underworld3/discretisation/discretisation_mesh_variables.py index 9bbcabbcf..5d8eb7d66 100644 --- a/src/underworld3/discretisation/discretisation_mesh_variables.py +++ b/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 diff --git a/tests/parallel/ptest_0010_snapshot_disk.py b/tests/parallel/ptest_0010_snapshot_disk.py index b96823135..14a6c482d 100644 --- a/tests/parallel/ptest_0010_snapshot_disk.py +++ b/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) diff --git a/tests/test_0010_snapshot_disk_format.py b/tests/test_0010_snapshot_disk_format.py index 012cce0a5..0b016e246 100644 --- a/tests/test_0010_snapshot_disk_format.py +++ b/tests/test_0010_snapshot_disk_format.py @@ -193,6 +193,7 @@ def _fresh_model_mesh_and_vars(): def test_write_snapshot_produces_wrapper_and_bulk_dir(tmp_path): """The two artifacts the convention promises: wrapper file + sibling .bulk/ directory containing PETSc HDF5 files.""" + import h5py import os import underworld3 as uw @@ -213,6 +214,16 @@ def test_write_snapshot_produces_wrapper_and_bulk_dir(tmp_path): assert any("T.00000.h5" in f for f in files) assert any("V.00000.h5" in f for f in files) + for variable_name in ("T", "V"): + variable_file = next( + filename + for filename in files + if filename.endswith(f".{variable_name}.00000.h5") + ) + with h5py.File(os.path.join(bulk, variable_file), "r") as h5: + assert variable_name in h5["topologies"]["uw_mesh"]["dms"] + assert variable_name in h5["uw_checkpoint"] + def test_snapshot_bulk_filenames_do_not_expand_loaded_mesh_name(tmp_path): """A source pathname remains metadata, not a PETSc bulk filename.""" @@ -289,8 +300,7 @@ def test_write_snapshot_populates_wrapper_layout(tmp_path): def test_write_read_snapshot_bit_exact_roundtrip(tmp_path): """The core phase-2 guarantee: write a snapshot, scribble all variables, read snapshot back, all variables match write-time - values bit-for-bit (#146's PETSc DMPlex same-rank reload, just - delivered via the wrapper).""" + values bit-for-bit through the exact same-layout PETSc vector payload.""" import underworld3 as uw uw, model, mesh, T, V = _fresh_model_mesh_and_vars() @@ -319,6 +329,49 @@ def test_write_read_snapshot_bit_exact_roundtrip(tmp_path): ) +def test_disk_restore_refreshes_packed_auxiliary_fields_before_solve(tmp_path): + """A solve after disk restore must use restored coefficient fields.""" + import underworld3 as uw + + uw.reset_default_model() + model = uw.get_default_model() + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=0.5, + ) + solution = uw.discretisation.MeshVariable("U", mesh, 1, degree=1) + source = uw.discretisation.MeshVariable("source", mesh, 1, degree=1) + + poisson = uw.systems.Poisson(mesh, u_Field=solution) + poisson.constitutive_model = uw.constitutive_models.DiffusionModel + poisson.constitutive_model.Parameters.diffusivity = 1.0 + poisson.f = source.sym[0] + poisson.add_dirichlet_bc(0.0, "All_Boundaries") + + source.array[...] = 1.0 + poisson.solve(zero_init_guess=True) + reference = np.asarray(solution.array[...]).copy() + + path = str(tmp_path / "auxiliary.snap.h5") + model.save_state(file=path) + + source.array[...] = 4.0 + poisson.solve(zero_init_guess=True) + assert not np.allclose(np.asarray(solution.array[...]), reference) + + model.load_state(path) + assert mesh._stale_lvec is True + poisson.solve(zero_init_guess=True) + + assert np.allclose( + np.asarray(solution.array[...]), + reference, + rtol=0.0, + atol=1.0e-5, + ) + + def test_read_snapshot_rejects_missing_bulk_dir(tmp_path): """If the user moves the wrapper without the bulk dir, read fails with a clear pointer rather than an obscure h5py error.""" @@ -338,6 +391,22 @@ def test_read_snapshot_rejects_missing_bulk_dir(tmp_path): model.load_state(path) +def test_read_snapshot_rejects_different_mpi_rank_count(tmp_path): + """Exact disk restart must not silently remap a different MPI layout.""" + import h5py + import underworld3 as uw + + uw, model, mesh, T, V = _fresh_model_mesh_and_vars() + path = str(tmp_path / "rank_count.snap.h5") + model.save_state(file=path) + + with h5py.File(path, "r+") as h5: + h5["metadata"].attrs["mpi_ranks_at_write"] = uw.mpi.size + 1 + + with pytest.raises(ValueError, match="same rank count"): + model.load_state(path) + + def test_read_snapshot_rejects_mismatched_mesh(tmp_path): """If the target model's meshes don't match the snapshot's, raise clearly — mesh-rebuild on read is v1.2 scope."""