From 44f3615d2123c87cd5e541c55fa8a2844b54426c Mon Sep 17 00:00:00 2001 From: lmoresi Date: Tue, 15 Sep 2026 10:07:59 +1000 Subject: [PATCH 01/12] =?UTF-8?q?docs:=20the=20timestepping=20pattern=20?= =?UTF-8?q?=E2=80=94=20the=20clock=20lives=20on=20model.tracker,=20units?= =?UTF-8?q?=20are=20chosen=20before=20the=20mesh?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every example hand-rolled its own `t += dt` in a loose variable that no snapshot could see, and set reference quantities after the mesh existed, when it was too late for them to take effect. The guide states the pattern once: the clock on `model.tracker`, reference quantities before the mesh, and why each of those is the only place the choice can be made. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- .../guides/HOW-TO-WRITE-UW3-SCRIPTS.md | 151 +++++++++++++++++- tests/test_0009_model_tracker.py | 86 ++++++++++ 2 files changed, 234 insertions(+), 3 deletions(-) diff --git a/docs/developer/guides/HOW-TO-WRITE-UW3-SCRIPTS.md b/docs/developer/guides/HOW-TO-WRITE-UW3-SCRIPTS.md index 1646bd499..5525dade1 100644 --- a/docs/developer/guides/HOW-TO-WRITE-UW3-SCRIPTS.md +++ b/docs/developer/guides/HOW-TO-WRITE-UW3-SCRIPTS.md @@ -14,9 +14,10 @@ This guide captures critical lessons learned from writing and debugging Underwor 3. [Units System Integration](#units-system-integration) 4. [Mesh and Variable Creation](#mesh-and-variable-creation) 5. [Solver Setup and Execution](#solver-setup-and-execution) -6. [Common Pitfalls and Anti-Patterns](#common-pitfalls-and-anti-patterns) -7. [Testing Best Practices](#testing-best-practices) -8. [Debugging Techniques](#debugging-techniques) +6. [The Timestepping Pattern](#the-timestepping-pattern) +7. [Common Pitfalls and Anti-Patterns](#common-pitfalls-and-anti-patterns) +8. [Testing Best Practices](#testing-best-practices) +9. [Debugging Techniques](#debugging-techniques) --- @@ -414,6 +415,135 @@ model = stokes.constitutive_model # Confusing with uw.Model --- +## The Timestepping Pattern + +### ⚠️ RULE: the clock lives on `model.tracker`, never in a loose variable + +Every time-dependent script needs a clock. Declare the model and its reference +quantities first (see [RULE #2](#rule-2-reference-quantities-before-mesh-creation) +— they must precede mesh creation), then keep the clock on the tracker: + +```python +uw.reset_default_model() +model = uw.get_default_model() +model.set_reference_quantities( + domain_depth=uw.quantity(500, "km"), + material_density=uw.quantity(3300, "kg/m**3"), + material_viscosity=uw.quantity(1e21, "Pa*s"), +) + +mesh = uw.meshing.UnstructuredSimplexBox(...) # inherits the reference quantities +# ... variables, solvers ... + +model.tracker.time = uw.quantity(0.0, "Myr") +model.tracker.step = 0 +model.tracker.dt = None + +while model.tracker.time < end_time: + dt = adv_diff.estimate_dt() # a dimensional quantity when units are active + + adv_diff.solve(timestep=dt) + stokes.solve(zero_init_guess=False) + + model.tracker.time = model.tracker.time + dt + model.tracker.step += 1 + model.tracker.dt = dt +``` + +**Start from the model, not from the mesh.** A default model is created for you +if you never ask for one, so it is easy to write a whole script without noticing +it exists — and then reference quantities, which must be set before the mesh, are +already too late. Declaring the model on the first line makes the units decision +explicit at the only point where it can still be made. `estimate_dt()` then +returns a dimensional quantity and the clock carries units with no extra work. + +Neither the model nor the units are enforced today; both arrived after much of +the surrounding code. Treat this ordering as the pattern regardless, because +retrofitting units to a script written without them means rebuilding the mesh. + +**Why this and not `t = 0.0; t += dt`.** The tracker is captured by +`model.save_state()` and reverted by `model.load_state()`. A loose Python +variable is not — the tracker's own docstring says so: + +> Everything on the tracker is captured by `snapshot` and reverted by +> `restore`; loose Python variables are not. + +So a script that keeps its clock in a local and then backsteps gets its +**fields** restored and its **time** left in the future. Nothing raises; the +run simply carries a clock that disagrees with the state it is describing, and +every output written from that point is mislabelled. Measured on a real +five-step run: after restoring a snapshot taken at t = 1140, a loose `t` still +read 2280 while the fields were correctly back at 1140. + +This holds for both snapshot flavours — the in-memory token and the on-disk +snapshot used for restart. + +`time`, `step` and `dt` are pre-seeded on the tracker as a convention. Anything +else you assign to it (`model.tracker.rms_velocity = ...`) is captured and +restored the same way, so a diagnostic you want to survive a backstep belongs +there too. + +### Backstepping + +The pattern above is what makes speculative stepping safe: + +```python +snap = model.save_state() # before the step, not after + +dt = big_dt +adv_diff.solve(timestep=dt) +stokes.solve(zero_init_guess=False) + +if courant_number() > courant_limit: + model.load_state(snap) # fields AND clock go back together + for _ in range(n_substeps): + ... # replay with smaller steps +else: + model.tracker.time += dt + model.tracker.step += 1 +``` + +Take the snapshot **before** the operator, not after. A `DDt` history plugin +shifts its history in its post-solve hook, so a snapshot taken after a solve +holds the shifted history, which is not the state that step ran from. + +Restoring a snapshot is bit-exact and repeatable. Re-running the same script is +not: warm starts and preconditioner reuse are solver history that is not part +of model state, so two independent runs of the same problem on the same solver +objects diverge at the 1e-13 level from the first step. If you need to look at a +step twice, restore it rather than re-run it. + +### Known gap: a dimensional clock does not survive a restart + +An in-memory snapshot round-trips a units-carrying tracker entry correctly. The +**on-disk** snapshot does not: a `pint` quantity falls through to the +"unserialisable type" branch, is recorded in the file as +`__skipped` and is simply **absent** after `load_state`, so reading +`model.tracker.time` afterwards raises. Nothing warns at save time. Plain +floats, ints and numpy arrays are unaffected. + +Until that is fixed, a script that needs to restart from disk should keep the +clock non-dimensional, or re-establish it explicitly after loading: + +```python +model.load_state(path) +model.tracker.time = uw.quantity(model.tracker.time_Myr, "Myr") # stored as a float +``` + +### Known gap: `mesh.t` is not this clock + +`mesh.t` is a separate, symbolic time atom bound to PETSc's `petsc_t`. The +high-level `solve()` wrappers never set it, so **an expression containing +`mesh.t` evaluates to zero inside a solve**, silently. A time-dependent +boundary condition written as `sympy.sin(omega * mesh.t)` is identically zero +and nothing warns. `solve(time=...)` is accepted and ignored. + +Until `mesh.t` is wired to the model clock, build time dependence from a +`uw.function.expression` you update yourself each step, and drive it from +`model.tracker.time`. + +--- + ## Common Pitfalls and Anti-Patterns ### ❌ Swarm Variable Creation After Population @@ -694,6 +824,13 @@ TypeError: unsupported operand type(s) for *: 'UnitAwareDerivativeMatrix' and 'N - [ ] If using units: Set reference quantities BEFORE mesh creation - [ ] If using units with [M]: Provide material_density or equivalent +### Writing a Timestepping Loop + +- [ ] Declare the model and its reference quantities BEFORE creating the mesh +- [ ] Keep `time`, `step` and `dt` on `model.tracker`, not in local variables +- [ ] Take snapshots BEFORE the operator you might want to undo +- [ ] Do not use `mesh.t` for time dependence — it is not the model clock + ### Creating a Swarm - [ ] Create mesh first @@ -727,6 +864,12 @@ TypeError: unsupported operand type(s) for *: 'UnitAwareDerivativeMatrix' and 'N ## Version History +- **2026-09-09**: The timestepping pattern + - Start from the model and its reference quantities, not from the mesh + - Clock on `model.tracker`, not loose variables (snapshot consistency) + - A dimensional clock is dropped by the on-disk snapshot + - Backstepping recipe; snapshot before the operator + - `mesh.t` is not the model clock and is silently zero in a solve - **2025-11-15**: Initial version - Swarm ordering rules from test_0850/0851 debugging - Units everywhere-or-nowhere principle @@ -744,3 +887,5 @@ TypeError: unsupported operand type(s) for *: 'UnitAwareDerivativeMatrix' and 'N - `docs/developer/COORDINATE-UNITS-TECHNICAL-NOTE.md`: Coordinate units implementation - `docs/beginner/tutorials/12-Units_System.ipynb`: Units system tutorial - `docs/beginner/tutorials/13-Non_Dimensional_Scaling.ipynb`: Dimensional analysis +- `docs/advanced/snapshot-restore.md`: Snapshot and restore semantics +- `tests/test_0009_model_tracker.py`: The pattern, enforced diff --git a/tests/test_0009_model_tracker.py b/tests/test_0009_model_tracker.py index 3f3f12ace..ca8e30da9 100644 --- a/tests/test_0009_model_tracker.py +++ b/tests/test_0009_model_tracker.py @@ -175,3 +175,89 @@ def do_step(dt): do_step(0.05) assert model.tracker.step == s_snap + 2 assert abs(model.tracker.time - (t_snap + 0.10)) < 1e-12 + + +@pytest.mark.xfail( + reason="mesh.t is a symbolic atom bound to PETSc's petsc_t, which the " + "high-level solve() wrappers never set, so an expression containing it is " + "silently ZERO inside a solve and solve(time=...) is accepted and ignored. " + "Remove this xfail when mesh.t resolves to the model clock (#410 ruling: " + "time is model-owned, mesh.t becomes a back-compat accessor onto it).", + strict=False, +) +def test_mesh_t_resolves_to_the_model_clock(): + """The other clock. `mesh.t` is what users reach for in a time-dependent + boundary condition, and it is NOT `model.tracker.time` — so a source term + proportional to it should scale with the clock, and today does not. + + The constant-source control is what makes the assertion meaningful: it + proves the Poisson problem produces a non-trivial solution at all, so a + zero answer with `mesh.t` is the clock's fault and not the setup's. + """ + uw, model = _fresh_model() + import sympy + + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1.0 / 8.0, qdegree=2 + ) + T = uw.discretisation.MeshVariable("T_clock", mesh, 1, degree=2) + poisson = uw.systems.Poisson(mesh, u_Field=T) + poisson.constitutive_model = uw.constitutive_models.DiffusionModel + poisson.constitutive_model.Parameters.diffusivity = 1.0 + for boundary in ("Top", "Bottom", "Left", "Right"): + poisson.add_dirichlet_bc(0.0, boundary) + poisson.petsc_options.delValue("ksp_monitor") + + poisson.f = sympy.sympify(1.0) + poisson.solve() + control = np.abs(np.asarray(T.array)).max() + assert control > 1.0e-3, "control failed: the Poisson setup itself is trivial" + + poisson.f = mesh.t + model.tracker.time = 5.0 + poisson.solve() + assert np.abs(np.asarray(T.array)).max() > 0.1 * control + + +@pytest.mark.xfail( + reason="A pint Quantity on the tracker falls through disk_snapshot's " + "'unserialisable type' branch: it is recorded as __skipped and is " + "ABSENT after load_state, with no warning at save time. A Quantity is " + "(magnitude, units) and is trivially serialisable. Remove this xfail when " + "the disk snapshot carries units.", + strict=False, +) +def test_a_dimensional_clock_survives_a_disk_snapshot(tmp_path): + """The pattern asks scripts to define units, which makes the clock + dimensional. That clock must survive a restart, and today it does not. + + The plain-float control is what makes this specific: it shows the disk + path works for ordinary values, so a dropped quantity is about units and + not about the tracker or the file. + """ + uw, model = _fresh_model() + + model.set_reference_quantities( + domain_depth=uw.quantity(500, "km"), + material_density=uw.quantity(3300, "kg/m**3"), + material_viscosity=uw.quantity(1e21, "Pa*s"), + ) + uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1.0 / 8.0, qdegree=2 + ) + + model.tracker.plain_control = 3.25 + model.tracker.time = uw.quantity(4.5, "Myr") + + path = str(tmp_path / "units.snap.h5") + model.save_state(file=path) + + model.tracker.plain_control = -1.0 + model.tracker.time = uw.quantity(-1.0, "Myr") + model.load_state(path) + + assert model.tracker.plain_control == pytest.approx(3.25), ( + "control failed: the disk snapshot lost an ordinary float too" + ) + assert model.tracker.time.magnitude == pytest.approx(4.5) + assert str(model.tracker.time.units) == "megayear" From 6f7421c5bac083b15b032dd7ec0bf5606a3c9cd6 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Tue, 15 Sep 2026 10:07:59 +1000 Subject: [PATCH 02/12] =?UTF-8?q?feat:=20model.step(dt)=20=E2=80=94=20the?= =?UTF-8?q?=20timestep=20as=20a=20transaction,=20with=20an=20ordered=20rec?= =?UTF-8?q?ord=20of=20what=20it=20did?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A run had no step boundary, so nothing could say what a timestep did, whether it finished, or what state it started from. `with model.step(dt):` owns a time interval: the clock reads its end for the whole block, because an implicit residual is centred at t + dt; the advance commits on clean exit and only then, so a rejected or failed step leaves the clock alone; and everything the block did — each solve, each history shift — is recorded in order. Three defects found by building it, each silent: mesh.t was zero in every solve (it now resolves to the model clock, #410); the on-disk snapshot dropped dimensional values; and a constants[] slot that stopped being constant packed a zero rather than saying so. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- .../guides/HOW-TO-WRITE-UW3-SCRIPTS.md | 98 +++++++---- src/underworld3/checkpoint/disk_snapshot.py | 59 ++++++- .../cython/petsc_generic_snes_solvers.pyx | 25 +++ .../discretisation/discretisation_mesh.py | 73 +++++--- src/underworld3/model.py | 165 ++++++++++++++++++ src/underworld3/systems/solvers.py | 10 +- src/underworld3/utilities/_jitextension.py | 29 ++- .../utilities/unit_aware_coordinates.py | 11 +- tests/test_0009_model_tracker.py | 131 +++++++++++--- tests/test_0011_model_step_journal.py | 127 ++++++++++++++ .../test_0104_constant_slot_still_constant.py | 102 +++++++++++ 11 files changed, 746 insertions(+), 84 deletions(-) create mode 100644 tests/test_0011_model_step_journal.py create mode 100644 tests/test_0104_constant_slot_still_constant.py diff --git a/docs/developer/guides/HOW-TO-WRITE-UW3-SCRIPTS.md b/docs/developer/guides/HOW-TO-WRITE-UW3-SCRIPTS.md index 5525dade1..ef719a8e1 100644 --- a/docs/developer/guides/HOW-TO-WRITE-UW3-SCRIPTS.md +++ b/docs/developer/guides/HOW-TO-WRITE-UW3-SCRIPTS.md @@ -483,6 +483,41 @@ else you assign to it (`model.tracker.rms_velocity = ...`) is captured and restored the same way, so a diagnostic you want to survive a backstep belongs there too. +### Wrap the step + +`model.step(dt)` makes the loop a transaction: + +```python +while model.tracker.time < end_time: + dt = adv_diff.estimate_dt() + + with model.step(dt): + adv_diff.solve(timestep=dt) + stokes.solve(zero_init_guess=False) +``` + +Three things follow, and none of them requires anything else in the script to +change. The clock reads the END of the interval for the whole block, which is +where an implicit scheme centres its residual, so a time-dependent coefficient +is evaluated at the right time. The advance commits only on clean exit, so a +step that raises — or one abandoned because the Courant number came out too +large — leaves the clock exactly as it was. And everything the block did is +recorded: + +```python +>>> for entry in model.journal[-3:]: +... print(entry) + solve:SNES_Stokes(V)> + solve:SNES_Stokes(V)> + solve:SNES_Stokes(V)> +``` + +That record is worth having on its own. It answers what a run actually did, +in order, without the script being instrumented for it — which is the question +you want to ask of someone else's model, or your own six months later. + +Opening a step is optional. A script that never does behaves exactly as before. + ### Backstepping The pattern above is what makes speculative stepping safe: @@ -490,17 +525,17 @@ The pattern above is what makes speculative stepping safe: ```python snap = model.save_state() # before the step, not after -dt = big_dt -adv_diff.solve(timestep=dt) -stokes.solve(zero_init_guess=False) - -if courant_number() > courant_limit: - model.load_state(snap) # fields AND clock go back together - for _ in range(n_substeps): - ... # replay with smaller steps -else: - model.tracker.time += dt - model.tracker.step += 1 +try: + with model.step(big_dt): + adv_diff.solve(timestep=big_dt) + stokes.solve(zero_init_guess=False) + if courant_number() > courant_limit: + raise StepRejected # abandons the step; the clock stays put +except StepRejected: + model.load_state(snap) # fields go back; the clock never moved + for sub_dt in substeps(big_dt): + with model.step(sub_dt): + ... ``` Take the snapshot **before** the operator, not after. A `DDt` history plugin @@ -513,34 +548,22 @@ of model state, so two independent runs of the same problem on the same solver objects diverge at the 1e-13 level from the first step. If you need to look at a step twice, restore it rather than re-run it. -### Known gap: a dimensional clock does not survive a restart - -An in-memory snapshot round-trips a units-carrying tracker entry correctly. The -**on-disk** snapshot does not: a `pint` quantity falls through to the -"unserialisable type" branch, is recorded in the file as -`__skipped` and is simply **absent** after `load_state`, so reading -`model.tracker.time` afterwards raises. Nothing warns at save time. Plain -floats, ints and numpy arrays are unaffected. +### Time-dependent expressions -Until that is fixed, a script that needs to restart from disk should keep the -clock non-dimensional, or re-establish it explicitly after loading: +`mesh.t` is the model clock as a symbol. It is repacked from +`model.tracker.time` before every solve, so a time-dependent source or +boundary condition follows the loop above with no recompilation per step: ```python -model.load_state(path) -model.tracker.time = uw.quantity(model.tracker.time_Myr, "Myr") # stored as a float +omega = 2 * sympy.pi / period +stokes.add_dirichlet_bc((V0 * sympy.sin(omega * mesh.t), 0.0), "Top") ``` -### Known gap: `mesh.t` is not this clock - -`mesh.t` is a separate, symbolic time atom bound to PETSc's `petsc_t`. The -high-level `solve()` wrappers never set it, so **an expression containing -`mesh.t` evaluates to zero inside a solve**, silently. A time-dependent -boundary condition written as `sympy.sin(omega * mesh.t)` is identically zero -and nothing warns. `solve(time=...)` is accepted and ignored. - -Until `mesh.t` is wired to the model clock, build time dependence from a -`uw.function.expression` you update yourself each step, and drive it from -`model.tracker.time`. +Two things to know. A script that never advances `model.tracker.time` leaves +`mesh.t` at zero, so the clock and the pattern above are the same subject. And +`mesh.t` should appear inside an expression rather than be handed bare to a +scalar setter — `poisson.f = mesh.t` stores a value, `poisson.f = 1.0 * mesh.t` +keeps the symbol. --- @@ -828,8 +851,9 @@ TypeError: unsupported operand type(s) for *: 'UnitAwareDerivativeMatrix' and 'N - [ ] Declare the model and its reference quantities BEFORE creating the mesh - [ ] Keep `time`, `step` and `dt` on `model.tracker`, not in local variables +- [ ] Wrap each step in `with model.step(dt):` - [ ] Take snapshots BEFORE the operator you might want to undo -- [ ] Do not use `mesh.t` for time dependence — it is not the model clock +- [ ] Use `mesh.t` inside an expression for time dependence, never bare ### Creating a Swarm @@ -867,7 +891,9 @@ TypeError: unsupported operand type(s) for *: 'UnitAwareDerivativeMatrix' and 'N - **2026-09-09**: The timestepping pattern - Start from the model and its reference quantities, not from the mesh - Clock on `model.tracker`, not loose variables (snapshot consistency) - - A dimensional clock is dropped by the on-disk snapshot + - Disk snapshots now carry dimensional values (magnitude + units) + - `mesh.t` now resolves to the model clock (#410) + - `model.step(dt)` — the step as a transaction, and the step journal - Backstepping recipe; snapshot before the operator - `mesh.t` is not the model clock and is silently zero in a solve - **2025-11-15**: Initial version diff --git a/src/underworld3/checkpoint/disk_snapshot.py b/src/underworld3/checkpoint/disk_snapshot.py index 9cbf7c22a..6a61a7a38 100644 --- a/src/underworld3/checkpoint/disk_snapshot.py +++ b/src/underworld3/checkpoint/disk_snapshot.py @@ -770,7 +770,8 @@ def _read_swarm_from_sidecar(swarm, sidecar_path: str) -> None: # # Serialisation is *generic over dataclass fields* — no per-class # special code. Handled value types: None, bool, int, float, str, -# numpy.ndarray, list (JSON-encoded), dict (recursive subgroup). Other +# numpy.ndarray, list (JSON-encoded), dict (recursive subgroup), +# dimensional quantities (magnitude + unit string). Other # types (notably sympy expressions in DDtSymbolicState.psi_star) are # marked with `_skipped` and not round-tripped — documented as # a v1.x limitation; consumers either use a non-Symbolic DDt flavor @@ -788,6 +789,40 @@ def _is_h5_attr_scalar(value: Any) -> bool: ) or isinstance(value, (bool, str)) +_MAGNITUDE_SUFFIX = "__magnitude" +_UNITS_SUFFIX = "__units" + + +def _is_quantity(value: Any) -> bool: + """A dimensional value, duck-typed. + + ``uw.quantity`` returns a ``UWQuantity``, which is NOT a + ``pint.Quantity`` subclass, so an isinstance test against either would + miss one of them. Both carry ``magnitude`` and ``units``. + """ + return hasattr(value, "magnitude") and hasattr(value, "units") + + +def _write_quantity(h5group, name: str, value: Any) -> None: + """Store a dimensional value as magnitude + unit string. + + The magnitude goes through the ordinary dispatch (attr for a scalar, + dataset for an array), so an array-valued quantity round-trips too. + """ + h5group.attrs[name + _UNITS_SUFFIX] = str(value.units) + _serialise_field(h5group, name + _MAGNITUDE_SUFFIX, np.asarray(value.magnitude).item() + if np.ndim(value.magnitude) == 0 else np.asarray(value.magnitude)) + + +def _read_quantity(h5group, name: str) -> Any: + """Inverse of :func:`_write_quantity`.""" + units = h5group.attrs[name + _UNITS_SUFFIX] + if isinstance(units, bytes): + units = units.decode() + magnitude = _deserialise_field(h5group, name + _MAGNITUDE_SUFFIX, None) + return uw.quantity(magnitude, str(units)) + + def _serialise_field(h5group, name: str, value: Any) -> None: """Write a Python value into an HDF5 group as attr/dataset/subgroup. @@ -797,6 +832,7 @@ def _serialise_field(h5group, name: str, value: Any) -> None: - attr `__json` for JSON-encodable lists / nested simple structures - dataset `` for numpy arrays - subgroup `` for dict values, recursing + - attrs `__magnitude` + `__units` for dimensional values - attr `__skipped` = '' for anything else """ if value is None: @@ -808,6 +844,9 @@ def _serialise_field(h5group, name: str, value: Any) -> None: if isinstance(value, str): h5group.attrs[name] = value return + if _is_quantity(value): + _write_quantity(h5group, name, value) + return if isinstance(value, np.ndarray): if name in h5group: del h5group[name] @@ -860,6 +899,21 @@ def _group_to_dict(h5group) -> dict: out[k] = _group_to_dict(item) else: out[k] = np.asarray(item[...]) + + # Fold `__magnitude` + `__units` back into one dimensional + # value. Done as a post-pass because the two halves may arrive from + # different loops above (a scalar magnitude is an attr, an array + # magnitude a dataset). + for units_key in [k for k in out if k.endswith(_UNITS_SUFFIX)]: + base = units_key[: -len(_UNITS_SUFFIX)] + magnitude_key = base + _MAGNITUDE_SUFFIX + if magnitude_key not in out: + continue + units = out.pop(units_key) + magnitude = out.pop(magnitude_key) + if isinstance(units, bytes): + units = units.decode() + out[base] = uw.quantity(magnitude, str(units)) return out @@ -886,6 +940,9 @@ def _deserialise_field(h5group, name: str, fallback: Any) -> Any: if (name + "__json") in h5group.attrs: return json.loads(h5group.attrs[name + "__json"]) + if (name + _UNITS_SUFFIX) in h5group.attrs: + return _read_quantity(h5group, name) + if (name + "__skipped") in h5group.attrs: # Skipped at write time — keep the current value rather than # clobber it with a placeholder. diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index aa22e699e..78a290861 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -2266,6 +2266,31 @@ class SolverBaseClass(uw_object): Called before each solve() to ensure constants are current without requiring JIT recompilation. """ + # Refresh mesh.t from the model clock first, so a time-dependent + # expression is repacked with the rest of the constants rather than + # needing its own hook (or a recompile) per timestep. + try: + self.mesh._sync_time_from_model() + except AttributeError: + pass + + # Note the solve in the model's step journal, if a step is open. This + # is the one place every solver passes through before solving, so one + # hook records them all, in order. A no-op outside a model.step block. + try: + # Name it by what it SOLVES, not by its auto-generated instance id: + # a journal reading "Stokes(V) -> AdvDiffusion(T)" is auditable, + # one reading "Solver_8_ -> Solver_14_" is not. + try: + unknown = self.u.name + except Exception: + unknown = "?" + uw.get_default_model()._record_step_event( + "solve", f"{type(self).__name__}({unknown})" + ) + except Exception: + pass + if not self.constants_manifest or self.dm is None: return diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index 7f045997e..16dcdd25a 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/src/underworld3/discretisation/discretisation_mesh.py @@ -1385,13 +1385,16 @@ def _setup_symbolic_coordinates(self, coordinate_system_type): self._Gamma.y._ccodestr = "petsc_n[1]" self._Gamma.z._ccodestr = "petsc_n[2]" - # Time coordinate — PETSc passes this as petsc_t to all pointwise - # functions. Solvers set dm.time before each solve via solve(time=t). - # Users reference it as mesh.t in expressions (e.g. V0 * sympy.sin(omega * mesh.t)) - from ..utilities.unit_aware_coordinates import TimeSymbol - - self._t = TimeSymbol("t") - self._t._units = None # patched below by _patch_time_units + # Time coordinate. This is a live-rampable ``constants[]`` atom, NOT + # PETSc's ``petsc_t``: the high-level solve() wrappers never set + # petsc_t, so an expression built on it evaluated to zero inside every + # solve (silently — a time-dependent BC was identically zero). Time is + # owned by the orchestration model; ``mesh.t`` reads that clock. + # ``_sync_time_from_model`` repacks it before each solve, from the + # solver's ``_update_constants``, so no kernel is recompiled per step. + self._t = uw.expression( + r"t", 0.0, "model time — the clock on uw.get_default_model().tracker" + ) # Add unit awareness to coordinate symbols if mesh has units or model has scales from ..utilities.unit_aware_coordinates import patch_coordinate_units @@ -4516,30 +4519,60 @@ def CoordinateSystem(self) -> CoordinateSystem: @property def t(self): - r"""Symbolic time coordinate. + r"""Symbolic model time. - PETSc passes a time value (``petsc_t``) to all pointwise residual - and Jacobian functions. Use ``mesh.t`` in expressions to reference - this time without forcing JIT recompilation each timestep. + A live-rampable ``constants[]`` atom carrying the clock owned by the + orchestration model, ``uw.get_default_model().tracker.time``. Every + solver repacks it from that clock immediately before solving, so an + expression built on ``mesh.t`` follows time with no JIT recompilation + per step. - The low-level PETSc solver accepts ``time=t`` to set the value - of ``petsc_t`` for pointwise functions. If not provided, ``petsc_t`` - defaults to 0. Note: the high-level Python ``solve()`` wrappers - do not yet pass ``time=`` through — set it directly via - ``UW_DMSetTime`` at the Cython level if needed. + Maintain the clock as part of the timestepping loop (see + ``docs/developer/guides/HOW-TO-WRITE-UW3-SCRIPTS.md``). A script that + never advances it leaves ``mesh.t`` at zero. - When the scaling system is active, ``mesh.t`` carries time units - (derived from the model's time scale) so that dimensional analysis - works correctly in expressions. + A dimensional clock is non-dimensionalised on the way in, so the value + the kernels see is always in solver units. + + .. note:: + Assign it as part of an expression rather than bare. A boundary + condition takes a Matrix / array form, and a bare atom handed to a + scalar setter is stored by value. Examples -------- >>> omega = 2 * np.pi / period >>> stokes.add_dirichlet_bc((V0 * sympy.sin(omega * mesh.t), 0.0), "Top") - >>> stokes.solve(time=current_time) # sets petsc_t before SNES + >>> model.tracker.time = 1.5 * uw.quantity(1, "Myr") + >>> stokes.solve() # mesh.t picks the clock up """ return self._t + def _sync_time_from_model(self): + """Repack ``mesh.t`` from the model clock. Called by every solver's + ``_update_constants`` immediately before a solve, so an expression + containing ``mesh.t`` sees the current time without a rebuild. + + Silent no-op when the model has no clock: ``mesh.t`` then stays at + whatever it was last set to (0.0 for a fresh mesh), which is the + behaviour a script that never advances a clock already expects. + """ + try: + model = uw.get_default_model() + time = model.tracker.time + except Exception: + return + if time is None: + return + try: + if hasattr(time, "magnitude") or hasattr(time, "_pint_qty"): + time = float(uw.non_dimensionalise(time)) + self._t.sym = sympy.sympify(float(time)) + except Exception: + # A clock we cannot reduce to a number is not worth failing a + # solve over; leave mesh.t as it stands. + return + @property def nullspace_rotations(self): """Symbolic velocity fields for rigid-body rotation null modes. diff --git a/src/underworld3/model.py b/src/underworld3/model.py index 1a5619ae8..c33eae69f 100644 --- a/src/underworld3/model.py +++ b/src/underworld3/model.py @@ -54,6 +54,40 @@ class ModelState(Enum): ERROR = "error" +class ModelStep: + """What one timestep did — the journal entry for a ``model.step`` block. + + Ordered, so ``[e.name for e in step.events]`` is the sequence of operators + the step actually applied. That sequence is what makes a step auditable + (did this run do what the write-up says?) and what a replay or an adjoint + needs in order to walk the run backwards. + """ + + __slots__ = ("index", "t0", "dt", "label", "events", "completed") + + def __init__(self, index, t0, dt, label=None): + self.index = index + self.t0 = t0 + self.dt = dt + self.label = label + self.events = [] + self.completed = False + + @property + def t1(self): + """The end of the interval this step covers.""" + return self.t0 + self.dt + + def _record(self, kind, name, **detail): + self.events.append({"kind": kind, "name": name, **detail}) + + def __repr__(self): + state = "" if self.completed else " ABANDONED" + seq = " -> ".join(f"{e['kind']}:{e['name']}" for e in self.events) or "(nothing)" + tag = f" {self.label!r}" if self.label else "" + return f"" + + class Model(PintNativeModelMixin, BaseModel): """ Central orchestrator for Underworld3 simulations. @@ -141,6 +175,14 @@ class Model(PintNativeModelMixin, BaseModel): # src/underworld3/checkpoint/tracker.py. _tracker: Any = PrivateAttr(default=None) + # The step journal: an ordered record of what each timestep actually did. + # ``_open_step`` is the ModelStep currently in progress (None outside a + # ``with model.step(dt):`` block); ``_journal`` is the bounded history of + # completed steps. See :meth:`step`. + _open_step: Any = PrivateAttr(default=None) + _journal: Any = PrivateAttr(default_factory=list) + _journal_limit: Any = PrivateAttr(default=512) + def __init__(self, name: Optional[str] = None, **kwargs): """ Initialize a new Model instance. @@ -604,6 +646,129 @@ def tracker(self): """ return self._tracker + # ------------------------------------------------------------------ + # The step journal + # ------------------------------------------------------------------ + + @property + def journal(self) -> List[Any]: + """Completed :class:`ModelStep` records, oldest first. + + An ordered account of what each timestep did — which solvers ran, in + what order, over which time interval. Answers "is this model doing the + thing I said it does" without instrumenting the script, and is the + record an adjoint or a replay needs. + + Bounded by ``model.journal_limit`` (default 512 steps); set it to + ``None`` to keep everything. + """ + return list(self._journal) + + @property + def journal_limit(self): + """How many completed steps to retain (None keeps all).""" + return self._journal_limit + + @journal_limit.setter + def journal_limit(self, value): + self._journal_limit = value + self._trim_journal() + + @property + def open_step(self): + """The step in progress, or None outside a ``model.step`` block.""" + return self._open_step + + def _trim_journal(self): + limit = self._journal_limit + if limit is not None and len(self._journal) > limit: + del self._journal[: len(self._journal) - limit] + + def _record_step_event(self, kind: str, name: str, **detail) -> None: + """Note that something happened inside the step in progress. + + Called by the machinery (solvers, history managers), not by users. + A no-op outside a ``model.step`` block, so nothing is required of a + script that does not use one. + """ + step = self._open_step + if step is not None: + step._record(kind, name, **detail) + + def step(self, dt, label: Optional[str] = None): + """One timestep, as a transaction. + + :: + + with model.step(dt): + adv_diff.solve(timestep=dt) + stokes.solve(zero_init_guess=False) + + The block owns a time INTERVAL. Three things follow: + + **The clock reads as the end of the interval for the whole block.** + An implicit scheme centres its residual at the new time, so a + time-dependent coefficient — a driven boundary above all — belongs at + ``t + dt``. Advancing only on exit would evaluate every implicit + coefficient one step late. + + **The advance commits on clean exit, and only then.** An exception, or + a step abandoned because the Courant number came out too large, leaves + ``model.tracker`` exactly as it was. Backstepping no longer has to + remember to unwind a counter. + + **Everything the block did is recorded** in :attr:`journal`, in order, + with the interval it ran over. + + Nothing is compulsory: a script that never opens a step behaves as + before, and the machinery's recording calls become no-ops. + + Parameters + ---------- + dt : float or dimensional quantity + The interval this step covers. + label : str, optional + A name for the step, carried into the journal. + """ + from contextlib import contextmanager + + @contextmanager + def _step_context(): + if self._open_step is not None: + raise RuntimeError( + "a model step is already open " + f"(step {self._open_step.index}, label {self._open_step.label!r}). " + "Steps do not nest — close the outer one first." + ) + + t0 = self.tracker.time if "time" in self.tracker else 0.0 + index = self.tracker.step if "step" in self.tracker else 0 + record = ModelStep(index=index, t0=t0, dt=dt, label=label) + self._open_step = record + + # Position the clock at the END of the interval for the duration of + # the block, so implicit coefficients (mesh.t) are evaluated there. + self.tracker.time = record.t1 + try: + yield record + except BaseException: + # Abandon: put the clock back and do not commit. + self.tracker.time = t0 + record.completed = False + self._open_step = None + raise + + # Commit. + self.tracker.time = record.t1 + self.tracker.step = index + 1 + self.tracker.dt = dt + record.completed = True + self._open_step = None + self._journal.append(record) + self._trim_journal() + + return _step_context() + def _register_state_bearer(self, obj) -> None: """Register a Snapshottable object with this model. diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index 24351f47b..019aafd43 100644 --- a/src/underworld3/systems/solvers.py +++ b/src/underworld3/systems/solvers.py @@ -602,8 +602,14 @@ def f(self, value): """Set the source term (handles units and scaling).""" self._needs_function_rewire = True - # Handle UWQuantity with units - enforce "units everywhere" principle - if hasattr(value, "value") and hasattr(value, "units"): + # Handle UWQuantity with units - enforce "units everywhere" principle. + # The `.value`/`.units` duck-test also matches a UWexpression, which is + # a SYMBOLIC atom, not a plain quantity — unwrapping one here baked a + # live-rampable constants[] atom to a C literal at assignment time + # (`poisson.f = mesh.t` became a constant zero). UWQuantity and pint + # Quantity are not sympy objects; UWexpression is, so that separates them. + if (hasattr(value, "value") and hasattr(value, "units") + and not isinstance(value, sympy.Basic)): # Extract the plain value plain_value = float(value.value) diff --git a/src/underworld3/utilities/_jitextension.py b/src/underworld3/utilities/_jitextension.py index 5301b43fb..46f9fa290 100644 --- a/src/underworld3/utilities/_jitextension.py +++ b/src/underworld3/utilities/_jitextension.py @@ -510,7 +510,34 @@ def _pack_constants(manifest): try: values[idx] = float(uw_expr.data) except Exception: - values[idx] = 0.0 + # Do NOT pack a zero here. A constants[] slot exists because + # this expression resolved to a single number when the kernel + # was COMPILED. If it no longer does, the compiled kernel is + # structurally wrong for the current model — it reads a scalar + # where the expression now varies in space — and packing 0.0 + # hands that kernel a zero coefficient. That is silent and + # catastrophic: a zero diffusivity or viscosity diverges, and + # nothing says why. + # + # The usual cause is a nested atom that has been ramped. + # `(1 + T**2)**(-m) + 1` is the NUMBER 2 while m is zero, so it + # banks as one constant; ramp m and it depends on T again, but + # the kernel still expects a scalar. + raise RuntimeError( + f"constants[] slot {idx} ({uw_expr.name!r}) no longer " + f"reduces to a number, so the compiled kernel — which " + f"treats it as a scalar constant — is out of date.\n" + f" current content: {str(getattr(uw_expr, '_sym', uw_expr))[:160]}\n" + f"This usually means an atom nested inside it has been " + f"ramped, and the expression has stopped being constant. " + f"Force a rebuild before solving again:\n" + f" solver.is_setup = False\n" + f" solver._needs_function_rewire = True\n" + f"To keep a coefficient rampable without this, give it its " + f"own atom rather than letting the enclosing expression " + f"collapse to a number at compile time — see " + f"uw.maths.functions.vanishing." + ) from None return values diff --git a/src/underworld3/utilities/unit_aware_coordinates.py b/src/underworld3/utilities/unit_aware_coordinates.py index 0073f2830..a25f3e7bf 100644 --- a/src/underworld3/utilities/unit_aware_coordinates.py +++ b/src/underworld3/utilities/unit_aware_coordinates.py @@ -268,9 +268,14 @@ def _patch_time_units(mesh): except Exception: pass - mesh._t._units = time_units - if not hasattr(mesh._t, "get_units"): - mesh._t.get_units = lambda: mesh._t._units + # mesh._t is a UWexpression (a rampable constants[] atom), which manages + # its own units; only the legacy symbol flavour needs patching. + try: + mesh._t._units = time_units + if not hasattr(mesh._t, "get_units"): + mesh._t.get_units = lambda: mesh._t._units + except AttributeError: + pass def get_coordinate_units(coord): diff --git a/tests/test_0009_model_tracker.py b/tests/test_0009_model_tracker.py index ca8e30da9..dfbe721d6 100644 --- a/tests/test_0009_model_tracker.py +++ b/tests/test_0009_model_tracker.py @@ -177,14 +177,6 @@ def do_step(dt): assert abs(model.tracker.time - (t_snap + 0.10)) < 1e-12 -@pytest.mark.xfail( - reason="mesh.t is a symbolic atom bound to PETSc's petsc_t, which the " - "high-level solve() wrappers never set, so an expression containing it is " - "silently ZERO inside a solve and solve(time=...) is accepted and ignored. " - "Remove this xfail when mesh.t resolves to the model clock (#410 ruling: " - "time is model-owned, mesh.t becomes a back-compat accessor onto it).", - strict=False, -) def test_mesh_t_resolves_to_the_model_clock(): """The other clock. `mesh.t` is what users reach for in a time-dependent boundary condition, and it is NOT `model.tracker.time` — so a source term @@ -213,27 +205,26 @@ def test_mesh_t_resolves_to_the_model_clock(): control = np.abs(np.asarray(T.array)).max() assert control > 1.0e-3, "control failed: the Poisson setup itself is trivial" - poisson.f = mesh.t + poisson.f = 1.0 * mesh.t model.tracker.time = 5.0 poisson.solve() - assert np.abs(np.asarray(T.array)).max() > 0.1 * control + at_five = np.abs(np.asarray(T.array)).max() + assert at_five > 0.1 * control + + # and it must TRACK the clock, not merely be non-zero once + model.tracker.time = 10.0 + poisson.solve() + at_ten = np.abs(np.asarray(T.array)).max() + assert at_ten == pytest.approx(2.0 * at_five, rel=1e-6) -@pytest.mark.xfail( - reason="A pint Quantity on the tracker falls through disk_snapshot's " - "'unserialisable type' branch: it is recorded as __skipped and is " - "ABSENT after load_state, with no warning at save time. A Quantity is " - "(magnitude, units) and is trivially serialisable. Remove this xfail when " - "the disk snapshot carries units.", - strict=False, -) def test_a_dimensional_clock_survives_a_disk_snapshot(tmp_path): """The pattern asks scripts to define units, which makes the clock - dimensional. That clock must survive a restart, and today it does not. + dimensional. That clock must survive a restart. The plain-float control is what makes this specific: it shows the disk - path works for ordinary values, so a dropped quantity is about units and - not about the tracker or the file. + path works for ordinary values, so a dropped quantity would be about + units and not about the tracker or the file. """ uw, model = _fresh_model() @@ -261,3 +252,101 @@ def test_a_dimensional_clock_survives_a_disk_snapshot(tmp_path): ) assert model.tracker.time.magnitude == pytest.approx(4.5) assert str(model.tracker.time.units) == "megayear" + + +def test_a_dimensional_array_survives_a_disk_snapshot(tmp_path): + """The magnitude may be an array, which is stored as a dataset rather + than an attribute — the other half of the quantity round-trip.""" + uw, model = _fresh_model() + + model.set_reference_quantities( + domain_depth=uw.quantity(500, "km"), + material_density=uw.quantity(3300, "kg/m**3"), + material_viscosity=uw.quantity(1e21, "Pa*s"), + ) + uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1.0 / 8.0, qdegree=2 + ) + + model.tracker.depths = uw.quantity(np.array([10.0, 20.0, 30.0]), "km") + + path = str(tmp_path / "arr.snap.h5") + model.save_state(file=path) + model.tracker.depths = uw.quantity(np.array([0.0]), "km") + model.load_state(path) + + assert np.allclose(model.tracker.depths.magnitude, [10.0, 20.0, 30.0]) + assert str(model.tracker.depths.units) == "kilometer" + + +def test_mesh_t_drives_a_time_dependent_boundary_condition(): + """The headline use case, and the one mesh.t's own docstring advertises: + a boundary value that varies with time. Boundary terms are assembled + through a different residual path from the source, so this is not implied + by the source-term test above. + """ + uw, model = _fresh_model() + import sympy + + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1.0 / 8.0, qdegree=2 + ) + T = uw.discretisation.MeshVariable("T_bc", mesh, 1, degree=2) + poisson = uw.systems.Poisson(mesh, u_Field=T) + poisson.constitutive_model = uw.constitutive_models.DiffusionModel + poisson.constitutive_model.Parameters.diffusivity = 1.0 + poisson.f = 0.0 + # A boundary condition takes a Matrix / array form, not a bare scalar + # expression, so the clock is wrapped rather than passed directly. + poisson.add_dirichlet_bc(sympy.Matrix([mesh.t]), "Top") + poisson.add_dirichlet_bc(0.0, "Bottom") + poisson.petsc_options.delValue("ksp_monitor") + + model.tracker.time = 1.0 + poisson.solve(zero_init_guess=True) + at_one = float(np.asarray(T.data)[:, 0].max()) + + model.tracker.time = 3.0 + poisson.solve(zero_init_guess=True) + at_three = float(np.asarray(T.data)[:, 0].max()) + + assert at_one > 0.5, "the driven boundary never reached the solution" + assert at_three == pytest.approx(3.0 * at_one, rel=1e-6) + + +def test_a_bare_rampable_atom_is_not_baked_by_the_source_setter(): + """`solver.f = ` must keep the atom symbolic. + + The setter's `.value`/`.units` duck-test for a dimensional quantity also + matches a UWexpression, which is a symbolic atom rather than a plain + quantity — so a bare assignment used to bake a live-rampable constant to a + literal at assignment time, and it never ramped again. + """ + uw, model = _fresh_model() + import sympy + + mesh = uw.meshing.StructuredQuadBox( + elementRes=(8, 8), minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0) + ) + T = uw.discretisation.MeshVariable("T_bare", mesh, 1, degree=2) + c = uw.expression(r"c_bare", 0.5, "a rampable atom") + + poisson = uw.systems.Poisson(mesh, u_Field=T) + poisson.constitutive_model = uw.constitutive_models.DiffusionModel + poisson.constitutive_model.Parameters.diffusivity = 1.0 + poisson.f = c + poisson.add_dirichlet_bc(0.0, "Top") + poisson.add_dirichlet_bc(0.0, "Bottom") + poisson.petsc_options.delValue("ksp_monitor") + + poisson.solve(zero_init_guess=True) + at_half = float(np.asarray(T.data)[:, 0].mean()) + + # the manifest is populated at setup, which happens on the first solve + assert "c_bare" in {e.name for _i, e in poisson.constants_manifest} + + c.sym = sympy.sympify(1.5) + poisson.solve(zero_init_guess=True) + at_one_and_a_half = float(np.asarray(T.data)[:, 0].mean()) + + assert at_one_and_a_half == pytest.approx(3.0 * at_half, rel=1e-6) diff --git a/tests/test_0011_model_step_journal.py b/tests/test_0011_model_step_journal.py new file mode 100644 index 000000000..a98e7ff8d --- /dev/null +++ b/tests/test_0011_model_step_journal.py @@ -0,0 +1,127 @@ +"""The step journal — ``with model.step(dt):``. + +One timestep as a transaction. Three guarantees, one test each: + + * the clock reads the END of the interval inside the block, because an + implicit scheme centres its residual there; + * the advance commits only on clean exit, so an abandoned step leaves the + clock alone; + * everything the block did is recorded, in order. + +Documented in ``docs/developer/guides/HOW-TO-WRITE-UW3-SCRIPTS.md``. +""" + +import numpy as np +import pytest +import sympy + +pytestmark = [pytest.mark.level_1, pytest.mark.tier_a] + + +def _fresh_model(): + import underworld3 as uw + + uw.reset_default_model() + return uw, uw.get_default_model() + + +def _poisson(uw, mesh, name): + T = uw.discretisation.MeshVariable(name, mesh, 1, degree=2) + solver = uw.systems.Poisson(mesh, u_Field=T) + solver.constitutive_model = uw.constitutive_models.DiffusionModel + solver.constitutive_model.Parameters.diffusivity = 1.0 + solver.f = 1.0 + solver.add_dirichlet_bc(0.0, "Top") + solver.add_dirichlet_bc(0.0, "Bottom") + solver.petsc_options.delValue("ksp_monitor") + return solver + + +def test_the_clock_reads_the_end_of_the_interval_inside_the_block(): + """An implicit residual is centred at t + dt, so that is where a + time-dependent coefficient must be evaluated. Committing only on exit + would evaluate every one of them a step late.""" + uw, model = _fresh_model() + model.tracker.time, model.tracker.step = 2.0, 7 + + with model.step(0.25) as step: + assert step.t0 == pytest.approx(2.0) + assert step.t1 == pytest.approx(2.25) + assert model.tracker.time == pytest.approx(2.25) + + assert model.tracker.time == pytest.approx(2.25) + assert model.tracker.step == 8 + + +def test_an_abandoned_step_does_not_commit(): + """A step rejected on a Courant check, or one that raises, must leave the + clock where it was — the caller should not have to unwind a counter.""" + uw, model = _fresh_model() + model.tracker.time, model.tracker.step = 1.0, 3 + + with pytest.raises(RuntimeError, match="Courant"): + with model.step(0.5): + raise RuntimeError("Courant too large") + + assert model.tracker.time == pytest.approx(1.0) + assert model.tracker.step == 3 + assert model.journal == [] + assert model.open_step is None + + +def test_the_journal_records_what_ran_and_in_what_order(): + """The point of the record: it answers what a step actually did, without + the script being instrumented.""" + uw, model = _fresh_model() + mesh = uw.meshing.StructuredQuadBox( + elementRes=(8, 8), minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0) + ) + first = _poisson(uw, mesh, "T_one") + second = _poisson(uw, mesh, "T_two") + model.tracker.time, model.tracker.step = 0.0, 0 + + with model.step(0.1, label="a step"): + first.solve() + second.solve() + + assert len(model.journal) == 1 + entry = model.journal[0] + assert entry.label == "a step" + assert entry.completed + names = [e["name"] for e in entry.events if e["kind"] == "solve"] + assert names == ["SNES_Poisson(T_one)", "SNES_Poisson(T_two)"], names + # named by what they solve, so the record is auditable + assert "T_one" in repr(entry) + + +def test_steps_do_not_nest(): + uw, model = _fresh_model() + model.tracker.time, model.tracker.step = 0.0, 0 + with model.step(0.1): + with pytest.raises(RuntimeError, match="already open"): + with model.step(0.1): + pass + + +def test_a_script_without_steps_is_unaffected(): + """Opening a step is optional; the recording hooks are no-ops without one.""" + uw, model = _fresh_model() + mesh = uw.meshing.StructuredQuadBox( + elementRes=(8, 8), minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0) + ) + solver = _poisson(uw, mesh, "T_free") + solver.solve() + assert model.open_step is None + assert model.journal == [] + assert np.abs(np.asarray(solver.u.data)).max() > 1.0e-3 + + +def test_the_journal_is_bounded(): + uw, model = _fresh_model() + model.tracker.time, model.tracker.step = 0.0, 0 + model.journal_limit = 3 + for _ in range(7): + with model.step(0.1): + pass + assert len(model.journal) == 3 + assert [e.index for e in model.journal] == [4, 5, 6] diff --git a/tests/test_0104_constant_slot_still_constant.py b/tests/test_0104_constant_slot_still_constant.py new file mode 100644 index 000000000..0b425f00e --- /dev/null +++ b/tests/test_0104_constant_slot_still_constant.py @@ -0,0 +1,102 @@ +"""A constants[] slot that stops being constant must say so, not pack a zero. + +An expression is given a ``constants[]`` slot because it resolved to a single +number when the kernel was compiled. Ramping an atom nested inside it can make +it depend on position again — the compiled kernel still reads a scalar, and +packing a zero into that slot hands the solve a zero coefficient. Silent, and +catastrophic: a zero diffusivity diverges and nothing says why. + +This is the failure behind the "rampable constant in exponent position does not +ramp" report. The atom is not compiled out; the ENCLOSING expression collapses +to a number while the atom is zero, banks as one constant, and then stops being +one. +""" + +import numpy as np +import pytest +import sympy + +pytestmark = [pytest.mark.level_1, pytest.mark.tier_a] + + +def _build(initial): + import underworld3 as uw + + uw.reset_default_model() + mesh = uw.meshing.StructuredQuadBox( + elementRes=(8, 8), minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0) + ) + T = uw.discretisation.MeshVariable("T_slot", mesh, 1, degree=2) + m = uw.expression(r"m_slot", initial, "rampable atom") + # constant while m == 0 (anything**0 is 1), field-dependent as soon as it isn't + kappa = uw.expression( + r"\kappa_slot", (1.0 + 0.5 * T.sym[0] ** 2) ** (-m) + 1.0, "collapsing" + ) + poisson = uw.systems.Poisson(mesh, u_Field=T) + poisson.constitutive_model = uw.constitutive_models.DiffusionModel + poisson.constitutive_model.Parameters.diffusivity = kappa + poisson.f = 1.0 + poisson.add_dirichlet_bc(0.0, "Top") + poisson.add_dirichlet_bc(0.0, "Bottom") + poisson.petsc_options.delValue("ksp_monitor") + return uw, poisson, T, m + + +def test_a_slot_that_stops_being_constant_raises(): + uw, poisson, T, m = _build(0.0) + poisson.solve(zero_init_guess=True) + + # Compiled while the whole expression was the number 2, so the diffusivity + # banks as a single scalar slot. (It is named for the parameter wrapper, + # not for the inner expression — the collector stops at the outermost thing + # that is truly constant and does not recurse past it.) + assert len(poisson.constants_manifest) == 1 + + m.sym = sympy.sympify(0.5) # now depends on T again + with pytest.raises(RuntimeError, match="no longer.*reduces to a number"): + poisson.solve(zero_init_guess=True) + + +def test_the_message_names_the_slot_and_says_how_to_recover(): + uw, poisson, T, m = _build(0.0) + poisson.solve(zero_init_guess=True) + m.sym = sympy.sympify(0.5) + with pytest.raises(RuntimeError) as excinfo: + poisson.solve(zero_init_guess=True) + message = str(excinfo.value) + assert "no longer" in message + assert "_needs_function_rewire" in message + assert "constants[] slot" in message + + +def test_forcing_a_rebuild_recovers_and_the_atom_then_ramps(): + """The recovery the message prescribes must actually work.""" + uw, poisson, T, m = _build(0.0) + poisson.solve(zero_init_guess=True) + + results = [] + for value in (0.5, 1.0): + m.sym = sympy.sympify(value) + poisson.is_setup = False + poisson._needs_function_rewire = True + poisson.constitutive_model._solver_is_setup = False + poisson.solve(zero_init_guess=True) + results.append(float(np.asarray(T.data)[:, 0].mean())) + + assert all(np.isfinite(results)) + assert results[0] != pytest.approx(results[1]), "the atom still does not ramp" + + +def test_an_expression_that_stays_constant_is_unaffected(): + """The negative control: a slot that remains a number keeps working, so the + guard is not just refusing every ramp.""" + uw, poisson, T, m = _build(0.0) + # a plainly constant coefficient, ramped in the ordinary way + c = uw.expression(r"c_plain", 1.0, "ordinary rampable constant") + poisson.constitutive_model.Parameters.diffusivity = c + poisson.solve(zero_init_guess=True) + first = float(np.asarray(T.data)[:, 0].mean()) + c.sym = sympy.sympify(2.0) + poisson.solve(zero_init_guess=True) + second = float(np.asarray(T.data)[:, 0].mean()) + assert second == pytest.approx(first / 2.0, rel=1e-6) From e0bcdbcf81eacec7802740c01b82993d058aa813 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Tue, 15 Sep 2026 10:07:59 +1000 Subject: [PATCH 03/12] feat: a run keeps the state each step started from, writes its log to disk, and draws it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `model.record_every` keeps a snapshot of the state a step started from, and `model.rewind()` returns to it — fields, histories and clock together. Replaying a step from its snapshot is bit-identical; re-running the script is not, because iterative-solver history is not model state. That is what lets a step that misbehaved be looked at twice. The record is written as it happens: one aligned line per step, flushed as each closes so a running job can be watched with `tail -f`, with every backtrack in it — a log that shows step 7 and then step 7 again with nothing between is not a log of what happened. JSON lines beside it for reading back. Rendered as a figure, time running down the page, SVG for the web and a PDF written without a plotting dependency. Two more silent defects: a snapshot rescaled the mesh under units (capture read metres, restore wrote model units); and a solver on a rotated boundary was recorded twice because a setup push re-announced it. Plus a guard for free-surface derived solvers whose copied configuration had drifted from the parent. An annulus convection case is the second worked example. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- .../guides/HOW-TO-WRITE-UW3-SCRIPTS.md | 249 ++++++ docs/examples/convection/README.md | 9 + .../Ex_Convection_Annulus_Recorded.py | 499 ++++++++++++ scripts/test.sh | 1 + src/underworld3/__init__.py | 2 + .../cython/petsc_generic_snes_solvers.pyx | 32 +- .../discretisation/discretisation_mesh.py | 11 +- src/underworld3/model.py | 650 ++++++++++++++- src/underworld3/systems/ddt.py | 51 ++ src/underworld3/systems/free_surface.py | 83 ++ src/underworld3/systems/solvers.py | 22 +- src/underworld3/utilities/journal_report.py | 741 ++++++++++++++++++ src/underworld3/utilities/rotated_bc.py | 5 +- tests/test_0011_model_step_journal.py | 179 +++++ tests/test_0012_snapshot_units_coords.py | 110 +++ tests/test_0013_step_record_fidelity.py | 173 ++++ tests/test_0014_journal_file.py | 334 ++++++++ tests/test_0015_journal_report.py | 332 ++++++++ tests/test_1074_free_surface_config_drift.py | 90 +++ 19 files changed, 3543 insertions(+), 30 deletions(-) create mode 100644 docs/examples/convection/intermediate/Ex_Convection_Annulus_Recorded.py create mode 100644 src/underworld3/utilities/journal_report.py create mode 100644 tests/test_0012_snapshot_units_coords.py create mode 100644 tests/test_0013_step_record_fidelity.py create mode 100644 tests/test_0014_journal_file.py create mode 100644 tests/test_0015_journal_report.py create mode 100644 tests/test_1074_free_surface_config_drift.py diff --git a/docs/developer/guides/HOW-TO-WRITE-UW3-SCRIPTS.md b/docs/developer/guides/HOW-TO-WRITE-UW3-SCRIPTS.md index ef719a8e1..685d0589f 100644 --- a/docs/developer/guides/HOW-TO-WRITE-UW3-SCRIPTS.md +++ b/docs/developer/guides/HOW-TO-WRITE-UW3-SCRIPTS.md @@ -518,6 +518,208 @@ you want to ask of someone else's model, or your own six months later. Opening a step is optional. A script that never does behaves exactly as before. +### Recording a run + +Ask the step to keep the state it started from and the journal becomes a +restorable record: + +```python +model.record_every = 1 # keep every step; None (default) keeps none +model.record_limit = 8 # how many snapshots to retain + +while model.tracker.time < end_time: + with model.step(dt): + adv_diff.solve(timestep=dt) + stokes.solve(zero_init_guess=False) + +model.rewind() # undo the last step: fields, history and clock +``` + +The snapshot is taken before the operators run, which is the only correct +point — a `DDt` shifts its history in its post-solve hook, so a snapshot taken +afterwards holds the shifted history rather than the step's input. + +Two things this buys beyond backstepping. Replaying a step from its own +snapshot reproduces it exactly, where re-running the script does not, so a step +that misbehaved can be looked at twice. And an adjoint needs precisely this: the +state at each step and the order the operators were applied in. + +Snapshots cost roughly 13 bytes per primary degree of freedom per step. Older +steps lose their snapshot and keep their journal record, so the account of what +happened outlives the state it happened to. + +A driver that runs the same model more than once — an inversion, a parameter +sweep, a restart — should start each run with a clean account: + +```python +model.clear_journal() +model.tracker.time = uw.quantity(0.0, "Myr") +model.tracker.step = 0 +``` + +Without it the journal is the concatenation of every run the process has done, +and `rewind()` will walk back into the previous one. + +On a mesh that deforms or adapts the snapshot cannot be taken yet; the run +warns once, keeps journalling, and `rewind()` will not reach those steps. + +### Writing the record down + +`model.journal` is what the run can still undo. It lives in memory, it is +bounded, and it dies with the process. `model.journal_file` is what the run +*did*: + +```python +model.journal_file = "output/run.log" +``` + +One aligned line per step, appended and flushed as it closes, so `tail -f` +follows a running job: + +``` +# underworld3 step log · model 'default' · started 2026-09-10T21:22:40+00:00 +# scales: length 2.2e+06 m | time 4.84e+18 s | mass 1.065e+47 kg | temperature 2500 K +# step t/Myr dt/Myr wall/s outcome operators, in order + 0 0.175907 0.175907 0.44 ok [convect] solve:SNES_AdvectionDiffusion_Composed(T) > history_shift:EulerianSUPG(T) > solve:SNES_Stokes(v) + 1 0.501546 0.325639 0.09 ok [convect] solve:SNES_AdvectionDiffusion_Composed(T) > history_shift:EulerianSUPG(T) > solve:SNES_Stokes(v) + 2 0.990939 0.489393 0.09 ok [convect] solve:SNES_AdvectionDiffusion_Composed(T) > history_shift:EulerianSUPG(T) > solve:SNES_Stokes(v) + 3 1.51459 0.523655 0.09 ok [convect] solve:SNES_AdvectionDiffusion_Composed(T) > history_shift:EulerianSUPG(T) > solve:SNES_Stokes(v) + 4 30.5663 29.0517 0.09 ABANDONED [too big] solve:SNES_AdvectionDiffusion_Composed(T) > history_shift:EulerianSUPG(T) > solve:SNES_Stokes(v) + -- restore from a snapshot; the clock now reads 1.51459 Myr + -- rewind to the start of step 3 (t = 0.990939 Myr); 1 step(s) undone + 3 1.51459 0.523655 0.42 ok [replay] solve:SNES_AdvectionDiffusion_Composed(T) > history_shift:EulerianSUPG(T) > solve:SNES_Stokes(v) +``` + +`wall/s` is how long the block took. It is not physics, but it is the number +you want when watching: a step that suddenly takes ten times as long is the +first sign of a solver in trouble. + +**A true log records the backtracks.** `rewind()` and a bare `load_state()` +each write their own line, because a log that shows step 3, then step 3 again +with nothing in between, is not a log of what happened. `rewind` writes the +more specific note and suppresses the generic one. + +Four other things go in the file that are not in `model.journal`, all +deliberate: + +- **An abandoned step.** A rejected step is the part of a run's history that is + otherwise invisible, and it is usually what you want when asking why a run + went the way it did. +- **A step aged out by `journal_limit`.** The account of what happened outlives + both the state and the bounded in-memory list. +- **An invariant complaint**, as an `invariant` event on the step, so it + survives the terminal the run happened to have. +- **Everything up to a kill.** The file is flushed per step. + +### For parsing: JSON lines + +A path ending `.jsonl`, `.ndjson` or `.json` — or `model.journal_format = +"jsonl"` — writes the same record as one JSON object per line: + +```json +{"kind": "run", "model": "default", "started": "2026-09-10T21:22:40+00:00", "scales": {"length": {"magnitude": 2200000.0, "units": "meter"}, ...}} +{"kind": "step", "index": 0, "label": "convect", + "t0": {"magnitude": 0.0, "units": "megayear"}, + "t1": {"magnitude": 0.1759, "units": "megayear"}, + "dt": {"magnitude": 5551210433127.7, "units": "second"}, + "completed": true, "restorable": true, "wall": 0.44, + "events": [{"kind": "solve", "name": "SNES_AdvectionDiffusion_Composed(T)"}, + {"kind": "history_shift", "name": "EulerianSUPG(T)", "dt": 1.1469e-06}, + {"kind": "solve", "name": "SNES_Stokes(v)"}]} +{"kind": "rewind", "message": "...", "to_step": 3, "steps_undone": 1, "t": {"magnitude": 0.9909, "units": "megayear"}} +``` + +Read it back with `uw.read_journal(path)`, which returns one entry per run — an +inversion driver that ran the forward model thirteen times leaves thirteen runs +in one file, delimited by the header `clear_journal()` writes. + +**Why JSON lines and not YAML.** One self-contained record per line is the +whole point. A killed run leaves a truncated final line that *fails* to parse, +so `read_journal` drops it and keeps everything before; a half-written YAML +mapping frequently still parses, as a real record with its last key missing. +Line-oriented also means `grep`, `wc -l` and `jq -c` work without a parser, and +`json` is stdlib with predictable float round-tripping. YAML is the right +format for a whole document written once and edited by hand — which is what +`Model.to_yaml` uses it for — but a log is a stream. + +The two formats differ in one more way. The text log is a **report**: the time +column is converted into one unit, named in the header. The JSON log is a +**record**: every value keeps the units the run actually held it in, which is +why `t0` may read in Myr beside a `dt` in seconds — the clock came from the +tracker and the interval from `estimate_dt()`. + +Rank 0 writes; the other ranks record in memory as usual. + +### The same account, as a figure + +A terminal is not where a run belongs in a paper. + +```python +uw.journal_diagram(model, out="figures/run.pdf") # or a .jsonl log +uw.journal_diagram(model, out="figures/run.svg") # same figure, SVG +uw.journal_flowchart(model) # Mermaid, for docs +``` + +`journal_diagram` puts **time down the page**: one row per step, A4 portrait, +paginated, so it drops into a document column and opens anywhere. Each row +carries the step index, the clock, `dt` as a number and as a bar, a wall-clock +tick, and one letter. Backtracks are drawn in the left gutter as an arrow from +the step that ended back up to the step it returned to. + +The PDF and the SVG are both written directly — no plotting library, no +rasterisation, nothing fetched at render time, and a print-safe palette that +separates in greyscale. The `dt` axis goes logarithmic when the range exceeds +20x and says so: a rejected step is often tens of times the accepted ones, +which is *why* it was rejected, and on a linear axis it flattens everything +else to nothing. + +**The letter is the layout.** Each distinct operator sequence gets one, defined +once at the foot of the figure: + +``` +step t/Myr dt/Myr seq dt + 9 6.472 1.176 A ▇▇▇▇▇▇▇▇▇▇▇▇ + 10 7.936 1.464 A ▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇ + 11 9.856 1.920 A ▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇ + 12 145 135.2 A ┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄ abandoned + 11 9.856 1.920 A ▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇ + 12 12.56 2.704 B ! ▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇ + +A AdvectionDiffusion(T) > shift EulerianSUPG(T) > Stokes(v) 14 steps +B AdvectionDiffusion(T) > shift EulerianSUPG(T) > Stokes(v) > + AdvectionDiffusion(T) > shift EulerianSUPG(T) > Stokes(v) 1 step +``` + +A column of `A` with a single `B` in it says at a glance that one step did +something different. A hundred spelled-out sequences say nothing and hide the +one that matters. + +`journal_flowchart` renders one step's operator flow as Mermaid. When a run has +more than one distinct sequence, each becomes its own subgraph labelled with +the steps that took it, so an anomalous step is visible rather than averaged +away. + +Both accept a live model, a `.jsonl` log, or the list `read_journal` returns. +Not a text log: that one is a report, and reading it back is refused with the +one line that fixes it. + +### What the record checks + +A step also checks that it can be what it claims to be. One invariant so far: +a history manager must advance exactly once per step. + +``` + history_shift:EulerianSUPG(T) -> solve:SNES_Stokes(V)> +``` + +Call a solver twice inside one step — a corrector, a Picard iteration on a +coupled system, a retry — and its history advances twice, so the physical step +is taken twice. The solve counter and the timestep history look identical to a +single step, so nothing else in the library can see it. The step warns. + +If a solver genuinely is called more than once within a step, only the last +call should carry the timestep. + ### Backstepping The pattern above is what makes speculative stepping safe: @@ -548,6 +750,26 @@ of model state, so two independent runs of the same problem on the same solver objects diverge at the 1e-13 level from the first step. If you need to look at a step twice, restore it rather than re-run it. +### Two worked cases + +**`docs/examples/convection/intermediate/Ex_Convection_Annulus_Recorded.py`** — +Boussinesq convection in an annulus. Four reference quantities, a body force +written as a force (Ra falls out of the nondimensionalisation rather than being +typed in), rotated free-slip on the curved boundaries, and a varying +`estimate_dt()`. It then demonstrates the four things the record buys, in +order: the journal, a rejected step, a bit-exact replay, and the invariant +catching a step that was taken twice. Compare +`../advanced/Ex_Convection_Cylinder.py`, which solves the same physics with a +bare `for step in range(n)` loop and no clock at all. + +**An adjoint driven from the journal.** The backward pass of a discrete adjoint +needs exactly what the record holds: the state at each step and the order the +operators were applied in. Walking `model.journal` backwards — +`load_state(entry.snapshot)`, replay, transpose-solve — replaces the +hand-written checkpoint dictionary that an adjoint normally carries, and +removes its dependence on knowing in advance which arrays the backward pass +will want. + ### Time-dependent expressions `mesh.t` is the model clock as a symbol. It is repacked from @@ -569,6 +791,29 @@ keeps the symbol. ## Common Pitfalls and Anti-Patterns +### ❌ Rebinding the name instead of setting `.sym` + +To change the value of an expression, set `.sym`. It is the only settable +property — `.value` and `.data` are derived, read-only views. + +```python +# ✅ CORRECT - a value change; the container keeps its identity +viscosity.sym = sympy.Integer(0) +solver._update_constants() # only if you are not about to solve + +# ❌ WRONG - rebinds a Python name and changes nothing +viscosity = 0 +``` + +The second line leaves every expression that already references the atom +pointing at the old object with its old value, and nothing complains. The +identity is the point: because the container is unchanged, a ramped value +reaches every residual that mentions it with no rebuild. + +`expr.copy(other)` does the same job from another expression, and assigning to +a constitutive parameter slot (`Parameters.diffusivity = 0.0`) is also a value +change rather than a replacement. + ### ❌ Swarm Variable Creation After Population ```python @@ -852,6 +1097,7 @@ TypeError: unsupported operand type(s) for *: 'UnitAwareDerivativeMatrix' and 'N - [ ] Declare the model and its reference quantities BEFORE creating the mesh - [ ] Keep `time`, `step` and `dt` on `model.tracker`, not in local variables - [ ] Wrap each step in `with model.step(dt):` +- [ ] To change an expression's value set `.sym`, never rebind the name - [ ] Take snapshots BEFORE the operator you might want to undo - [ ] Use `mesh.t` inside an expression for time dependence, never bare @@ -894,6 +1140,9 @@ TypeError: unsupported operand type(s) for *: 'UnitAwareDerivativeMatrix' and 'N - Disk snapshots now carry dimensional values (magnitude + units) - `mesh.t` now resolves to the model clock (#410) - `model.step(dt)` — the step as a transaction, and the step journal + - `model.record_every` / `model.rewind()` — the journal as a restorable record + - A step warns when a history advances more than once + - Set `.sym` to change a value; rebinding the name changes nothing - Backstepping recipe; snapshot before the operator - `mesh.t` is not the model clock and is silently zero in a solve - **2025-11-15**: Initial version diff --git a/docs/examples/convection/README.md b/docs/examples/convection/README.md index 72717925f..966b95fd1 100644 --- a/docs/examples/convection/README.md +++ b/docs/examples/convection/README.md @@ -42,6 +42,15 @@ Thermal convection combines heat transfer and fluid mechanics to model buoyancy- - Multiple convection cells and interactions - Introduces: domain geometry effects, cell interactions +7. **Annulus Convection, Recorded** - `Ex_Convection_Annulus_Recorded.py` + - Boussinesq convection in a 2D annulus, written in the timestepping pattern + - Reference quantities first, so the buoyancy is written as a force and the + Rayleigh number falls out of the nondimensionalisation + - Rotated free-slip on the curved boundaries; a varying `estimate_dt()` + - Demonstrates the run's own record: the journal, a rejected step, a + bit-exact replay, and the step invariant that catches a doubled step + - See `docs/developer/guides/HOW-TO-WRITE-UW3-SCRIPTS.md` + ### 🎓 Advanced Examples (`advanced/`) **Complex convection systems.** diff --git a/docs/examples/convection/intermediate/Ex_Convection_Annulus_Recorded.py b/docs/examples/convection/intermediate/Ex_Convection_Annulus_Recorded.py new file mode 100644 index 000000000..e76438340 --- /dev/null +++ b/docs/examples/convection/intermediate/Ex_Convection_Annulus_Recorded.py @@ -0,0 +1,499 @@ +# --- +# jupyter: +# jupytext: +# formats: py:percent +# text_representation: +# extension: .py +# format_name: percent +# format_version: '1.3' +# kernelspec: +# display_name: Python 3 +# language: python +# name: python3 +# --- + +# %% [markdown] +""" +# Convection in an Annulus — a recorded run + +**PHYSICS:** convection +**DIFFICULTY:** intermediate + +## Description + +Boussinesq thermal convection in a 2D annulus, written in the timestepping +pattern: the model and its reference quantities come first, the clock lives on +`model.tracker`, and each timestep is a `model.step(dt)` block. + +Compare `../advanced/Ex_Convection_Cylinder.py`, which solves the same physics +with a bare `for step in range(n)` loop and no clock at all. The physics here is +unchanged. What the pattern adds is that the run keeps an account of itself, and +that account is worth four things this script demonstrates in turn: + +1. **what ran** — an ordered journal, named by what each solver solves +2. **a rejected step** — the clock does not move when a step is abandoned +3. **playback** — a recorded step replays bit-for-bit, where a re-run does not +4. **an invariant** — a step that took the physical step twice says so +5. **a log on disk** — the same account in aligned columns, flushed as each + step closes, so a run that dies keeps its history and a run in progress can + be watched with `tail -f` +6. **a figure** — the same account as a portrait PDF (or SVG), and the step's + operator flow as Mermaid + +## Key concepts + +- **Units first.** Four reference quantities fix the scaling, and the body + force is then written as the physics — `-rho0 alpha T g rhat` — rather than + as a Rayleigh number. Ra falls out of the nondimensionalisation; the script + prints it so you can check. +- **Rotated free-slip on curved boundaries.** `add_rotated_freeslip_bc` + enforces `v.n = 0` to machine precision on a circle, where a penalty or + Nitsche condition leaks at ~1e-3. +- **A varying timestep.** `estimate_dt()` returns a dimensional quantity that + goes straight into `model.step(dt)` and `adv.solve(timestep=dt)`. With an + implicit SUPG transport this is an accuracy choice, not a stability limit. + +## Parameters + +Override from the command line, e.g. `-uw_n_steps 20 -uw_cell_size 0.075`. +""" + +# %% +import warnings + +import numpy as np +import sympy + +import underworld3 as uw + + +def say(*args): + """Rank-safe print that keeps its own formatting. + + `uw.pprint`'s default `clean_display=True` rewrites the string it is given + — it strips braces and collapses runs of whitespace — so an aligned table + printed through it loses its columns. Pass `clean_display=False` whenever + the layout is yours rather than SymPy's. + """ + uw.pprint(*args, clean_display=False) + + +params = uw.Params( + uw_cell_size=0.1, # mesh resolution, as a fraction of the outer radius + uw_n_steps=8, # timesteps in the recorded run + uw_dt_fraction=0.5, # accuracy factor on estimate_dt() + uw_demos=1, # run the journal demonstrations after the loop + uw_journal_file="output/annulus_convection.log", +) + +# %% [markdown] +""" +## The model comes first + +Reference quantities must be set BEFORE the mesh is created, so the model is +the first thing the script declares rather than something the mesh conjures for +you. Four quantities fix all four dimensions this problem uses: + +| quantity | fixes | +|---|---| +| `shell_thickness` | length | +| `thermal_diffusivity` | time, as `d^2 / kappa` | +| `mantle_viscosity` | mass | +| `temperature_contrast` | temperature | +""" + +# %% +uw.reset_default_model() +model = uw.get_default_model() + +SHELL_THICKNESS = uw.quantity(2200, "km") +KAPPA = uw.quantity(1e-6, "m**2/s") +ETA = uw.quantity(1e22, "Pa*s") +DELTA_T = uw.quantity(2500, "K") + +RHO0 = uw.quantity(3300, "kg/m**3") +ALPHA = uw.quantity(3e-5, "1/K") +GRAVITY = uw.quantity(9.81, "m/s**2") + +model.set_reference_quantities( + shell_thickness=SHELL_THICKNESS, + thermal_diffusivity=KAPPA, + mantle_viscosity=ETA, + temperature_contrast=DELTA_T, +) + +RAYLEIGH = (RHO0 * ALPHA * DELTA_T * GRAVITY * SHELL_THICKNESS**3 / (KAPPA * ETA)) +say(f"Ra = {float(RAYLEIGH.to('dimensionless').magnitude):.3e}") +say(f"diffusion time d^2/kappa = " + f"{model.get_fundamental_scales()['time'].to('Gyr')}") + +# %% [markdown] +""" +## Mesh and variables + +An annulus with `radiusInner / radiusOuter = 0.55`, roughly Earth's +core-mantle ratio. The mesh is built in model units; `mesh.X.coords` reads back +in metres because the model declares a length scale. +""" + +# %% +R_OUTER = 1.0 +R_INNER = 0.55 + +mesh = uw.meshing.Annulus( + radiusInner=R_INNER, + radiusOuter=R_OUTER, + cellSize=params.uw_cell_size, + degree=1, + qdegree=3, +) + +v = uw.discretisation.MeshVariable("v", mesh, mesh.dim, degree=2) +p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1) +T = uw.discretisation.MeshVariable("T", mesh, 1, degree=3) + +# %% [markdown] +""" +## Stokes: rotated free-slip, and the body force as physics + +`add_rotated_freeslip_bc(0.0, boundary)` rotates each boundary node into its +own normal / tangent frame and constrains the normal component strongly. On a +circle that is exact to machine precision; a penalty or Nitsche condition +leaks at around 1e-3, which on a convection run shows up as spurious radial +flow at the boundary. + +The buoyancy is written as the force it is, in the units it has. The +nondimensionalisation turns it into `Ra T rhat` — that is where the Rayleigh +number printed above comes from, and writing it this way means the script +never has to be told what Ra is. +""" + +# %% +stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) +stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel +stokes.constitutive_model.Parameters.shear_viscosity_0 = ETA +stokes.tolerance = 1.0e-8 +stokes.petsc_options.delValue("ksp_monitor") + +stokes.add_rotated_freeslip_bc(0.0, "Upper") +stokes.add_rotated_freeslip_bc(0.0, "Lower") + +radius = sympy.sqrt(mesh.X.dot(mesh.X)) +rhat = mesh.X / radius +stokes.bodyforce = -RHO0 * ALPHA * GRAVITY * T.sym[0] * rhat + +# %% [markdown] +""" +## Transport + +`uw.systems.AdvDiffusion` composes an implicit Eulerian SUPG transport step +(Crank-Nicolson by default). Hot inner boundary, cold outer. +""" + +# %% +adv = uw.systems.AdvDiffusion(mesh, u_Field=T, V_fn=v.sym) +adv.constitutive_model = uw.constitutive_models.DiffusionModel +adv.constitutive_model.Parameters.diffusivity = KAPPA +adv.add_dirichlet_bc(1.0, "Lower") +adv.add_dirichlet_bc(0.0, "Upper") +adv.tolerance = 1.0e-8 +adv.petsc_options.delValue("ksp_monitor") + +# %% [markdown] +""" +## Initial condition + +A conductive profile with a mode-5 perturbation. Built from the nodal +coordinates in model units, which is what `beta`-style level sets and initial +conditions generally want: `T.coords` reads in metres, so divide by the length +scale once and work in the box's own units. +""" + +# %% +LENGTH_SCALE = float(model.get_fundamental_scales()["length"].to("m").magnitude) + + +def myr(q): + """A time quantity as a Myr string. `UWQuantity.__format__` delegates to + the bare float, so pint's `~` format specs do not apply to it.""" + return f"{float(q.to('Myr').magnitude):.4f} Myr" + +Xn = np.asarray(T.coords)[:, :2] / LENGTH_SCALE +rn = np.sqrt((Xn**2).sum(axis=1)) +thn = np.arctan2(Xn[:, 1], Xn[:, 0]) +shell = (rn - R_INNER) / (R_OUTER - R_INNER) + +T.array[:, 0, 0] = (1.0 - shell) + 0.1 * np.sin(5.0 * thn) * np.sin(np.pi * shell) + +adv.Unknowns.DuDt.initialise_history() +stokes.solve(zero_init_guess=True) + +# %% [markdown] +""" +## Diagnostics on the tracker + +`model.tracker.time`, `.step` and `.dt` are pre-seeded by convention. Anything +else you assign to it is captured by a snapshot and restored by a rewind, in +the same breath as the fields — which is exactly what a diagnostic wants, and +what a loose Python variable cannot give you. +""" + +# %% +v_rms_fn = sympy.sqrt(v.sym.dot(v.sym)) +area = float(uw.maths.Integral(mesh, sympy.sympify(1.0)).evaluate()) + + +def v_rms(): + return float(uw.maths.Integral(mesh, v_rms_fn).evaluate()) / area + + +# %% [markdown] +""" +## The loop + +Everything above is ordinary. This is the pattern: + +```python +while ...: + dt = adv.estimate_dt() + with model.step(dt, label="convect"): + adv.solve(timestep=dt, zero_init_guess=False) + stokes.solve(zero_init_guess=False) +``` + +`record_every = 1` asks each step to keep the state it started from — every +field, the transport history, the clock and the tracker diagnostics together — +captured before the operators run, which is the only correct point. +""" + +# %% +model.tracker.time = uw.quantity(0.0, "Myr") +model.tracker.step = 0 +model.tracker.dt = None +model.tracker.v_rms = v_rms() + +model.record_every = 1 +model.record_limit = params.uw_n_steps + +# The in-memory journal is what the run can still UNDO; it is bounded and it +# dies with the process. The log is what the run DID: one aligned line per step, +# appended and flushed as each step closes, including the steps that were +# abandoned and the backtracks. Setting it is optional and costs a line per +# step. A `.jsonl` suffix (or `model.journal_format = "jsonl"`) writes the same +# record as JSON objects instead, for parsing rather than reading. +model.journal_file = str(params.uw_journal_file) + +for _ in range(int(params.uw_n_steps)): + dt = params.uw_dt_fraction * adv.estimate_dt() + + with model.step(dt, label="convect"): + adv.solve(timestep=dt, zero_init_guess=False) + stokes.solve(zero_init_guess=False) + model.tracker.v_rms = v_rms() + + say(f"step {model.tracker.step:>3d} " + f"t = {myr(model.tracker.time)} " + f"dt = {myr(dt)} " + f"v_rms = {model.tracker.v_rms:.4e}") + +# %% [markdown] +""" +## 1. What ran + +The journal is an ordered account of each step: the interval it covered and +the operators it applied, named by what they solve. It answers "is this model +doing the thing the write-up says it does" without the script being +instrumented for it — which is the question you want to ask of someone else's +model, or of your own six months later. + +Note the `history_shift` between the two solves. That is the transport history +advancing, and it is what the step's invariant checks. +""" + +# %% +if params.uw_demos: + say("") + say("--- 1. the journal " + "-" * 55) + for entry in model.journal: + say(f" step {entry.index:>2d} dt = {myr(entry.dt):>12s} " + + " -> ".join(f"{e['kind']}:{e['name']}" for e in entry.events)) + say(f" {len(model.restore_points)} of {len(model.journal)} steps " + f"are restorable") + +# %% [markdown] +""" +## 2. A rejected step + +A `model.step` block is a transaction. If it does not exit cleanly — an +exception, or a step abandoned because a diagnostic came out wrong — the clock +and the step counter are left exactly as they were, and nothing is added to the +journal. Backstepping no longer has to remember to unwind a counter. + +The fields are yours to restore: take a snapshot before the block, and load it +in the handler. The clock never moved, so the two stay consistent. +""" + + +# %% +class StepRejected(Exception): + """Raised inside a step block to abandon it.""" + + +if params.uw_demos: + say("") + say("--- 2. a rejected step " + "-" * 51) + + before = (myr(model.tracker.time), model.tracker.step, len(model.journal)) + snap = model.save_state() # BEFORE the step, not after + + reckless_dt = 50.0 * params.uw_dt_fraction * adv.estimate_dt() + try: + with model.step(reckless_dt, label="too big"): + adv.solve(timestep=reckless_dt, zero_init_guess=False) + stokes.solve(zero_init_guess=False) + if v_rms() > 4.0 * model.tracker.v_rms: + raise StepRejected("v_rms jumped; the step is not resolved") + except StepRejected as why: + model.load_state(snap) + say(f" rejected: {why}") + + after = (myr(model.tracker.time), model.tracker.step, len(model.journal)) + say(f" clock/step/journal before : {before}") + say(f" clock/step/journal after : {after}") + say(f" unchanged: {before == after}") + +# %% [markdown] +""" +## 3. Playback + +`model.rewind()` puts the run back to the start of a completed step — fields, +transport history, clock and tracker diagnostics together — and truncates the +journal to match, so it continues to describe the run that actually happened. + +Replaying the step from there reproduces it exactly. Re-*running* the script +does not: warm starts and preconditioner reuse are solver history rather than +model state, so two independent runs of the same problem diverge at the 1e-13 +level from the first step. If you need to look at a step twice, restore it +rather than re-run it. +""" + +# %% +if params.uw_demos: + say("") + say("--- 3. playback " + "-" * 58) + + T_end = np.asarray(T.array)[:, 0, 0].copy() + v_rms_end = model.tracker.v_rms + + target = model.rewind() + say(f" rewound to the start of step {target.index}: " + f"t = {myr(model.tracker.time)}, " + f"v_rms = {model.tracker.v_rms:.4e} " + f"(was {v_rms_end:.4e})") + + with model.step(target.dt, label="replay"): + adv.solve(timestep=target.dt, zero_init_guess=False) + stokes.solve(zero_init_guess=False) + model.tracker.v_rms = v_rms() + + T_replay = np.asarray(T.array)[:, 0, 0] + say(f" replayed: identical to the original step: " + f"{np.array_equal(T_replay, T_end)} " + f"max |dT| = {np.abs(T_replay - T_end).max():.3e}") + +# %% [markdown] +""" +## 4. An invariant + +A history manager must advance exactly once per step. Advancing twice means +the step was taken twice — a corrector, a Picard iteration on the coupled +system, or a retry that called the solver again — and the temperature moves +two intervals while the timestep history and the solve counter look identical +to a single step. Nothing else in the library can see that. + +The step says so. If a solver genuinely is called more than once within a step, +only the last call should carry the timestep. +""" + +# %% +if params.uw_demos: + say("") + say("--- 4. the invariant " + "-" * 53) + + dt = params.uw_dt_fraction * adv.estimate_dt() + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + with model.step(dt, label="taken twice"): + adv.solve(timestep=dt, zero_init_guess=False) # a "predictor" + stokes.solve(zero_init_guess=False) + adv.solve(timestep=dt, zero_init_guess=False) # and a "corrector" + stokes.solve(zero_init_guess=False) + + for w in caught: + if issubclass(w.category, RuntimeWarning): + say(" " + " ".join(str(w.message).split())[:200]) + say(f" the step as recorded: {model.journal[-1]}") + +# %% [markdown] +""" +## 5. The log on disk + +`model.journal_file` writes the same account to a file, one line per step, +flushed as it closes — so `tail -f` on it follows a running job, and a run that +is killed keeps everything up to the moment it died. + +Three differences from `model.journal`, all deliberate. An **abandoned** step +appears in the file and not in memory. A step aged out by `journal_limit` +leaves memory but stays in the file. And a **backtrack** — `rewind()` or a +bare `load_state()` — writes its own line, because a log that shows step 7 and +then step 7 again, with nothing in between, is not a log of what happened. +""" + +# %% +if params.uw_demos: + say("") + say("--- 5. the log on disk " + "-" * 51) + say(f" {model.journal_file}") + + with open(model.journal_file, encoding="utf-8") as handle: + for line in handle.read().splitlines(): + say(" " + line) + +# %% [markdown] +""" +## 6. The same account, as a figure + +A terminal is not where a run belongs in a paper. `uw.journal_diagram` renders +the record with **time running down the page** — one row per step, A4 portrait, +paginated — and writes it as a PDF, which opens anywhere, or an SVG if the +suffix says so. Both are written directly: no plotting library, no +rasterisation, no theme to fight with. `uw.journal_flowchart` renders one +step's operator flow as Mermaid, for dropping into documentation. + +The layout decision worth knowing about: each distinct operator sequence gets a +**letter**, defined once at the foot of the figure. A column of `A` with a +single `B` in it says at a glance that one step did something different, where +a hundred spelled-out sequences say nothing and hide the one that matters. The +doubled step below is found that way rather than by reading. + +Both take a live model, which matters here because the text log is a report and +cannot be read back — pass a `.jsonl` log or the model itself. +""" + +# %% +if params.uw_demos: + say("") + say("--- 6. the figure " + "-" * 56) + + stem = model.journal_file.rsplit(".", 1)[0] + say(f" {uw.journal_diagram(model, out=stem + '.pdf', title='Annulus convection - run log')}") + say(f" {uw.journal_diagram(model, out=stem + '.svg', title='Annulus convection - run log')}") + say("") + for line in uw.journal_flowchart(model).splitlines(): + say(" " + line) + +# %% +say("") +say(f"final: t = {myr(model.tracker.time)}, " + f"{model.tracker.step} steps, " + f"v_rms = {model.tracker.v_rms:.4e}") diff --git a/scripts/test.sh b/scripts/test.sh index f58ca36f9..eff3eebc6 100755 --- a/scripts/test.sh +++ b/scripts/test.sh @@ -134,6 +134,7 @@ if [ $PARALLEL_ONLY -eq 0 ]; then # unbatched file would have closed the issue without closing the gap. # level_2/tier_b, ~55s serial; passes at np=1 and np=2. $PYTEST tests/test_1072_free_surface_spherical.py || status=1 + $PYTEST tests/test_1074_free_surface_config_drift.py || status=1 # Diffusion / Advection tests $PYTEST tests/test_1100*py || status=1 diff --git a/src/underworld3/__init__.py b/src/underworld3/__init__.py index 4ed69e815..176283e79 100644 --- a/src/underworld3/__init__.py +++ b/src/underworld3/__init__.py @@ -218,9 +218,11 @@ def view(): create_model, get_default_model, reset_default_model, + read_journal, ThermalConvectionConfig, create_thermal_convection_model, ) +from .utilities.journal_report import journal_diagram, journal_flowchart from .parameters import ParameterRegistry, ParameterType from .materials import MaterialRegistry, MaterialProperty from .constitutive_models import MultiMaterialConstitutiveModel diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index 78a290861..c796b07f1 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -2260,11 +2260,18 @@ class SolverBaseClass(uw_object): cdef double[::1] vals_view = np.ascontiguousarray(values, dtype=np.float64) CHKERRQ(PetscDSSetConstants(cds.ds, n_constants, &vals_view[0])) - def _update_constants(self): + def _update_constants(self, record=True): """Re-pack current UWexpression values and call PetscDSSetConstants. Called before each solve() to ensure constants are current without requiring JIT recompilation. + + ``record=False`` suppresses the step-journal entry. Pass it from any + site that pushes constants for its OWN assembly rather than to + dispatch a solve — otherwise the journal reports one operator as two. + The rotated free-slip loop is such a site: it re-attaches the + auxiliary vector and re-packs before running its own manual Krylov + loop, after the public ``solve()`` has already announced itself. """ # Refresh mesh.t from the model clock first, so a time-dependent # expression is repacked with the rest of the constants rather than @@ -2277,19 +2284,20 @@ class SolverBaseClass(uw_object): # Note the solve in the model's step journal, if a step is open. This # is the one place every solver passes through before solving, so one # hook records them all, in order. A no-op outside a model.step block. - try: - # Name it by what it SOLVES, not by its auto-generated instance id: - # a journal reading "Stokes(V) -> AdvDiffusion(T)" is auditable, - # one reading "Solver_8_ -> Solver_14_" is not. + if record: try: - unknown = self.u.name + # Name it by what it SOLVES, not by its auto-generated instance + # id: a journal reading "Stokes(V) -> AdvDiffusion(T)" is + # auditable, one reading "Solver_8_ -> Solver_14_" is not. + try: + unknown = self.u.name + except Exception: + unknown = "?" + uw.get_default_model()._record_step_event( + "solve", f"{type(self).__name__}({unknown})" + ) except Exception: - unknown = "?" - uw.get_default_model()._record_step_event( - "solve", f"{type(self).__name__}({unknown})" - ) - except Exception: - pass + pass if not self.constants_manifest or self.dm is None: return diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index 16dcdd25a..4ee75bf84 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/src/underworld3/discretisation/discretisation_mesh.py @@ -5241,7 +5241,12 @@ def snapshot_payload(self) -> dict: - ``name``: stable string identifier for the mesh. - ``mesh_version``: current ``_mesh_version`` integer. - - ``coords``: deformed mesh coordinates (numpy array). + - ``coords``: deformed mesh coordinates, in MODEL UNITS — the + representation :meth:`_deform_mesh` writes back. ``mesh.X.coords`` + is the unit-aware view and returns metres when a model declares a + length scale; capturing that and restoring it through + ``_deform_mesh`` would multiply the mesh by the length scale on + every restore, silently and without changing any array's shape. - ``vars``: ``{var.clean_name: gvec_array.copy()}`` for every mesh variable on this mesh. @@ -5249,7 +5254,7 @@ def snapshot_payload(self) -> dict: section / DM-topology data sufficient to rebuild the DM on restore. """ - coords = numpy.asarray(self.X.coords).copy() + coords = numpy.asarray(self._coords).copy() var_arrays: dict[str, numpy.ndarray] = {} for var in self.vars.values(): var._sync_lvec_to_gvec() @@ -5294,7 +5299,7 @@ def apply_snapshot_payload(self, payload: dict) -> None: ) coords = numpy.asarray(payload["coords"]) - expected_shape = numpy.asarray(self.X.coords).shape + expected_shape = numpy.asarray(self._coords).shape if coords.shape != expected_shape: raise SnapshotInvalidatedError( f"mesh {self.name!r}: coordinate shape changed " diff --git a/src/underworld3/model.py b/src/underworld3/model.py index c33eae69f..178784bc2 100644 --- a/src/underworld3/model.py +++ b/src/underworld3/model.py @@ -63,7 +63,8 @@ class ModelStep: needs in order to walk the run backwards. """ - __slots__ = ("index", "t0", "dt", "label", "events", "completed") + __slots__ = ("index", "t0", "dt", "label", "events", "completed", "snapshot", + "wall") def __init__(self, index, t0, dt, label=None): self.index = index @@ -72,6 +73,20 @@ def __init__(self, index, t0, dt, label=None): self.label = label self.events = [] self.completed = False + # Seconds of wall clock the block took. Not physics, but the number you + # want when watching a run: a step that suddenly takes ten times as + # long is the first sign of a solver in trouble. + self.wall = None + # The state this step STARTED from, when the recording policy kept one. + # Taken before the operators ran, which is the only correct point: a + # DDt shifts its history in its post-solve hook, so a snapshot taken + # afterwards holds the shifted history rather than the step's input. + self.snapshot = None + + @property + def restorable(self): + """Whether this step kept the state it started from.""" + return self.snapshot is not None @property def t1(self): @@ -81,6 +96,71 @@ def t1(self): def _record(self, kind, name, **detail): self.events.append({"kind": kind, "name": name, **detail}) + def _check_invariants(self): + """Complain about a step that cannot be what it claims to be. + + One invariant so far, and it catches a mistake that is otherwise + invisible: a history manager must advance EXACTLY ONCE per step. Twice + means the step was taken twice — a corrector, a Picard iteration or a + retry that called the solver again — and the field advances twice while + the solve counter and the timestep history look identical to a single + step. + """ + import warnings + from collections import Counter + + shifts = Counter( + e["name"] for e in self.events if e["kind"] == "history_shift" + ) + repeated = {name: n for name, n in shifts.items() if n > 1} + if repeated: + detail = ", ".join(f"{name} x{n}" for name, n in sorted(repeated.items())) + # Also record it against the step, so the log and the journal carry + # the complaint and not just the terminal the run happened to have. + self.events.append({ + "kind": "invariant", + "name": "history advanced more than once", + "detail": detail, + }) + warnings.warn( + f"step {self.index}: history advanced more than once ({detail}). " + f"The step has been taken more than once, so the field is " + f"further ahead than dt says. If a solver is called twice " + f"within one step deliberately — a corrector or a Picard " + f"iteration — only the last call should carry the timestep.", + RuntimeWarning, + stacklevel=3, + ) + + def as_dict(self): + """This step as plain JSON-able data — the on-disk log's line format. + + Dimensional values become ``{"magnitude": ..., "units": ...}``, the + same split the on-disk snapshot uses, so a log written by a run with + units is readable without a live model to interpret it. + + ``snapshot`` is deliberately absent: it is megabytes of field data and + does not survive the process. ``restorable`` records whether one was + held, which is what a reader of the log can act on. + + Each value keeps the units the run actually held it in, which is why + ``t0`` may read in Myr beside a ``dt`` in seconds: the clock came from + the tracker and the interval from ``estimate_dt()``. The log is a + record of the run, not a tidied report of it — convert on the way out. + """ + return { + "kind": "step", + "index": self.index, + "label": self.label, + "t0": _jsonable_quantity(self.t0), + "t1": _jsonable_quantity(self.t1), + "dt": _jsonable_quantity(self.dt), + "completed": bool(self.completed), + "restorable": bool(self.restorable), + "wall": None if self.wall is None else float(self.wall), + "events": [dict(e) for e in self.events], + } + def __repr__(self): state = "" if self.completed else " ABANDONED" seq = " -> ".join(f"{e['kind']}:{e['name']}" for e in self.events) or "(nothing)" @@ -88,6 +168,104 @@ def __repr__(self): return f"" +def _quantity_parts(value): + """``(magnitude, unit string or None)`` for a value that may be dimensional.""" + if hasattr(value, "magnitude") and hasattr(value, "units"): + try: + return float(value.magnitude), str(value.units) + except (TypeError, ValueError): + return None, str(value.units) + try: + return float(value), None + except (TypeError, ValueError): + return None, None + + +def _abbreviate_unit(unit): + """A short unit name for a column header. Falls back to the full name.""" + if unit is None: + return "" + return { + "second": "s", "minute": "min", "hour": "hr", "day": "d", + "year": "yr", "kiloyear": "kyr", "megayear": "Myr", "gigayear": "Gyr", + "meter": "m", "kilometer": "km", "kelvin": "K", "kilogram": "kg", + }.get(str(unit), str(unit)) + + +def _in_units_of(value, unit): + """``value`` as a bare number in ``unit``, or its own magnitude if it cannot + be converted. A text log is a report: one time column, one unit.""" + if unit is None: + magnitude, _ = _quantity_parts(value) + return magnitude + try: + return float(value.to(unit).magnitude) + except Exception: + magnitude, _ = _quantity_parts(value) + return magnitude + + +def _bare(value, unit): + """A serialised value as a bare number, converted to ``unit`` if it can be. + + ``value`` is what :meth:`ModelStep.as_dict` produced: a float, or a + ``{"magnitude", "units"}`` pair. The text log shows one time column in one + unit, so a ``dt`` in seconds beside a clock in Myr is converted rather than + printed as it stands. + """ + if not isinstance(value, dict): + try: + return float(value) + except (TypeError, ValueError): + return float("nan") + magnitude = value.get("magnitude") + units = value.get("units") + if unit is None or units is None or str(units) == str(unit): + try: + return float(magnitude) + except (TypeError, ValueError): + return float("nan") + try: + import underworld3 as uw + + return float(uw.quantity(float(magnitude), str(units)).to(unit).magnitude) + except Exception: + try: + return float(magnitude) + except (TypeError, ValueError): + return float("nan") + + +def _pretty_time(value): + """A compact, readable rendering of a clock value for a log note.""" + magnitude, unit = _quantity_parts(value) + if magnitude is None: + return str(value) + if unit is None: + return f"{magnitude:.6g}" + return f"{magnitude:.6g} {_abbreviate_unit(unit)}" + + +def _jsonable_quantity(value): + """A number, or a dimensional value split into magnitude and units. + + Duck-typed, because ``uw.quantity`` returns a ``UWQuantity``, which is not + a ``pint.Quantity`` subclass — an isinstance test against either would + miss one of them. Both carry ``magnitude`` and ``units``. + """ + if hasattr(value, "magnitude") and hasattr(value, "units"): + magnitude = value.magnitude + try: + magnitude = float(magnitude) + except (TypeError, ValueError): + magnitude = str(magnitude) + return {"magnitude": magnitude, "units": str(value.units)} + try: + return float(value) + except (TypeError, ValueError): + return str(value) + + class Model(PintNativeModelMixin, BaseModel): """ Central orchestrator for Underworld3 simulations. @@ -183,6 +361,22 @@ class Model(PintNativeModelMixin, BaseModel): _journal: Any = PrivateAttr(default_factory=list) _journal_limit: Any = PrivateAttr(default=512) + # Optional on-disk log of the journal: one JSON object per line, appended + # and flushed as each step closes. See :attr:`journal_file`. + _journal_path: Any = PrivateAttr(default=None) + _journal_fh: Any = PrivateAttr(default=None) + _journal_format: Any = PrivateAttr(default=None) + _journal_columns: Any = PrivateAttr(default=None) + # Set while rewind() is doing its own restore, so load_state does not log a + # second, less informative note for the same backtrack. + _restoring: Any = PrivateAttr(default=False) + + # Recording policy: how often a step keeps a restorable snapshot of the + # state it started from, and how many of those to retain. See :meth:`step`. + _record_every: Any = PrivateAttr(default=None) + _record_limit: Any = PrivateAttr(default=8) + _record_warned: Any = PrivateAttr(default=False) + def __init__(self, name: Optional[str] = None, **kwargs): """ Initialize a new Model instance. @@ -679,11 +873,321 @@ def open_step(self): """The step in progress, or None outside a ``model.step`` block.""" return self._open_step + def clear_journal(self): + """Start a new run's journal, discarding the records and snapshots in it. + + A driver that runs the same model many times — an inversion, a + parameter sweep, a restart from a saved state — needs each run to have + its own account. Without this the journal is a concatenation of every + run the process has done, and ``rewind()`` will happily walk back into + the previous one. + + Does not touch the clock: reset ``model.tracker.time`` / ``step`` + yourself if the new run starts from zero. + """ + if self._open_step is not None: + raise RuntimeError( + "cannot clear the journal from inside a model.step block " + f"(step {self._open_step.index} is open)." + ) + self._journal.clear() + self._record_warned = False + # A new run gets a new section in the log rather than a new file, so + # one file holds the whole process — thirteen forward runs of an + # inversion, say — delimited by their headers. + self._write_journal_line(self._run_header()) + + @property + def journal_file(self): + """Path of the on-disk step log, or None (the default: memory only). + + Assign a path and every step that closes — completed OR abandoned — + is appended as one JSON object on its own line, and flushed. A run + that crashes keeps the log up to the crash, which is when it is worth + most. + + :: + + model.journal_file = "output/run.journal.jsonl" + + The file records what the run DID; ``model.journal`` is what it can + still UNDO. They differ in two ways, both deliberate: an abandoned step + appears in the file and not in memory, and a step trimmed by + ``journal_limit`` leaves memory but stays in the file. + + Read one back with :func:`underworld3.read_journal`. Rank 0 writes; + other ranks record in memory as usual. + """ + return self._journal_path + + @journal_file.setter + def journal_file(self, path): + if self._journal_fh is not None: + self._journal_fh.close() + self._journal_fh = None + self._journal_path = None if path is None else str(path) + self._journal_columns = None + if self._journal_path is None: + return + import underworld3 as uw + + if uw.mpi.rank != 0: + return + directory = os.path.dirname(self._journal_path) + if directory: + os.makedirs(directory, exist_ok=True) + self._journal_fh = open(self._journal_path, "w", encoding="utf-8") + self._write_journal_line(self._run_header()) + + @property + def journal_format(self): + """``"text"`` (default) or ``"jsonl"``. + + Text is for reading — aligned columns, one line per step, designed to + be watched with ``tail -f`` while a run is going. It is a report: the + time column is converted to a single unit named in the header. + + ``"jsonl"`` is for parsing — one JSON object per line, every value in + the units the run actually held it in. Chosen automatically when the + path ends ``.jsonl``, ``.ndjson`` or ``.json``; set this explicitly to + override. + """ + if self._journal_format is not None: + return self._journal_format + if self._journal_path and self._journal_path.lower().endswith( + (".jsonl", ".ndjson", ".json")): + return "jsonl" + return "text" + + @journal_format.setter + def journal_format(self, value): + if value not in (None, "text", "jsonl"): + raise ValueError( + f"journal_format must be 'text', 'jsonl' or None, not {value!r}") + self._journal_format = value + + def _run_header(self): + """The record that opens a run in the log, so the file is self-describing.""" + from datetime import datetime, timezone + + scales = {} + try: + for name, scale in (self.get_fundamental_scales() or {}).items(): + scales[str(name)] = _jsonable_quantity(scale) + except Exception: + scales = {} + return { + "kind": "run", + "model": getattr(self, "name", None), + "started": datetime.now(timezone.utc).isoformat(timespec="seconds"), + "scales": scales, + } + + # ------------------------------------------------------------------ + # Rendering + # ------------------------------------------------------------------ + + def _render_journal_text(self, payload): + """One record as human-readable text. Returns a string, possibly + several lines, or None for a record this format does not show.""" + kind = payload.get("kind") + + if kind == "run": + scales = payload.get("scales") or {} + summary = " | ".join( + f"{name} {value['magnitude']:.4g} {_abbreviate_unit(value['units'])}" + for name, value in scales.items() + if isinstance(value, dict) + ) + lines = [ + "", + f"# underworld3 step log · model {payload.get('model')!r} " + f"· started {payload.get('started')}", + ] + if summary: + lines.append(f"# scales: {summary}") + else: + lines.append("# scales: none declared (nondimensional run)") + # Column names are written lazily, with the first step, because the + # time unit is not known until a step carries one. + self._journal_columns = None + return "\n".join(lines) + + if kind == "step": + prefix = "" + unit = (payload["t1"] or {}).get("units") if isinstance( + payload.get("t1"), dict) else None + short = _abbreviate_unit(unit) + if self._journal_columns is None: + self._journal_columns = short + t_col = f"t/{short}" if short else "t" + dt_col = f"dt/{short}" if short else "dt" + prefix = ( + f"#{'step':>5s} {t_col:>14s} {dt_col:>14s} {'wall/s':>8s} " + f"{'outcome':<9s} operators, in order\n" + ) + + t1 = _bare(payload["t1"], unit) + dt = _bare(payload["dt"], unit) + + wall = payload.get("wall") + wall_text = "-" if wall is None else f"{wall:.2f}" + outcome = "ok" if payload.get("completed") else "ABANDONED" + label = payload.get("label") + tag = f"[{label}] " if label else "" + # An invariant is a flag on the step, not an operator it applied — + # it belongs beside the outcome, not in the sequence. + flagged = any(e.get("kind") == "invariant" + for e in payload.get("events", [])) + operators = " > ".join( + f"{e['kind']}:{e['name']}" for e in payload.get("events", []) + if e.get("kind") != "invariant" + ) or "(nothing)" + if flagged: + outcome = f"{outcome} !" + return ( + f"{prefix}" + f" {payload['index']:>5d} {t1:>14.6g} {dt:>14.6g} " + f"{wall_text:>8s} {outcome:<9s} {tag}{operators}" + ) + + # Everything else — rewind, restore — is a note about the run rather + # than a row of the table, so it breaks the columns deliberately. + return f" -- {payload.get('message', kind)}" + + def _write_journal_line(self, payload): + """Append one record and flush, so a killed run keeps its log.""" + if self._journal_fh is None: + return + import json + + try: + if self.journal_format == "jsonl": + text = json.dumps(payload, default=str) + else: + text = self._render_journal_text(payload) + if text is None: + return + self._journal_fh.write(text + "\n") + self._journal_fh.flush() + except Exception: + # A log is a convenience: never take a run down for it. Drop the + # handle so the failure is reported once rather than per step. + try: + self._journal_fh.close() + except Exception: + pass + self._journal_fh = None + import warnings + + warnings.warn( + f"could not append to the journal file {self._journal_path!r}; " + f"logging is off for the rest of this run. The in-memory " + f"model.journal is unaffected.", + RuntimeWarning, + ) + + def _write_journal_note(self, kind, message, **fields): + """Log something that happened to the run but is not a step. + + A backtrack above all: a log that shows step 7, then step 7 again, with + nothing in between, is not a log of what happened. + """ + payload = {"kind": kind, "message": message} + payload.update(fields) + self._write_journal_line(payload) + def _trim_journal(self): limit = self._journal_limit if limit is not None and len(self._journal) > limit: del self._journal[: len(self._journal) - limit] + @property + def record_every(self): + """Keep a restorable snapshot every N steps (None keeps none). + + ``1`` records every step, which is what replay and an adjoint want. + Snapshots cost roughly 13 bytes per primary degree of freedom each, so + a long run on a large mesh should either raise :attr:`record_limit` + with care or record less often and recompute between. + """ + return self._record_every + + @record_every.setter + def record_every(self, value): + self._record_every = value + + @property + def record_limit(self): + """How many snapshots to retain (None retains all). + + Older steps keep their journal record and lose their snapshot, so the + account of what happened survives even where the state does not. + """ + return self._record_limit + + @record_limit.setter + def record_limit(self, value): + self._record_limit = value + self._trim_records() + + @property + def restore_points(self): + """Completed steps that can still be restored, oldest first.""" + return [entry for entry in self._journal if entry.restorable] + + def _trim_records(self): + limit = self._record_limit + if limit is None: + return + restorable = [e for e in self._journal if e.restorable] + for entry in restorable[: max(0, len(restorable) - limit)]: + entry.snapshot = None + + def rewind(self, steps: int = 1): + """Go back to the state at the start of a completed step. + + ``steps=1`` returns to the beginning of the most recent completed step, + undoing it. Fields, histories and the clock all come back together, + because the clock lives on the tracker and the tracker is captured with + everything else. + + The journal is truncated to match, so it continues to describe the run + that actually happened. + """ + restorable = [e for e in self._journal if e.restorable] + if not restorable: + raise RuntimeError( + "nothing to rewind to: no completed step kept a snapshot. " + "Set model.record_every = 1 before the loop to record every step." + ) + if steps < 1 or steps > len(restorable): + raise ValueError( + f"cannot rewind {steps} step(s); {len(restorable)} restorable " + f"step(s) are retained (see model.record_limit)." + ) + target = restorable[-steps] + self._restoring = True + try: + self.load_state(target.snapshot) + finally: + self._restoring = False + cut = self._journal.index(target) + dropped = len(self._journal) - cut + del self._journal[cut:] + + # A log that shows step 7, then step 7 again with nothing in between is + # not a log of what happened. Say where the run went back to. + self._write_journal_note( + "rewind", + f"rewind to the start of step {target.index} " + f"(t = {_pretty_time(self.tracker.time)}); {dropped} step(s) undone", + to_step=int(target.index), + steps_undone=int(dropped), + t=_jsonable_quantity(self.tracker.time), + ) + return target + def _record_step_event(self, kind: str, name: str, **detail) -> None: """Note that something happened inside the step in progress. @@ -744,20 +1248,62 @@ def _step_context(): t0 = self.tracker.time if "time" in self.tracker else 0.0 index = self.tracker.step if "step" in self.tracker else 0 record = ModelStep(index=index, t0=t0, dt=dt, label=label) + + # Record the state this step starts FROM, before any operator runs. + every = self._record_every + if every and index % every == 0: + try: + record.snapshot = self.save_state() + except Exception as exc: + # A snapshot is a convenience here, not a precondition — a + # deforming or adapted mesh cannot be captured yet, and the + # run should carry on with a journal but no restore + # point rather than fail. Say so once. + if not self._record_warned: + self._record_warned = True + import warnings + + warnings.warn( + f"step {index}: could not record the starting state " + f"({type(exc).__name__}: {exc}). The journal still " + f"holds what ran, but model.rewind() will not " + f"reach this step. This is expected on a mesh that " + f"deforms or adapts.", + RuntimeWarning, + ) + self._open_step = record # Position the clock at the END of the interval for the duration of # the block, so implicit coefficients (mesh.t) are evaluated there. self.tracker.time = record.t1 + import time as _time + + wall0 = _time.monotonic() try: yield record except BaseException: + record.wall = _time.monotonic() - wall0 # Abandon: put the clock back and do not commit. self.tracker.time = t0 record.completed = False self._open_step = None + # The abandoned record never joins the journal, so the state it + # captured is unreachable — drop it rather than hold a field- + # sized object until the exception's traceback is collected. + # The idiom for going back is the caller's own save_state() + # taken before the block. + record.snapshot = None + # The log keeps the record itself. A rejected step is the part + # of a run's history that is otherwise invisible, and it is + # usually the part you want when asking why a run went the way + # it did. + self._write_journal_line(record.as_dict()) raise + record.wall = _time.monotonic() - wall0 + record._check_invariants() + # Commit. self.tracker.time = record.t1 self.tracker.step = index + 1 @@ -765,7 +1311,9 @@ def _step_context(): record.completed = True self._open_step = None self._journal.append(record) + self._write_journal_line(record.as_dict()) self._trim_journal() + self._trim_records() return _step_context() @@ -845,13 +1393,29 @@ def load_state(self, source) -> None: from underworld3.checkpoint import read_snapshot as _read_snapshot if isinstance(source, Snapshot): - return _restore(self, source) - if isinstance(source, (str, os.PathLike)): - return _read_snapshot(self, str(source)) - raise TypeError( - f"load_state expects a Snapshot token or a path string, " - f"got {type(source).__name__}" - ) + result = _restore(self, source) + elif isinstance(source, (str, os.PathLike)): + result = _read_snapshot(self, str(source)) + else: + raise TypeError( + f"load_state expects a Snapshot token or a path string, " + f"got {type(source).__name__}" + ) + + # A restore moves the run backwards. It belongs in the log for the same + # reason a rewind does: without it the log shows a step, then an + # earlier step, with nothing to say why. ``rewind`` writes its own, + # more specific, note and suppresses this one. + if not self._restoring: + where = "a file" if isinstance(source, (str, os.PathLike)) else "a snapshot" + self._write_journal_note( + "restore", + f"restore from {where}; the clock now reads " + f"{_pretty_time(self.tracker.time)}", + source=str(source) if isinstance(source, (str, os.PathLike)) else "memory", + t=_jsonable_quantity(self.tracker.time), + ) + return result def define_parameter(self, name: str, ptype=None, **kwargs): """ @@ -4834,6 +5398,76 @@ def view(self): _default_model = None +def read_journal(path): + """Read a journal file back as a list of runs. + + Each entry is ``{"run":
, "steps": [, ...]}``, in the order + the process produced them — an inversion driver that ran the forward model + thirteen times leaves thirteen runs in one file. + + The file is JSON lines, so it is also readable with ``jq`` and survives a + run that was killed part way: a truncated final line is dropped and + everything before it is returned. + + Parameters + ---------- + path : str + A file written by a model with :attr:`Model.journal_file` set. + + Returns + ------- + list of dict + """ + runs = [] + first = True + with open(path, encoding="utf-8") as handle: + for line in handle: + line = line.strip() + if not line: + continue + if first: + first = False + if line.startswith("#"): + raise ValueError( + f"{path} is the TEXT journal format, which is a report " + f"rather than a record — it converts the time column to " + f"one unit and drops each event's detail, so it cannot " + f"be read back. Write JSON lines instead: give the path " + f"a .jsonl suffix, or set model.journal_format = 'jsonl'." + ) + try: + entry = json.loads(line) + except json.JSONDecodeError: + # A run killed mid-write leaves a partial last line. Everything + # before it is intact, which is the point of one object per line. + break + kind = entry.get("kind") + if kind == "run": + runs.append({"run": entry, "steps": [], "notes": []}) + continue + if not runs: + runs.append({"run": None, "steps": [], "notes": []}) + if kind == "step": + runs[-1]["steps"].append(entry) + else: + # A backtrack, or anything else that is not a step. Record WHERE + # in the sequence it happened — a rewind means nothing without + # the step it interrupted and the step it went back to. + entry = dict(entry) + entry["after_position"] = len(runs[-1]["steps"]) - 1 + if kind == "rewind" and entry.get("to_step") is not None: + target = entry["to_step"] + entry["to_position"] = next( + (i for i, step in enumerate(runs[-1]["steps"]) + if step.get("index") == target), None) + entry.setdefault("short", f"rewind {entry.get('steps_undone', 1)}") + else: + entry["to_position"] = entry["after_position"] + entry.setdefault("short", kind) + runs[-1]["notes"].append(entry) + return runs + + def get_default_model(): """ Get or create the default model for this UW3 session. diff --git a/src/underworld3/systems/ddt.py b/src/underworld3/systems/ddt.py index 180e43463..671b8051d 100644 --- a/src/underworld3/systems/ddt.py +++ b/src/underworld3/systems/ddt.py @@ -619,6 +619,51 @@ def _init_coefficient_expressions(self, order, theta, with_exp): if with_exp: _update_exp_values(self._exp_coeffs, None, None) + def _note_history_shift(self, dt): + """Tell the model's open step that this history advanced. + + A history manager should shift EXACTLY ONCE per model step. Shifting + twice means the step was taken twice — a Picard iteration, a corrector + or a retry that called the solver again — and the field advances twice + while ``n_solves_completed`` (capped at ``order``) and ``dt_history`` + look identical. Recording the shift lets ``model.step`` say so; without + it the mistake is invisible. + + A no-op outside a ``model.step`` block. + """ + try: + import underworld3 as uw + + uw.get_default_model()._record_step_event( + "history_shift", self._history_label(), dt=float(dt) + ) + except Exception: + pass + + def _history_label(self): + """Name this history by the field it TRACKS, for the step record. + + Not by its ``psi_star`` slot, whose name is generated from the instance + number and tells a reader nothing. + """ + tracked = None + try: + psi = self.psi_fn + tracked = getattr(psi, "name", None) + if tracked is None: + # a MeshVariable's .sym prints as "{name}(N.x, N.y)", possibly + # wrapped in a Matrix for a vector or tensor unknown + import re + + match = re.search(r"\{([^{}]+)\}", str(psi)) + if match: + tracked = match.group(1) + except Exception: + pass + if tracked is None: + tracked = getattr(self, "instance_number", "?") + return f"{type(self).__name__}({tracked})" + def _register_with_default_model(self): """Register with the active default model as a snapshot state-bearer. @@ -1133,6 +1178,7 @@ def update_post_solve( for i in range(self.order - 1, 0, -1): self._dt_history[i] = self._dt_history[i - 1] self._dt_history[0] = dt + self._note_history_shift(dt) # Shift history: copy each element down the chain. for i in range(self.order - 1, 0, -1): @@ -1576,6 +1622,7 @@ def update_post_solve( for i in range(self.order - 1, 0, -1): self._dt_history[i] = self._dt_history[i - 1] self._dt_history[0] = dt + self._note_history_shift(dt) ### copy values down the chain for i in range(self.order - 1, 0, -1): @@ -3020,6 +3067,7 @@ def update_post_solve( for i in range(self.order - 1, 0, -1): self._dt_history[i] = self._dt_history[i - 1] self._dt_history[0] = dt + self._note_history_shift(dt) if self._n_solves_completed < self.order: self._n_solves_completed += 1 @@ -3819,6 +3867,7 @@ def update_post_solve( for i in range(self.order - 1, 0, -1): self._dt_history[i] = self._dt_history[i - 1] self._dt_history[0] = dt + self._note_history_shift(dt) for h in range(self.order - 1): i = self.order - (h + 1) @@ -4168,6 +4217,7 @@ def update_post_solve( for i in range(self.order - 1, 0, -1): self._dt_history[i] = self._dt_history[i - 1] self._dt_history[0] = dt + self._note_history_shift(dt) for h in range(self.order - 1): i = self.order - (h + 1) @@ -4502,5 +4552,6 @@ def update_post_solve(self, dt, evalf=False, verbose=False, **_ignored): for i in range(self.order - 1, 0, -1): self._dt_history[i] = self._dt_history[i - 1] self._dt_history[0] = dt + self._note_history_shift(dt) if self._n_solves_completed < self.order: self._n_solves_completed += 1 diff --git a/src/underworld3/systems/free_surface.py b/src/underworld3/systems/free_surface.py index 6fd131772..d75bba677 100644 --- a/src/underworld3/systems/free_surface.py +++ b/src/underworld3/systems/free_surface.py @@ -558,10 +558,92 @@ def _velocity_rebind_map(self, solver): rebind[source_L[i, j]] = target_L[i, j] return rebind + def _derived_solvers(self): + """The solves this manager owns, with what each is meant to share. + + ``held_bg`` is omitted: it exists precisely to carry a different body + force, and its rheology is checked through ``held``. + """ + pairs = [("held", self.held), ("consistent", self.consistent)] + return [(name, solver) for name, solver in pairs if solver is not None] + + def _check_derived_solvers_match(self): + """Refuse to solve if a derived lid has drifted from the free solve. + + ``held`` and ``consistent`` are separate Stokes solvers, and the free + solve's configuration was copied into them ONCE, when this manager was + built. Change the free solve afterwards — a different rheology, a + retuned tolerance, a new body force — and they keep the old values. + That is a wrong answer rather than a crash: ``h_inf``, the equilibrium + the surface relaxes toward, is recovered from the HELD solve, so the + surface would relax toward an equilibrium computed with stale physics + while the free solve uses the new. + + A parameter whose value is an expression over mesh variables — a + temperature-dependent viscosity — is shared symbolically and tracks + correctly, as does a rampable constant. Only re-assignment drifts. + """ + from underworld3.utilities._api_tools import ExpressionDescriptor + from underworld3.function.expressions import unwrap + + drift = [] + + def note(name, what, mine, theirs): + if str(mine) != str(theirs): + drift.append(f"{what} — free: {str(mine)[:60]} | {name}: {str(theirs)[:60]}") + + free_params = self.free.constitutive_model.Parameters + for name, solver in self._derived_solvers(): + for setting in ("penalty", "tolerance", "consistent_jacobian"): + note(name, setting, getattr(self.free, setting, None), + getattr(solver, setting, None)) + + # Compare against what copying the parameter TODAY would produce, + # applying the same velocity rebinding :meth:`_copy_constitutive_model` + # applies. A nonlinear rheology is deliberately rebound onto each + # derived solver's own unknowns, so the expressions are MEANT to + # differ textually; only a change of substance is drift. + rebind = self._velocity_rebind_map(solver) + their_params = solver.constitutive_model.Parameters + for cls in type(free_params).__mro__: + for attr, descriptor in cls.__dict__.items(): + if not isinstance(descriptor, ExpressionDescriptor): + continue + try: + expected = getattr(free_params, attr) + if hasattr(expected, "subs"): + expected = unwrap( + expected, keep_constants=True, return_self=True + ).subs(rebind) + actual = getattr(their_params, attr) + if hasattr(actual, "subs"): + actual = unwrap( + actual, keep_constants=True, return_self=True + ) + note(name, f"Parameters.{attr}", expected, actual) + except Exception: + pass + # `held` carries its own body force when driving_buoyancy was given; + # otherwise both are meant to be the free solve's. + if not (name == "held" and self._driving_buoyancy_given): + note(name, "bodyforce", self.free.bodyforce, solver.bodyforce) + + if drift: + raise RuntimeError( + "the free surface's derived solves no longer match the free " + "solve:\n " + "\n ".join(sorted(set(drift))) + "\n" + "They were configured when the FreeSurface was built. Configure " + "the Stokes solver fully BEFORE constructing the FreeSurface, or " + "rebuild the manager after changing it. Solving now would relax " + "the surface toward an equilibrium computed with the stale " + "values, silently." + ) + def _build_held(self, driving_buoyancy): r"""The held free-slip lid: rotated ``u.n = 0`` on every wall and the surface, driving body force only. Its constraint reaction is :math:`\sigma_{nn}`, handed to ``dynamic_topography`` as :math:`h_\infty`.""" + self._driving_buoyancy_given = driving_buoyancy is not None self.held = self._new_stokes("held") self.held.bodyforce = ( self.free.bodyforce if driving_buoyancy is None else driving_buoyancy @@ -861,6 +943,7 @@ def solve(self): rotated free-slip solve gives :math:`\sigma_{nn}` and hence :math:`h_\infty`. Call once per step before :meth:`estimate_dt` / :meth:`advance`. """ + self._check_derived_solvers_match() self.free.solve(zero_init_guess=True) self.held.solve(zero_init_guess=True) self.held.dynamic_topography( diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index 019aafd43..936c0d014 100644 --- a/src/underworld3/systems/solvers.py +++ b/src/underworld3/systems/solvers.py @@ -409,14 +409,24 @@ def _reduce_dt(per_elem): def _dimensionalise_dt(dt_estimate): """Return a timestep estimate with physical time units when a model with - reference scales is active, otherwise as a plain nondimensional scalar.""" + reference scales is active, otherwise as a plain nondimensional scalar. + + ``_as_scalar`` is applied BEFORE dimensionalising, not only in the + no-units fallback. ``np.squeeze`` promotes a Python float to a 0-d array, + and ``uw.dimensionalise`` maps an array to a ``UnitAwareArray`` — which + follows the transparent-container principle and drops its units under + arithmetic. A timestep is a scalar quantity, not a field, so the estimate + must come back as a ``UWQuantity``: the pattern's own idiom + ``dt = fraction * solver.estimate_dt()`` silently loses the units + otherwise, and the loss only surfaces later, wherever the bare number + meets the dimensional clock. + """ + scalar = _as_scalar(np.squeeze(dt_estimate)) try: - return uw.dimensionalise(np.squeeze(dt_estimate), {'[time]': 1}) + return uw.dimensionalise(scalar, {'[time]': 1}) except Exception: - # Sanctioned fallback: no active scaling model. _as_scalar because - # np.squeeze promotes a Python float to a 0-d array, which is not a - # number any caller expects (see _apply_unit_aware_scaling). - return _as_scalar(np.squeeze(dt_estimate)) + # Sanctioned fallback: no active scaling model. + return scalar def _invalidate_solution_cache(u): diff --git a/src/underworld3/utilities/journal_report.py b/src/underworld3/utilities/journal_report.py new file mode 100644 index 000000000..a73598d9d --- /dev/null +++ b/src/underworld3/utilities/journal_report.py @@ -0,0 +1,741 @@ +"""Turn a run's step log into a figure. + +The log is written to be watched (:attr:`underworld3.Model.journal_file`); this +module turns it into something to put in a paper or read on a page. + +``journal_diagram`` + What the run DID, as SVG or PDF. Time runs DOWN the page, one row per step, + so the figure is portrait, paginates, and drops into a document column. + +``journal_flowchart`` + What ONE step does, as Mermaid, for dropping into documentation. + +The layout decision that makes a long run legible: each distinct operator +sequence gets a LETTER, and the letters are defined once at the foot of the +figure. A column of ``A`` with a single ``B`` in it says, at a glance, that one +step did something different — where a hundred repeated sequences say nothing +and hide the one that matters. + +Neither renderer needs a plotting library. The SVG and the PDF are both written +directly, so there is no rasterisation, no dependency, and no theme to fight. +""" + +from __future__ import annotations + +import html +import math +import os +import zlib + +__all__ = ["journal_diagram", "journal_flowchart"] + + +# --- palette --------------------------------------------------------------- +# Print-safe: the two step colours differ in value as well as hue, so they +# survive a greyscale photocopier and a colour-blind reader alike. +_INK = (0.11, 0.11, 0.12) +_MUTED = (0.42, 0.45, 0.50) +_RULE = (0.84, 0.83, 0.81) +_PAPER = (0.992, 0.988, 0.980) +_ACCEPTED = (0.29, 0.44, 0.65) +_ABANDONED = (0.71, 0.33, 0.29) +_WALL = (0.60, 0.65, 0.69) +_FLAG = (0.76, 0.46, 0.12) + +# A4 portrait in points, which is also close enough to US Letter that the +# figure sits inside either with margins to spare. +PAGE_W, PAGE_H = 595.0, 842.0 + + +# --------------------------------------------------------------------------- +# Reading +# --------------------------------------------------------------------------- + +def _as_runs(source): + """Accept a path, the list ``read_journal`` returns, or a live model.""" + if isinstance(source, (str, os.PathLike)): + import underworld3 as uw + + return uw.read_journal(str(source)) + if hasattr(source, "journal") and hasattr(source, "tracker"): + return [{"run": source._run_header(), + "steps": [entry.as_dict() for entry in source.journal], + "notes": []}] + if isinstance(source, list): + if source and isinstance(source[0], dict) and "steps" in source[0]: + return source + return [{"run": None, "steps": list(source), "notes": []}] + raise TypeError( + f"expected a journal path, the list read_journal returns, or a Model; " + f"got {type(source).__name__}" + ) + + +def _pick_run(runs, index): + populated = [r for r in runs if r.get("steps")] + if not populated: + raise ValueError("this journal holds no steps") + return populated[index] + + +def _magnitude(value): + if isinstance(value, dict): + return float(value.get("magnitude", float("nan"))) + try: + return float(value) + except (TypeError, ValueError): + return float("nan") + + +def _unit(value): + return value.get("units") if isinstance(value, dict) else None + + +def _short_unit(unit): + if unit is None: + return "" + return { + "second": "s", "minute": "min", "hour": "hr", "day": "d", "year": "yr", + "kiloyear": "kyr", "megayear": "Myr", "gigayear": "Gyr", + "meter": "m", "kilometer": "km", "kelvin": "K", "kilogram": "kg", + }.get(str(unit), str(unit)) + + +def _converted(value, unit): + """``value`` as a bare number in ``unit`` — a figure has one time axis.""" + if not isinstance(value, dict) or unit is None: + return _magnitude(value) + if str(value.get("units")) == str(unit): + return _magnitude(value) + try: + import underworld3 as uw + + return float( + uw.quantity(_magnitude(value), str(value["units"])).to(unit).magnitude + ) + except Exception: + return _magnitude(value) + + +def _short_operator(name): + """``SNES_AdvectionDiffusion_Composed(T)`` -> ``AdvectionDiffusion(T)``. + + The prefix says which base class implemented it, which is never the thing + the reader is checking.""" + text = str(name) + for prefix in ("SNES_", "uw_"): + if text.startswith(prefix): + text = text[len(prefix):] + return text.replace("_Composed", "") + + +def _signature(step): + """The operator sequence of a step, as a comparable tuple.""" + return tuple( + (event["kind"], _short_operator(event["name"])) + for event in step.get("events", []) + if event.get("kind") in ("solve", "history_shift") + ) + + +def _describe(signature): + return " ".join( + name if kind == "solve" else f"shift {name}" for kind, name in signature + ) or "(nothing)" + + +def _sequence_text(signature): + parts = [name if kind == "solve" else f"shift {name}" + for kind, name in signature] + return " > ".join(parts) or "(nothing)" + + +def _wrap(text, budget): + """Break a sequence between operators, never inside one.""" + if len(text) <= budget: + return [text] + lines, current = [], "" + for token in text.split(" > "): + piece = token if not current else f" > {token}" + if current and len(current) + len(piece) > budget: + lines.append(current + " >") + current = token + else: + current += piece + lines.append(current) + return lines + + +# --------------------------------------------------------------------------- +# A tiny drawing model, so one layout can be written to two formats +# --------------------------------------------------------------------------- + +class _Canvas: + """Ops in a top-left origin, y increasing downward (SVG's convention). + + The PDF writer flips y on the way out; nothing in the layout code has to + know which format it is being drawn into. + """ + + def __init__(self): + self.pages = [[]] + + @property + def ops(self): + return self.pages[-1] + + def new_page(self): + self.pages.append([]) + + def rect(self, x, y, w, h, fill=None, stroke=None, dash=None, width=1.0): + self.ops.append(("rect", x, y, max(w, 0.4), max(h, 0.4), fill, stroke, + dash, width)) + + def line(self, x1, y1, x2, y2, stroke=_RULE, width=0.7, dash=None): + self.ops.append(("line", x1, y1, x2, y2, stroke, width, dash)) + + def curve(self, points, stroke=_ABANDONED, width=1.0, dash=None): + self.ops.append(("curve", list(points), stroke, width, dash)) + + def text(self, x, y, content, size=9.0, fill=_INK, anchor="start", + bold=False, mono=False): + self.ops.append(("text", x, y, str(content), size, fill, anchor, + bold, mono)) + + +_HELV_EM, _COUR_EM = 0.53, 0.60 + + +def _text_width(content, size, mono): + return len(content) * (_COUR_EM if mono else _HELV_EM) * size + + +# --------------------------------------------------------------------------- +# Layout — time runs down the page +# --------------------------------------------------------------------------- + +_MARGIN = 46.0 +_ROW = 14.0 + + +def _layout(header, steps, notes, title=None, width=PAGE_W, page_height=None): + """Draw the run onto a canvas. Returns ``(canvas, width, height)``.""" + canvas = _Canvas() + right = width - _MARGIN + + unit = None + for step in steps: + unit = _unit(step.get("t1")) or _unit(step.get("dt")) + if unit: + break + short = _short_unit(unit) + + dts = [_converted(step.get("dt"), unit) for step in steps] + t1s = [_converted(step.get("t1"), unit) for step in steps] + walls = [step.get("wall") or 0.0 for step in steps] + finite = [d for d in dts if math.isfinite(d) and d > 0] + dt_max = max(finite) if finite else 1.0 + dt_min = min(finite) if finite else 1.0 + wall_max = max(walls) if any(walls) else 1.0 + + # A rejected step is often tens of times the accepted ones — that IS why it + # was rejected — and on a linear axis it flattens everything else to + # nothing. Switch to log and say so, rather than quietly clipping the bar + # that carries the story. + log_scale = dt_max / max(dt_min, 1e-300) > 20.0 + + # --- one letter per distinct operator sequence ------------------------- + order, letters = [], {} + for step in steps: + signature = _signature(step) + if signature not in letters: + letters[signature] = chr(ord("A") + len(order)) if len(order) < 26 \ + else f"#{len(order)}" + order.append(signature) + counts = {} + for step in steps: + counts[_signature(step)] = counts.get(_signature(step), 0) + 1 + + # --- columns ----------------------------------------------------------- + x_gutter = _MARGIN + 14.0 # backtrack arrows live to the left + x_index = x_gutter + 26.0 # step number, right aligned + x_time = x_index + 54.0 # t, right aligned + x_dt = x_time + 52.0 # dt, right aligned + x_letter = x_dt + 16.0 # sequence letter + x_bar = x_letter + 16.0 + x_wall = right - 34.0 + bar_max = x_wall - x_bar - 12.0 + + def bar_length(value): + if not math.isfinite(value) or value <= 0: + return 0.6 + if log_scale: + lo, hi = math.log10(dt_min), math.log10(dt_max) + frac = 0.06 + 0.94 * ((math.log10(value) - lo) / (hi - lo) + if hi > lo else 1.0) + else: + frac = value / dt_max + return max(1.0, frac * bar_max) + + def draw_header(y, first): + if first: + canvas.text(_MARGIN, y + 12, + title or f"Run log — {header.get('model', 'model')!r}", + size=14, bold=True) + y += 20 + bits = [] + if header.get("started"): + bits.append(f"started {header['started']}") + accepted = [s for s in steps if s.get("completed")] + if accepted: + t_from = _converted(accepted[0].get("t0"), unit) + t_to = _converted(accepted[-1].get("t1"), unit) + if math.isfinite(t_from) and math.isfinite(t_to): + bits.append(f"t = {t_from:.4g} to {t_to:.4g}" + + (f" {short}" if short else "")) + bits.append(f"{len(steps)} steps") + abandoned = sum(1 for s in steps if not s.get("completed")) + if abandoned: + bits.append(f"{abandoned} abandoned") + if notes: + bits.append(f"{len(notes)} backtrack(s)") + canvas.text(_MARGIN, y + 8, " · ".join(bits), size=8.5, fill=_MUTED) + y += 13 + scales = header.get("scales") or {} + if scales: + canvas.text(_MARGIN, y + 8, "scales: " + " ".join( + f"{name} {value['magnitude']:.4g} {_short_unit(value['units'])}" + for name, value in scales.items() if isinstance(value, dict) + ), size=8, fill=_MUTED) + y += 12 + y += 10 + # column captions + canvas.text(x_index, y + 8, "step", size=8, fill=_MUTED, anchor="end") + canvas.text(x_time, y + 8, f"t/{short}" if short else "t", size=8, + fill=_MUTED, anchor="end") + canvas.text(x_dt, y + 8, f"dt/{short}" if short else "dt", size=8, + fill=_MUTED, anchor="end") + canvas.text(x_letter, y + 8, "seq", size=8, fill=_MUTED) + caption = "dt" + (" (log scale)" if log_scale else "") + canvas.text(x_bar + 2, y + 8, caption, size=8, fill=_MUTED) + canvas.text(right, y + 8, "wall", size=8, fill=_MUTED, anchor="end") + y += 12 + canvas.line(_MARGIN, y, right, y, _RULE, 0.7) + return y + 4 + + # --- rows -------------------------------------------------------------- + y = _MARGIN + y = draw_header(y, first=True) + row_y = {} + row_page = {} + + for i, step in enumerate(steps): + if page_height is not None and y + _ROW > page_height - _MARGIN - 20: + canvas.text(_MARGIN, page_height - _MARGIN + 4, "continued", size=7.5, + fill=_MUTED) + canvas.new_page() + y = _MARGIN + y = draw_header(y, first=False) + + row_y[i] = y + row_page[i] = len(canvas.pages) - 1 + completed = bool(step.get("completed")) + colour = _ACCEPTED if completed else _ABANDONED + text_colour = _INK if completed else _ABANDONED + base = y + _ROW - 4 + + canvas.text(x_index, base, step.get("index", i), size=8.5, + fill=text_colour, anchor="end") + if math.isfinite(t1s[i]): + canvas.text(x_time, base, f"{t1s[i]:.4g}", size=8.5, + fill=text_colour, anchor="end") + if math.isfinite(dts[i]): + canvas.text(x_dt, base, f"{dts[i]:.4g}", size=8.5, + fill=text_colour, anchor="end") + canvas.text(x_letter, base, letters[_signature(step)], size=8.5, + fill=colour, bold=True) + + # An abandoned bar is usually the longest on the page — that is why it + # was abandoned — so it has to leave room for the word that says so. + room = bar_max - (0.0 if completed else 46.0) + length = min(bar_length(dts[i]), room) + canvas.rect(x_bar, y + 2.5, length, _ROW - 6.5, + fill=colour if completed else None, + stroke=None if completed else _ABANDONED, + dash=None if completed else (2.0, 1.5)) + if not completed: + canvas.text(x_bar + length + 4, base, "abandoned", size=7.5, + fill=_ABANDONED) + + if any(e.get("kind") == "invariant" for e in step.get("events", [])): + canvas.text(x_bar - 8, base, "!", size=10, fill=_FLAG, bold=True) + + if any(walls): + w = max(0.6, (walls[i] / wall_max) * 30.0) + canvas.rect(right - w, y + 4.0, w, _ROW - 9.0, fill=_WALL) + + y += _ROW + + canvas.line(_MARGIN, y + 2, right, y + 2, _RULE, 0.7) + y += 6 + + # --- backtracks, in the left gutter ------------------------------------ + # Backtracks go in the left gutter, each on its own track so two that land + # on adjacent rows do not draw over one another. They are drawn last but + # must land on the PAGE THEIR ROWS ARE ON, not on whichever page the + # cursor happens to have reached. + for track, note in enumerate(notes): + after = note.get("after_position") + target = note.get("to_position") + if after is None or after not in row_y: + continue + if target is None or target not in row_y: + target = after + page = row_page[after] + ops = canvas.pages[page] + x = _MARGIN + 10 - 4.0 * (track % 3) + y_from = row_y[after] + _ROW - 3 + label = note.get("short", "back") + + if row_page[target] != page: + # It reached back past a page break; say so where it happened + # rather than draw a line to a row that is not on this page. + ops.append(("line", x, y_from, x + 4, y_from, _ABANDONED, 0.8, None)) + ops.append(("line", x, y_from, x, y_from - _ROW * 0.8, _ABANDONED, + 0.8, (2.0, 1.5))) + ops.append(("tri", x, y_from - _ROW * 0.8, 2.5, _ABANDONED)) + ops.append(("text", _MARGIN - 2, y_from, f"{label} \u2191", 7, + _ABANDONED, "end", False, False)) + continue + + y_to = row_y[target] + 2 + if y_to > y_from: + continue + ops.append(("line", x, y_from, x, y_to, _ABANDONED, 0.8, (2.0, 1.5))) + ops.append(("line", x, y_from, x + 4, y_from, _ABANDONED, 0.8, None)) + ops.append(("tri", x, y_to, 2.5, _ABANDONED)) + # Only label a backtrack that spans enough rows to hold the word. + if y_from - y_to >= _ROW * 1.5: + ops.append(("text", _MARGIN - 2, (y_from + y_to) / 2 + 3, label, 7, + _ABANDONED, "end", False, False)) + + # --- the legend: what each letter means -------------------------------- + if page_height is not None and y + 30 + 14 * len(order) > page_height - _MARGIN: + canvas.new_page() + y = _MARGIN + + y += 10 + canvas.text(_MARGIN, y + 8, "Operator sequences", size=9.5, bold=True) + y += 18 + budget = int((right - _MARGIN - 34) / (_COUR_EM * 8.0)) + for signature in order: + canvas.text(_MARGIN + 2, y + 8, letters[signature], size=9, bold=True, + fill=_ACCEPTED if signature == order[0] else _FLAG) + count = counts[signature] + canvas.text(right, y + 8, f"{count} step{'' if count == 1 else 's'}", + size=8, fill=_MUTED, anchor="end") + for line in _wrap(_sequence_text(signature), budget): + canvas.text(_MARGIN + 18, y + 8, line, size=8, mono=True, + fill=_ACCEPTED if signature == order[0] else _FLAG) + y += 11 + y += 5 + + flagged = [ + (step.get("index", i), event.get("detail", "")) + for i, step in enumerate(steps) + for event in step.get("events", []) + if event.get("kind") == "invariant" + ] + if flagged: + y += 6 + canvas.text(_MARGIN, y + 8, "! Invariant", size=9.5, bold=True, fill=_FLAG) + y += 15 + for index, detail in flagged: + canvas.text(_MARGIN + 12, y + 8, + f"step {index}: history advanced more than once " + f"({detail}) — the step was taken twice", + size=8, fill=_INK) + y += 12 + + height = page_height if page_height is not None else y + _MARGIN + return canvas, width, height + + +# --------------------------------------------------------------------------- +# SVG +# --------------------------------------------------------------------------- + +def _hex(colour): + return "#" + "".join(f"{int(round(c * 255)):02x}" for c in colour) + + +def _svg_ops(ops, width, height): + out = [f''] + for op in ops: + kind = op[0] + if kind == "rect": + _, x, y, w, h, fill, stroke, dash, lw = op + attrs = f'x="{x:.1f}" y="{y:.1f}" width="{w:.1f}" height="{h:.1f}" rx="1"' + attrs += f' fill="{_hex(fill)}"' if fill else ' fill="none"' + if stroke: + attrs += f' stroke="{_hex(stroke)}" stroke-width="{lw}"' + if dash: + attrs += f' stroke-dasharray="{dash[0]} {dash[1]}"' + out.append(f"") + elif kind == "line": + _, x1, y1, x2, y2, stroke, lw, dash = op + attrs = (f'x1="{x1:.1f}" y1="{y1:.1f}" x2="{x2:.1f}" y2="{y2:.1f}" ' + f'stroke="{_hex(stroke)}" stroke-width="{lw}"') + if dash: + attrs += f' stroke-dasharray="{dash[0]} {dash[1]}"' + out.append(f"") + elif kind == "tri": + _, x, y, r, colour = op + out.append(f'') + elif kind == "text": + _, x, y, content, size, fill, anchor, bold, mono = op + family = ("'SF Mono', Menlo, monospace" if mono + else "'Helvetica Neue', Helvetica, Arial, sans-serif") + out.append( + f'' + f'{html.escape(content)}' + ) + body = "\n".join(out) + return (f'\n' + f'{body}\n\n') + + +# --------------------------------------------------------------------------- +# PDF — written directly, so the figure needs nothing installed to become one +# --------------------------------------------------------------------------- + +_PDF_SUBSTITUTIONS = { + "→": "->", "·": "-", "⚠": "!", "—": "-", "–": "-", + "'": "'", "'": "'", """: '"', """: '"', "≥": ">=", "≤": "<=", +} + + +def _pdf_text(content): + """WinAnsi-safe, with the escapes PDF strings need.""" + for source, target in _PDF_SUBSTITUTIONS.items(): + content = content.replace(source, target) + content = content.encode("latin-1", "replace").decode("latin-1") + return content.replace("\\", r"\\").replace("(", r"\(").replace(")", r"\)") + + +def _pdf_page_stream(ops, width, height): + """One page's content stream. PDF's origin is bottom-left, so y flips.""" + def fy(y): + return height - y + + out = [f"{_PAPER[0]:.3f} {_PAPER[1]:.3f} {_PAPER[2]:.3f} rg", + f"0 0 {width:.1f} {height:.1f} re f"] + for op in ops: + kind = op[0] + if kind == "rect": + _, x, y, w, h, fill, stroke, dash, lw = op + out.append("q") + if dash: + out.append(f"[{dash[0]} {dash[1]}] 0 d") + if fill: + out.append(f"{fill[0]:.3f} {fill[1]:.3f} {fill[2]:.3f} rg") + if stroke: + out.append(f"{stroke[0]:.3f} {stroke[1]:.3f} {stroke[2]:.3f} RG " + f"{lw} w") + out.append(f"{x:.2f} {fy(y + h):.2f} {w:.2f} {h:.2f} re") + out.append("B" if (fill and stroke) else ("f" if fill else "S")) + out.append("Q") + elif kind == "line": + _, x1, y1, x2, y2, stroke, lw, dash = op + out.append("q") + if dash: + out.append(f"[{dash[0]} {dash[1]}] 0 d") + out.append(f"{stroke[0]:.3f} {stroke[1]:.3f} {stroke[2]:.3f} RG {lw} w") + out.append(f"{x1:.2f} {fy(y1):.2f} m {x2:.2f} {fy(y2):.2f} l S") + out.append("Q") + elif kind == "tri": + _, x, y, r, colour = op + out.append(f"q {colour[0]:.3f} {colour[1]:.3f} {colour[2]:.3f} rg") + out.append(f"{x:.2f} {fy(y):.2f} m {x - r:.2f} {fy(y + r * 1.6):.2f} l " + f"{x + r:.2f} {fy(y + r * 1.6):.2f} l f Q") + elif kind == "text": + _, x, y, content, size, fill, anchor, bold, mono = op + font = "/F3" if mono else ("/F2" if bold else "/F1") + w = _text_width(content, size, mono) + if anchor == "end": + x -= w + elif anchor == "middle": + x -= w / 2 + out.append(f"BT {font} {size:.1f} Tf " + f"{fill[0]:.3f} {fill[1]:.3f} {fill[2]:.3f} rg " + f"{x:.2f} {fy(y):.2f} Td ({_pdf_text(content)}) Tj ET") + return "\n".join(out).encode("latin-1", "replace") + + +def _pdf_document(pages, width, height): + objects = {} + n_pages = len(pages) + font_ids = {"F1": 3, "F2": 4, "F3": 5} + first_page = 6 + + objects[1] = b"<< /Type /Catalog /Pages 2 0 R >>" + kids = " ".join(f"{first_page + 2 * i} 0 R" for i in range(n_pages)) + objects[2] = (f"<< /Type /Pages /Count {n_pages} /Kids [{kids}] >>" + ).encode("latin-1") + objects[3] = b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica /Encoding /WinAnsiEncoding >>" + objects[4] = b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica-Bold /Encoding /WinAnsiEncoding >>" + objects[5] = b"<< /Type /Font /Subtype /Type1 /BaseFont /Courier /Encoding /WinAnsiEncoding >>" + + for i, ops in enumerate(pages): + page_id = first_page + 2 * i + stream_id = page_id + 1 + resources = ("<< /Font << " + " ".join( + f"/{name} {oid} 0 R" for name, oid in font_ids.items()) + " >> >>") + objects[page_id] = ( + f"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 {width:.1f} {height:.1f}] " + f"/Resources {resources} /Contents {stream_id} 0 R >>" + ).encode("latin-1") + raw = _pdf_page_stream(ops, width, height) + packed = zlib.compress(raw) + objects[stream_id] = ( + f"<< /Length {len(packed)} /Filter /FlateDecode >>\nstream\n" + ).encode("latin-1") + packed + b"\nendstream" + + out = bytearray(b"%PDF-1.4\n%\xe2\xe3\xcf\xd3\n") + offsets = {} + for oid in sorted(objects): + offsets[oid] = len(out) + out += f"{oid} 0 obj\n".encode("latin-1") + objects[oid] + b"\nendobj\n" + + xref_at = len(out) + top = max(objects) + 1 + out += f"xref\n0 {top}\n".encode("latin-1") + out += b"0000000000 65535 f \n" + for oid in range(1, top): + out += f"{offsets.get(oid, 0):010d} 00000 n \n".encode("latin-1") + out += (f"trailer\n<< /Size {top} /Root 1 0 R >>\nstartxref\n{xref_at}\n" + f"%%EOF\n").encode("latin-1") + return bytes(out) + + +# --------------------------------------------------------------------------- +# Entry points +# --------------------------------------------------------------------------- + +def journal_diagram(source, out=None, run=-1, title=None, format=None, + width=None): + """Render a run's log as a figure, with time running DOWN the page. + + Parameters + ---------- + source : str, list or Model + A ``.jsonl`` journal file, the list :func:`underworld3.read_journal` + returns, or a live model. Not a text log — that format is a report and + cannot be read back. + out : str, optional + Where to write. Defaults to the source path with the format's suffix, + else ``journal.pdf``. + run : int, default -1 + Which run in the file. A file holds one per ``clear_journal()``. + title : str, optional + Overrides the heading taken from the run header. + format : {"pdf", "svg"}, optional + Inferred from ``out``'s suffix; PDF by default. PDF paginates onto A4 + portrait; SVG is one continuous page. + width : float, optional + Page width in points. Defaults to A4 portrait. + + Returns + ------- + str + The path written. + """ + runs = _as_runs(source) + entry = _pick_run(runs, run) + + if format is None: + if out and str(out).lower().endswith(".svg"): + format = "svg" + else: + format = "pdf" + if format not in ("pdf", "svg"): + raise ValueError(f"format must be 'pdf' or 'svg', not {format!r}") + + if out is None: + suffix = ".svg" if format == "svg" else ".pdf" + out = (os.path.splitext(str(source))[0] + suffix + if isinstance(source, (str, os.PathLike)) else "journal" + suffix) + + page_width = width or PAGE_W + canvas, page_width, height = _layout( + entry.get("run") or {}, entry["steps"], entry.get("notes", []), + title=title, width=page_width, + page_height=PAGE_H if format == "pdf" else None, + ) + + directory = os.path.dirname(out) + if directory: + os.makedirs(directory, exist_ok=True) + + if format == "svg": + with open(out, "w", encoding="utf-8") as handle: + handle.write(_svg_ops(canvas.pages[0], page_width, height)) + else: + with open(out, "wb") as handle: + handle.write(_pdf_document(canvas.pages, page_width, PAGE_H)) + return out + + +def journal_flowchart(source, run=-1, out=None): + """The operator flow of a step, as Mermaid, for dropping into documentation. + + When every step ran the same sequence — the usual case — that is one + flowchart. When they did not, each distinct sequence becomes its own + subgraph, labelled with the steps that took it, which is what makes an + anomalous step visible rather than averaged away. + """ + runs = _as_runs(source) + steps = _pick_run(runs, run)["steps"] + + signatures = {} + for i, step in enumerate(steps): + signatures.setdefault(_signature(step), []).append(step.get("index", i)) + + lines = ["flowchart LR"] + for group, (signature, indices) in enumerate(signatures.items()): + if len(signatures) > 1: + named = ", ".join(str(i) for i in indices[:6]) + if len(indices) > 6: + named += f", +{len(indices) - 6}" + lines.append(f' subgraph g{group}["step {named}"]') + lines.append(" direction LR") + indent = " " if len(signatures) > 1 else " " + if not signature: + lines.append(f'{indent}n{group}_0["(nothing)"]') + previous = None + for i, (kind, name) in enumerate(signature): + node = f"n{group}_{i}" + if kind == "history_shift": + lines.append(f'{indent}{node}[/"shift {name}"/]') + else: + lines.append(f'{indent}{node}["{name}"]') + if previous is not None: + lines.append(f"{indent}{previous} --> {node}") + previous = node + if len(signatures) > 1: + lines.append(" end") + + text = "\n".join(lines) + "\n" + if out: + directory = os.path.dirname(out) + if directory: + os.makedirs(directory, exist_ok=True) + with open(out, "w", encoding="utf-8") as handle: + handle.write(text) + return text diff --git a/src/underworld3/utilities/rotated_bc.py b/src/underworld3/utilities/rotated_bc.py index 992d7b2b2..e82d13a95 100644 --- a/src/underworld3/utilities/rotated_bc.py +++ b/src/underworld3/utilities/rotated_bc.py @@ -1043,7 +1043,10 @@ def solve_rotated_freeslip(solver, boundaries, remove_rotation_gauge=True, # already attached it. solver.mesh.update_lvec() solver.dm.setAuxiliaryVec(solver.mesh.lvec, None) - solver._update_constants() + # record=False: the public solve() that dispatched here has already + # announced this solver to the step journal. This push is for THIS + # function's own assembly; recording it again reports one operator as two. + solver._update_constants(record=False) if rtol is None: rtol = float(solver.tolerance) diff --git a/tests/test_0011_model_step_journal.py b/tests/test_0011_model_step_journal.py index a98e7ff8d..af969290f 100644 --- a/tests/test_0011_model_step_journal.py +++ b/tests/test_0011_model_step_journal.py @@ -125,3 +125,182 @@ def test_the_journal_is_bounded(): pass assert len(model.journal) == 3 assert [e.index for e in model.journal] == [4, 5, 6] + + +# --------------------------------------------------------------------------- +# Recording: a step keeps the state it started from, so the run can be replayed +# --------------------------------------------------------------------------- + + +def _advdiff(uw, mesh): + import sympy + + T = uw.discretisation.MeshVariable("T_record", mesh, 1, degree=2) + V = uw.discretisation.MeshVariable("V_record", mesh, 2, degree=2) + x, y = mesh.X + V.array[:, 0, :] = np.asarray( + uw.function.evaluate(sympy.Matrix([[-(y - 0.5), (x - 0.5)]]), V.coords) + ).reshape(-1, 2) + T.array[:, 0, 0] = np.asarray( + uw.function.evaluate(sympy.exp(-(((x - 0.3) ** 2 + (y - 0.5) ** 2) / 0.02)), T.coords) + ).ravel() + solver = uw.systems.AdvDiffusion(mesh, u_Field=T, V_fn=V.sym) + solver.constitutive_model = uw.constitutive_models.DiffusionModel + solver.constitutive_model.Parameters.diffusivity = 1.0e-4 + solver.petsc_options.delValue("ksp_monitor") + return solver, T + + +def test_recording_is_off_by_default(): + uw, model = _fresh_model() + model.tracker.time, model.tracker.step = 0.0, 0 + with model.step(0.1): + pass + assert model.journal[0].restorable is False + assert model.restore_points == [] + + +def test_rewind_undoes_a_step_exactly(): + """Fields, history and clock all come back, and the step is undone.""" + uw, model = _fresh_model() + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1.0 / 8, qdegree=3 + ) + solver, T = _advdiff(uw, mesh) + model.tracker.time, model.tracker.step = 0.0, 0 + model.record_every = 1 + + for _ in range(2): + with model.step(0.02): + solver.solve(timestep=0.02) + + at_two = np.array(T.array) + assert model.tracker.step == 2 + + # take a third step, then undo it + with model.step(0.02): + solver.solve(timestep=0.02) + assert model.tracker.step == 3 + assert not np.allclose(np.array(T.array), at_two) + + model.rewind() + + assert np.array_equal(np.array(T.array), at_two), "fields did not come back" + assert model.tracker.step == 2, "the clock did not come back" + assert model.tracker.time == pytest.approx(0.04) + assert len(model.journal) == 2, "the journal still claims the undone step" + + +def test_replaying_a_rewound_step_reproduces_it(): + """The property replay debugging rests on: the same step, taken twice from + the same state, gives the same answer.""" + uw, model = _fresh_model() + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1.0 / 8, qdegree=3 + ) + solver, T = _advdiff(uw, mesh) + model.tracker.time, model.tracker.step = 0.0, 0 + model.record_every = 1 + + with model.step(0.02): + solver.solve(timestep=0.02) + first = np.array(T.array) + + model.rewind() + with model.step(0.02): + solver.solve(timestep=0.02) + second = np.array(T.array) + + assert np.array_equal(first, second), ( + "replaying a step from its own snapshot did not reproduce it" + ) + + +def test_the_record_is_bounded_but_the_journal_survives(): + """Old steps lose their snapshot and keep their record, so the account of + what happened outlives the state.""" + uw, model = _fresh_model() + model.tracker.time, model.tracker.step = 0.0, 0 + model.record_every = 1 + model.record_limit = 2 + + for _ in range(5): + with model.step(0.1): + pass + + assert len(model.journal) == 5 + assert [e.index for e in model.restore_points] == [3, 4] + + +def test_rewind_without_a_record_says_what_to_do(): + uw, model = _fresh_model() + model.tracker.time, model.tracker.step = 0.0, 0 + with model.step(0.1): + pass + with pytest.raises(RuntimeError, match="record_every"): + model.rewind() + + +# --------------------------------------------------------------------------- +# Invariants: a step that cannot be what it claims to be +# --------------------------------------------------------------------------- + + +def test_a_history_that_advances_twice_in_one_step_is_reported(): + """Two solves inside one step take the physical step twice. The solve + counter and the timestep history look identical to a single step, so + without this the mistake is invisible.""" + uw, model = _fresh_model() + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1.0 / 8, qdegree=3 + ) + solver, T = _advdiff(uw, mesh) + model.tracker.time, model.tracker.step = 0.0, 0 + + with pytest.warns(RuntimeWarning, match="advanced more than once"): + with model.step(0.02): + solver.solve(timestep=0.02) + solver.solve(timestep=0.02) # the same step, taken twice + + entry = model.journal[0] + shifts = [e for e in entry.events if e["kind"] == "history_shift"] + assert len(shifts) == 2 + + +def test_one_solve_per_step_is_quiet(): + """The negative control: the ordinary loop must not warn.""" + uw, model = _fresh_model() + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1.0 / 8, qdegree=3 + ) + solver, T = _advdiff(uw, mesh) + model.tracker.time, model.tracker.step = 0.0, 0 + + import warnings + + with warnings.catch_warnings(): + warnings.simplefilter("error", RuntimeWarning) + for _ in range(3): + with model.step(0.02): + solver.solve(timestep=0.02) + + assert len(model.journal) == 3 + + +def test_the_journal_shows_the_history_that_moved(): + """The record names which history advanced, not just that a solve ran.""" + uw, model = _fresh_model() + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1.0 / 8, qdegree=3 + ) + solver, T = _advdiff(uw, mesh) + model.tracker.time, model.tracker.step = 0.0, 0 + + with model.step(0.02): + solver.solve(timestep=0.02) + + kinds = [e["kind"] for e in model.journal[0].events] + assert "solve" in kinds and "history_shift" in kinds + shift = next(e for e in model.journal[0].events if e["kind"] == "history_shift") + assert shift["dt"] == pytest.approx(0.02) + assert "T_record" in shift["name"], shift["name"] diff --git a/tests/test_0012_snapshot_units_coords.py b/tests/test_0012_snapshot_units_coords.py new file mode 100644 index 000000000..428b5be99 --- /dev/null +++ b/tests/test_0012_snapshot_units_coords.py @@ -0,0 +1,110 @@ +"""A snapshot must not rescale the mesh. + +``mesh.X.coords`` is the UNIT-AWARE view: with a model that declares a length +scale it returns metres, while the DM coordinate vector the restore path writes +back into holds model units. Capturing one and restoring the other multiplies +the mesh by the length scale — silently, because every array keeps its shape +and every field is restored correctly. Only the geometry is wrong, so what +fails afterwards is every integral, every evaluate, and every subsequent solve. + +``model.rewind()`` goes straight through this path, which is how it was found. +""" + +import pytest + +pytestmark = [pytest.mark.level_1, pytest.mark.tier_a] + +import numpy as np + + +LENGTH_SCALE_M = 500e3 + + +def _model_with_units(): + import underworld3 as uw + + uw.reset_default_model() + model = uw.get_default_model() + model.set_reference_quantities( + domain_depth=uw.quantity(500, "km"), + material_viscosity=uw.quantity(1e21, "Pa*s"), + lithostatic_pressure=uw.quantity(3300 * 9.81 * 500e3, "Pa"), + ) + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1.0 / 8.0 + ) + return uw, model, mesh + + +def test_restore_leaves_the_mesh_at_its_own_size(): + """Round-tripping a snapshot must not scale the coordinates.""" + uw, model, mesh = _model_with_units() + + dimensional_before = np.asarray(mesh.X.coords).copy() + raw_before = np.asarray(mesh._coords).copy() + assert dimensional_before.max() == pytest.approx(LENGTH_SCALE_M, rel=1e-6), ( + "the fixture is not exercising the unit-aware view" + ) + assert raw_before.max() == pytest.approx(1.0, rel=1e-6) + + snap = model.save_state() + model.load_state(snap) + + assert np.asarray(mesh._coords).max() == pytest.approx(1.0, rel=1e-12), ( + "restore rescaled the mesh: the captured coordinates were dimensional " + "but were written back as model units" + ) + assert np.allclose(np.asarray(mesh.X.coords), dimensional_before, rtol=0, atol=0) + + +def test_repeated_restores_do_not_drift(): + """The scaling error compounds, so check more than one round trip.""" + uw, model, mesh = _model_with_units() + raw_before = np.asarray(mesh._coords).copy() + + for _ in range(3): + snap = model.save_state() + model.load_state(snap) + + assert np.allclose(np.asarray(mesh._coords), raw_before, rtol=0, atol=0) + + +def test_evaluate_still_works_after_a_restore(): + """The symptom, not the mechanism: a rescaled mesh puts every sample point + outside the domain, and evaluate quietly returns the value at one corner.""" + uw, model, mesh = _model_with_units() + T = uw.discretisation.MeshVariable("T", mesh, 1, degree=2) + x, y = mesh.X + T.array[:, 0, 0] = np.asarray(T.coords)[:, 1] / LENGTH_SCALE_M + + sample = np.column_stack([np.full(9, 0.5), np.linspace(0.05, 0.95, 9)]) + before = np.asarray(uw.function.evaluate(T.sym[0], sample)).ravel() + assert np.ptp(before) > 0.5, "the fixture should vary across the sample line" + + model.load_state(model.save_state()) + + after = np.asarray(uw.function.evaluate(T.sym[0], sample)).ravel() + assert np.allclose(after, before, rtol=1e-10, atol=1e-12), ( + "evaluate disagrees with itself across a snapshot round trip" + ) + + +def test_rewind_reaches_the_state_the_step_started_from(): + """The path this was found on: record a step, rewind, and check the mesh.""" + uw, model, mesh = _model_with_units() + T = uw.discretisation.MeshVariable("T", mesh, 1, degree=2) + T.array[:, 0, 0] = 1.0 + + model.tracker.time = 0.0 + model.tracker.step = 0 + model.record_every = 1 + + raw_before = np.asarray(mesh._coords).copy() + with model.step(0.5): + T.array[:, 0, 0] = 2.0 + + model.rewind() + + assert np.allclose(np.asarray(mesh._coords), raw_before, rtol=0, atol=0) + assert np.allclose(np.asarray(T.array)[:, 0, 0], 1.0) + assert model.tracker.time == 0.0 diff --git a/tests/test_0013_step_record_fidelity.py b/tests/test_0013_step_record_fidelity.py new file mode 100644 index 000000000..60cca00cc --- /dev/null +++ b/tests/test_0013_step_record_fidelity.py @@ -0,0 +1,173 @@ +"""What the step journal claims must be what happened. + +Two ways it was over- or under-reporting, both found by writing a real +annulus convection run in the timestepping pattern. + +1. The journal counted one operator as two. The hook lives in + ``_update_constants``, which is the single point every solver passes on its + way to a solve — except that the rotated free-slip loop pushes constants a + second time for its own assembly, after the public ``solve()`` has already + announced the solver. So a Stokes solve on a curved boundary recorded twice. + +2. ``estimate_dt()`` lost its units under the pattern's own idiom. + ``dt = fraction * solver.estimate_dt()`` came back as a bare float whenever + the estimate happened to be a Python float rather than a numpy scalar, + because ``np.squeeze`` promoted it to a 0-d array and a dimensionalised + array is a ``UnitAwareArray``, which drops units under arithmetic. +""" + +import pytest + +pytestmark = [pytest.mark.level_1, pytest.mark.tier_a] + +import numpy as np +import sympy + + +def _annulus_model(units=True): + import underworld3 as uw + + uw.reset_default_model() + model = uw.get_default_model() + if units: + model.set_reference_quantities( + shell_thickness=uw.quantity(2200, "km"), + thermal_diffusivity=uw.quantity(1e-6, "m**2/s"), + mantle_viscosity=uw.quantity(1e22, "Pa*s"), + temperature_contrast=uw.quantity(2500, "K"), + ) + mesh = uw.meshing.Annulus( + radiusInner=0.55, radiusOuter=1.0, cellSize=0.25, degree=1, qdegree=3 + ) + v = uw.discretisation.MeshVariable("v", mesh, mesh.dim, degree=2) + p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1) + T = uw.discretisation.MeshVariable("T", mesh, 1, degree=2) + + stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = ( + uw.quantity(1e22, "Pa*s") if units else 1.0 + ) + stokes.tolerance = 1.0e-6 + stokes.petsc_options.delValue("ksp_monitor") + stokes.add_rotated_freeslip_bc(0.0, "Upper") + stokes.add_rotated_freeslip_bc(0.0, "Lower") + + radius = sympy.sqrt(mesh.X.dot(mesh.X)) + if units: + stokes.bodyforce = ( + -uw.quantity(3300, "kg/m**3") + * uw.quantity(3e-5, "1/K") + * uw.quantity(9.81, "m/s**2") + * T.sym[0] + * mesh.X + / radius + ) + else: + stokes.bodyforce = -1.0e5 * T.sym[0] * mesh.X / radius + + adv = uw.systems.AdvDiffusion(mesh, u_Field=T, V_fn=v.sym) + adv.constitutive_model = uw.constitutive_models.DiffusionModel + adv.constitutive_model.Parameters.diffusivity = ( + uw.quantity(1e-6, "m**2/s") if units else 1.0 + ) + adv.add_dirichlet_bc(1.0, "Lower") + adv.add_dirichlet_bc(0.0, "Upper") + adv.tolerance = 1.0e-6 + adv.petsc_options.delValue("ksp_monitor") + + scale = 1.0 + if units: + scale = float(model.get_fundamental_scales()["length"].to("m").magnitude) + X = np.asarray(T.coords)[:, :2] / scale + r = np.sqrt((X**2).sum(axis=1)) + th = np.arctan2(X[:, 1], X[:, 0]) + shell = (r - 0.55) / (1.0 - 0.55) + T.array[:, 0, 0] = (1.0 - shell) + 0.1 * np.sin(5.0 * th) * np.sin(np.pi * shell) + adv.Unknowns.DuDt.initialise_history() + + return uw, model, mesh, stokes, adv, T + + +def _names(entry, kind="solve"): + return [e["name"] for e in entry.events if e["kind"] == kind] + + +def test_rotated_freeslip_solve_is_recorded_once(): + """A curved-boundary Stokes solve is one operator, not two.""" + uw, model, mesh, stokes, adv, T = _annulus_model(units=False) + stokes.solve(zero_init_guess=True) + + model.tracker.time = 0.0 + model.tracker.step = 0 + + with model.step(0.01, label="convect"): + adv.solve(timestep=0.01, zero_init_guess=False) + stokes.solve(zero_init_guess=False) + + entry = model.journal[-1] + solves = _names(entry) + assert sum(1 for n in solves if "Stokes" in n) == 1, ( + f"the rotated free-slip dispatch recorded more than one Stokes solve: {solves}" + ) + assert sum(1 for n in solves if "AdvectionDiffusion" in n) == 1, solves + assert len(_names(entry, "history_shift")) == 1 + + +def test_a_solver_called_twice_is_still_recorded_twice(): + """The de-duplication must not hide a genuinely repeated solve.""" + uw, model, mesh, stokes, adv, T = _annulus_model(units=False) + stokes.solve(zero_init_guess=True) + + model.tracker.time = 0.0 + model.tracker.step = 0 + + with pytest.warns(RuntimeWarning, match="history advanced more than once"): + with model.step(0.01): + adv.solve(timestep=0.01, zero_init_guess=False) + stokes.solve(zero_init_guess=False) + adv.solve(timestep=0.01, zero_init_guess=False) + stokes.solve(zero_init_guess=False) + + solves = _names(model.journal[-1]) + assert sum(1 for n in solves if "Stokes" in n) == 2, solves + + +@pytest.mark.parametrize("solver_name", ["stokes", "adv"]) +def test_estimate_dt_survives_being_scaled(solver_name): + """`dt = fraction * solver.estimate_dt()` must keep its units.""" + uw, model, mesh, stokes, adv, T = _annulus_model(units=True) + stokes.solve(zero_init_guess=True) + + solver = stokes if solver_name == "stokes" else adv + dt = solver.estimate_dt() + assert hasattr(dt, "to"), f"{solver_name}.estimate_dt() returned {type(dt).__name__}" + + scaled = 0.5 * dt + assert hasattr(scaled, "to"), ( + f"0.5 * {solver_name}.estimate_dt() dropped its units " + f"({type(dt).__name__} -> {type(scaled).__name__})" + ) + assert float(scaled.to("s").magnitude) == pytest.approx( + 0.5 * float(dt.to("s").magnitude), rel=1e-12 + ) + + +def test_a_scaled_estimate_drives_the_clock(): + """The whole idiom, end to end: the scaled estimate must reach the clock.""" + uw, model, mesh, stokes, adv, T = _annulus_model(units=True) + stokes.solve(zero_init_guess=True) + + model.tracker.time = uw.quantity(0.0, "Myr") + model.tracker.step = 0 + + dt = 0.5 * adv.estimate_dt() + with model.step(dt, label="convect"): + adv.solve(timestep=dt, zero_init_guess=False) + stokes.solve(zero_init_guess=False) + + elapsed = model.tracker.time + assert hasattr(elapsed, "to") + assert float(elapsed.to("s").magnitude) == pytest.approx( + float(dt.to("s").magnitude), rel=1e-12 + ) diff --git a/tests/test_0014_journal_file.py b/tests/test_0014_journal_file.py new file mode 100644 index 000000000..20f34be77 --- /dev/null +++ b/tests/test_0014_journal_file.py @@ -0,0 +1,334 @@ +"""The journal, written down. + +``model.journal`` is what a run can still undo: bounded, in memory, gone with +the process. ``model.journal_file`` is what the run did: one JSON object per +line, appended and flushed as each step closes. + +The two differ deliberately, and the differences are what the tests below pin. +An abandoned step appears in the file and not in memory — it is the part of a +run's history that is otherwise invisible. A step aged out by ``journal_limit`` +leaves memory and stays in the file. And one object per line means a run that +is killed keeps everything up to the moment it died. +""" + +import pytest + +pytestmark = [pytest.mark.level_1, pytest.mark.tier_a] + +import json + + +def _model(tmp_path, units=False, name="run.journal.jsonl", fmt=None): + import underworld3 as uw + + uw.reset_default_model() + model = uw.get_default_model() + if units: + model.set_reference_quantities( + domain_depth=uw.quantity(500, "km"), + material_viscosity=uw.quantity(1e21, "Pa*s"), + lithostatic_pressure=uw.quantity(3300 * 9.81 * 500e3, "Pa"), + ) + path = tmp_path / name + model.journal_file = str(path) + if fmt is not None: + model.journal_format = fmt + model.tracker.time = uw.quantity(0.0, "Myr") if units else 0.0 + model.tracker.step = 0 + return uw, model, path + + +def test_a_completed_step_is_one_line(tmp_path): + uw, model, path = _model(tmp_path) + + with model.step(0.25, label="convect"): + pass + + lines = path.read_text().splitlines() + assert len(lines) == 2, "expected a run header and one step" + + header, step = (json.loads(line) for line in lines) + assert header["kind"] == "run" + + wall = step.pop("wall") + assert wall >= 0.0, "a step should record how long it took" + assert step == { + "kind": "step", + "index": 0, + "label": "convect", + "t0": 0.0, + "t1": 0.25, + "dt": 0.25, + "completed": True, + "restorable": False, + "events": [], + } + + +def test_an_abandoned_step_is_in_the_file_and_not_in_memory(tmp_path): + uw, model, path = _model(tmp_path) + + with model.step(0.25, label="fine"): + pass + with pytest.raises(RuntimeError): + with model.step(9.0, label="too big"): + raise RuntimeError("courant") + + assert [e.label for e in model.journal] == ["fine"] + + steps = [json.loads(line) for line in path.read_text().splitlines()][1:] + assert [s["label"] for s in steps] == ["fine", "too big"] + assert [s["completed"] for s in steps] == [True, False] + + # The clock did not move for the abandoned step, so its t0 is the previous + # step's t1 and nothing after it is shifted. + assert steps[1]["t0"] == pytest.approx(0.25) + assert model.tracker.time == pytest.approx(0.25) + + +def test_an_abandoned_step_does_not_retain_its_snapshot(tmp_path): + """It is unreachable — the journal never holds it — so it must not be kept.""" + uw, model, path = _model(tmp_path) + model.record_every = 1 + + captured = {} + with pytest.raises(RuntimeError): + with model.step(0.25): + captured["open"] = model.open_step + raise RuntimeError("nope") + + assert captured["open"].snapshot is None + steps = [json.loads(line) for line in path.read_text().splitlines()][1:] + assert steps[0]["restorable"] is False + + +def test_a_step_aged_out_of_memory_stays_in_the_file(tmp_path): + uw, model, path = _model(tmp_path) + model.journal_limit = 2 + + for _ in range(5): + with model.step(0.1, label="convect"): + pass + + assert len(model.journal) == 2 + steps = [json.loads(line) for line in path.read_text().splitlines()][1:] + assert len(steps) == 5 + assert [s["index"] for s in steps] == [0, 1, 2, 3, 4] + + +def test_events_are_recorded_in_order(tmp_path): + uw, model, path = _model(tmp_path) + + with model.step(0.1): + model._record_step_event("solve", "SNES_Stokes(v)") + model._record_step_event("history_shift", "EulerianSUPG(T)", dt=0.1) + model._record_step_event("solve", "SNES_AdvectionDiffusion(T)") + + step = json.loads(path.read_text().splitlines()[-1]) + assert [(e["kind"], e["name"]) for e in step["events"]] == [ + ("solve", "SNES_Stokes(v)"), + ("history_shift", "EulerianSUPG(T)"), + ("solve", "SNES_AdvectionDiffusion(T)"), + ] + assert step["events"][1]["dt"] == pytest.approx(0.1) + + +def test_dimensional_values_survive_the_round_trip(tmp_path): + uw, model, path = _model(tmp_path, units=True) + + dt = uw.quantity(1.5, "Myr") + with model.step(dt, label="sink"): + pass + + header, step = (json.loads(line) for line in path.read_text().splitlines()) + assert header["scales"]["length"]["units"] == "meter" + assert header["scales"]["length"]["magnitude"] == pytest.approx(500e3, rel=1e-9) + assert step["dt"] == {"magnitude": pytest.approx(1.5), "units": "megayear"} + assert step["t1"] == {"magnitude": pytest.approx(1.5), "units": "megayear"} + + +def test_clear_journal_opens_a_new_run_in_the_same_file(tmp_path): + """An inversion runs the forward model many times; one file, many runs.""" + uw, model, path = _model(tmp_path) + + for run in range(3): + model.clear_journal() + model.tracker.time = 0.0 + model.tracker.step = 0 + for _ in range(run + 1): + with model.step(0.1, label=f"run{run}"): + pass + + runs = uw.read_journal(path) + # The first header is written when journal_file is set; clear_journal adds + # one per run, so the leading empty section is expected. + populated = [r for r in runs if r["steps"]] + assert [len(r["steps"]) for r in populated] == [1, 2, 3] + assert [r["steps"][0]["label"] for r in populated] == ["run0", "run1", "run2"] + + +def test_a_truncated_final_line_does_not_lose_the_rest(tmp_path): + """A run killed mid-write: everything before the partial line is intact.""" + uw, model, path = _model(tmp_path) + + for _ in range(3): + with model.step(0.1, label="convect"): + pass + + with open(path, "a", encoding="utf-8") as handle: + handle.write('{"kind": "step", "index": 3, "lab') + + runs = uw.read_journal(path) + assert len(runs) == 1 + assert [s["index"] for s in runs[0]["steps"]] == [0, 1, 2] + + +def test_logging_is_off_by_default(tmp_path): + uw, model, path = _model(tmp_path) + model.journal_file = None + + assert model.journal_file is None + before = path.read_text() + with model.step(0.1): + pass + assert path.read_text() == before, "writing continued after logging was off" + assert len(model.journal) == 1, "the in-memory journal must be unaffected" + + +# --------------------------------------------------------------------------- +# The text format — what a run is watched through +# --------------------------------------------------------------------------- + + +def test_the_default_format_is_text_and_the_suffix_chooses_json(tmp_path): + uw, model, path = _model(tmp_path, name="run.log") + assert model.journal_format == "text" + + model.journal_file = str(tmp_path / "run.jsonl") + assert model.journal_format == "jsonl" + + model.journal_format = "text" + assert model.journal_format == "text", "an explicit format must win" + + +def test_text_log_is_one_aligned_line_per_step(tmp_path): + uw, model, path = _model(tmp_path, units=True, name="run.log") + + dt = uw.quantity(0.5, "Myr") + for _ in range(3): + with model.step(dt, label="convect"): + model._record_step_event("solve", "SNES_Stokes(v)") + + lines = path.read_text().splitlines() + comments = [l for l in lines if l.startswith("#")] + rows = [l for l in lines if l.strip() and not l.startswith("#")] + + assert any("underworld3 step log" in c for c in comments) + assert any("scales:" in c for c in comments) + assert any("t/Myr" in c and "dt/Myr" in c for c in comments), ( + "the column header must name the unit the time column is in" + ) + assert len(rows) == 3 + for index, row in enumerate(rows): + assert row.split()[0] == str(index) + assert "solve:SNES_Stokes(v)" in row + assert "ok" in row + + +def test_text_log_converts_dt_into_the_clock_unit(tmp_path): + """A dt in seconds beside a clock in Myr is converted; the table has one unit.""" + uw, model, path = _model(tmp_path, units=True, name="run.log") + + dt = uw.quantity(0.5, "Myr").to("s") # same interval, other unit + with model.step(dt, label="convect"): + pass + + row = [l for l in path.read_text().splitlines() + if l.strip() and not l.startswith("#")][0] + fields = row.split() + assert float(fields[1]) == pytest.approx(0.5, rel=1e-6), fields + assert float(fields[2]) == pytest.approx(0.5, rel=1e-6), ( + f"dt was not converted into the clock's unit: {fields}" + ) + + +def test_a_backtrack_is_in_the_log(tmp_path): + """A log that shows step 2, then step 2 again, must say what happened.""" + uw, model, path = _model(tmp_path, name="run.log") + model.record_every = 1 + + for _ in range(3): + with model.step(0.1, label="convect"): + pass + + model.rewind() + + lines = path.read_text().splitlines() + notes = [l for l in lines if l.strip().startswith("--")] + assert len(notes) == 1, lines + assert "rewind to the start of step 2" in notes[0] + assert "1 step(s) undone" in notes[0] + + +def test_a_bare_restore_is_in_the_log_too(tmp_path): + """The backstepping idiom is save_state / load_state, not rewind.""" + uw, model, path = _model(tmp_path, name="run.log") + + with model.step(0.1, label="convect"): + pass + snap = model.save_state() + with model.step(0.1, label="convect"): + pass + model.load_state(snap) + + notes = [l for l in path.read_text().splitlines() if l.strip().startswith("--")] + assert len(notes) == 1, notes + assert "restore from a snapshot" in notes[0] + + +def test_rewind_logs_one_note_not_two(tmp_path): + """rewind() restores internally; only its own, more specific, note is written.""" + uw, model, path = _model(tmp_path, name="run.log") + model.record_every = 1 + + with model.step(0.1): + pass + model.rewind() + + notes = [l for l in path.read_text().splitlines() if l.strip().startswith("--")] + assert len(notes) == 1, notes + assert "rewind" in notes[0] + assert "restore from" not in notes[0] + + +def test_backtracks_are_records_in_the_json_format(tmp_path): + uw, model, path = _model(tmp_path, name="run.jsonl") + model.record_every = 1 + + for _ in range(2): + with model.step(0.1): + pass + model.rewind() + + kinds = [json.loads(line)["kind"] for line in path.read_text().splitlines()] + assert kinds == ["run", "step", "step", "rewind"] + + rewind = json.loads(path.read_text().splitlines()[-1]) + assert rewind["to_step"] == 1 + assert rewind["steps_undone"] == 1 + + +def test_the_invariant_is_recorded_against_the_step(tmp_path): + """The complaint belongs in the log, not only in whatever terminal ran it.""" + uw, model, path = _model(tmp_path, name="run.jsonl") + + with pytest.warns(RuntimeWarning, match="history advanced more than once"): + with model.step(0.1): + model._record_step_event("history_shift", "EulerianSUPG(T)", dt=0.1) + model._record_step_event("history_shift", "EulerianSUPG(T)", dt=0.1) + + step = json.loads(path.read_text().splitlines()[-1]) + flags = [e for e in step["events"] if e["kind"] == "invariant"] + assert len(flags) == 1 + assert "more than once" in flags[0]["name"] + assert "EulerianSUPG(T) x2" in flags[0]["detail"] diff --git a/tests/test_0015_journal_report.py b/tests/test_0015_journal_report.py new file mode 100644 index 000000000..ea6df57f0 --- /dev/null +++ b/tests/test_0015_journal_report.py @@ -0,0 +1,332 @@ +"""A run's log, as a figure. + +The log is written to be watched; these renderers turn it into something to +put in a paper. Time runs DOWN the page, one row per step, so the figure is +portrait and paginates. + +The layout decision under test is the one that makes a long run legible: each +distinct operator sequence gets a LETTER, defined once at the foot. A column +of ``A`` with a single ``B`` in it says at a glance that one step did something +different, where a hundred repeated sequences say nothing and hide the one that +matters. +""" + +import pytest + +pytestmark = [pytest.mark.level_1, pytest.mark.tier_a] + +import json +import xml.dom.minidom + + +def _run(steps, notes=None, model="test", scales=None): + return [{ + "run": {"kind": "run", "model": model, "started": "2026-01-01T00:00:00+00:00", + "scales": scales or {}}, + "steps": steps, + "notes": notes or [], + }] + + +def _step(index, dt=0.5, unit="megayear", t0=None, completed=True, wall=0.1, + events=None, label="convect"): + t0 = index * dt if t0 is None else t0 + quantity = (lambda v: {"magnitude": v, "units": unit}) if unit else (lambda v: v) + return { + "kind": "step", "index": index, "label": label, + "t0": quantity(t0), "t1": quantity(t0 + dt), "dt": quantity(dt), + "completed": completed, "restorable": True, "wall": wall, + "events": events if events is not None else [ + {"kind": "solve", "name": "SNES_AdvectionDiffusion_Composed(T)"}, + {"kind": "history_shift", "name": "EulerianSUPG(T)", "dt": dt}, + {"kind": "solve", "name": "SNES_Stokes(v)"}, + ], + } + + +def _svg(tmp_path, runs, **kwargs): + import underworld3 as uw + + out = str(tmp_path / "run.svg") + uw.journal_diagram(runs, out=out, **kwargs) + text = open(out, encoding="utf-8").read() + xml.dom.minidom.parseString(text) # must be well-formed + return text + + +def test_the_figure_is_self_contained_svg(tmp_path): + text = _svg(tmp_path, _run([_step(i) for i in range(4)])) + assert text.startswith("]*width="([-\d.]+)"', text)] + assert max(edges) <= width + 0.5 + assert min(float(x) for x in re.findall(r'= -0.5 + + +def test_a_shared_sequence_is_written_once(tmp_path): + text = _svg(tmp_path, _run([_step(i) for i in range(12)])) + assert "Operator sequences" in text + assert text.count("AdvectionDiffusion(T)") == 1, ( + "an identical sequence must not be spelled out per step" + ) + assert ">12 steps<" in text.replace(" ", " ") or "12 steps" in text + + +def test_a_step_that_differs_gets_its_own_letter(tmp_path): + odd = _step(5, events=[ + {"kind": "solve", "name": "SNES_AdvectionDiffusion_Composed(T)"}, + {"kind": "history_shift", "name": "EulerianSUPG(T)", "dt": 0.5}, + {"kind": "solve", "name": "SNES_Stokes(v)"}, + {"kind": "solve", "name": "SNES_Stokes(v)"}, + ]) + steps = [_step(i) for i in range(5)] + [odd] + [_step(i) for i in range(6, 10)] + + text = _svg(tmp_path, _run(steps)) + assert ">A<" in text and ">B<" in text, "two sequences, two letters" + assert "9 steps" in text and "1 step<" in text + # The A column is one letter per row, so the legend must be the only place + # the sequences are spelled out. + assert text.count("shift EulerianSUPG(T)") == 2 + + +def test_one_sequence_means_one_letter(tmp_path): + text = _svg(tmp_path, _run([_step(i) for i in range(6)])) + assert ">B<" not in text + + +def test_an_abandoned_step_is_marked(tmp_path): + steps = [_step(0), _step(1, dt=9.0, completed=False), _step(1)] + text = _svg(tmp_path, _run(steps)) + assert "abandoned" in text + assert "1 abandoned" in text + + +def test_a_backtrack_is_drawn(tmp_path): + steps = [_step(i) for i in range(6)] + notes = [{"kind": "rewind", "after_position": 5, "to_position": 1, + "short": "rewind 4", "steps_undone": 4}] + text = _svg(tmp_path, _run(steps, notes)) + assert "rewind 4" in text + assert "backtrack(s)" in text + assert "]*monospace" font-size="([\d.]+)"()[^>]*' + r'text-anchor="(\w+)"[^>]*>([^<]*)<', text): + plain = content.replace(">", ">").replace("&", "&") + assert float(x) + len(plain) * 0.60 * float(size) <= width - 20, plain + + +def test_the_page_fits_a_document_column(tmp_path): + """Time runs down the page, so width is fixed and height grows.""" + import re + + def size(n): + text = _svg(tmp_path, _run([_step(i) for i in range(n)])) + return tuple(int(v) for v in + re.search(r'width="(\d+)" height="(\d+)"', text).groups()) + + small_w, small_h = size(5) + big_w, big_h = size(40) + assert small_w == big_w <= 620, "the page must not widen with the run" + assert big_h > small_h + 30 * 12, "height must grow one row per step" + + +def test_the_time_span_is_on_the_figure(tmp_path): + """dt is what is plotted; without this the figure never says WHEN.""" + text = _svg(tmp_path, _run([_step(i) for i in range(4)])) + assert "t = 0 to" in text + assert "Myr" in text + + +def test_a_nondimensional_run_still_renders(tmp_path): + text = _svg(tmp_path, _run([_step(i, unit=None) for i in range(4)])) + assert "dt" in text + assert "None" not in text + + +# --------------------------------------------------------------------------- +# Mermaid +# --------------------------------------------------------------------------- + +def test_flowchart_is_one_chain_when_every_step_agrees(): + import underworld3 as uw + + text = uw.journal_flowchart(_run([_step(i) for i in range(6)])) + assert text.startswith("flowchart LR") + assert "subgraph" not in text + assert text.count("-->") == 2 + assert "shift EulerianSUPG(T)" in text + + +def test_flowchart_separates_the_step_that_differs(): + import underworld3 as uw + + odd = _step(3, events=[{"kind": "solve", "name": "SNES_Stokes(v)"}]) + text = uw.journal_flowchart(_run([_step(0), _step(1), _step(2), odd])) + assert text.count("subgraph") == 2 + assert 'step 3' in text + assert 'step 0, 1, 2' in text + + +# --------------------------------------------------------------------------- +# Sources +# --------------------------------------------------------------------------- + +def test_a_live_model_can_be_drawn_without_a_file(tmp_path): + """The text log cannot be read back, so a live model must be a source.""" + import underworld3 as uw + + uw.reset_default_model() + model = uw.get_default_model() + model.tracker.time = 0.0 + model.tracker.step = 0 + for _ in range(3): + with model.step(0.25, label="convect"): + model._record_step_event("solve", "SNES_Stokes(v)") + + out = str(tmp_path / "live.svg") + uw.journal_diagram(model, out=out) + text = open(out, encoding="utf-8").read() + xml.dom.minidom.parseString(text) + assert "3 steps" in text + + +def test_reading_a_text_log_says_what_to_do_instead(tmp_path): + import underworld3 as uw + + uw.reset_default_model() + model = uw.get_default_model() + path = tmp_path / "run.log" + model.journal_file = str(path) + model.tracker.time = 0.0 + model.tracker.step = 0 + with model.step(0.1): + pass + + with pytest.raises(ValueError, match="journal_format = 'jsonl'"): + uw.read_journal(str(path)) + + +# --------------------------------------------------------------------------- +# PDF +# --------------------------------------------------------------------------- + + +def _pdf(tmp_path, runs, name="run.pdf", **kwargs): + import underworld3 as uw + + out = str(tmp_path / name) + uw.journal_diagram(runs, out=out, **kwargs) + return open(out, "rb").read() + + +def test_pdf_is_the_default_and_is_a_real_pdf(tmp_path): + import underworld3 as uw + + out = str(tmp_path / "run") + written = uw.journal_diagram(_run([_step(i) for i in range(5)]), out=out) + assert written == out + data = open(out, "rb").read() + assert data.startswith(b"%PDF-1.4") + assert data.rstrip().endswith(b"%%EOF") + assert b"/Type /Catalog" in data and b"xref" in data + assert b"/BaseFont /Helvetica" in data + + +def test_pdf_paginates_a_long_run(tmp_path): + steps = [_step(i) for i in range(200)] + data = _pdf(tmp_path, _run(steps)) + pages = data.count(b"/Type /Page ") + assert pages >= 4, f"200 steps should not fit on {pages} page(s)" + assert data.count(b"/Type /Pages") == 1 + + +def test_pdf_offsets_point_at_their_objects(tmp_path): + """A cross-reference table that lies makes an unopenable file.""" + import re + + data = _pdf(tmp_path, _run([_step(i) for i in range(30)])) + start = int(re.search(rb"startxref\s+(\d+)", data).group(1)) + assert data[start:start + 4] == b"xref" + body = data[start:].split(b"trailer")[0].splitlines() + entries = [line for line in body[2:] if line.strip().endswith(b"n")] + for index, line in enumerate(entries, start=1): + offset = int(line.split()[0]) + assert data[offset:offset + len(f"{index} 0 obj")] == \ + f"{index} 0 obj".encode(), f"object {index} is not at its offset" + + +def test_pdf_is_written_without_a_plotting_library(tmp_path): + """No dependency is imported to produce the figure.""" + import sys + + for module in ("matplotlib", "cairosvg", "reportlab", "PIL"): + sys.modules.pop(module, None) + _pdf(tmp_path, _run([_step(i) for i in range(5)])) + for module in ("matplotlib", "cairosvg", "reportlab"): + assert module not in sys.modules, f"{module} was imported to draw" + + +def test_a_jsonl_log_round_trips_into_a_figure(tmp_path): + import underworld3 as uw + + uw.reset_default_model() + model = uw.get_default_model() + path = tmp_path / "run.jsonl" + model.journal_file = str(path) + model.record_every = 1 + model.tracker.time = 0.0 + model.tracker.step = 0 + + for _ in range(4): + with model.step(0.25, label="convect"): + model._record_step_event("solve", "SNES_Stokes(v)") + model.rewind() + + out = uw.journal_diagram(str(path)) + assert out.endswith(".pdf") + assert open(out, "rb").read().startswith(b"%PDF") + + out_svg = uw.journal_diagram(str(path), out=str(tmp_path / "run.svg")) + text = open(out_svg, encoding="utf-8").read() + xml.dom.minidom.parseString(text) + assert "1 backtrack(s)" in text + assert "stroke-dasharray" in text, "the backtrack should be drawn" diff --git a/tests/test_1074_free_surface_config_drift.py b/tests/test_1074_free_surface_config_drift.py new file mode 100644 index 000000000..bbc6adf61 --- /dev/null +++ b/tests/test_1074_free_surface_config_drift.py @@ -0,0 +1,90 @@ +"""FreeSurface's derived lids must not drift from the free solve. + +`held` and `consistent` are separate Stokes solvers, and the free solve's +configuration is copied into them ONCE, when the manager is built. Changing the +free solve afterwards leaves them stale — and `h_inf`, the equilibrium the +surface relaxes toward, is recovered from the HELD solve. So the surface would +relax toward an equilibrium computed with the old rheology while the free solve +uses the new one: a wrong answer, not a crash. +""" + +import numpy as np +import pytest +import sympy +import underworld3 as uw + +pytestmark = [pytest.mark.level_1, pytest.mark.tier_a] + + +def _model(): + uw.reset_default_model() + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1.0 / 8, qdegree=3 + ) + v = uw.discretisation.MeshVariable("V_fs", mesh, 2, degree=2) + p = uw.discretisation.MeshVariable("P_fs", mesh, 1, degree=1) + stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = 1.0 + stokes.bodyforce = sympy.Matrix([0, -1.0]) + stokes.add_dirichlet_bc((0.0, sympy.oo), "Left") + stokes.add_dirichlet_bc((0.0, sympy.oo), "Right") + stokes.add_dirichlet_bc((sympy.oo, 0.0), "Bottom") + stokes.petsc_options.delValue("ksp_monitor") + return mesh, stokes + + +def test_changing_the_rheology_after_construction_is_refused(): + mesh, stokes = _model() + fs = uw.systems.FreeSurface(stokes, "Top", buoyancy_scale=1.0) + + stokes.constitutive_model.Parameters.shear_viscosity_0 = 1000.0 + + with pytest.raises(RuntimeError, match="no longer match the free solve"): + fs.solve() + + +def test_the_message_names_what_drifted_and_how_to_recover(): + mesh, stokes = _model() + fs = uw.systems.FreeSurface(stokes, "Top", buoyancy_scale=1.0) + stokes.bodyforce = sympy.Matrix([0, -5.0]) + + with pytest.raises(RuntimeError) as excinfo: + fs.solve() + message = str(excinfo.value) + assert "bodyforce" in message + assert "BEFORE constructing" in message + + +def test_an_untouched_manager_solves(): + """The negative control: the guard must not refuse an ordinary run.""" + mesh, stokes = _model() + fs = uw.systems.FreeSurface(stokes, "Top", buoyancy_scale=1.0) + fs.solve() + assert np.all(np.isfinite(np.asarray(stokes.u.data))) + + +def test_a_shared_symbolic_parameter_still_tracks(): + """A parameter whose value is an expression over mesh variables is shared + symbolically and must NOT trip the guard — only re-assignment drifts.""" + uw.reset_default_model() + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1.0 / 8, qdegree=3 + ) + v = uw.discretisation.MeshVariable("V_sym", mesh, 2, degree=2) + p = uw.discretisation.MeshVariable("P_sym", mesh, 1, degree=1) + T = uw.discretisation.MeshVariable("T_sym", mesh, 1, degree=2) + T.array[:, 0, 0] = 0.5 + stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = sympy.exp(-T.sym[0]) + stokes.bodyforce = sympy.Matrix([0, -T.sym[0]]) + stokes.add_dirichlet_bc((0.0, sympy.oo), "Left") + stokes.add_dirichlet_bc((0.0, sympy.oo), "Right") + stokes.add_dirichlet_bc((sympy.oo, 0.0), "Bottom") + stokes.petsc_options.delValue("ksp_monitor") + + fs = uw.systems.FreeSurface(stokes, "Top", buoyancy_scale=1.0) + T.array[:, 0, 0] = 0.9 # the FIELD changes, not the expression + fs.solve() # must not raise + assert np.all(np.isfinite(np.asarray(stokes.u.data))) From 9e82d403484a5ad4bf6d4b69091ada1ed3fa875c Mon Sep 17 00:00:00 2001 From: lmoresi Date: Tue, 15 Sep 2026 10:07:59 +1000 Subject: [PATCH 04/12] =?UTF-8?q?feat:=20the=20transcript=20=E2=80=94=20wh?= =?UTF-8?q?at=20a=20run=20did,=20judged=20by=20a=20later=20pass=20rather?= =?UTF-8?q?=20than=20in=20the=20loop,=20and=20left=20by=20default?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Named: journal, record (noun) and log become the transcript — what the run actually did, false starts and re-takes included. Record stays the verb. An earlier revision asserted a rule inside the loop — a history must advance exactly once per step — and it was wrong twice over: it fired on legitimate sub-cycling, and it could not see the checks that matter, which are growth rates across steps. It is gone. The step records; a pass over a finished transcript judges. The design note sets out what such a pass can read from the record and what it cannot yet. Every run now leaves its transcript without being asked, in ./transcripts/-