From 3eba6aa49e9f8863970f99caa3cb6829917f6340 Mon Sep 17 00:00:00 2001 From: Tyagi Date: Sat, 25 Jul 2026 11:36:05 +0530 Subject: [PATCH 1/4] Reuse the rotated free-slip solver workspace across solves (#417) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Repeated transient rotated free-slip Stokes solves rebuilt the rotation Q, the PtAP'd operator, the fieldsplit Schur KSP/PC and the GAMG hierarchy on every timestep — an allocator high-water problem (RSS growth to OOM, issue #417) and roughly half the per-step cost of the Zhong A1 production runs. Cache the workspace on the solver between solves, re-derived from the original PR #418 (bec76bbd) at the seam of the rewritten unified Newton loop: - geometry tier (Q/Qt, constrained rows, fault pair blocks, custom-FMG prolongation) reused while the boundary/fault registration and DM are unchanged; - structure tier (Ahat, Schur pmat, nullspaces, KSP/PC) reused as objects with values refreshed in place — the loop's own between-iteration pattern; - an iteration-0 fast path skips Jacobian assembly / ptap / PCSetUp entirely when the operator coefficient state counters match (RHS-only timesteps); - the cache is forfeited for direct-LU, prescribed-datum and fault interface-law solves, and torn down by _reset() and the _build() full rebuild before the SNES/DM are destroyed; - expose the existing time= argument through the Stokes.solve wrapper; an explicit time vetoes the fast path (petsc_t bypasses every counter). Unlike the original one-shot linear path, a wrong fast-path verdict here cannot return a stale solution: the loop measures the true residual at every iterate and reassembles from iteration 1 on. Regression: RHS-only reuse (same Q/Ahat/KSP handles), viscosity-field invalidation with in-place refresh, and the time= veto, in test_1018_rotated_freeslip. Production evidence on the original mechanism: 310 guarded Zhong A1 steps, flat RSS (PR #418 thread). Underworld development team with AI support from Claude Code --- .../design/ROTATED_FREESLIP_LINEAR_REUSE.md | 112 +++++++ .../cython/petsc_generic_snes_solvers.pyx | 35 ++- src/underworld3/systems/solvers.py | 6 + src/underworld3/utilities/rotated_bc.py | 291 ++++++++++++++++-- tests/test_1018_rotated_freeslip.py | 98 ++++++ 5 files changed, 516 insertions(+), 26 deletions(-) create mode 100644 docs/developer/design/ROTATED_FREESLIP_LINEAR_REUSE.md diff --git a/docs/developer/design/ROTATED_FREESLIP_LINEAR_REUSE.md b/docs/developer/design/ROTATED_FREESLIP_LINEAR_REUSE.md new file mode 100644 index 00000000..3935a116 --- /dev/null +++ b/docs/developer/design/ROTATED_FREESLIP_LINEAR_REUSE.md @@ -0,0 +1,112 @@ +# Rotated Free-Slip Workspace Reuse Across Solves + +## Scope + +This note documents the cross-solve workspace cache for rotated strong +free-slip Stokes solves. It addresses +[issue #417](https://github.com/underworldcode/underworld3/issues/417): +repeated transient solves rebuilt the rotated operator, the fieldsplit Schur +preconditioner and the GAMG hierarchy on every timestep, growing resident +memory until out-of-memory and paying distributed assembly and `PCSetUp` +every step. + +The change does not alter the rotated boundary condition or its discrete +equations. It changes PETSc object ownership across solves, and decides when +per-solve work can be skipped. + +## Where the reuse lives now + +The rotated solve is a single manual Newton/Picard loop +(`rotated_bc.solve_rotated_freeslip`) — there is no separate linear path. +Within one solve the loop already reuses its operator and KSP context between +Newton iterations (ptap-with-result, `setOperators` refresh). The cache +extends exactly that pattern across solves, split by what each piece depends +on: + +- **Geometry tier** — the rotation `Q`/`Q.T`, the constrained normal rows, + the fault contact pair blocks, and the custom-FMG prolongation depend only + on the mesh, the boundary specs and the fault registration. They are reused + whenever the geometry signature matches (boundary names + normals, fault + registration, DM identity). +- **Structure tier** — the transformed operator `Ahat`, the pressure-mass + Schur block `Mp`, the nullspaces and the fieldsplit KSP/PC are reused as + *objects* with their values refreshed in place — the same operation + sequence the Newton loop performs between its own iterations, so it carries + the same (production-validated) risk profile. This tier is always correct + regardless of any change detection, because values are reassembled. +- **Iteration-0 fast path** — when the operator key proves nothing + operator-relevant changed, the first Newton increment additionally skips + Jacobian assembly, the ptap and `PCSetUp` entirely. This is the timestep + fast path: a body-force (temperature) change alters only the residual, so + repeated constant-viscosity Stokes solves pay one residual assembly and one + Krylov solve per step. + +## The structural safety net + +A wrong fast-path decision cannot produce a wrong answer. The loop measures +the **true residual** (fresh kernels, current constants) at every iterate, +convergence is declared only on that measurement, and every iteration after +the first always reassembles the operator. A stale cached operator therefore +costs one extra Newton increment; it can never return a stale solution. + +## Invalidation + +- Operator coefficient **mesh variables** are collected from the constitutive + parameters, constraint term, penalty and saddle preconditioner (expressions + unwrapped first, unknowns excluded). Their base `MeshVariable._state` + counters must match for the fast path. +- If the coefficient enumeration fails, the fast path is forfeited and every + solve reassembles (structure-tier reuse only) — correctness first. +- An explicit `solve(time=...)` vetoes the fast path: `petsc_t` reaches the + kernels through the DM, invisible to any counter. +- Mesh deform, field-layout, boundary-condition or forced-setup changes route + through the solver's full-rebuild teardown, which destroys the whole + workspace (`_reset_rotated_solver_cache`, called from `_reset()` and from + the `_build()` full-rebuild branch, before the SNES/DM are destroyed). +- A geometry-signature mismatch (different boundaries or fault registration + on the same solver) destroys and rebuilds the workspace. + +The cache is forfeited entirely for direct-LU solves, prescribed-datum solves +(the datum is re-evaluated from possibly-changed fields inside +`build_rotation`), and fault **interface-law** solves (the interface tangent +and the Picard-lagged normal stress are solution-dependent). Frictionless +split-fault contacts cache normally: their pair blocks are geometry. + +The high-level `uw.systems.Stokes.solve()` wrapper exposes and forwards +`time=` for the veto. It remains distinct from the viscoelastic integration +`timestep=`. + +## PETSc ownership + +`_destroy_rotated_ksp_ctx` destroys the KSP, the pressure-mass block and the +owned nullspaces. `_destroy_rotated_linear_cache` additionally destroys the +transformed operator and the rotation matrices, and *dereferences* (does not +destroy) the custom-FMG prolongation list, whose coarse matrices are shared +with the solver's registered multigrid hierarchy. The result dict shares the +`Q`/`Q.T` Python wrappers with the cache, so garbage collection and explicit +teardown compose without double-destroys. + +## Validation + +The focused regression (`test_1018_rotated_freeslip.py`) verifies that: + +- a body-force-only change preserves the `Q`, `Ahat` and KSP handles and + rides the fast path (`workspace_reused`), with the velocity scaling exactly + with the right-hand side; +- a viscosity mesh-variable change is detected and refreshes the operator + values in place on the same objects; +- an explicit `time=` solve vetoes the fast path. + +Production validation for the original mechanism: 310 guarded Zhong A1 +steps (8 ranks, `cellsize=1/8`) with flat RSS and a continuous physical +trajectory — see the PR #418 thread. + +Always launch MPI tests with the worktree MPI executable: + +```bash +.pixi/envs/amr-dev/bin/mpirun -np 2 \ + .pixi/envs/amr-dev/bin/python +``` + +Using a system `mpirun` from a different Open MPI installation can stall +during initialization and is not a solver failure. diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index 8d07446b..ba5250e5 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -1291,8 +1291,24 @@ class SolverBaseClass(uw_object): return + def _reset_rotated_solver_cache(self): + """Release the rotated-free-slip cross-solve workspace (rotated_bc + cache: Q/Qt, the PtAP'd operator, the fieldsplit KSP/PC) before any + solver/DM teardown — those PETSc objects reference the current DM row + layout and must not survive it. No-op for solvers without the cache + (getattr guard: only SNES_Stokes_SaddlePt ever populates it). The + last solve's result dict (``_rotated_freeslip_info``) is NOT dropped: + its reaction vector is independent of the cache and the σ_nn / + dynamic-topography recoveries may still need it.""" + cache = getattr(self, "_rotated_linear_cache", None) + if cache is not None: + from underworld3.utilities.rotated_bc import _destroy_rotated_linear_cache + _destroy_rotated_linear_cache(cache) + self._rotated_linear_cache = None + def _reset(self): + self._reset_rotated_solver_cache() self.natural_bcs = [] self.essential_bcs = [] # A teardown means the next solve is a different discrete problem: a resume @@ -2065,6 +2081,15 @@ class SolverBaseClass(uw_object): # sequence; this brings _build() into line with it. # NB self.snes / self.dm_hierarchy may not exist yet on the first # build, so use getattr/hasattr-style guards rather than `is not None`. + + # The rotated free-slip cross-solve workspace (rotation Q, PtAP'd + # operator, fieldsplit KSP/PC) was built against the SNES/DM we are + # about to destroy — release it first. Fast paths 1 and 2 above keep + # the DM/SNES, so the workspace legitimately survives them (a rewire's + # new kernels are caught by the workspace's own JIT-key/constants + # invalidation signature). + self._reset_rotated_solver_cache() + if getattr(self, "snes", None) is not None: if verbose and uw.mpi.rank == 0: print(f"Destroy solver SNES", flush=True) @@ -5904,6 +5929,10 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): self._rotated_freeslip_bcs = [] self._rotated_freeslip_datum = {} self._rotated_freeslip_info = None + # Cross-solve rotated workspace (rotated_bc cache, issue #417): + # populated/keyed/invalidated entirely inside solve_rotated_freeslip; + # torn down here by _reset_rotated_solver_cache on any DM rebuild. + self._rotated_linear_cache = None # Split-fault interface conditions (add_fault_bc): fault names whose # coincident DOF pairs carry a contact (the laws themselves live in # _fault_interface_laws, set lazily by utilities/fault_contact.py). @@ -9366,7 +9395,11 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): from underworld3.utilities.rotated_bc import solve_rotated_freeslip self._rotated_freeslip_info = solve_rotated_freeslip( self, self._rotated_freeslip_bcs, verbose=verbose, - zero_init_guess=zero_init_guess, picard=picard) + zero_init_guess=zero_init_guess, picard=picard, + # An explicit `time=` reaches the kernels through petsc_t on + # the DM — invisible to every state counter and constant + # value, so it must veto the cached-operator fast path. + force_operator_refresh=time is not None) # This path solves via ksp.solve on the rotated operator (not self.snes), # so give it a report from the rotated result rather than leaving a stale one. _rotated_report = self._capture_rotated_report(self._rotated_freeslip_info) diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index 2f129b8e..fe997d8f 100644 --- a/src/underworld3/systems/solvers.py +++ b/src/underworld3/systems/solvers.py @@ -1415,6 +1415,7 @@ def solve( self, zero_init_guess: bool = None, timestep: float = None, + time=None, _force_setup: bool = False, verbose: bool = False, debug: bool = False, @@ -1445,6 +1446,9 @@ def solve( ``has_solution``. timestep : float, optional Advection timestep. Required when stress history is active. + time : float or Quantity, optional + Physical evaluation time for expressions using ``mesh.t``. This is + distinct from the viscoelastic integration ``timestep``. _force_setup : bool Force rebuild of pointwise functions. verbose : bool @@ -1565,6 +1569,7 @@ def solve( _force_setup=_force_setup, verbose=verbose, picard=picard, + time=time, divergence_retries=divergence_retries, ) @@ -1622,6 +1627,7 @@ def solve( _force_setup=_force_setup, verbose=verbose, picard=picard, + time=time, divergence_retries=divergence_retries, ) diff --git a/src/underworld3/utilities/rotated_bc.py b/src/underworld3/utilities/rotated_bc.py index adb502a8..925aef68 100644 --- a/src/underworld3/utilities/rotated_bc.py +++ b/src/underworld3/utilities/rotated_bc.py @@ -402,6 +402,130 @@ def _set_rows_local(vec, row_val_map): vec.setArray(a) +# --------------------------------------------------------------------------- # +# Cross-solve workspace reuse (issue #417) +# +# Repeated rotated free-slip solves used to rebuild the whole linear +# workspace every time: the rotation Q, the PtAP'd operator, the fieldsplit +# Schur KSP/PC and its GAMG (or FMG) hierarchy. In a time-stepping loop that +# is both slow (distributed assembly + PCSetUp every step) and an allocator +# high-water problem (a fresh KSP/PC object graph per step). The cache below +# keeps the workspace alive on the solver between solves, split by what each +# piece actually depends on: +# +# * GEOMETRY tier — Q/Qt, the constrained rows, the fault pair blocks and +# the custom-FMG prolongation depend only on the mesh, the boundary specs +# and the fault registration. They are reused whenever the geometry +# signature matches, and torn down by the solver's own rebuild paths +# (_reset / the _build full-rebuild teardown), which fire on mesh.deform, +# BC changes and every is_setup=False invalidation. +# * STRUCTURE tier — Ahat, the Schur pmat Mp and the KSP/PC context are +# reused as OBJECTS with their values refreshed in place (ptap-with-result +# / createSubMatrix-with-submat / setOperators), exactly the pattern the +# Newton loop already uses between its own iterations. This is always +# correct regardless of any change-detection: values are reassembled. +# * FAST PATH — iteration 0 of the Newton loop may additionally skip the +# Jacobian assembly / ptap / PCSetUp entirely when the operator key proves +# nothing operator-relevant changed. A wrong skip cannot produce a wrong +# answer: the loop measures the TRUE residual at every iterate and every +# iteration after the first always reassembles, so a stale operator costs +# one extra increment, never a stale solution (the structural safety net +# that the state-counter-keyed original lacked). +# +# The cache is forfeited entirely (correctness first) for the paths where +# reuse is not provably sound: direct-LU solves, prescribed-datum solves +# (the datum is re-evaluated from possibly-changed fields inside +# build_rotation), and fault interface laws (the interface tangent and the +# Picard-lagged normal stress are solution-dependent state). +# --------------------------------------------------------------------------- # + +def _operator_coefficient_variables(solver): + """Mesh variables that can change the assembled Stokes OPERATOR (not just + the RHS): everything reachable from the constitutive parameters, the + constraint term, the penalty and the saddle preconditioner. The solver's + own unknowns are excluded (the solve writes them every time). Expressions + are unwrapped first so variables hidden inside nested UWexpressions are + seen (the "unwrap before extracting atoms" rule). + + Returns a tuple of variables, or ``None`` when the enumeration fails — + the caller must then treat the operator as always-changed (no fast path; + the structural refresh path is used every solve).""" + from underworld3.function.expressions import mesh_vars_in_expression, unwrap + + expressions = [] + try: + expressions.append(solver.constraints) + expressions.append(solver.penalty) + saddle_preconditioner = getattr(solver, "saddle_preconditioner", None) + if saddle_preconditioner is not None: + expressions.append(saddle_preconditioner) + parameters = getattr(solver.constitutive_model, "Parameters", None) + if parameters is not None: + for name in parameters._list_valid_parameters(type(parameters)): + expressions.append(getattr(parameters, name)) + except Exception: + return None + + variables = set() + for expression in expressions: + expression = getattr(expression, "sym", expression) + if expression is None or not hasattr(expression, "args"): + continue + try: + expression = unwrap(expression, keep_constants=False, + return_self=False) + _, regular, derivatives = mesh_vars_in_expression( + sympy.sympify(expression)) + except Exception: + return None + variables.update(fn.meshvar() for fn in regular) + variables.update(derivatives) + + unknowns = {getattr(variable, "_base_var", variable) + for variable in solver.fields.values()} + variables = {getattr(variable, "_base_var", variable) + for variable in variables} + variables.difference_update(unknowns) + return tuple(sorted(variables, key=lambda variable: variable._uw_id)) + + +def _coefficient_states(variables): + """UW data-version counters of the operator coefficient variables.""" + return tuple(variable._state for variable in variables) + + +def _rotated_geometry_signature(solver, boundaries): + """Signature of everything the GEOMETRY tier of the workspace was built + from: the boundary specs (name + normal), the fault-contact registration + (names + analytic-normal overrides) and the DM identity. All entries are + registration state, identical across ranks.""" + dm = solver.dm + sig_boundaries = tuple((name, repr(normal)) + for name, normal in map(_boundary_spec, boundaries)) + sig_faults = tuple(getattr(solver, "_fault_contact_faults", []) or []) + sig_fault_normals = repr(sorted( + (k, repr(v)) + for k, v in getattr(solver, "_fault_normal_overrides", {}).items())) + return (sig_boundaries, sig_faults, sig_fault_normals, + dm.handle if dm is not None else 0) + + +def _destroy_rotated_linear_cache(cache): + """Release the persistent rotated-free-slip workspace. The custom-FMG + prolongation list is only DEREFERENCED (its coarse Mats are shared with + the solver's registered multigrid hierarchy).""" + if not cache: + return + _destroy_rotated_ksp_ctx(cache.get("ctx")) + cache["ctx"] = None + for key in ("Ahat", "Q", "Qt"): + obj = cache.get(key) + if obj is not None: + obj.destroy() + cache[key] = None + cache["custom_Pl"] = None + + # --------------------------------------------------------------------------- # # The rotated solve # --------------------------------------------------------------------------- # @@ -540,7 +664,8 @@ def _backtracking_line_search(u, d, rnorm, rotated_residual, Q, Qt, normal_rows, def solve_rotated_freeslip(solver, boundaries, remove_rotation_gauge=True, verbose=False, zero_init_guess=True, picard=0, - rtol=None, atol=1.0e-11, stol=1.0e-8, max_it=50): + rtol=None, atol=1.0e-11, stol=1.0e-8, max_it=50, + force_operator_refresh=False): """THE rotated strong-free-slip solve (linear and nonlinear models alike): a manual outer Newton/Picard loop that rotates the residual F(u), the Jacobian J(u) and the strong ``v_n = ũ_n`` constraint (``ũ_n = 0`` for pure free-slip; @@ -621,7 +746,16 @@ def solve_rotated_freeslip(solver, boundaries, remove_rotation_gauge=True, of Newton increments solved, ``== len(ksp_its)``) and status; * ``"rnorm"``, ``"rnorm0"`` — final and initial rotated residual norms ‖F̂‖ (feed the solve report); - * ``"continuation_switched"`` — whether the Picard→Newton tangent switch fired. + * ``"continuation_switched"`` — whether the Picard→Newton tangent switch fired; + * ``"rotation_reused"`` — the GEOMETRY tier (Q/Qt/constrained rows/ + prolongation) came from the cross-solve cache; + * ``"workspace_reused"`` — the iteration-0 operator fast path fired + (Jacobian assembly, ptap and PCSetUp all skipped — see the workspace + reuse block comment above ``_operator_coefficient_variables``). + + ``force_operator_refresh=True`` disables the iteration-0 fast path for this + solve (used by ``time=`` solves: the DM time reaches the kernels through + ``petsc_t``, which no state counter or constants value can see). """ if getattr(solver, "snes", None) is None: solver._setup_pointwise_functions() @@ -676,11 +810,65 @@ def solve_rotated_freeslip(solver, boundaries, remove_rotation_gauge=True, # with no lift, tangent transparency, custom_Pl, nullspace) is unchanged from # pure free-slip. datum_specs = getattr(solver, "_rotated_freeslip_datum", None) - Q, Qt, normal_rows, datum_map = build_rotation(solver, boundaries, datum_specs) # Direct LU per increment: no FMG prolongation / null space to build — the # gauge is fixed by the naive pressure pin instead (see _naive_pressure_pin). use_lu = bool(getattr(solver, "_rotated_use_lu", False)) - custom_Pl = None if use_lu else _build_rotated_custom_Pl(solver, Q, normal_rows) + interface_laws = bool(getattr(solver, "_fault_interface_laws", {})) + + # ---- cross-solve workspace cache (issue #417; see the block comment + # above _operator_coefficient_variables for the tier design) ---- + # Forfeited for LU (per-solve pin/factorisation), prescribed datum (the + # datum values are re-evaluated from possibly-changed fields inside + # build_rotation) and fault interface laws (solution-dependent interface + # tangent + Picard-lagged normal stress) — correctness first. + cache_allowed = not (use_lu or datum_specs or interface_laws) + cache = getattr(solver, "_rotated_linear_cache", None) + if cache is not None and not cache_allowed: + _destroy_rotated_linear_cache(cache) + solver._rotated_linear_cache = None + cache = None + + geometry_sig = _rotated_geometry_signature(solver, boundaries) \ + if cache_allowed else None + coefficient_variables = None + coefficient_states = None + if cache_allowed: + if cache is not None and cache.get("geometry_sig") == geometry_sig \ + and cache.get("coefficient_variables") is not None: + coefficient_variables = cache["coefficient_variables"] + else: + coefficient_variables = _operator_coefficient_variables(solver) + if coefficient_variables is not None: + coefficient_states = _coefficient_states(coefficient_variables) + + op_ok = False + if cache is not None: + geom_ok = cache.get("geometry_sig") == geometry_sig + op_ok = (geom_ok + and not force_operator_refresh + and coefficient_states is not None + and cache.get("coefficient_states") == coefficient_states) + if not geom_ok: + _destroy_rotated_linear_cache(cache) + solver._rotated_linear_cache = None + cache = None + + if cache is not None: + Q, Qt = cache["Q"], cache["Qt"] + normal_rows = cache["normal_rows"] + datum_map = {} # cache_allowed excludes datum solves + custom_Pl = cache.get("custom_Pl") + rotation_reused = True + else: + Q, Qt, normal_rows, datum_map = build_rotation(solver, boundaries, datum_specs) + custom_Pl = None if use_lu else _build_rotated_custom_Pl(solver, Q, normal_rows) + rotation_reused = False + + # The iteration-0 fast path: skip Jacobian assembly / ptap / PCSetUp when + # the operator key proves nothing operator-relevant changed. + reuse_operator = bool(cache is not None and op_ok) + workspace_reused = False + # Interface constitutive laws (fault_contact.add_viscous_fault_bc / # add_coulomb_fault_bc): the assembler caches the fault-trace geometry # once; each iterate it adds the interface force ∫τ(V)δV to the rotated @@ -688,13 +876,15 @@ def solve_rotated_freeslip(solver, boundaries, remove_rotation_gauge=True, # operator — full Newton in V, the stiff direction. (Only a future # reaction-fed σ_n argument would be Picard-lagged, by choice.) interface = None - if getattr(solver, "_fault_interface_laws", {}): + if interface_laws: from underworld3.utilities import fault_contact interface = fault_contact._InterfaceAssembler(solver) # The null space is built INSIDE the loop, after the first Jacobian assembly: # _mode_satisfies_constraints verifies each candidate mode against the # ASSEMBLED operator (‖J·m‖ ≈ 0), which an unassembled J cannot support. - nsp = None + # A cached ctx carries the (geometry-tier) null space it was built with. + nsp = cache["ctx"].get("nsp") if (cache is not None + and cache.get("ctx")) else None # initial guess (cartesian, composite): warm-start from the fields or zero. # A prescribed datum on a COLD start is imposed through the FIRST increment's @@ -738,11 +928,13 @@ def solve_rotated_freeslip(solver, boundaries, remove_rotation_gauge=True, # 1/mu pressure-mass Schur pmat (values refreshed in place), the constraint # diagonal scale (frozen at the first tangent — only the magnitude matters), # and the KSP/PC context (fieldsplit ISs, FMG hierarchy, GAMG setup survive). - Ahat = None + # A cross-solve cache seeds all of them — the STRUCTURE tier: same objects, + # values refreshed in place unless the iteration-0 fast path fires. + Ahat = cache["Ahat"] if cache is not None else None Atot = None - Mp = None - ctx = None - diag_scale = None + ctx = cache["ctx"] if cache is not None else None + Mp = ctx["Mp"] if ctx is not None else None + diag_scale = cache["diag_scale"] if cache is not None else None lin_its = [] # velocity / pressure sub-KSP LAST-APPLICATION counts, one per Newton increment # (iterative path only — the direct-LU path has no sub-KSPs and records None) @@ -843,11 +1035,21 @@ def rotated_residual(uvec, keep_cartesian=False): if verbose: mpi.pprint(f"[rotated_bc] continuation: Picard→Newton at iter {iters} " f"(rel |F̂| {rnorm/(r0+1e-300):.2e})") - snes.computeJacobian(u, J, Jp) # Jp carries the 1/mu mass (Schur pmat) - if Ahat is None: - Ahat = J.ptap(Qt) + # Iteration 0 may ride the cross-solve fast path: the cached Ahat still + # holds the rotated, constraint-eliminated operator and the cached + # KSP/PC is set up on it. Every LATER iteration always reassembles — + # that, plus the exact residual measured at every iterate, is the + # structural safety net: a wrong skip costs one extra increment, never + # a stale solution. + assemble = not (reuse_operator and iters == 0) + if assemble: + snes.computeJacobian(u, J, Jp) # Jp carries the 1/mu mass (Schur pmat) + if Ahat is None: + Ahat = J.ptap(Qt) + else: + J.ptap(Qt, result=Ahat) # same nonzero pattern → in-place refresh else: - J.ptap(Qt, result=Ahat) # same nonzero pattern → in-place refresh + workspace_reused = True # The interface tangent is REASSEMBLED at every iterate (that is # what makes it consistent Newton, dtau/dV at the current slip # rates) and added to a COPY: the ptap-with-result refresh above @@ -874,7 +1076,7 @@ def rotated_residual(uvec, keep_cartesian=False): if ctx is None: Mp = _pressure_mass_schur_pmat(solver) nsp = _rotated_nullspace(solver, Q, normal_rows) # J assembled above - elif Mp is not None: + elif assemble and Mp is not None: Jp.createSubMatrix(pres_is, pres_is, submat=Mp) # viscosity may be u-dependent if diag_scale is None: diag_scale = _velocity_diag_scale(Aop, solver) @@ -888,8 +1090,11 @@ def rotated_residual(uvec, keep_cartesian=False): Aop.zeroRowsColumns(normal_rows, diag=diag_scale, x=xhat, b=bhat) xhat.destroy() xhat = None - else: + elif assemble: Aop.zeroRowsColumns(normal_rows, diag=diag_scale) + # else: the cached Aop already carries the eliminated constraint + # rows/cols from the solve that built it — re-zeroing would bump the + # Mat state and force PCSetUp, defeating the fast path. if use_lu: if ctx is None: ksp_lu = PETSc.KSP().create(comm=dm.comm) @@ -915,7 +1120,8 @@ def rotated_residual(uvec, keep_cartesian=False): else: dhat, last_reason, ctx = _solve_rotated_iterative( solver, Aop, bhat, Q, Qt, normal_rows, - custom_Pl=custom_Pl, nsp=nsp, Mp=Mp, verbose=False, ctx=ctx) + custom_Pl=custom_Pl, nsp=nsp, Mp=Mp, verbose=False, ctx=ctx, + refresh_operator=assemble) lin_its.append(ctx["ksp"].getIterationNumber()) vel_its_last.append(ctx.get("vel_its_last")) pres_its_last.append(ctx.get("pres_its_last")) @@ -1004,9 +1210,26 @@ def rotated_residual(uvec, keep_cartesian=False): f"this number).") Fc.destroy() # residual output buffer (reaction persists in the result dict) - _destroy_rotated_ksp_ctx(ctx) # KSP/PC + the owned Schur pmat - if Ahat is not None: - Ahat.destroy() # the reused rotated operator + if cache_allowed and ctx is not None and Ahat is not None and not use_lu: + # Persist the workspace for the next solve. The stored key describes + # the operator values Ahat now holds. Q/Qt are SHARED with the result + # dict below — one Python wrapper each, so teardown and GC compose. + solver._rotated_linear_cache = { + "geometry_sig": geometry_sig, + "Q": Q, "Qt": Qt, "normal_rows": normal_rows, + "custom_Pl": custom_Pl, + "Ahat": Ahat, "diag_scale": diag_scale, "ctx": ctx, + "coefficient_variables": coefficient_variables, + "coefficient_states": coefficient_states, + } + else: + if cache_allowed: + # nothing solvable was built (e.g. an already-converged warm start + # on a fresh solver) — leave no partial cache behind. + solver._rotated_linear_cache = None + _destroy_rotated_ksp_ctx(ctx) # KSP/PC + the owned Schur pmat + if Ahat is not None: + Ahat.destroy() # the reused rotated operator if Atot is not None: Atot.destroy() # rotated operator + interface term if xhat is not None: @@ -1024,7 +1247,9 @@ def rotated_residual(uvec, keep_cartesian=False): "velocity_pc": "direct-LU" if use_lu else (ctx or {}).get("velocity_pc"), "schur_pre": "none" if use_lu else (ctx or {}).get("schur_pre"), "velocity_pc_type": None if use_lu else (ctx or {}).get("velocity_pc_type"), - "continuation_switched": continuation and phase == "newton"} + "continuation_switched": continuation and phase == "newton", + "rotation_reused": rotation_reused, + "workspace_reused": workspace_reused} def _build_rotated_custom_Pl(solver, Q, normal_rows): @@ -1095,17 +1320,26 @@ def _velocity_diag_scale(Ahat, solver): def _destroy_rotated_ksp_ctx(ctx): - """Release the reusable rotated-KSP context (KSP and the owned Schur pmat).""" + """Release a reusable rotated-KSP context and its owned PETSc objects.""" if ctx is None: return if ctx.get("ksp") is not None: ctx["ksp"].destroy() + ctx["ksp"] = None if ctx.get("Mp") is not None: ctx["Mp"].destroy() + ctx["Mp"] = None + if ctx.get("nsp") is not None: + ctx["nsp"].destroy() + ctx["nsp"] = None + if ctx.get("cns") is not None: + ctx["cns"].destroy() + ctx["cns"] = None def _solve_rotated_iterative(solver, Ahat, bhat, Q, Qt, normal_rows, verbose=False, - custom_Pl=None, nsp=None, Mp=None, ctx=None): + custom_Pl=None, nsp=None, Mp=None, ctx=None, + refresh_operator=True): """Solve the rotated saddle with a SELF-CONTAINED fieldsplit-Schur KSP on the rotated operator. The velocity block is geometric FMG on the CUSTOM prolongation (PR#290, rotated) whenever the solver has a hierarchy — ``set_custom_fmg`` or a @@ -1138,7 +1372,11 @@ def _solve_rotated_iterative(solver, Ahat, bhat, Q, Qt, normal_rows, verbose=Fal across Newton iterations — the fieldsplit ISs, Schur USER pmat and FMG prolongations survive; only the operator-values refresh is paid. The caller must keep ``Ahat``/``Mp`` the SAME Mat objects (values updated in place) and - release the context with ``_destroy_rotated_ksp_ctx`` when done.""" + release the context with ``_destroy_rotated_ksp_ctx`` when done. + + ``refresh_operator=False`` is the cross-solve fast path for an exactly + unchanged matrix: the KSP/PC/GAMG hierarchy is left untouched and only the + right-hand side changes.""" from underworld3.utilities import custom_mg dm = solver.dm vel_is = solver._subdict["velocity"][0] @@ -1320,7 +1558,10 @@ def _solve_rotated_iterative(solver, Ahat, bhat, Q, Qt, normal_rows, verbose=Fal nsp = ctx["nsp"] # Same Mat objects, new values (ptap-with-result / createSubMatrix-with- # submat) — poke the KSP so PCSetUp refreshes on the changed operator. - ksp.setOperators(Ahat) + # refresh_operator=False is the cross-solve fast path for a provably + # unchanged matrix: leave the KSP/PC/GAMG setup untouched entirely. + if refresh_operator: + ksp.setOperators(Ahat) if nsp is not None: nsp.remove(bhat) # project EVERY rhs diff --git a/tests/test_1018_rotated_freeslip.py b/tests/test_1018_rotated_freeslip.py index 96b82241..d702208c 100644 --- a/tests/test_1018_rotated_freeslip.py +++ b/tests/test_1018_rotated_freeslip.py @@ -67,6 +67,104 @@ def test_rotated_freeslip_box_reproduces_essential(): assert sol.velocity_error(v) < 1e-3 +def test_rotated_linear_workspace_reuses_unchanged_operator(): + """Repeated linear solves reuse the rotated workspace across solves. + + RHS-only changes ride the iteration-0 fast path (no Jacobian assembly, no + ptap, no PCSetUp — ``workspace_reused``); an operator-coefficient FIELD + change is detected by the state-counter key and refreshes the operator + values IN PLACE on the same objects (same Q/Ahat/KSP handles — the + structure tier); an explicit ``time=`` solve vetoes the fast path.""" + mesh = uw.meshing.StructuredQuadBox( + elementRes=(8, 8), minCoords=(0, 0), maxCoords=(1, 1), qdegree=3 + ) + temperature = uw.discretisation.MeshVariable( + "Tcache", mesh, 1, degree=1, continuous=True + ) + viscosity = uw.discretisation.MeshVariable( + "Etacache", mesh, 1, degree=1, continuous=True + ) + velocity = uw.discretisation.MeshVariable( + "Vcache", mesh, mesh.dim, degree=2, continuous=True + ) + pressure = uw.discretisation.MeshVariable( + "Pcache", mesh, 1, degree=1, continuous=False + ) + temperature.data[:, 0] = 1.0 + temperature.coords[:, 0] + viscosity.data[:, 0] = 1.0 + + stokes = uw.systems.Stokes( + mesh, velocityField=velocity, pressureField=pressure + ) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = viscosity.sym[0] + stokes.bodyforce = sympy.Matrix([[0.0, -temperature.sym[0]]]) + for wall in ("Top", "Bottom", "Left", "Right"): + stokes.add_rotated_freeslip_bc(0, wall) + stokes.petsc_use_pressure_nullspace = True + stokes.petsc_options["snes_type"] = "ksponly" + stokes.tolerance = 1.0e-8 + + stokes.solve() + velocity_1 = velocity.data.copy() + cache_1 = stokes._rotated_linear_cache + assert cache_1 is not None, "no workspace cached after a linear rotated solve" + handles_1 = ( + cache_1["Q"].handle, + cache_1["Ahat"].handle, + cache_1["ctx"]["ksp"].handle, + ) + assert not stokes._rotated_freeslip_info["workspace_reused"] + assert not stokes._rotated_freeslip_info["rotation_reused"] + + # RHS-only change: fast path (same handles, no reassembly), and the + # linear-in-forcing solution exactly doubles. + temperature.data[:, 0] *= 2.0 + stokes.solve(zero_init_guess=False) + + cache_2 = stokes._rotated_linear_cache + handles_2 = ( + cache_2["Q"].handle, + cache_2["Ahat"].handle, + cache_2["ctx"]["ksp"].handle, + ) + assert handles_2 == handles_1 + assert stokes._rotated_freeslip_info["workspace_reused"] + assert stokes._rotated_freeslip_info["rotation_reused"] + relative_scaling_error = ( + np.linalg.norm(velocity.data - 2.0 * velocity_1) + / np.linalg.norm(2.0 * velocity_1) + ) + assert relative_scaling_error < 1.0e-6 + + # Operator-coefficient FIELD change: the state-counter key catches it, the + # operator values are refreshed in place (same objects, no fast path). + velocity_2 = velocity.data.copy() + viscosity.data[:, 0] *= 2.0 + stokes.solve(zero_init_guess=False) + + assert not stokes._rotated_freeslip_info["workspace_reused"] + assert stokes._rotated_freeslip_info["rotation_reused"] + assert stokes._rotated_linear_cache["ctx"]["ksp"].handle == handles_1[2] + viscosity_scaling_error = ( + np.linalg.norm(velocity.data - 0.5 * velocity_2) + / np.linalg.norm(0.5 * velocity_2) + ) + assert viscosity_scaling_error < 1.0e-6 + + # An explicit time= solve must veto the fast path (petsc_t reaches the + # kernels outside any state counter), while the workspace objects persist. + refreshed_velocity = velocity.data.copy() + stokes.solve(zero_init_guess=False, time=0.5) + assert not stokes._rotated_freeslip_info["workspace_reused"] + assert stokes._rotated_linear_cache["ctx"]["ksp"].handle == handles_1[2] + time_refresh_error = ( + np.linalg.norm(velocity.data - refreshed_velocity) + / np.linalg.norm(refreshed_velocity) + ) + assert time_refresh_error < 1.0e-6 + + @pytest.mark.level_2 def test_rotated_freeslip_spherical_shell_3d(): """3D spherical shell, free-slip inner+outer (the Zhong #248 configuration): From 783382113813688f33d8c196cc2c1018de87dfde Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 13 Aug 2026 11:37:01 +1000 Subject: [PATCH 2/4] Close the rotated-workspace blind spot: rampable constants join the key, the verdict goes collective MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The PR #418 review's unresolved finding: the reuse key was built from MeshVariable._state counters, which are BLIND to rampable UWexpression constants — the #416 contract lets a constant change value with no state bump, so a 2x viscosity ramp between solves reported "unchanged" and (on the original one-shot path) returned a bit-identical stale solution while the matrix-probe safety net was disabled by the very verdict it was meant to check. The reworked loop already made a stale verdict non-fatal (the true residual is measured every iterate and iteration 1 onward always reassembles), but the verdict itself must still be honest: - the operator key now includes the packed constants[] values the kernels will actually assemble with, plus the JIT bundle key (covers an in-place kernel rewire). Measured on the ramp probe: the naive key reported workspace_reused=True for a 2x constant ramp; this key reports False and reassembles. Over-invalidation on RHS-only constant changes is accepted — reassembly is the safe default; - if the constants manifest or the coefficient enumeration cannot be read, the fast path is forfeited outright — correctness first; - the fast path additionally requires a self-measured linear hint (last solve converged in <= 1 increment): for a nonlinear model the cached operator is the previous solve's tangent, and the skip would only trade an assembly for a wasted increment; - the match verdict is allgathered and must be unanimous before it gates any collective PETSc call — state counters follow rank-local writes, and a rank-divergent verdict is a deadlock; - on a detected change the stored signature is poisoned before the in-place refresh, so an exception mid-refresh cannot leave a stale key that later matches half-updated values. Regression: test_rotated_workspace_constant_ramp_invalidates (fail-before validated on the naive key: the reused flag lies there) with its own armed- fast-path negative control and a fresh-solver control at the ramped viscosity; test_rotated_workspace_deform_invalidates re-proves the mesh.deform teardown on the reworked cache. Underworld development team with AI support from Claude Code --- .../design/ROTATED_FREESLIP_LINEAR_REUSE.md | 34 +++++- src/underworld3/utilities/rotated_bc.py | 69 ++++++++++- tests/test_1018_rotated_freeslip.py | 107 ++++++++++++++++++ 3 files changed, 199 insertions(+), 11 deletions(-) diff --git a/docs/developer/design/ROTATED_FREESLIP_LINEAR_REUSE.md b/docs/developer/design/ROTATED_FREESLIP_LINEAR_REUSE.md index 3935a116..964d95c8 100644 --- a/docs/developer/design/ROTATED_FREESLIP_LINEAR_REUSE.md +++ b/docs/developer/design/ROTATED_FREESLIP_LINEAR_REUSE.md @@ -47,18 +47,40 @@ A wrong fast-path decision cannot produce a wrong answer. The loop measures the **true residual** (fresh kernels, current constants) at every iterate, convergence is declared only on that measurement, and every iteration after the first always reassembles the operator. A stale cached operator therefore -costs one extra Newton increment; it can never return a stale solution. +costs one extra Newton increment; it can never return a stale solution. This +is deliberately independent of the invalidation key — a safety net gated by +the trigger it guards is not a safety net, which was the unresolved finding +of the original PR #418 review (the state-counter key was blind to rampable +constants, and the same "unchanged" verdict disabled the matrix probe that +was supposed to catch it, returning a bit-identical stale solution flagged +as legitimately reused). ## Invalidation -- Operator coefficient **mesh variables** are collected from the constitutive +The iteration-0 fast path requires ALL of: + +- Operator coefficient **mesh variables** — collected from the constitutive parameters, constraint term, penalty and saddle preconditioner (expressions unwrapped first, unknowns excluded). Their base `MeshVariable._state` - counters must match for the fast path. -- If the coefficient enumeration fails, the fast path is forfeited and every - solve reassembles (structure-tier reuse only) — correctness first. + counters must match. +- The **constants signature** — the packed `constants[]` values the compiled + kernels will assemble with, plus the JIT bundle key. This is what sees a + rampable UWexpression constant (the #416 contract: a value change bumps no + state counter — the original PR's blind spot) and an in-place kernel + rewire. It deliberately over-invalidates on RHS-only constant changes: + reassembly is the safe default. +- A **linear hint** — the previous solve on this workspace converged in at + most one increment. A nonlinear model's cached operator is last solve's + tangent, not this iterate's; the hint is self-measured by the loop, so no + up-front nonlinearity probe is paid. +- If the coefficient enumeration or the constants manifest cannot be read, + the fast path is forfeited and every solve reassembles (structure-tier + reuse only) — correctness first. - An explicit `solve(time=...)` vetoes the fast path: `petsc_t` reaches the - kernels through the DM, invisible to any counter. + kernels through the DM, invisible to any counter or constant value. +- The match verdict is made **collective** (allgather + unanimity) before it + gates any collective PETSc call — state counters follow rank-local writes, + and a rank-divergent verdict would be a deadlock, not a wrong answer. - Mesh deform, field-layout, boundary-condition or forced-setup changes route through the solver's full-rebuild teardown, which destroys the whole workspace (`_reset_rotated_solver_cache`, called from `_reset()` and from diff --git a/src/underworld3/utilities/rotated_bc.py b/src/underworld3/utilities/rotated_bc.py index 925aef68..54dc00e1 100644 --- a/src/underworld3/utilities/rotated_bc.py +++ b/src/underworld3/utilities/rotated_bc.py @@ -494,6 +494,35 @@ def _coefficient_states(variables): return tuple(variable._state for variable in variables) +def _operator_constants_signature(solver): + """Signature of everything that reaches the compiled kernels OUTSIDE + mesh-variable data: the packed ``constants[]`` values and the JIT bundle + key. + + This closes the blind spot the original PR #418 cache had: a rampable + UWexpression constant changes VALUE with no ``_state`` bump anywhere (the + #416 live-constants contract — every δ-continuation and viscosity-ramp + driver relies on it), so a key built from MeshVariable state counters + alone reported "unchanged" for a 2x viscosity ramp. The packed constants + ARE the values the kernels will assemble with, so comparing them closes + the gap exactly — at the cost of also invalidating on RHS-only constant + changes, which errs toward reassembly (robust generality). The JIT key + covers a function rewire in place (new kernels on the same DM), where the + manifest itself is replaced. If the manifest cannot be read, return None: + the caller must forfeit the fast path.""" + from underworld3.utilities._jitextension import _pack_constants + try: + manifest = getattr(solver, "constants_manifest", None) + values = tuple(float(v) for v in _pack_constants(manifest)) \ + if manifest else () + except Exception: + return None + jit_key = getattr(solver, "_current_jit_cache_key", None) + if jit_key is None: + jit_key = getattr(solver, "_last_jit_cache_key", None) + return (str(jit_key), values) + + def _rotated_geometry_signature(solver, boundaries): """Signature of everything the GEOMETRY tier of the workspace was built from: the boundary specs (name + normal), the fault-contact registration @@ -832,6 +861,7 @@ def solve_rotated_freeslip(solver, boundaries, remove_rotation_gauge=True, if cache_allowed else None coefficient_variables = None coefficient_states = None + operator_sig = None if cache_allowed: if cache is not None and cache.get("geometry_sig") == geometry_sig \ and cache.get("coefficient_variables") is not None: @@ -840,6 +870,10 @@ def solve_rotated_freeslip(solver, boundaries, remove_rotation_gauge=True, coefficient_variables = _operator_coefficient_variables(solver) if coefficient_variables is not None: coefficient_states = _coefficient_states(coefficient_variables) + # The state counters alone are BLIND to rampable UWexpression + # constants (#416: value changes bump nothing) — the constants + # signature is the other half of the key. + operator_sig = _operator_constants_signature(solver) op_ok = False if cache is not None: @@ -847,11 +881,25 @@ def solve_rotated_freeslip(solver, boundaries, remove_rotation_gauge=True, op_ok = (geom_ok and not force_operator_refresh and coefficient_states is not None - and cache.get("coefficient_states") == coefficient_states) + and operator_sig is not None + and cache.get("coefficient_states") == coefficient_states + and cache.get("operator_sig") == operator_sig) + # COLLECTIVE verdict: state counters are bumped by rank-local data + # writes, so a rank-divergent match here would desync the collective + # assembly/PCSetUp calls the verdict gates (the np>1 deadlock class). + if mpi.size > 1: + verdicts = mpi.comm.allgather((bool(geom_ok), bool(op_ok))) + geom_ok = all(v[0] for v in verdicts) + op_ok = all(v[1] for v in verdicts) if not geom_ok: _destroy_rotated_linear_cache(cache) solver._rotated_linear_cache = None cache = None + elif not op_ok: + # The operator values are about to be refreshed in place; poison + # the stored key NOW so an exception mid-refresh cannot leave a + # stale signature that later matches half-updated values. + cache["operator_sig"] = None if cache is not None: Q, Qt = cache["Q"], cache["Qt"] @@ -865,8 +913,14 @@ def solve_rotated_freeslip(solver, boundaries, remove_rotation_gauge=True, rotation_reused = False # The iteration-0 fast path: skip Jacobian assembly / ptap / PCSetUp when - # the operator key proves nothing operator-relevant changed. - reuse_operator = bool(cache is not None and op_ok) + # the operator key proves nothing operator-relevant changed AND the + # previous solve behaved linearly (converged in <= 1 increment). For a + # genuinely nonlinear model the cached operator is the LAST solve's + # converged tangent, not this iterate's, so the skip would only trade one + # assembly for one wasted increment — the linear hint is self-measured, + # no up-front nonlinearity probe. + reuse_operator = bool(cache is not None and op_ok + and cache.get("linear_hint", False)) workspace_reused = False # Interface constitutive laws (fault_contact.add_viscous_fault_bc / @@ -1212,8 +1266,11 @@ def rotated_residual(uvec, keep_cartesian=False): Fc.destroy() # residual output buffer (reaction persists in the result dict) if cache_allowed and ctx is not None and Ahat is not None and not use_lu: # Persist the workspace for the next solve. The stored key describes - # the operator values Ahat now holds. Q/Qt are SHARED with the result - # dict below — one Python wrapper each, so teardown and GC compose. + # the operator values Ahat now holds; linear_hint records whether + # this solve behaved linearly (<= 1 increment), which is what + # licenses the iteration-0 fast path next time. Q/Qt are SHARED with + # the result dict below — one Python wrapper each, so teardown and + # GC compose. solver._rotated_linear_cache = { "geometry_sig": geometry_sig, "Q": Q, "Qt": Qt, "normal_rows": normal_rows, @@ -1221,6 +1278,8 @@ def rotated_residual(uvec, keep_cartesian=False): "Ahat": Ahat, "diag_scale": diag_scale, "ctx": ctx, "coefficient_variables": coefficient_variables, "coefficient_states": coefficient_states, + "operator_sig": operator_sig, + "linear_hint": bool(converged and newton_its <= 1), } else: if cache_allowed: diff --git a/tests/test_1018_rotated_freeslip.py b/tests/test_1018_rotated_freeslip.py index d702208c..c7fba4b7 100644 --- a/tests/test_1018_rotated_freeslip.py +++ b/tests/test_1018_rotated_freeslip.py @@ -165,6 +165,113 @@ def test_rotated_linear_workspace_reuses_unchanged_operator(): assert time_refresh_error < 1.0e-6 +def _rampable_rotated_stokes(mesh, k_expr, tag, forcing=None): + """Rotated free-slip Stokes with viscosity given by a rampable + UWexpression constant (the #416 idiom used by every continuation + driver). ``forcing`` optionally supplies a scalar field for the body + force (an RHS-only mesh variable); default is a fixed analytic load.""" + v = uw.discretisation.MeshVariable(f"v{tag}", mesh, mesh.dim, degree=2) + p = uw.discretisation.MeshVariable(f"p{tag}", mesh, 1, degree=1, + continuous=False) + s = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + s.constitutive_model = uw.constitutive_models.ViscousFlowModel + s.constitutive_model.Parameters.shear_viscosity_0 = k_expr + x, y = mesh.X + load = forcing.sym[0] if forcing is not None \ + else sympy.sin(sympy.pi * x) * sympy.cos(sympy.pi * y) + s.bodyforce = sympy.Matrix([[0.0, load]]) + s.penalty = 0.0 + s.tolerance = 1e-9 + for wall in ("Top", "Bottom", "Left", "Right"): + s.add_rotated_freeslip_bc(0, wall) + s.petsc_use_pressure_nullspace = True + return s, v + + +def test_rotated_workspace_constant_ramp_invalidates(): + """THE blind-spot regression (PR #418 review, unresolved finding): a + rampable UWexpression constant changes value with NO state-counter bump + (#416 contract). The workspace key must see it through the packed + constants[] signature: the fast path must NOT fire, and the second + solution must match a fresh-solver control at the ramped viscosity. + + Fail-before validated: on the state-counter-only key (the ported + original), solve 2 reports workspace_reused=True — the flag lies about a + stale operator (the answer is still rescued by the loop's structural + safety net, unlike the original one-shot path which returned a + bit-identical stale solution).""" + mesh = uw.meshing.StructuredQuadBox( + elementRes=(12, 12), minCoords=(0, 0), maxCoords=(1, 1), qdegree=3) + T = uw.discretisation.MeshVariable("Trmp", mesh, 1, degree=1) + xy = T.coords + T.data[:, 0] = np.sin(np.pi * xy[:, 0]) * np.cos(np.pi * xy[:, 1]) + + k = uw.function.expression(r"k_\eta", 1.0, "rampable viscosity") + s1, v1 = _rampable_rotated_stokes(mesh, k, "Rmp", forcing=T) + s1.solve() + + # negative control for the test itself: an RHS-only field change must + # ride the fast path, proving the skip is ARMED before we assert the + # ramp defeats it. + T.data[:, 0] *= 2.0 + s1.solve(zero_init_guess=False) + assert s1._rotated_freeslip_info["workspace_reused"], ( + "fast path did not fire on an RHS-only change — the ramp assertion " + "below would pass vacuously") + u1 = v1.data.copy() + + # THE RAMP: value change only — no .sym rebuild, no state bump anywhere + k.sym = 2.0 + s1.solve(zero_init_guess=False) + info = s1._rotated_freeslip_info + assert not info["workspace_reused"], ( + "constant ramp rode the fast path — the operator key is blind to " + "rampable constants again (#416 / PR #418 review finding)") + assert info["rotation_reused"], "geometry tier should survive a ramp" + u2 = v1.data.copy() + + # fresh-solver control at the ramped viscosity, same forcing + k_c = uw.function.expression(r"k_c", 2.0, "control viscosity") + s_c, v_c = _rampable_rotated_stokes(mesh, k_c, "Ctl", forcing=T) + s_c.solve() + err = np.linalg.norm(u2 - v_c.data) / np.linalg.norm(v_c.data) + assert err < 1e-6, f"ramped solve differs from fresh control by {err:.2e}" + # and the linear model's exact halved-velocity scaling + half = np.linalg.norm(u2 - 0.5 * u1) / np.linalg.norm(0.5 * u1) + assert half < 1e-6, f"ramped solve is not the halved velocity ({half:.2e})" + + +def test_rotated_workspace_deform_invalidates(): + """mesh.deform between solves: geometry changed, so the whole workspace + must be rebuilt (rotation_reused False) and the answer must match a fresh + solver on the deformed mesh.""" + mesh = uw.meshing.StructuredQuadBox( + elementRes=(10, 10), minCoords=(0, 0), maxCoords=(1, 1), qdegree=3) + + k = uw.function.expression(r"k_d", 1.0, "viscosity") + s1, v1 = _rampable_rotated_stokes(mesh, k, "Dfm") + s1.solve() + assert s1._rotated_linear_cache is not None + + # bump the top boundary (the free-surface pattern) + coords = mesh.X.coords.copy() + coords[:, 1] += 0.02 * coords[:, 1] * np.sin(np.pi * coords[:, 0]) + mesh.deform(coords) + + s1.solve() + info = s1._rotated_freeslip_info + assert not info["rotation_reused"], ( + "workspace survived a mesh.deform — stale rotation Q in use") + assert not info["workspace_reused"] + + k_c = uw.function.expression(r"k_dc", 1.0, "control viscosity") + s_c, v_c = _rampable_rotated_stokes(mesh, k_c, "DfC") + s_c.solve() + err = np.linalg.norm(v1.data - v_c.data) / np.linalg.norm(v_c.data) + assert err < 1e-6, ( + f"post-deform solve differs from fresh control by {err:.2e}") + + @pytest.mark.level_2 def test_rotated_freeslip_spherical_shell_3d(): """3D spherical shell, free-slip inner+outer (the Zhong #248 configuration): From f8c2b726f58cb91ef6cabfebe39578b8f0f8bee5 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 13 Aug 2026 11:38:16 +1000 Subject: [PATCH 3/4] Fault contact composes with the rotated workspace cache: pair blocks cache, interface laws opt out MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit solve_with_fault drives the same rotated Newton loop, so the cross-solve workspace decision had to be made explicitly for the fault machinery: - FRICTIONLESS pair blocks are geometry (coincident-node pairing + fault normals live in Q, keyed by the fault registration in the geometry signature) — they cache. A warm repeat reuses the rotation; a cold re-solve rides the iteration-0 fast path; both match a fresh-solver control on the same mesh. - INTERFACE-LAW solvers (viscous / Coulomb / rate-state) opt out entirely: the interface tangent is reassembled per iterate at the current slip rates and the reaction-fed normal stress is Picard-lagged solver state — neither is keyable registration state, so cache_allowed excludes them at the top of solve_rotated_freeslip. Regression: test_fault_repeat_solve_composes_with_workspace_reuse covers both arms, with a fresh-solver control for the cached arm and the absent- cache assertion for the opt-out. Underworld development team with AI support from Claude Code --- tests/test_0846_fault_contact.py | 75 ++++++++++++++++++++++++++++++++ 1 file changed, 75 insertions(+) diff --git a/tests/test_0846_fault_contact.py b/tests/test_0846_fault_contact.py index db0de6ef..5d9c560a 100644 --- a/tests/test_0846_fault_contact.py +++ b/tests/test_0846_fault_contact.py @@ -116,6 +116,81 @@ def test_frictionless_fault_slips_like_a_crack(): f"peak slip {Vmag.max():.4f} vs crack value {HALF:.4f}") +def test_fault_repeat_solve_composes_with_workspace_reuse(): + """The cross-solve rotated workspace (issue #417) composes with the + split-node fault machinery, which enters the SAME rotated Newton loop. + + Frictionless pair blocks are GEOMETRY (coincident-node pairing + fault + normals), so they cache: a repeat solve reuses the rotation, a cold + re-solve rides the iteration-0 fast path, and both must match a fresh + solver on the same mesh. Interface-LAW solvers must opt out entirely + (the interface tangent and the Picard-lagged normal stress are + solution-dependent) — asserted via the absent cache.""" + mesh = _split_box() + x, y = mesh.X + + def build(tag): + v = uw.discretisation.MeshVariable(f"vRe{tag}", mesh, 2, degree=2) + p = uw.discretisation.MeshVariable(f"pRe{tag}", mesh, 1, degree=0, + continuous=False) + 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.tolerance = 1e-8 + stokes.petsc_use_pressure_nullspace = True + for side in ("Top", "Bottom", "Left", "Right"): + stokes.add_dirichlet_bc((y - 0.5, 0.0), side) + return stokes, v + + stokes, v = build("A") + add_frictionless_fault_bc(stokes, "Flt") + info1 = solve_with_fault(stokes) + assert info1["converged"] + assert not info1["rotation_reused"] + assert stokes._rotated_linear_cache is not None, ( + "frictionless fault solve did not cache — pair blocks are geometry " + "and should reuse") + v_first = v.data.copy() + + # warm repeat, nothing changed: geometry tier reused, answer unchanged + info2 = solve_with_fault(stokes, zero_init_guess=False) + assert info2["converged"] and info2["rotation_reused"] + drift = np.linalg.norm(v.data - v_first) / np.linalg.norm(v_first) + assert drift < 1e-8, f"repeat fault solve drifted by {drift:.2e}" + + # cold re-solve: a genuine increment on the cached operator (fast path) + info3 = solve_with_fault(stokes, zero_init_guess=True) + assert info3["converged"] and info3["rotation_reused"] + assert info3["workspace_reused"], ( + "cold fault re-solve did not ride the iteration-0 fast path") + err = np.linalg.norm(v.data - v_first) / np.linalg.norm(v_first) + assert err < 1e-6, f"fast-path fault solve differs by {err:.2e}" + + # fresh-solver control on the same mesh + control, vc = build("B") + add_frictionless_fault_bc(control, "Flt") + infoc = solve_with_fault(control) + assert infoc["converged"] + ctrl = np.linalg.norm(v.data - vc.data) / np.linalg.norm(vc.data) + assert ctrl < 1e-6, f"repeat fault solve differs from fresh control by {ctrl:.2e}" + + # interface LAWS opt out: solution-dependent tangent + Picard-lagged + # normal stress must never ride a cross-solve cache + lawful, vl = build("C") + add_viscous_fault_bc(lawful, 1.0 / HALF, "Flt") + il1 = solve_with_fault(lawful) + assert il1["converged"] + assert lawful._rotated_linear_cache is None, ( + "interface-law solve left a workspace cache — the opt-out regressed") + vl_first = vl.data.copy() + il2 = solve_with_fault(lawful, zero_init_guess=False) + assert il2["converged"] + assert not il2["rotation_reused"] and not il2["workspace_reused"] + assert lawful._rotated_linear_cache is None + ldrift = np.linalg.norm(vl.data - vl_first) / np.linalg.norm(vl_first) + assert ldrift < 1e-6, f"repeat interface-law solve drifted by {ldrift:.2e}" + + def test_viscous_fault_bridges_welded_to_free(): """The linear interface law spans its two exact limits monotonically. From 2f17743d27a8c327567e74723c74f357873213db Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 13 Aug 2026 15:23:49 +1000 Subject: [PATCH 4/4] Review response: persist the operator key only for values Ahat holds; null before destroy Two minors from the #543 review. (m2) The persisted operator_sig now requires that this solve either reassembled Ahat or rode a key-matched fast path - a poisoned solve exiting at iteration 0 without assembling can no longer store a fresh key against unrefreshed values (the window was bounded by the always-reassemble net; now it is closed). (m4) Both teardown sites null the cache attribute before destroying, so an exception mid-destroy leaves objects unreachable rather than arming a double-destroy. Underworld development team with AI support from Claude Code --- .../cython/petsc_generic_snes_solvers.pyx | 5 ++++- src/underworld3/utilities/rotated_bc.py | 16 +++++++++++++--- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index ba5250e5..20d7f309 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -1301,10 +1301,13 @@ class SolverBaseClass(uw_object): its reaction vector is independent of the cache and the σ_nn / dynamic-topography recoveries may still need it.""" cache = getattr(self, "_rotated_linear_cache", None) + # Null BEFORE destroying: an exception mid-destroy must leave objects + # unreachable (leaked-but-safe), never a half-destroyed cache a later + # reset would double-destroy (#543 review, m4). + self._rotated_linear_cache = None if cache is not None: from underworld3.utilities.rotated_bc import _destroy_rotated_linear_cache _destroy_rotated_linear_cache(cache) - self._rotated_linear_cache = None def _reset(self): diff --git a/src/underworld3/utilities/rotated_bc.py b/src/underworld3/utilities/rotated_bc.py index 54dc00e1..61ae92a3 100644 --- a/src/underworld3/utilities/rotated_bc.py +++ b/src/underworld3/utilities/rotated_bc.py @@ -892,9 +892,11 @@ def solve_rotated_freeslip(solver, boundaries, remove_rotation_gauge=True, geom_ok = all(v[0] for v in verdicts) op_ok = all(v[1] for v in verdicts) if not geom_ok: - _destroy_rotated_linear_cache(cache) + # Null before destroying — same double-destroy discipline as + # _reset_rotated_solver_cache (#543 review, m4). solver._rotated_linear_cache = None - cache = None + stale, cache = cache, None + _destroy_rotated_linear_cache(stale) elif not op_ok: # The operator values are about to be refreshed in place; poison # the stored key NOW so an exception mid-refresh cannot leave a @@ -1063,6 +1065,7 @@ def rotated_residual(uvec, keep_cartesian=False): ref = None last_reason = 0 iters = 0 + did_assemble = False converged = False phase = "picard" if continuation else "newton" for iters in range(max_it): @@ -1097,6 +1100,7 @@ def rotated_residual(uvec, keep_cartesian=False): # a stale solution. assemble = not (reuse_operator and iters == 0) if assemble: + did_assemble = True snes.computeJacobian(u, J, Jp) # Jp carries the 1/mu mass (Schur pmat) if Ahat is None: Ahat = J.ptap(Qt) @@ -1278,7 +1282,13 @@ def rotated_residual(uvec, keep_cartesian=False): "Ahat": Ahat, "diag_scale": diag_scale, "ctx": ctx, "coefficient_variables": coefficient_variables, "coefficient_states": coefficient_states, - "operator_sig": operator_sig, + # The stored key must describe what Ahat HOLDS: valid if this + # solve reassembled (values current), or if the fast path rode a + # key-matched operator untouched. A poisoned solve that exited at + # iteration 0 without assembling must not persist a fresh key + # against unrefreshed values (#543 review, m2). + "operator_sig": (operator_sig if (did_assemble or reuse_operator) + else None), "linear_hint": bool(converged and newton_its <= 1), } else: