Reuse linear rotated free-slip solver workspace

docs/developer/design/ROTATED_FREESLIP_LINEAR_REUSE.md

@@ -0,0 +1,108 @@
+# Rotated Free-Slip Linear Solver Reuse
+
+## Scope
+
+This note documents the repeated-solve workspace introduced for linear Stokes
+problems using rotated strong free slip. It addresses
+[issue #417](https://github.com/underworldcode/underworld3/issues/417).
+
+The change does not alter the rotated boundary condition or its discrete
+equations. It changes PETSc object ownership and determines when the assembled
+operator and its preconditioner can be reused.
+
+## Previous Lifecycle
+
+Every linear rotated solve previously rebuilt:
+
+1. the nodal rotation matrices `Q` and `Q.T`;
+2. the transformed saddle operator `Ahat = Q A Q.T`;
+3. the pressure mass block;
+4. the coupled and pressure nullspaces;
+5. the fieldsplit Schur KSP/PC;
+6. the velocity GAMG hierarchy.
+
+The linear/nonlinear detector also assembled and copied two trial Jacobians on
+every call. Long transient runs therefore accumulated a large PETSc and system
+allocator high-water mark even when only the body-force field changed.
+
+## Reused Workspace
+
+The Stokes solver now owns one `_rotated_linear_cache` containing:
+
+- `Q`, `Q.T`, and constrained normal rows;
+- the transformed matrix `Ahat`;
+- two deterministic operator-probe vector triples;
+- operator coefficient mesh variables and their UW data-version counters;
+- the pressure mass block, nullspaces, fieldsplit KSP/PC, and GAMG hierarchy;
+- the constraint-row diagonal scale.
+
+For an unchanged linear operator, a subsequent solve:
+
+1. assembles only `F(0)` and forms the new right-hand side;
+2. rotates and constrains that right-hand side;
+3. solves with the existing KSP/PC hierarchy;
+4. rotates the solution back;
+5. forms and retains the Cartesian reaction needed by topography recovery.
+
+The prior per-solve solution and reaction vectors are explicitly destroyed
+before replacement.
+
+## Invalidation
+
+The cache is destroyed before a solver or DM rebuild. The cached
+linear/nonlinear classification is invalidated at the same time.
+
+Mutable operator coefficients are collected from the constitutive parameters,
+constraint, penalty, and saddle preconditioner. Their base
+`MeshVariable._state` counters decide whether Jacobian assembly is required.
+
+- A body-force field change does not refresh the operator.
+- A viscosity field change refreshes `J`, `Jp`, and the rotated operator.
+- An explicit `solve(time=...)` refreshes and compares the operator, covering
+  expressions that depend on `mesh.t`.
+- Mesh, field-layout, boundary-condition, constitutive, or forced setup changes
+  destroy the full workspace.
+
+After a coefficient-triggered assembly, two deterministic matrix-vector
+products decide whether matrix values actually changed. This avoids exact
+`Mat.equal`, which rejects harmless distributed-assembly roundoff, and avoids
+copying CSR arrays, which recreated the memory high-water problem.
+
+The high-level `uw.systems.Stokes.solve()` wrapper exposes and forwards
+`time=` for this purpose. It remains distinct from the viscoelastic integration
+`timestep=`.
+
+## PETSc Ownership
+
+`_destroy_rotated_ksp_ctx` explicitly destroys the KSP, pressure mass matrix,
+and owned nullspaces. `_destroy_rotated_linear_cache` additionally destroys the
+operator probes, transformed operator, and rotation matrices.
+
+Linear topography no longer requires retaining a copy of the native operator
+and right-hand side. The Cartesian reaction `J U - b` is computed immediately
+after the solve and stored in the solve result.
+
+## Validation
+
+The focused regression verifies that:
+
+- changing only a temperature/body-force field preserves the `Q`, `Ahat`, and
+  KSP handles;
+- the velocity scales with the changed right-hand side;
+- changing a viscosity mesh variable destroys the old KSP and refreshes the
+  operator;
+- the refreshed velocity has the expected inverse-viscosity scaling.
+
+Additional serial and eight-rank tests cover spherical velocity/leakage,
+reaction-derived topography, boundary traction, and dynamic-topography field
+recovery.
+
+Always launch MPI tests with the worktree MPI executable:
+
+```bash
+.pixi/envs/amr-dev/bin/mpirun -np 8 \
+  .pixi/envs/amr-dev/bin/python <script.py>
+```
+
+Using a system `mpirun` from a different Open MPI installation can stall during
+initialization and is not a solver failure.

src/underworld3/cython/petsc_generic_snes_solvers.pyx

@@ -898,6 +898,7 @@ class SolverBaseClass(uw_object):
 
     def _reset(self):
 
+        self._reset_rotated_solver_cache()
         self.natural_bcs = []
         self.essential_bcs = []
 
@@ -916,6 +917,21 @@ class SolverBaseClass(uw_object):
 
         return
 
+    def _reset_rotated_solver_cache(self):
+        """Release rotated-free-slip state before solver/DM teardown."""
+        result = getattr(self, "_rotated_freeslip_info", None)
+        cache = getattr(self, "_rotated_linear_cache", None)
+        if result is not None or cache is not None:
+            from underworld3.utilities.rotated_bc import (
+                _destroy_rotated_linear_cache,
+                _destroy_rotated_solve_result,
+            )
+            _destroy_rotated_solve_result(result)
+            _destroy_rotated_linear_cache(cache)
+        self._rotated_freeslip_info = None
+        self._rotated_linear_cache = None
+        self._rotated_residual_is_nonlinear = None
+
     def get_snes_diagnostics(self):
         """
         Extract comprehensive SNES convergence diagnostics with string representations.
@@ -1229,6 +1245,11 @@ class SolverBaseClass(uw_object):
         if self.is_setup:
             return
 
+        # Rotated matrices and KSP/PC objects reference the current DM/SNES
+        # layout. Any requested solver rebuild invalidates the whole workspace
+        # and the cached linear/nonlinear classification.
+        self._reset_rotated_solver_cache()
+
         # Resolve the preconditioner choice (auto/fmg/gamg) against the current
         # mesh hierarchy and push the option bundle before the per-class
         # _setup_solver runs snes.setFromOptions(). No-op unless the solver
@@ -5098,6 +5119,8 @@ class SNES_Stokes_SaddlePt(SolverBaseClass):
         # reaction = sigma_nn). Empty by default → the solve path is unchanged.
         self._rotated_freeslip_bcs = []
         self._rotated_freeslip_info = None
+        self._rotated_linear_cache = None
+        self._rotated_residual_is_nonlinear = None
         # Give the Lagrange-multiplier (lambda) block its own viscosity-scaled
         # Schur preconditioner. The constraint Schur complement S_lambda = C A^-1 C^T
         # scales as 1/mu (since A ~ mu K), exactly like the pressure Schur S_p ~ mu^-1 M_p
@@ -8066,9 +8089,12 @@ class SNES_Stokes_SaddlePt(SolverBaseClass):
             #    consistent_jacobian tangent (Picard / Newton / continuation).
             from underworld3.utilities.rotated_bc import (
                 solve_rotated_freeslip, solve_rotated_freeslip_nonlinear)
-            if not self._residual_is_nonlinear():
+            if self._rotated_residual_is_nonlinear is None:
+                self._rotated_residual_is_nonlinear = self._residual_is_nonlinear()
+            if not self._rotated_residual_is_nonlinear:
                 self._rotated_freeslip_info = solve_rotated_freeslip(
-                    self, self._rotated_freeslip_bcs, verbose=verbose)
+                    self, self._rotated_freeslip_bcs, verbose=verbose,
+                    force_operator_refresh=time is not None)
             else:
                 self._rotated_freeslip_info = solve_rotated_freeslip_nonlinear(
                     self, self._rotated_freeslip_bcs, verbose=verbose,

src/underworld3/systems/solvers.py

@@ -1411,6 +1411,7 @@ def solve(
         self,
         zero_init_guess: bool = True,
         timestep: float = None,
+        time=None,
         _force_setup: bool = False,
         verbose: bool = False,
         debug: bool = False,
@@ -1432,6 +1433,9 @@ def solve(
             If True, use zero initial guess. Otherwise use current field values.
         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
@@ -1513,6 +1517,7 @@ def solve(
                 _force_setup=_force_setup,
                 verbose=verbose,
                 picard=picard,
+                time=time,
                 divergence_retries=divergence_retries,
             )
 
@@ -1570,6 +1575,7 @@ def solve(
                 _force_setup=_force_setup,
                 verbose=verbose,
                 picard=picard,
+                time=time,
                 divergence_retries=divergence_retries,
             )