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..8978ef31 --- /dev/null +++ b/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 +``` + +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 362e647e..ffcad313 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/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, diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index f91f8ea1..32c3751b 100644 --- a/src/underworld3/systems/solvers.py +++ b/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, ) diff --git a/src/underworld3/utilities/rotated_bc.py b/src/underworld3/utilities/rotated_bc.py index 1bb9c9c9..381b5dc3 100644 --- a/src/underworld3/utilities/rotated_bc.py +++ b/src/underworld3/utilities/rotated_bc.py @@ -253,10 +253,111 @@ def _zero_rows_local(vec, normal_rows): vec.setArray(a) +def _build_matrix_probe(matrix): + """Build two deterministic matrix-vector products for reuse detection.""" + probes = [] + for phase in (0.6180339887498948, 1.4142135623730951): + x = matrix.createVecRight() + row_start, row_end = x.getOwnershipRange() + indices = np.arange(row_start, row_end, dtype=float) + 1.0 + values = x.getArray() + values[:] = np.sin(phase * indices) + np.cos( + (phase + 0.5) * indices + ) + y = matrix.createVecLeft() + work = matrix.createVecLeft() + matrix.mult(x, y) + probes.append((x, y, work)) + return probes + + +def _matrix_matches_probe(matrix, probes, rtol=5.0e-13): + """Test whether an assembled operator matches two cached products. + + Exact ``Mat.equal`` rejects harmless distributed-assembly roundoff, while + copying CSR values on every solve creates its own allocator high-water + growth. Two deterministic products provide a tight, allocation-free check. + """ + for x, reference, work in probes: + matrix.mult(x, work) + work.axpy(-1.0, reference) + relative_error = work.norm() / (reference.norm() + 1.0e-300) + if relative_error > rtol: + return False + return True + + +def _refresh_matrix_probe(matrix, probes): + """Replace cached matrix-vector products after an operator change.""" + for x, reference, _ in probes: + matrix.mult(x, reference) + + +def _destroy_matrix_probe(probes): + """Release vectors owned by an operator-reuse probe.""" + if probes is None: + return + for vectors in probes: + for vector in vectors: + vector.destroy() + + +def _operator_coefficient_variables(solver): + """Return mutable mesh variables that can change the linear Stokes operator.""" + from underworld3.function.expressions import mesh_vars_in_expression + + expressions = [solver.constraints, 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)): + try: + expressions.append(getattr(parameters, name)) + except (AttributeError, TypeError): + pass + + variables = set() + for expression in expressions: + expression = getattr(expression, "sym", expression) + if expression is None or not hasattr(expression, "args"): + continue + try: + _, regular, derivatives = mesh_vars_in_expression(expression) + except (AttributeError, TypeError): + continue + 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): + """Return UW data-version counters for operator coefficient variables.""" + return tuple(variable._state for variable in variables) + + # --------------------------------------------------------------------------- # # The rotated solve # --------------------------------------------------------------------------- # -def solve_rotated_freeslip(solver, boundaries, remove_rotation_gauge=True, verbose=False): +def solve_rotated_freeslip( + solver, + boundaries, + remove_rotation_gauge=True, + verbose=False, + force_operator_refresh=False, +): """Assemble + solve the rotated strong-free-slip Stokes saddle. Fills the solver's velocity/pressure fields with the (rotated-back, gauge-removed) solution. @@ -271,14 +372,16 @@ def solve_rotated_freeslip(solver, boundaries, remove_rotation_gauge=True, verbo ``dynamic_topography_field`` (as their ``solve_result`` argument), with keys: * ``"Q"``, ``"Qt"`` — the rotation and its transpose (PETSc Mats); - * ``"A"``, ``"b"`` — the unrotated operator and RHS, kept so the reaction - ``A·u − b`` can be reconstructed (linear path only); + * ``"reaction"`` — the Cartesian nodal reaction ``A·u − b``; * ``"U"``, ``"Uhat"`` — the Cartesian and rotated-frame solutions; * ``"normal_rows"`` — global rows of the constrained normal components; * ``"boundaries"`` — the boundary specs the rotation was built from; * ``"rotation_gauge_removed"`` — whether a rigid-rotation gauge was projected out; * ``"ksp_reason"``, ``"ksp_its"`` — outer KSP converged-reason and iteration count. """ + _destroy_rotated_solve_result(getattr(solver, "_rotated_freeslip_info", None)) + solver._rotated_freeslip_info = None + if getattr(solver, "snes", None) is None: solver._setup_pointwise_functions() solver._setup_discretisation() @@ -286,18 +389,36 @@ def solve_rotated_freeslip(solver, boundaries, remove_rotation_gauge=True, verbo dm = solver.dm snes = solver.snes - # Assemble the operator FIRST so its parallel row layout is final before we build - # Q against it (A = exact Jacobian at 0 — linear; b = -F(0)). The Pmat is - # assembled alongside: its p-p block is the native 1/mu pressure mass (the - # DS JacobianPreconditioner term) that preconditions the Schur complement — - # petsc4py's computeJacobian(x, J) would silently pass J as its own Pmat and - # the mass block would never be assembled. snes.setUp() U0 = dm.getGlobalVec() U0.set(0.0) J, Jp = snes.getJacobian()[:2] - snes.computeJacobian(U0, J, Jp) - Aorig = J.copy() + + use_lu = bool(getattr(solver, "_rotated_use_lu", False)) + cache = getattr(solver, "_rotated_linear_cache", None) + if use_lu and cache is not None: + _destroy_rotated_linear_cache(cache) + solver._rotated_linear_cache = None + cache = None + + if cache is None: + coefficient_variables = _operator_coefficient_variables(solver) + else: + coefficient_variables = cache["coefficient_variables"] + coefficient_states = _coefficient_states(coefficient_variables) + coefficients_unchanged = ( + cache is not None + and coefficient_states == cache["coefficient_states"] + and not force_operator_refresh + ) + + # Assemble J/Jp only when a coefficient used by the linear operator changed. + # A time-dependent body force changes F(0), not J, so repeated constant- + # viscosity Stokes solves can reuse the transformed matrix and GAMG hierarchy + # without paying distributed finite-element assembly every timestep. + if not coefficients_unchanged: + snes.computeJacobian(U0, J, Jp) + F0 = dm.getGlobalVec() snes.computeFunction(U0, F0) b = F0.copy() @@ -306,12 +427,27 @@ def solve_rotated_freeslip(solver, boundaries, remove_rotation_gauge=True, verbo dm.restoreGlobalVec(U0) dm.restoreGlobalVec(F0) - Q, Qt, normal_rows = build_rotation(solver, boundaries) - - # rotate: Â = Q A Qᵀ, b̂ = Q b - Ahat = Aorig.ptap(Qt) - bhat = b.duplicate() - Q.mult(b, bhat) + operator_reused = False + if cache is None: + Q, Qt, normal_rows = build_rotation(solver, boundaries) + Ahat = J.ptap(Qt) + operator_probe = _build_matrix_probe(J) + diag_scale = _velocity_diag_scale(Ahat, solver) + else: + Q = cache["Q"] + Qt = cache["Qt"] + normal_rows = cache["normal_rows"] + Ahat = cache["Ahat"] + operator_probe = cache["operator_probe"] + operator_reused = coefficients_unchanged or _matrix_matches_probe( + J, operator_probe + ) + if operator_reused: + diag_scale = cache["diag_scale"] + else: + J.ptap(Qt, result=Ahat) + _refresh_matrix_probe(J, operator_probe) + diag_scale = _velocity_diag_scale(Ahat, solver) # constrain rotated normal rows (v_n=0): zero the matrix rows/cols AND the RHS # at those rows — zeroRowsColumns does NOT touch the RHS, so a nonzero b there @@ -325,14 +461,19 @@ def solve_rotated_freeslip(solver, boundaries, remove_rotation_gauge=True, verbo # zeroed after the solve. # zeroRowsColumns takes GLOBAL row indices (correct); the RHS write goes through # _zero_rows_local (ownership-relative indexing — the np>1 crash class). - Ahat.zeroRowsColumns(normal_rows, diag=_velocity_diag_scale(Ahat, solver)) + if not operator_reused: + Ahat.zeroRowsColumns(normal_rows, diag=diag_scale) + + bhat = b.duplicate() + Q.mult(b, bhat) _zero_rows_local(bhat, normal_rows) # ITERATIVE by default (LU is almost never right): a self-contained fieldsplit- # Schur solve whose velocity block is geometric FMG on the custom prolongation # when a hierarchy is registered (set_custom_fmg), else GAMG. Direct LU only when # explicitly opted in via solver._rotated_use_lu. - if getattr(solver, "_rotated_use_lu", False): + if use_lu: + Aorig = J.copy() # Pin one pressure DOF (datum) — row only, keeps the B^T coupling. Only the # direct solve needs it; the iterative path fixes the pressure gauge with the # constant-pressure null space instead. @@ -362,12 +503,37 @@ def solve_rotated_freeslip(solver, boundaries, remove_rotation_gauge=True, verbo ksp_reason = ksp.getConvergedReason() ksp_its = ksp.getIterationNumber() _warn_if_ksp_diverged(ksp, kind="rotated direct-LU") + ksp.destroy() + Ahat.destroy() + _destroy_matrix_probe(operator_probe) + owns_rotation = True else: - Mp = _pressure_mass_schur_pmat(solver) + if cache is None or not operator_reused: + if cache is not None: + _destroy_rotated_ksp_ctx(cache["ctx"]) + Mp = _pressure_mass_schur_pmat(solver) + ctx = None + else: + Mp = cache["ctx"]["Mp"] + ctx = cache["ctx"] Uhat, ksp_reason, ctx = _solve_rotated_iterative( - solver, Ahat, bhat, Q, Qt, normal_rows, verbose=verbose, Mp=Mp) + solver, Ahat, bhat, Q, Qt, normal_rows, verbose=verbose, Mp=Mp, + ctx=ctx, refresh_operator=not operator_reused) ksp_its = ctx["ksp"].getIterationNumber() - _destroy_rotated_ksp_ctx(ctx) + cache = { + "Q": Q, + "Qt": Qt, + "normal_rows": normal_rows, + "Ahat": Ahat, + "operator_probe": operator_probe, + "coefficient_variables": coefficient_variables, + "coefficient_states": coefficient_states, + "diag_scale": diag_scale, + "ctx": ctx, + "boundaries": list(boundaries), + } + solver._rotated_linear_cache = cache + owns_rotation = False # rotate back u = Qᵀ û (U is returned in the result dict → create, don't # borrow from the pool) @@ -376,10 +542,21 @@ def solve_rotated_freeslip(solver, boundaries, remove_rotation_gauge=True, verbo removed = _finalize_rotated_solution(solver, U, Q, normal_rows, remove_rotation_gauge) - return {"Q": Q, "Qt": Qt, "A": Aorig, "b": b, "U": U, "Uhat": Uhat, + reaction = J.createVecLeft() + J.mult(U, reaction) + reaction.axpy(-1.0, b) + bhat.destroy() + + result = {"Q": Q, "Qt": Qt, "reaction": reaction, "U": U, "Uhat": Uhat, "normal_rows": normal_rows, "boundaries": list(boundaries), "rotation_gauge_removed": removed, "ksp_reason": ksp_reason, - "ksp_its": ksp_its} + "ksp_its": ksp_its, "workspace_reused": operator_reused, + "_owns_rotation": owns_rotation} + if use_lu: + result.update({"A": Aorig, "b": b}) + else: + b.destroy() + return result def _finalize_rotated_solution(solver, U, Q, normal_rows, remove_rotation_gauge): @@ -560,6 +737,12 @@ def solve_rotated_freeslip_nonlinear(solver, boundaries, remove_rotation_gauge=T of Newton increments solved, ``== len(ksp_its)``) and status; * ``"continuation_switched"`` — whether the Picard→Newton tangent switch fired. """ + _destroy_rotated_solve_result(getattr(solver, "_rotated_freeslip_info", None)) + solver._rotated_freeslip_info = None + if getattr(solver, "_rotated_linear_cache", None) is not None: + _destroy_rotated_linear_cache(solver._rotated_linear_cache) + solver._rotated_linear_cache = None + if getattr(solver, "snes", None) is None: solver._setup_pointwise_functions() solver._setup_discretisation() @@ -725,7 +908,8 @@ def rotated_residual(uvec, keep_cartesian=False): "rotation_gauge_removed": removed, "ksp_reason": last_reason, "nonlinear_iterations": newton_its, "converged": converged, "ksp_its": lin_its, - "continuation_switched": continuation and phase == "newton"} + "continuation_switched": continuation and phase == "newton", + "_owns_rotation": True} def _build_rotated_custom_Pl(solver, Q, normal_rows): @@ -779,17 +963,57 @@ 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 _destroy_rotated_solve_result(result): + """Release PETSc objects owned by one completed rotated solve result.""" + if not result: + return + for key in ("reaction", "U", "Uhat", "A", "b"): + obj = result.get(key) + if obj is not None: + obj.destroy() + result[key] = None + if result.get("_owns_rotation", True): + for key in ("Q", "Qt"): + obj = result.get(key) + if obj is not None: + obj.destroy() + result[key] = None + + +def _destroy_rotated_linear_cache(cache): + """Release the persistent linear rotated-free-slip workspace.""" + if not cache: + return + _destroy_rotated_ksp_ctx(cache.get("ctx")) + _destroy_matrix_probe(cache.get("operator_probe")) + cache["operator_probe"] = None + for key in ("Ahat", "Q", "Qt"): + obj = cache.get(key) + if obj is not None: + obj.destroy() + cache[key] = 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) when a hierarchy is registered (``set_custom_fmg``), else GAMG @@ -820,7 +1044,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. It leaves the KSP/PC/GAMG hierarchy untouched and updates + only the right-hand side.""" from underworld3.utilities import custom_mg dm = solver.dm vel_is = solver._subdict["velocity"][0] @@ -935,7 +1163,8 @@ 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) + 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 f6db5294..7e9db3f3 100644 --- a/tests/test_1018_rotated_freeslip.py +++ b/tests/test_1018_rotated_freeslip.py @@ -67,6 +67,96 @@ 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 update the RHS without rebuilding GAMG.""" + 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 + ) + with mesh.access(temperature, viscosity): + 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 + handles_1 = ( + cache_1["Q"].handle, + cache_1["Ahat"].handle, + cache_1["ctx"]["ksp"].handle, + ) + assert not stokes._rotated_freeslip_info["workspace_reused"] + + with mesh.access(temperature): + 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"] + 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 + + velocity_2 = velocity.data.copy() + old_cache = stokes._rotated_linear_cache + with mesh.access(viscosity): + viscosity.data[:, 0] *= 2.0 + stokes.solve(zero_init_guess=False) + + assert not stokes._rotated_freeslip_info["workspace_reused"] + assert old_cache["ctx"]["ksp"] is None + 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 + + refreshed_cache = stokes._rotated_linear_cache + refreshed_velocity = velocity.data.copy() + stokes.solve(zero_init_guess=False, time=0.5) + assert stokes._rotated_freeslip_info["workspace_reused"] + assert ( + stokes._rotated_linear_cache["ctx"]["ksp"].handle + == refreshed_cache["ctx"]["ksp"].handle + ) + 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):