From 1b97569469e01bbc1a2bdcd282a6ccb82f5d8d6c Mon Sep 17 00:00:00 2001 From: lmoresi Date: Tue, 22 Sep 2026 06:56:09 +1000 Subject: [PATCH 1/2] fix: capture expression contents in the snapshot, so parameters rewind with the fields (#774) snapshot() captured each 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 was not captured at all. A restore therefore put the FIELDS back and left every parameter wherever the run had since moved it, with nothing raised. A replayed step then solved a different problem from the one the transcript recorded. Ramping a parameter between steps - kappa.sym = 7.0 - is the ordinary way to write such a run, so this was not an exotic path: after rewind to step 1: kappa=7.0 (should be 1.0) T=1.0 (correct) Every live expression's contents are now captured alongside the fields and put back on restore. Two details that decide whether this is correct: * stored by REFERENCE, not deep-copied. _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. A test asserts the restored value keeps its coordinate symbols. * 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 directly above, which raises - a missing state-bearer means the snapshot came from a different Model, whereas expressions are made and dropped freely during a normal run. Scope is wider than rewind(): checkpoint/snapshot.py is on development and only rewind() is confined to the timestepping branch, so save_state()/load_state() users have the same gap today. It also reaches the discrete adjoint, whose driver restores a step's snapshot and replays the solve - replaying at the wrong parameter values gives a gradient for a problem the run never solved, and a Taylor test would not catch it because both sides move together. Verified by negative control: with the restore disabled, the two contract tests fail and the other 25 pass. Noted, 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 capture registry is a WeakSet and so adds no lifetime, but it is not what bounds the set today - the docstring says so rather than claiming a weakness the code does not currently have. Worth its own issue. tests/test_0007_snapshot_inmemory.py 27 passed; tests/test_00[0-4]*py 145 passed. Underworld development team with AI support from Claude Code --- src/underworld3/checkpoint/snapshot.py | 71 ++++++++++++++++++ src/underworld3/function/expressions.py | 39 ++++++++++ tests/test_0007_snapshot_inmemory.py | 95 +++++++++++++++++++++++++ 3 files changed, 205 insertions(+) diff --git a/src/underworld3/checkpoint/snapshot.py b/src/underworld3/checkpoint/snapshot.py index eff45baca..878280caa 100644 --- a/src/underworld3/checkpoint/snapshot.py +++ b/src/underworld3/checkpoint/snapshot.py @@ -134,6 +134,11 @@ class Snapshot: # state-bearers. List preserves capture order — informational only, # since lookup is by key. state_bearers: list = field(default_factory=list) + # Expression captures: list of (stable_key, _sym, _wrapped). A parameter is + # not a mesh, a swarm or a state-bearer, so before this it was not captured + # at all and a restore returned the fields of one step with the parameters of + # another, silently. See _capture_expressions. + expressions: list = field(default_factory=list) metadata: dict[str, Any] = field(default_factory=dict) @@ -194,6 +199,7 @@ def snapshot(model, *, path: Optional[str] = None) -> Snapshot: _capture_swarm(snap, swarm) for obj in list(model._state_bearers): _capture_state_bearer(snap, obj) + _capture_expressions(snap) return snap @@ -243,6 +249,69 @@ def _capture_state_bearer(snap: Snapshot, obj) -> None: snap.state_bearers.append((_state_bearer_key(obj), copy.deepcopy(state))) +def _expression_key(expr) -> str: + """Stable per-process key for an expression, matching the state-bearer + convention: the class name and ``uw_object.instance_number``.""" + return f"{type(expr).__name__}_{expr.instance_number}" + + +def _capture_expressions(snap: Snapshot) -> None: + """Capture every live expression's contents. + + A run that ramps a parameter (``kappa.sym = 7.0`` between steps) used to be + restorable in its FIELDS only: the mesh variables came back and the + parameters stayed wherever the run had left them, so a replayed step solved + a different problem from the one the transcript recorded — with no error, in + the one place a wrong answer is hardest to notice. + + Contents are stored by REFERENCE, not deep-copied. ``_sym`` is a sympy + object and sympy objects are immutable, so the reference cannot go stale + under a later ``sym =`` assignment: that assignment REPLACES the object + rather than mutating it. Deep-copying instead would be actively wrong here — + an expression's ``_sym`` can carry mesh-variable symbols, and cloning that + graph would detach the restored parameter from the live mesh. + """ + from underworld3.function.expressions import live_expressions + + for expr in live_expressions(): + try: + snap.expressions.append( + (_expression_key(expr), expr._sym, expr._wrapped) + ) + except AttributeError: + # Not every expression subclass carries both slots; one that does + # not is one whose value is derived, and deriving it again is right. + continue + + +def _restore_expressions(snap: Snapshot) -> None: + """Put captured expression contents back. + + Matched by key against the expressions alive NOW. Two cases are deliberately + quiet rather than fatal, because neither means the restore is wrong: + + * captured but no longer alive — the expression was dropped since; there is + nothing to write to; + * alive but not captured — it was created after the snapshot, so the + snapshot has no opinion about what it should hold. + + Both differ from the state-bearer rule above, which raises: a missing + state-bearer means the snapshot came from a different Model, whereas + expressions are created and dropped freely throughout a normal run. + """ + if not snap.expressions: + return + from underworld3.function.expressions import live_expressions + + live = {_expression_key(e): e for e in live_expressions()} + for key, captured_sym, captured_wrapped in snap.expressions: + expr = live.get(key) + if expr is None: + continue + expr._sym = captured_sym + expr._wrapped = captured_wrapped + + def _capture_swarm(snap: Snapshot, swarm) -> None: payload = swarm.snapshot_payload() name = payload["name"] @@ -352,6 +421,8 @@ def restore(model, snap: Snapshot) -> None: ) obj.state = copy.deepcopy(captured_state) + _restore_expressions(snap) + def _build_mesh_payload(snap: Snapshot, mesh_name: str) -> dict: return { diff --git a/src/underworld3/function/expressions.py b/src/underworld3/function/expressions.py index e34c094b1..d85ba547d 100644 --- a/src/underworld3/function/expressions.py +++ b/src/underworld3/function/expressions.py @@ -14,6 +14,7 @@ - UWexpression.to() simply calls uw.convert_units(self, target) """ +import weakref import sympy import numpy as np from sympy import Symbol, simplify, Number @@ -602,6 +603,35 @@ def substitute_expr(fn, sub_expr, keep_constants=True, return_self=True): # UWexpression Class - Simplified (no UWQuantity inheritance) # ============================================================================ +# Every live UWexpression, so a snapshot can capture what the parameters WERE. +# +# Field data rewinds and parameters did not: `snapshot()` captured meshes, +# swarms and registered state-bearers, and an expression is none of those. A run +# that ramps a parameter — `kappa.sym = 7.0` between steps, which is the ordinary +# way to do it — could be rewound to an earlier step and come back with the +# fields of that step and the PARAMETERS OF THE LATEST ONE, silently. Replaying +# the step then solves a different problem from the one the transcript records. +# +# WeakSet so this registry never becomes the thing that keeps an expression +# alive: transient expressions are made constantly (unit wrapping, template +# substitution, derivative lowering). Note that it is NOT what bounds the set +# today — expressions are already held strongly elsewhere in UW3 by a +# name-keyed cache, so dropping the last user reference does not collect one. +# The WeakSet is therefore defensive rather than load-bearing: it means capture +# does not ADD a lifetime, and it starts working the day that cache is fixed. +_LIVE_EXPRESSIONS: "weakref.WeakSet" = weakref.WeakSet() + + +def live_expressions(): + """Every UWexpression still alive, as a list with a stable order. + + A list, not the WeakSet: entries can vanish mid-iteration as they are + collected, and capture needs a traversal that cannot change under it. + Sorted by instance number so a capture reads the same way twice. + """ + return sorted(_LIVE_EXPRESSIONS, key=lambda e: getattr(e, "instance_number", 0)) + + class UWexpression(MathematicalMixin, uw_object, Symbol): """ A SymPy Symbol that wraps a value for lazy evaluation. @@ -883,6 +913,15 @@ def __init__( self._uw_id = uw_object._obj_count uw_object._obj_count += 1 + # Visible to snapshot capture, so a rewind puts parameters back with the + # fields. Registered last: a half-built expression must never be + # captured. Failure here must not take out construction — an expression + # that cannot be registered is merely one a snapshot will not restore. + try: + _LIVE_EXPRESSIONS.add(self) + except TypeError: + pass + # ========================================================================= # Core Properties # ========================================================================= diff --git a/tests/test_0007_snapshot_inmemory.py b/tests/test_0007_snapshot_inmemory.py index d52fee33a..91278339f 100644 --- a/tests/test_0007_snapshot_inmemory.py +++ b/tests/test_0007_snapshot_inmemory.py @@ -767,3 +767,98 @@ def test_continuation_bit_identical_across_stash_and_recover(): _assert_bit_identical(ctrl, stash, "stash-and-recover") + + +# --------------------------------------------------------------------------- # +# Expressions: parameters must rewind with the fields +# --------------------------------------------------------------------------- # +# A snapshot captured meshes, swarms and registered state-bearers. A parameter +# is none of those, so it was not captured at all: restore returned the FIELDS +# of the captured step and left every `uw.expression` wherever the run had since +# moved it. A replayed step then solved a different problem from the one the +# transcript recorded, with no error raised — and ramping a parameter between +# steps (`kappa.sym = ...`) is the ordinary way to write such a run. + + +def test_expression_contents_are_restored_with_the_fields(): + """The contract. Scribble on a field AND a parameter; both come back. + + The field half already held; it is asserted alongside so a regression that + breaks both cannot pass by breaking them symmetrically. + """ + import sympy + + uw, model, mesh = _fresh_model_and_mesh() + T = uw.discretisation.MeshVariable("T_expr_snap", mesh, 1, degree=2) + kappa = uw.expression(r"\kappa", 1.0, "diffusivity") + + T.array[...] = 3.0 + snap = model.save_state() + + kappa.sym = sympy.Float(7.0) + T.array[...] = -42.0 + + model.load_state(snap) + + assert float(kappa.sym) == pytest.approx(1.0), ( + f"the parameter did not rewind with the field; kappa is {kappa.sym}") + assert np.allclose(np.asarray(T.array[...]), 3.0) + + +def test_a_symbolic_expression_value_is_restored_not_just_a_number(): + """A parameter's contents can be an expression, not only a float. + + Restoring must put the whole symbolic value back. Stored by reference + rather than deep-copied, because `_sym` can carry mesh-variable symbols and + cloning that graph would detach the restored parameter from the live mesh — + so this also checks the reference has not gone stale. + """ + import sympy + + uw, model, mesh = _fresh_model_and_mesh() + x, y = mesh.X + eta = uw.expression(r"\eta", sympy.sympify(1), "viscosity") + + eta.sym = 2 + x * y + captured = eta.sym + snap = model.save_state() + + eta.sym = sympy.Float(99.0) + model.load_state(snap) + + assert eta.sym == captured + assert eta.sym.free_symbols == captured.free_symbols, ( + "the restored value lost its coordinate symbols — it was cloned rather " + "than referenced") + + +def test_an_expression_created_after_the_snapshot_is_left_alone(): + """A snapshot has no opinion about something that did not exist when it was + taken. Leaving it is right; raising would break the ordinary case of a run + that builds a new parameter after a restore point.""" + uw, model, mesh = _fresh_model_and_mesh() + snap = model.save_state() + + latecomer = uw.expression(r"\alpha", 5.0, "made after the snapshot") + model.load_state(snap) + + assert float(latecomer.sym) == pytest.approx(5.0) + + +def test_capture_covers_every_live_expression(): + """Capture is driven by the registry, so the registry must see expressions + made anywhere — including the ones UW3 builds for itself (strain-rate + invariants, unit wrappers), not just the ones a user names.""" + import sympy + from underworld3.function.expressions import live_expressions + + uw, model, mesh = _fresh_model_and_mesh() + before = {e.instance_number for e in live_expressions()} + mine = uw.expression(r"\beta", 2.0, "made here") + after = {e.instance_number for e in live_expressions()} + + assert mine.instance_number in after - before + + snap = model.save_state() + captured = {key for key, _sym, _wrapped in snap.expressions} + assert f"{type(mine).__name__}_{mine.instance_number}" in captured From 30b21300a4eb13dbe0c9d3f0e673d06c91ed05b5 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Tue, 22 Sep 2026 07:05:33 +1000 Subject: [PATCH 2/2] Capture from the existing container registry, not a parallel one The first version added a WeakSet of every UWexpression. That was a second registry for a job the first one already does, and it described expression lifetime wrongly in the bargain. UWexpression._expr_names IS the definition of a UW expression: a container with identity by NAME, looked up rather than rebuilt, so uw.expression(r"\eta", ...) reaches the same object twice and a formula written early keeps seeing later edits to its contents. The strong hold is the point, not a leak. That registry is exactly the set whose contents a snapshot must capture - the parameters of the run - so capture now reads it directly. _ephemeral_expr_names is deliberately excluded. Those are the _unique_name_generation=True expressions built for derivative lowering and template substitution: contents derived, held weakly, keyed (name, uw_id), and rebuilt from the persistent ones. Restoring them would write over a value the machinery is about to recompute. Two tests follow the corrected model: capture reads the container registry and skips the ephemerals, and a name reused twice is ONE container captured ONCE - which is the identity rule the design rests on. tests/test_0007_snapshot_inmemory.py 28 passed; tests/test_00[0-4]*py 146 passed. Underworld development team with AI support from Claude Code --- src/underworld3/function/expressions.py | 53 +++++++++---------------- tests/test_0007_snapshot_inmemory.py | 35 +++++++++++----- 2 files changed, 45 insertions(+), 43 deletions(-) diff --git a/src/underworld3/function/expressions.py b/src/underworld3/function/expressions.py index d85ba547d..6ca0c29c1 100644 --- a/src/underworld3/function/expressions.py +++ b/src/underworld3/function/expressions.py @@ -603,33 +603,27 @@ def substitute_expr(fn, sub_expr, keep_constants=True, return_self=True): # UWexpression Class - Simplified (no UWQuantity inheritance) # ============================================================================ -# Every live UWexpression, so a snapshot can capture what the parameters WERE. -# -# Field data rewinds and parameters did not: `snapshot()` captured meshes, -# swarms and registered state-bearers, and an expression is none of those. A run -# that ramps a parameter — `kappa.sym = 7.0` between steps, which is the ordinary -# way to do it — could be rewound to an earlier step and come back with the -# fields of that step and the PARAMETERS OF THE LATEST ONE, silently. Replaying -# the step then solves a different problem from the one the transcript records. -# -# WeakSet so this registry never becomes the thing that keeps an expression -# alive: transient expressions are made constantly (unit wrapping, template -# substitution, derivative lowering). Note that it is NOT what bounds the set -# today — expressions are already held strongly elsewhere in UW3 by a -# name-keyed cache, so dropping the last user reference does not collect one. -# The WeakSet is therefore defensive rather than load-bearing: it means capture -# does not ADD a lifetime, and it starts working the day that cache is fixed. -_LIVE_EXPRESSIONS: "weakref.WeakSet" = weakref.WeakSet() - - def live_expressions(): - """Every UWexpression still alive, as a list with a stable order. - - A list, not the WeakSet: entries can vanish mid-iteration as they are - collected, and capture needs a traversal that cannot change under it. - Sorted by instance number so a capture reads the same way twice. + """The persistent expression containers, in a stable order. + + Reads ``UWexpression._expr_names``, which is the registry that already + defines what a UW expression IS: a container with identity by name, looked + up rather than rebuilt, so the same ``uw.expression(r"\\eta", ...)`` reaches + the same object and a formula written against it keeps seeing later edits to + its contents. That is exactly the set whose CONTENTS a snapshot has to + capture — the parameters of the run. + + ``_ephemeral_expr_names`` is deliberately NOT included. Those are the + ``_unique_name_generation=True`` expressions made for derivative lowering + and template substitution: their contents are derived, they are held weakly, + and they are rebuilt from the persistent ones. Restoring them would write + over a value the machinery is about to recompute. + + A list, not the live dict: capture must not iterate a mapping that other + construction can mutate underneath it. Sorted by name so two captures of the + same state read the same way. """ - return sorted(_LIVE_EXPRESSIONS, key=lambda e: getattr(e, "instance_number", 0)) + return [UWexpression._expr_names[k] for k in sorted(UWexpression._expr_names)] class UWexpression(MathematicalMixin, uw_object, Symbol): @@ -913,15 +907,6 @@ def __init__( self._uw_id = uw_object._obj_count uw_object._obj_count += 1 - # Visible to snapshot capture, so a rewind puts parameters back with the - # fields. Registered last: a half-built expression must never be - # captured. Failure here must not take out construction — an expression - # that cannot be registered is merely one a snapshot will not restore. - try: - _LIVE_EXPRESSIONS.add(self) - except TypeError: - pass - # ========================================================================= # Core Properties # ========================================================================= diff --git a/tests/test_0007_snapshot_inmemory.py b/tests/test_0007_snapshot_inmemory.py index 91278339f..68af29ed4 100644 --- a/tests/test_0007_snapshot_inmemory.py +++ b/tests/test_0007_snapshot_inmemory.py @@ -845,20 +845,37 @@ def test_an_expression_created_after_the_snapshot_is_left_alone(): assert float(latecomer.sym) == pytest.approx(5.0) -def test_capture_covers_every_live_expression(): - """Capture is driven by the registry, so the registry must see expressions - made anywhere — including the ones UW3 builds for itself (strain-rate - invariants, unit wrappers), not just the ones a user names.""" - import sympy - from underworld3.function.expressions import live_expressions +def test_capture_reads_the_persistent_container_registry(): + """Capture is driven by the registry that defines what an expression IS — + ``_expr_names``, the by-name container store — so anything a user can reach + later by name is captured, and the ephemeral derivative/template expressions + are not.""" + from underworld3.function.expressions import UWexpression, live_expressions uw, model, mesh = _fresh_model_and_mesh() - before = {e.instance_number for e in live_expressions()} mine = uw.expression(r"\beta", 2.0, "made here") - after = {e.instance_number for e in live_expressions()} - assert mine.instance_number in after - before + assert any(e is mine for e in live_expressions()) snap = model.save_state() captured = {key for key, _sym, _wrapped in snap.expressions} assert f"{type(mine).__name__}_{mine.instance_number}" in captured + + # ephemerals are held in their own weak registry and deliberately skipped + ephemeral_names = {k[0] for k in list(UWexpression._ephemeral_expr_names)} + persistent_names = set(UWexpression._expr_names) + assert not (ephemeral_names & persistent_names & {r"\beta"}) + + +def test_a_reused_name_is_one_container_and_is_captured_once(): + """Identity is the NAME: asking for the same name returns the same object, + which is what lets a formula written early keep seeing later edits. Capture + must therefore record it once, not once per construction site.""" + uw, model, mesh = _fresh_model_and_mesh() + first = uw.expression(r"\gamma_{shared}", 1.0, "first use") + second = uw.expression(r"\gamma_{shared}", 2.0, "second use") + assert first is second + + snap = model.save_state() + key = f"{type(first).__name__}_{first.instance_number}" + assert [k for k, _s, _w in snap.expressions].count(key) == 1