Snapshot/restore does not capture expression contents: a rewind returns the fields of one step with the parameters of another #774

snapshot() captures each registered mesh's coordinates and mesh-variable DOFs, each swarm's particles, and every registered state-bearer. A uw.expression is none of those, so it is not captured at all.

A restore therefore puts the fields back and leaves every parameter wherever the run has since moved it. Nothing is raised. A replayed step solves a different problem from the one the transcript records.

Ramping a parameter between steps is the ordinary way to write such a run, so this is not an exotic path.

Reproduction

model.record_every = 1
kappa = uw.expression(r"\kappa", 1.0, "diffusivity")

T.array[...] = 1.0
with model.step(0.1): pass          # step 1 at kappa = 1

kappa.sym = sympy.Float(7.0)
T.array[...] = 2.0
with model.step(0.1): pass          # step 2 at kappa = 7

model.rewind(1)                     # back to the start of step 2
model.rewind(1)                     # back to the start of step 1

Observed:

after rewind to step 1: kappa=7.0 (should be 1.0)   T=1.0 (correct)

T rewinds. kappa does not.

Scope

Wider than rewind(). src/underworld3/checkpoint/snapshot.py is on development; only rewind() is confined to the docs/timestepping-pattern branch. Anyone using model.save_state() / model.load_state() on development today has the same gap — the field half of the state round-trips and the parameter half does not.

It also reaches the discrete adjoint, whose driver restores a step's snapshot and replays the solve: replaying with the wrong parameter values gives a gradient for a problem the run never solved, and the Taylor test would not catch it because both sides move together.

Fix

Capture every live expression's contents alongside the fields, and restore them.

Two details that matter:

• Store by reference, not deep copy. _sym is a sympy object and sympy objects are immutable, so a later sym = assignment replaces rather than mutates and the reference cannot go stale. Deep-copying would be actively wrong: an expression's _sym can carry mesh-variable symbols, and cloning that graph detaches the restored parameter from the live mesh.
• Missing on either side is quiet, not fatal. Captured-but-collected has nothing to write to; alive-but-not-captured was created after the snapshot, so the snapshot has no opinion about it. That differs from the state-bearer rule, which raises — a missing state-bearer means the snapshot came from a different Model, whereas expressions are created and dropped freely during a normal run.

Noted while fixing, not fixed

Expressions are already held strongly by a name-keyed cache elsewhere in UW3, so dropping the last user reference does not collect one. The new capture registry is a WeakSet and so does not add a lifetime, but it is not what bounds the set today. Worth its own issue: it means every expression ever constructed — including the transient ones from unit wrapping, template substitution and derivative lowering — stays alive for the length of the run.

Underworld development team with AI support from Claude Code