Fix invalid symbolic atoms in DDt snapshot copies

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

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:

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):

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)

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")

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")
 
-

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}")

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}")