diff --git a/docs/developer/CHANGELOG.md b/docs/developer/CHANGELOG.md index 5aaa56964..4aad91431 100644 --- a/docs/developer/CHANGELOG.md +++ b/docs/developer/CHANGELOG.md @@ -6,6 +6,29 @@ 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 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 180e43463..3fb7fd3b6 100644 --- a/src/underworld3/systems/ddt.py +++ b/src/underworld3/systems/ddt.py @@ -111,11 +111,40 @@ 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) + 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/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/parallel/ptest_1079_diffusion_disk_restart.py b/tests/parallel/ptest_1079_diffusion_disk_restart.py new file mode 100644 index 000000000..6b3e58501 --- /dev/null +++ b/tests/parallel/ptest_1079_diffusion_disk_restart.py @@ -0,0 +1,31 @@ +"""Fresh-process worker: write a symbolic snapshot or restore and continue.""" + +import underworld3 as uw + +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( + 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 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") + +for step, dt in enumerate((0.01, 0.015)): + solver.solve(timestep=dt, zero_init_guess=False) + model.save_state(file=f"{params.uw_phase}_{step}.h5") diff --git a/tests/test_0007_snapshot_inmemory.py b/tests/test_0007_snapshot_inmemory.py index d52fee33a..44398a574 100644 --- a/tests/test_0007_snapshot_inmemory.py +++ b/tests/test_0007_snapshot_inmemory.py @@ -402,6 +402,58 @@ 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 + ) + temperature.array[:, 0, 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.array)) + + + def test_semilagrangian_ddt_roundtrip(): import underworld3 as uw from underworld3.systems.ddt import DDtSemiLagrangianState @@ -766,4 +818,3 @@ def test_continuation_bit_identical_across_stash_and_recover(): _assert_bit_identical(ctrl, stash, "stash-and-recover") - diff --git a/tests/test_1074_diffusion_restart.py b/tests/test_1074_diffusion_restart.py new file mode 100644 index 000000000..97a02afae --- /dev/null +++ b/tests/test_1074_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}") 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}")