From ef8abbb366214b0ca6f522c9f024c253a6e4c469 Mon Sep 17 00:00:00 2001 From: Tyagi Date: Wed, 9 Sep 2026 11:27:31 +1000 Subject: [PATCH 1/4] Preserve live symbolic atoms when copying DDt snapshot history Generic deepcopy reconstructs UWexpression atoms without their required internal state. Copy the symbolic history matrices and mutable timestep containers without reconstructing their symbolic atoms. Reproduced the invalid-expression failure on development 87091138. The focused clean branch passes all 25 in-memory snapshot tests and the two new checks on eight MPI ranks. This fixes expression validity, not the separately observed legacy Diffusion uninterrupted-continuation discrepancy. --- src/underworld3/systems/ddt.py | 24 +++++++++++++ tests/test_0007_snapshot_inmemory.py | 54 +++++++++++++++++++++++++++- 2 files changed, 77 insertions(+), 1 deletion(-) diff --git a/src/underworld3/systems/ddt.py b/src/underworld3/systems/ddt.py index 180e43463..25d9b799c 100644 --- a/src/underworld3/systems/ddt.py +++ b/src/underworld3/systems/ddt.py @@ -116,6 +116,30 @@ class DDtSymbolicState(_DDtCoreState): psi_star: list = field(default_factory=list) + def __deepcopy__(self, memo): + """Copy history containers without reconstructing symbolic atoms. + + SymPy matrices are value containers but their expression atoms include + UWexpression objects whose identity binds them to live parameter and + coefficient registries. Generic ``copy.deepcopy`` reconstructs those + Symbol subclasses without their wrapped value, producing invalid atoms + after snapshot restore. Matrix ``copy()`` keeps the immutable symbolic + atoms while separating the mutable history list and matrices. + """ + import copy + + duplicate = type(self)( + _schema_version=self._schema_version, + dt_history=copy.deepcopy(self.dt_history, memo), + history_initialised=self.history_initialised, + n_solves_completed=self.n_solves_completed, + dt=copy.deepcopy(self.dt, memo), + psi_star=[value.copy() for value in self.psi_star], + ) + memo[id(self)] = duplicate + return duplicate + + @dataclass class DDtEulerianState(_DDtCoreState): diff --git a/tests/test_0007_snapshot_inmemory.py b/tests/test_0007_snapshot_inmemory.py index d52fee33a..8dce22386 100644 --- a/tests/test_0007_snapshot_inmemory.py +++ b/tests/test_0007_snapshot_inmemory.py @@ -402,6 +402,59 @@ def test_eulerian_ddt_roundtrip(): assert ddt.state.psi_star_var_names == state_pre.psi_star_var_names +def test_symbolic_history_copy_preserves_atoms_and_separates_containers(): + import copy + import sympy + import underworld3 as uw + from underworld3.systems.ddt import DDtSymbolicState + + coefficient = uw.function.expression("snapshot_coefficient", 2.0) + original = DDtSymbolicState( + dt_history=[0.1, None], psi_star=[sympy.Matrix([[coefficient]])] + ) + captured = copy.deepcopy(original) + + assert captured.psi_star[0][0] is coefficient + assert captured.psi_star is not original.psi_star + assert captured.psi_star[0] is not original.psi_star[0] + captured.psi_star[0][0] = 0 + captured.dt_history[0] = 0.2 + assert original.psi_star[0][0] is coefficient + assert original.dt_history == [0.1, None] + + +def test_symbolic_flux_history_remains_valid_after_solver_restore(): + """Deep-copying symbolic history must preserve live UWexpression atoms.""" + import numpy as np + 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.3 + ) + temperature = uw.discretisation.MeshVariable( + "T_symbolic_restore", mesh, 1, degree=1 + ) + with mesh.access(temperature): + temperature.data[:, 0] = temperature.coords[:, 0] + + diffusion = uw.systems.Diffusion( + mesh, u_Field=temperature, order=2, theta=1.0 + ) + diffusion.constitutive_model = uw.constitutive_models.DiffusionModel + diffusion.constitutive_model.Parameters.diffusivity = 0.05 + for _ in range(3): + diffusion.solve(timestep=0.01, zero_init_guess=False) + + snapshot = model.save_state() + model.load_state(snapshot) + diffusion.solve(timestep=0.01, zero_init_guess=False) + + assert np.all(np.isfinite(temperature.data)) + + + def test_semilagrangian_ddt_roundtrip(): import underworld3 as uw from underworld3.systems.ddt import DDtSemiLagrangianState @@ -766,4 +819,3 @@ def test_continuation_bit_identical_across_stash_and_recover(): _assert_bit_identical(ctrl, stash, "stash-and-recover") - From d551b6fd6a8a16c08a292768d5bb51262c6aebef Mon Sep 17 00:00:00 2001 From: Tyagi Date: Wed, 9 Sep 2026 12:12:41 +1000 Subject: [PATCH 2/4] Fix legacy Diffusion symbolic flux initialization before kernel compilation Populate symbolic history before first residual construction so cold and restored solvers embed identical Adams-Moulton terms. Previously zero placeholders disappeared from the initial operator, producing different continuation after restore despite exact field recovery. Add eight small Level 1 operator and replay regressions covering startup, orders 1-3 and variable timesteps. Validation: 33 serial tests and 10 tests per rank on eight MPI ranks pass; tight-tolerance replay errors remain at roundoff. --- docs/developer/CHANGELOG.md | 12 ++++++ src/underworld3/systems/solvers.py | 5 +++ tests/test_1058_diffusion_restart.py | 63 ++++++++++++++++++++++++++++ 3 files changed, 80 insertions(+) create mode 100644 tests/test_1058_diffusion_restart.py diff --git a/docs/developer/CHANGELOG.md b/docs/developer/CHANGELOG.md index 5aaa56964..c1205357e 100644 --- a/docs/developer/CHANGELOG.md +++ b/docs/developer/CHANGELOG.md @@ -6,6 +6,18 @@ This log tracks significant development work at a conceptual level, suitable for ## 2026 Q3 (July – September) +### Legacy Diffusion Restart Consistency (September 2026, #708) + +`Diffusion` now initializes symbolic flux history before compiling its +first residual. Previously the initial kernel embedded zero history slots; +subsequent coefficient updates could not restore the missing symbolic terms. +A snapshot restore rebuilt the kernel with populated slots and therefore +changed the operator despite restoring all fields exactly. Initializing the +slots before compilation makes cold and rebuilt operators agree. This also +corrects affected uninterrupted legacy Diffusion trajectories; it does not +change the composed `AdvDiffusion` or SLCN implementation. Small serial/MPI +tests cover startup and established histories, orders 1-3 and varying steps. + ### The Multiplier Was Not the Whole Traction (August 2026) **`Stokes_Constrained.topography()` now returns the traction the boundary is diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index 24351f47b..56c8a37a5 100644 --- a/src/underworld3/systems/solvers.py +++ b/src/underworld3/systems/solvers.py @@ -4847,6 +4847,11 @@ def solve( # self._flux = self.constitutive_model.flux.T # self._flux_star = self._flux.copy() + # Symbolic slots are embedded in F1 at compilation, unlike nodal + # histories. Populate them before building, not in the later hook. + if isinstance(self.DFDt, Symbolic_DDt) and not self.DFDt._history_initialised: + self.DFDt.initialise_history() + if not self.is_setup: self._setup_pointwise_functions(verbose) self._setup_discretisation(verbose) diff --git a/tests/test_1058_diffusion_restart.py b/tests/test_1058_diffusion_restart.py new file mode 100644 index 000000000..97a02afae --- /dev/null +++ b/tests/test_1058_diffusion_restart.py @@ -0,0 +1,63 @@ +"""Legacy diffusion must not change its operator when restore rebuilds it.""" + +import numpy as np +import pytest +import sympy +import underworld3 as uw + +pytestmark = [pytest.mark.level_1, pytest.mark.tier_b] + + +def _problem(order, theta): + uw.reset_default_model() + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=0.25, + ) + temperature = uw.discretisation.MeshVariable("T", mesh, 1, degree=1) + temperature.array[:, 0, 0] = temperature.coords[:, 0] + solver = uw.systems.Diffusion(mesh, temperature, order=order, theta=theta) + solver.constitutive_model = uw.constitutive_models.DiffusionModel + solver.constitutive_model.Parameters.diffusivity = 0.05 + solver.tolerance = 1e-12 + return uw.get_default_model(), mesh, solver, temperature + + +@pytest.mark.parametrize("theta", [0.5, 1.0]) +def test_compiled_flux_contains_initial_symbolic_history(theta): + _model, _mesh, solver, _temperature = _problem(2, theta) + for _ in range(3): + solver.solve(timestep=0.01, zero_init_guess=False) + unwrap = uw.function.expressions.unwrap + compiled = unwrap(solver._f1.sym, keep_constants=False) + live = unwrap(solver.DFDt.adams_moulton_flux(), keep_constants=False) + assert all(sympy.simplify(term) == 0 for term in compiled - live) + + +@pytest.mark.parametrize("order,theta,warm_steps", [ + (1, 1.0, 3), (1, 0.5, 3), (2, 1.0, 0), + (2, 1.0, 1), (2, 1.0, 3), (3, 1.0, 4), +]) +def test_diffusion_snapshot_replays_continuation(order, theta, warm_steps): + model, mesh, solver, temperature = _problem(order, theta) + for _ in range(warm_steps): + solver.solve(timestep=0.01, zero_init_guess=False) + initial_fields = {name: np.array(var.array) for name, var in mesh.vars.items()} + snapshot = model.save_state() + timesteps = [0.01, 0.01, 0.015, 0.0075] + reference = [] + for dt in timesteps: + solver.solve(timestep=dt, zero_init_guess=False) + reference.append(np.array(temperature.array)) + + model.load_state(snapshot) + for name, values in initial_fields.items(): + np.testing.assert_array_equal(mesh.vars[name].array, values) + max_error = 0.0 + for dt, expected in zip(timesteps, reference): + solver.solve(timestep=dt, zero_init_guess=False) + actual = np.asarray(temperature.array) + max_error = max(max_error, float(np.max(np.abs(actual - expected)))) + np.testing.assert_allclose(actual, expected, rtol=1e-11, atol=1e-12) + max_error = max(uw.mpi.comm.allgather(max_error)) + uw.pprint(f"DIFFUSION_REPLAY order={order} theta={theta} warm={warm_steps} " + f"ranks={uw.mpi.size} max_error={max_error:.12g}") From 21a84519196f6d20e1593ed8cb3ca1f87e0abbfa Mon Sep 17 00:00:00 2001 From: Tyagi Date: Thu, 10 Sep 2026 10:59:58 +1000 Subject: [PATCH 3/4] Address PR 708 review: expose skipped disk state and test fresh-process Diffusion replay Warn when snapshot fields cannot be serialized and document live symbolic atom reference limitations. Add bounded fresh-process serial/MPI replay regression with a strict, narrowly recognized expected numerical failure: all fields and metadata restore exactly, but omitted symbolic history still changes continuation. Do not claim disk fidelity or reconstruct arbitrary symbolic history in this patch. Modernize array access and document the general folded-zero kernel hazard. Validation: 62 serial passes plus one known disk xfail; eight-rank disk workers reproduce the same 0.00350097 mismatch with exact field restoration. --- docs/developer/CHANGELOG.md | 11 +++ src/underworld3/checkpoint/disk_snapshot.py | 12 ++++ src/underworld3/systems/ddt.py | 7 +- .../ptest_1058_diffusion_disk_restart.py | 34 +++++++++ tests/test_0007_snapshot_inmemory.py | 5 +- tests/test_1059_diffusion_disk_restart.py | 71 +++++++++++++++++++ 6 files changed, 136 insertions(+), 4 deletions(-) create mode 100644 tests/parallel/ptest_1058_diffusion_disk_restart.py create mode 100644 tests/test_1059_diffusion_disk_restart.py diff --git a/docs/developer/CHANGELOG.md b/docs/developer/CHANGELOG.md index c1205357e..4aad91431 100644 --- a/docs/developer/CHANGELOG.md +++ b/docs/developer/CHANGELOG.md @@ -18,6 +18,17 @@ corrects affected uninterrupted legacy Diffusion trajectories; it does not change the composed `AdvDiffusion` or SLCN implementation. Small serial/MPI tests cover startup and established histories, orders 1-3 and varying steps. +The general hazard is a symbolic slot compiled while it holds zero: +simplification can remove it permanently from the kernel (for example, +`x**0` folds to one), so later ramping cannot recover the omitted term. + +PR #708 review follow-up adds a fresh-process disk regression. Disk snapshots +still skip live symbolic matrices; their replay limitation is tracked as a +strict expected failure, not claimed fixed by the in-memory correction. +Saving any unsupported state field now emits a warning naming the field. +Symbolic snapshot docstrings explicitly document shared live atom references +and their lack of isolation from subsequent mutation or disk portability. + ### The Multiplier Was Not the Whole Traction (August 2026) **`Stokes_Constrained.topography()` now returns the traction the boundary is diff --git a/src/underworld3/checkpoint/disk_snapshot.py b/src/underworld3/checkpoint/disk_snapshot.py index 9cbf7c22a..ea20d49b9 100644 --- a/src/underworld3/checkpoint/disk_snapshot.py +++ b/src/underworld3/checkpoint/disk_snapshot.py @@ -829,10 +829,22 @@ def _serialise_field(h5group, name: str, value: Any) -> None: f"unserialisable list (len={len(value)}, " f"first-type={type(value[0]).__name__ if value else 'empty'})" ) + warnings.warn( + f"Snapshot skipped state field {h5group.name}/{name}: " + f"{h5group.attrs[name + '__skipped']}. " + "Disk restart may not reproduce continuation.", + RuntimeWarning, stacklevel=2, + ) return h5group.attrs[name + "__skipped"] = ( f"unserialisable type {type(value).__name__}" ) + warnings.warn( + f"Snapshot skipped state field {h5group.name}/{name}: " + f"{h5group.attrs[name + '__skipped']}. " + "Disk restart may not reproduce continuation.", + RuntimeWarning, stacklevel=2, + ) def _group_to_dict(h5group) -> dict: diff --git a/src/underworld3/systems/ddt.py b/src/underworld3/systems/ddt.py index 25d9b799c..3fb7fd3b6 100644 --- a/src/underworld3/systems/ddt.py +++ b/src/underworld3/systems/ddt.py @@ -111,7 +111,12 @@ class DDtSymbolicState(_DDtCoreState): """Snapshot of a :class:`Symbolic` DDt instance's evolution state. ``Symbolic`` is the pure-symbolic flavor — ``psi_star`` history - slots hold sympy expressions (immutable), captured by value. + slots hold symbolic matrices whose containers are copied, but whose live + UWexpression atoms are retained by reference. This supports in-memory + backstepping, not isolation from subsequent changes to those atoms. + These references cannot be transported through disk snapshots; symbolic + history is currently skipped there. Reconstructing symbolic forms from + the solver would be needed for a general disk-restart guarantee. """ psi_star: list = field(default_factory=list) diff --git a/tests/parallel/ptest_1058_diffusion_disk_restart.py b/tests/parallel/ptest_1058_diffusion_disk_restart.py new file mode 100644 index 000000000..d6b333e61 --- /dev/null +++ b/tests/parallel/ptest_1058_diffusion_disk_restart.py @@ -0,0 +1,34 @@ +"""Fresh-process worker: write a symbolic snapshot or restore and continue.""" + +import sys +import numpy as np +import underworld3 as uw + +phase = sys.argv[1] +uw.reset_default_model() +model = uw.get_default_model() +mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0., 0.), maxCoords=(1., 1.), cellSize=0.3, +) +temperature = uw.discretisation.MeshVariable("T", mesh, 1, degree=1) +temperature.array[:, 0, 0] = temperature.coords[:, 0] +solver = uw.systems.Diffusion(mesh, temperature, order=2, theta=1.) +solver.constitutive_model = uw.constitutive_models.DiffusionModel +solver.constitutive_model.Parameters.diffusivity = 0.05 +solver.tolerance = 1e-12 +if phase == "write": + for _ in range(3): + solver.solve(timestep=0.01, zero_init_guess=False) + model.save_state(file="restart.h5") +else: + model.load_state("restart.h5") + +values = {f"restored_{name}": np.array(var.array) + for name, var in mesh.vars.items()} +values["dt_history"] = np.array(solver.DFDt._dt_history) +values["n_solves"] = solver.DFDt._n_solves_completed +values["initialized"] = solver.DFDt._history_initialised +for step, dt in enumerate((0.01, 0.015)): + solver.solve(timestep=dt, zero_init_guess=False) + values[f"T_{step}"] = np.array(temperature.array) +np.savez(f"{phase}_{uw.mpi.rank}.npz", **values) diff --git a/tests/test_0007_snapshot_inmemory.py b/tests/test_0007_snapshot_inmemory.py index 8dce22386..44398a574 100644 --- a/tests/test_0007_snapshot_inmemory.py +++ b/tests/test_0007_snapshot_inmemory.py @@ -436,8 +436,7 @@ def test_symbolic_flux_history_remains_valid_after_solver_restore(): temperature = uw.discretisation.MeshVariable( "T_symbolic_restore", mesh, 1, degree=1 ) - with mesh.access(temperature): - temperature.data[:, 0] = temperature.coords[:, 0] + temperature.array[:, 0, 0] = temperature.coords[:, 0] diffusion = uw.systems.Diffusion( mesh, u_Field=temperature, order=2, theta=1.0 @@ -451,7 +450,7 @@ def test_symbolic_flux_history_remains_valid_after_solver_restore(): model.load_state(snapshot) diffusion.solve(timestep=0.01, zero_init_guess=False) - assert np.all(np.isfinite(temperature.data)) + assert np.all(np.isfinite(temperature.array)) diff --git a/tests/test_1059_diffusion_disk_restart.py b/tests/test_1059_diffusion_disk_restart.py new file mode 100644 index 000000000..6a8e3ecc1 --- /dev/null +++ b/tests/test_1059_diffusion_disk_restart.py @@ -0,0 +1,71 @@ +"""Run parent in serial; UW_DIFFUSION_TEST_RANKS selects fresh worker ranks.""" + +import os +from pathlib import Path +import subprocess +import sys + +import h5py +import numpy as np +import pytest +import sympy +import underworld3 as uw +from parallel.serial_reference import _MPI_ENV_PREFIXES + +pytestmark = [pytest.mark.tier_b] + + +class SymbolicDiskContinuationMismatch(AssertionError): + """Known disk limitation; setup and exact-restore failures are not xfailed.""" + + +@pytest.mark.parametrize("value", [object(), [sympy.Matrix([1])]]) +@pytest.mark.level_1 +def test_skipped_snapshot_state_warns(tmp_path, value): + from underworld3.checkpoint.disk_snapshot import _serialise_field + + with h5py.File(tmp_path / "state.h5", "w") as f: + group = f.create_group("state") + with pytest.warns(RuntimeWarning, match="Snapshot skipped state field /state/history"): + _serialise_field(group, "history", value) + assert "history__skipped" in group.attrs + + +@pytest.mark.level_2 +@pytest.mark.xfail( + strict=True, raises=SymbolicDiskContinuationMismatch, + reason="PR #708: disk snapshots skip live symbolic flux history", +) +def test_symbolic_diffusion_fresh_process_disk_restart(tmp_path): + if uw.mpi.size != 1: + pytest.skip("Run parent in serial; UW_DIFFUSION_TEST_RANKS selects worker ranks.") + ranks = int(os.environ.get("UW_DIFFUSION_TEST_RANKS", "1")) + root = Path(__file__).resolve().parents[1] + worker = root / "tests/parallel/ptest_1058_diffusion_disk_restart.py" + env = {k: v for k, v in os.environ.items() + if not k.startswith(_MPI_ENV_PREFIXES)} + launcher = [] if ranks == 1 else [str(Path(sys.executable).with_name("mpirun")), "-np", str(ranks)] + for phase in ("write", "resume"): + command = [sys.executable, str(root / "scripts/mpi_supervisor.py"), + "--silence", "45", "--hard-cap", "90", "--", + *launcher, sys.executable, "-m", "mpi4py", str(worker), phase] + with (tmp_path / f"{phase}.log").open("w") as log: + result = subprocess.run(command, cwd=tmp_path, env=env, + stdout=log, stderr=subprocess.STDOUT, timeout=100) + assert result.returncode == 0, (tmp_path / f"{phase}.log").read_text() + assert "Snapshot skipped state field" in (tmp_path / "write.log").read_text() + maximum = 0.0 + for rank in range(ranks): + with np.load(tmp_path / f"write_{rank}.npz") as expected, np.load( + tmp_path / f"resume_{rank}.npz") as actual: + assert set(actual) == set(expected) + for name in expected: + if not name.startswith("T_"): + np.testing.assert_array_equal(actual[name], expected[name]) + for name in ("T_0", "T_1"): + maximum = max(maximum, float(np.max(np.abs(actual[name] - expected[name])))) + # All exact field/metadata restoration checks on every rank run before + # recognizing this known numerical limitation. Other failures stay failures. + print(f"DIFFUSION_DISK_REPLAY ranks={ranks} max_error={maximum:.12g}") + if maximum > 1e-11: + raise SymbolicDiskContinuationMismatch(f"Disk replay max difference {maximum:.12g}") From 71c8a12b9f9e076f7bb4436dd734947d50671f44 Mon Sep 17 00:00:00 2001 From: Tyagi Date: Fri, 11 Sep 2026 13:05:58 +1000 Subject: [PATCH 4/4] Use real checkpoint output for Diffusion restart validation in PR 708 Replace npz sidecar comparisons with flushed checkpoint wrapper and PETSc bulk-vector checks. Assert exact restored field layouts and values, saved history metadata, and explicit skipped symbolic state before checking continuation. Use uw.Params for worker phases. Assign unused 1074 and 1079 test numbers to remove unrelated 1058 collisions. Serial and eight-rank checkpoint workers reproduce the documented 0.00350097 expected failure; 33 in-memory tests pass. No solver or checkpoint implementation changes. --- ...y => ptest_1079_diffusion_disk_restart.py} | 19 ++- tests/test_1059_diffusion_disk_restart.py | 71 ----------- ...tart.py => test_1074_diffusion_restart.py} | 0 tests/test_1079_diffusion_disk_restart.py | 117 ++++++++++++++++++ 4 files changed, 125 insertions(+), 82 deletions(-) rename tests/parallel/{ptest_1058_diffusion_disk_restart.py => ptest_1079_diffusion_disk_restart.py} (66%) delete mode 100644 tests/test_1059_diffusion_disk_restart.py rename tests/{test_1058_diffusion_restart.py => test_1074_diffusion_restart.py} (100%) create mode 100644 tests/test_1079_diffusion_disk_restart.py diff --git a/tests/parallel/ptest_1058_diffusion_disk_restart.py b/tests/parallel/ptest_1079_diffusion_disk_restart.py similarity index 66% rename from tests/parallel/ptest_1058_diffusion_disk_restart.py rename to tests/parallel/ptest_1079_diffusion_disk_restart.py index d6b333e61..6b3e58501 100644 --- a/tests/parallel/ptest_1058_diffusion_disk_restart.py +++ b/tests/parallel/ptest_1079_diffusion_disk_restart.py @@ -1,10 +1,12 @@ """Fresh-process worker: write a symbolic snapshot or restore and continue.""" -import sys -import numpy as np import underworld3 as uw -phase = sys.argv[1] +params = uw.Params( + uw_phase=uw.Param("write", type=uw.ParamType.STRING, + description="Checkpoint phase: write or resume."), +) +assert params.uw_phase in ("write", "resume") uw.reset_default_model() model = uw.get_default_model() mesh = uw.meshing.UnstructuredSimplexBox( @@ -16,19 +18,14 @@ solver.constitutive_model = uw.constitutive_models.DiffusionModel solver.constitutive_model.Parameters.diffusivity = 0.05 solver.tolerance = 1e-12 -if phase == "write": +if params.uw_phase == "write": for _ in range(3): solver.solve(timestep=0.01, zero_init_guess=False) model.save_state(file="restart.h5") else: model.load_state("restart.h5") + model.save_state(file="restored.h5") -values = {f"restored_{name}": np.array(var.array) - for name, var in mesh.vars.items()} -values["dt_history"] = np.array(solver.DFDt._dt_history) -values["n_solves"] = solver.DFDt._n_solves_completed -values["initialized"] = solver.DFDt._history_initialised for step, dt in enumerate((0.01, 0.015)): solver.solve(timestep=dt, zero_init_guess=False) - values[f"T_{step}"] = np.array(temperature.array) -np.savez(f"{phase}_{uw.mpi.rank}.npz", **values) + model.save_state(file=f"{params.uw_phase}_{step}.h5") diff --git a/tests/test_1059_diffusion_disk_restart.py b/tests/test_1059_diffusion_disk_restart.py deleted file mode 100644 index 6a8e3ecc1..000000000 --- a/tests/test_1059_diffusion_disk_restart.py +++ /dev/null @@ -1,71 +0,0 @@ -"""Run parent in serial; UW_DIFFUSION_TEST_RANKS selects fresh worker ranks.""" - -import os -from pathlib import Path -import subprocess -import sys - -import h5py -import numpy as np -import pytest -import sympy -import underworld3 as uw -from parallel.serial_reference import _MPI_ENV_PREFIXES - -pytestmark = [pytest.mark.tier_b] - - -class SymbolicDiskContinuationMismatch(AssertionError): - """Known disk limitation; setup and exact-restore failures are not xfailed.""" - - -@pytest.mark.parametrize("value", [object(), [sympy.Matrix([1])]]) -@pytest.mark.level_1 -def test_skipped_snapshot_state_warns(tmp_path, value): - from underworld3.checkpoint.disk_snapshot import _serialise_field - - with h5py.File(tmp_path / "state.h5", "w") as f: - group = f.create_group("state") - with pytest.warns(RuntimeWarning, match="Snapshot skipped state field /state/history"): - _serialise_field(group, "history", value) - assert "history__skipped" in group.attrs - - -@pytest.mark.level_2 -@pytest.mark.xfail( - strict=True, raises=SymbolicDiskContinuationMismatch, - reason="PR #708: disk snapshots skip live symbolic flux history", -) -def test_symbolic_diffusion_fresh_process_disk_restart(tmp_path): - if uw.mpi.size != 1: - pytest.skip("Run parent in serial; UW_DIFFUSION_TEST_RANKS selects worker ranks.") - ranks = int(os.environ.get("UW_DIFFUSION_TEST_RANKS", "1")) - root = Path(__file__).resolve().parents[1] - worker = root / "tests/parallel/ptest_1058_diffusion_disk_restart.py" - env = {k: v for k, v in os.environ.items() - if not k.startswith(_MPI_ENV_PREFIXES)} - launcher = [] if ranks == 1 else [str(Path(sys.executable).with_name("mpirun")), "-np", str(ranks)] - for phase in ("write", "resume"): - command = [sys.executable, str(root / "scripts/mpi_supervisor.py"), - "--silence", "45", "--hard-cap", "90", "--", - *launcher, sys.executable, "-m", "mpi4py", str(worker), phase] - with (tmp_path / f"{phase}.log").open("w") as log: - result = subprocess.run(command, cwd=tmp_path, env=env, - stdout=log, stderr=subprocess.STDOUT, timeout=100) - assert result.returncode == 0, (tmp_path / f"{phase}.log").read_text() - assert "Snapshot skipped state field" in (tmp_path / "write.log").read_text() - maximum = 0.0 - for rank in range(ranks): - with np.load(tmp_path / f"write_{rank}.npz") as expected, np.load( - tmp_path / f"resume_{rank}.npz") as actual: - assert set(actual) == set(expected) - for name in expected: - if not name.startswith("T_"): - np.testing.assert_array_equal(actual[name], expected[name]) - for name in ("T_0", "T_1"): - maximum = max(maximum, float(np.max(np.abs(actual[name] - expected[name])))) - # All exact field/metadata restoration checks on every rank run before - # recognizing this known numerical limitation. Other failures stay failures. - print(f"DIFFUSION_DISK_REPLAY ranks={ranks} max_error={maximum:.12g}") - if maximum > 1e-11: - raise SymbolicDiskContinuationMismatch(f"Disk replay max difference {maximum:.12g}") diff --git a/tests/test_1058_diffusion_restart.py b/tests/test_1074_diffusion_restart.py similarity index 100% rename from tests/test_1058_diffusion_restart.py rename to tests/test_1074_diffusion_restart.py diff --git a/tests/test_1079_diffusion_disk_restart.py b/tests/test_1079_diffusion_disk_restart.py new file mode 100644 index 000000000..2bce0e560 --- /dev/null +++ b/tests/test_1079_diffusion_disk_restart.py @@ -0,0 +1,117 @@ +"""Run parent in serial; UW_DIFFUSION_TEST_RANKS selects fresh worker ranks.""" + +import os +from pathlib import Path +import subprocess +import sys + +import h5py +import numpy as np +import pytest +import sympy +import underworld3 as uw +from parallel.serial_reference import _MPI_ENV_PREFIXES + +pytestmark = [pytest.mark.tier_b] + + +class SymbolicDiskContinuationMismatch(AssertionError): + """Known disk limitation; setup and exact-restore failures are not xfailed.""" + + +def _checkpoint_fields(wrapper): + """Read the writer's flushed PETSc vectors and layouts, not sidecar logs.""" + fields = {} + with h5py.File(wrapper) as snapshot: + bulk = wrapper.parent / snapshot["meshes"].attrs["bulk_dir"] + for mesh_name, mesh in snapshot["meshes"].items(): + for name, variable in mesh["variables"].items(): + with h5py.File(bulk / variable.attrs["external_file"]) as data: + arrays = {} + data.visititems(lambda key, obj: arrays.update({key: obj[()]}) + if isinstance(obj, h5py.Dataset) else None) + assert arrays, f"Missing checkpoint datasets for {name}" + assert f"uw_checkpoint/{name}" in arrays + fields[mesh_name, name] = arrays + assert fields + return fields + + +def _assert_saved_metadata_equal(reference, restored): + with h5py.File(reference) as expected, h5py.File(restored) as actual: + assert set(expected["python_state"]) == set(actual["python_state"]) + for name, group in expected["python_state"].items(): + other = actual["python_state"][name] + assert set(group.attrs) == set(other.attrs) + for key in group.attrs: + np.testing.assert_array_equal(other.attrs[key], group.attrs[key]) + symbolic = [group for name, group in expected["python_state"].items() + if name.startswith("Symbolic_")] + assert len(symbolic) == 1 + assert "psi_star__skipped" in symbolic[0].attrs + assert "MutableDenseMatrix" in symbolic[0].attrs["psi_star__skipped"] + + +@pytest.mark.parametrize("value", [object(), [sympy.Matrix([1])]]) +@pytest.mark.level_1 +def test_skipped_snapshot_state_warns(tmp_path, value): + from underworld3.checkpoint.disk_snapshot import _serialise_field + + with h5py.File(tmp_path / "state.h5", "w") as f: + group = f.create_group("state") + with pytest.warns(RuntimeWarning, match="Snapshot skipped state field /state/history"): + _serialise_field(group, "history", value) + assert "history__skipped" in group.attrs + + +@pytest.mark.level_2 +@pytest.mark.xfail( + strict=True, raises=SymbolicDiskContinuationMismatch, + reason="PR #708: disk snapshots skip live symbolic flux history", +) +def test_symbolic_diffusion_fresh_process_disk_restart(tmp_path): + if uw.mpi.size != 1: + pytest.skip("Run parent in serial; UW_DIFFUSION_TEST_RANKS selects worker ranks.") + ranks = int(os.environ.get("UW_DIFFUSION_TEST_RANKS", "1")) + root = Path(__file__).resolve().parents[1] + worker = root / "tests/parallel/ptest_1079_diffusion_disk_restart.py" + env = {k: v for k, v in os.environ.items() + if not k.startswith(_MPI_ENV_PREFIXES)} + launcher = [] if ranks == 1 else [str(Path(sys.executable).with_name("mpirun")), "-np", str(ranks)] + for phase in ("write", "resume"): + command = [sys.executable, str(root / "scripts/mpi_supervisor.py"), + "--silence", "45", "--hard-cap", "90", "--", + *launcher, sys.executable, "-m", "mpi4py", str(worker), "-uw_phase", phase] + with (tmp_path / f"{phase}.log").open("w") as log: + result = subprocess.run(command, cwd=tmp_path, env=env, + stdout=log, stderr=subprocess.STDOUT, timeout=100) + assert result.returncode == 0, (tmp_path / f"{phase}.log").read_text() + assert "Snapshot skipped state field" in (tmp_path / "write.log").read_text() + _assert_saved_metadata_equal(tmp_path / "restart.h5", tmp_path / "restored.h5") + expected = _checkpoint_fields(tmp_path / "restart.h5") + actual = _checkpoint_fields(tmp_path / "restored.h5") + assert set(actual) == set(expected) + for field in expected: + assert set(actual[field]) == set(expected[field]) + for dataset in expected[field]: + np.testing.assert_array_equal(actual[field][dataset], expected[field][dataset]) + + maximum = 0.0 + for step in range(2): + expected = _checkpoint_fields(tmp_path / f"write_{step}.h5") + actual = _checkpoint_fields(tmp_path / f"resume_{step}.h5") + assert set(actual) == set(expected) + for field in expected: + assert set(actual[field]) == set(expected[field]) + for dataset, values in expected[field].items(): + if "/vecs/" in dataset or dataset.startswith("uw_checkpoint/"): + assert actual[field][dataset].shape == values.shape + assert np.all(np.isfinite(actual[field][dataset])) + maximum = max(maximum, float(np.max(np.abs(actual[field][dataset] - values)))) + else: + np.testing.assert_array_equal(actual[field][dataset], values) + # Exact checkpoint restoration and layout checks precede the known + # numerical limitation. Missing state or corrupt layouts fail normally. + print(f"DIFFUSION_DISK_REPLAY ranks={ranks} max_error={maximum:.12g}") + if maximum > 1e-11: + raise SymbolicDiskContinuationMismatch(f"Disk replay max difference {maximum:.12g}")