diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index 70592a4a..bdc41d4e 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -1,3 +1,5 @@ +import collections + import numpy as np import sympy @@ -109,6 +111,23 @@ class SolverBaseClass(uw_object): # add_update_callback(). self._snes_update_callbacks = [] + # Solver difficulty / convergence reporting (default-on; see the + # solve_report / solve_history properties and _capture_solve_report). + # A SolveReport is recorded after every solve; solve_history keeps a + # short bounded trail so continuation / time-stepping loops can read + # trends without logging each step themselves. + self._solve_report = None + self._solve_history = collections.deque(maxlen=32) + # Bounded/resumable difficulty probe state (see estimate_difficulty). + # _difficulty_probe gates probe-only behaviour (iteration-cap solves, + # suppressed divergence warnings, restart anchoring); _difficulty_max_it + # is the active cap; _resume_abs_target anchors a chunked start/stop/restart + # chain to the ORIGINAL ||F0|| so it terminates at the same point an + # uninterrupted solve would. None = not in a resumable chain. + self._difficulty_probe = False + self._difficulty_max_it = None + self._resume_abs_target = None + # Preconditioner selection — see the `preconditioner` property. # `_pc_option_prefix` is set by subclasses that participate in the easy # FMG/GAMG switch ("" for scalar/vector, "fieldsplit_velocity_" for @@ -916,6 +935,166 @@ class SolverBaseClass(uw_object): return + @property + def solve_report(self): + """Difficulty / convergence record of the most recent solve (read-only). + + A :class:`~underworld3.systems.solve_report.SolveReport` populated by default after + every ``solve()`` — no opt-in. ``None`` before the first solve. Exposes the converged + reason, nonlinear/linear iteration counts, final and initial ‖F‖, the residual + reduction and contraction ρ, and the residual ladder. See also ``solve_history`` and + ``estimate_difficulty``. + """ + return self._solve_report + + @solve_report.setter + def solve_report(self, value): + raise AttributeError("solve_report is read-only (set by solve()).") + + @property + def solve_history(self): + """Bounded trail of recent :class:`SolveReport`s (read-only, most recent last). + + A ``collections.deque`` (``maxlen=32``) so continuation / time-stepping loops can read + difficulty trends without logging each solve themselves. + """ + return self._solve_history + + @solve_history.setter + def solve_history(self, value): + raise AttributeError("solve_history is read-only (set by solve()).") + + def _capture_solve_report(self, *, bounded=False): + """Record a SolveReport for the just-completed solve into ``self._solve_report`` and + append it to ``self._solve_history``. Reads SNES state only (side-effect free). + + INVARIANT: every physical ``self.snes.solve(...)`` site must be followed by a call to + this method, so that reporting covers every solve path — including any that returns + early and bypasses ``_snes_solve_with_retries`` / ``_warn_on_divergence`` (e.g. the + rotated-free-slip handoff in ``utilities/rotated_bc.py``). Do NOT anchor capture on + ``_warn_on_divergence`` (some paths skip it). + """ + from underworld3.systems.solve_report import ( + SolveReport, reason_string, contraction, + ) + + if getattr(self, "snes", None) is None: + self._solve_report = None + return None + + reason = int(self.snes.getConvergedReason()) + nl_its = int(self.snes.getIterationNumber()) + ksp_its = int(self.snes.getLinearSolveIterations()) + try: + fnorm = float(self.snes.getFunctionNorm()) + except Exception: + fnorm = float("nan") + try: + fev = int(self.snes.getFunctionEvaluations()) + except Exception: + fev = None + try: + hist = tuple(float(h) for h in self.snes.getConvergenceHistory()[0]) + except Exception: + hist = () + fnorm0 = hist[0] if hist else None + reduction = (fnorm / fnorm0) if (fnorm0 not in (None, 0.0)) else None + rho = contraction(hist) + + report = SolveReport( + reason=reason, + reason_str=reason_string(reason), + converged=reason > 0, + nl_its=nl_its, + ksp_its=ksp_its, + fnorm=fnorm, + fnorm0=fnorm0, + reduction=reduction, + rho=rho, + fev=fev, + history=hist, + bounded=bool(bounded), + ) + self._solve_report = report + self._solve_history.append(report) + return report + + def _capture_rotated_report(self, info): + """Record a SolveReport for a rotated-free-slip solve, which runs its own + ``ksp.solve()`` on the rotated operator (``utilities/rotated_bc.py``) and never + touches ``self.snes`` — so the generic reader would see stale SNES state. Build the + report from the rotated result dict (``ksp_reason`` / ``ksp_its``) so a rotated solve + also leaves a (read-only) report rather than a stale one. Best-effort, never raises.""" + from underworld3.systems.solve_report import SolveReport, reason_string + try: + info = info or {} + reason = int(info.get("ksp_reason", 0)) + ksp_its = int(info.get("ksp_its", 0)) + nl_its = int(info.get("newton_its", 1)) # linear path = one outer solve + fnorm = float(info.get("rnorm", float("nan"))) + report = SolveReport( + reason=reason, reason_str=reason_string(reason), converged=reason > 0, + nl_its=nl_its, ksp_its=ksp_its, fnorm=fnorm, + fnorm0=None, reduction=None, rho=None, fev=None, history=(), bounded=False, + ) + self._solve_report = report + self._solve_history.append(report) + return report + except Exception: + self._solve_report = None + return None + + def estimate_difficulty(self, max_nl_its, *, warm=True, **solve_kwargs): + """Run a bounded, resumable solve to *estimate solver difficulty* and return its report. + + Caps the solve at ``max_nl_its`` nonlinear iterations (for THIS call only), runs the + normal ``solve()`` — which leaves the partial iterate in the fields — and returns the + :class:`SolveReport` (``bounded=True``). The reported effort (iterations, residual + reduction, contraction ρ) IS the difficulty estimate. Continue losslessly with another + ``estimate_difficulty(warm=True)`` (a large cap runs it to completion); a chunked + start/stop/restart chain terminates at the same point an uninterrupted solve would, + because the chain anchors convergence to the original ‖F0‖. + + This bounds *solver work predictably* (an iteration count) — it is NOT a wall-time limit. + + Parameters + ---------- + max_nl_its : int + Cap on nonlinear (SNES) iterations for this call. + warm : bool, default True + ``True`` continues from the current fields (a restart within a chain); + ``False`` starts a fresh chain from a zero initial guess. + **solve_kwargs + Forwarded to ``solve()`` (e.g. ``timestep``, ``picard`` for Stokes/VE). + + Returns + ------- + SolveReport + The (``bounded=True``) report for this capped solve; also available as + ``self.solve_report``. + """ + if not warm: + self._resume_abs_target = None # fresh chain: forget any prior anchor + self._difficulty_probe = True + self._difficulty_max_it = int(max_nl_its) + try: + self.solve(zero_init_guess=(not warm), **solve_kwargs) + finally: + self._difficulty_probe = False + self._difficulty_max_it = None + + # Anchor a NEW chain to the original ||F0||. Subsequent warm restarts via + # estimate_difficulty(warm=True) — including a large-cap call to run to + # completion — terminate at tolerance*||F0||, the same point an uninterrupted + # solve would, regardless of where the caps fell. The anchor is consulted ONLY + # under _difficulty_probe (see _snes_solve_with_retries), so plain solves are + # unaffected and cannot inherit a stale anchor. + if self._resume_abs_target is None and self._solve_report is not None: + f0 = self._solve_report.fnorm0 + if f0: + self._resume_abs_target = self.tolerance * float(f0) + return self._solve_report + def get_snes_diagnostics(self): """ Extract comprehensive SNES convergence diagnostics with string representations. @@ -1129,6 +1308,11 @@ class SolverBaseClass(uw_object): if phase == "picard" and reason == -5: return + # Bounded difficulty probe: hitting the iteration cap (DIVERGED_MAX_IT) is + # the intended stop, not a failure — the effort is reported, not warned. + if self._difficulty_probe and reason == -5: + return + its = self.snes.getIterationNumber() _entry = self._convergence_reasons.get(reason) reason_str = _entry[0] if _entry is not None else f"UNKNOWN({reason})" @@ -1169,23 +1353,61 @@ class SolverBaseClass(uw_object): # Attach the per-iteration callback dispatcher here (after all # setFromOptions in the solve path). No-op when no callbacks registered. self._attach_snes_update_hook() - if self.consistent_jacobian == "continuation": - self._continuation_solve(gvec, verbose=verbose) - else: - self.snes.solve(None, gvec) - if divergence_retries <= 0: - return - for _r in range(divergence_retries): - reason = self.snes.getConvergedReason() - if reason >= 0: - return - if verbose and uw.mpi.rank == 0: - print( - f"SNES DIVERGED (reason={reason}); " - f"warm-start retry {_r + 1}/{divergence_retries}", - flush=True, - ) - self.snes.solve(None, gvec) + # Arm the residual-history buffer for THIS solve (reset clears stale data so the + # reported contraction is computed on the current ladder only). Done here — every + # solve funnels through this method and reads self.snes freshly — so a recreated + # SNES (new object after a setup-dirtying _build) is armed automatically. Guarded: + # not all PETSc builds expose it identically. + try: + self.snes.setConvergenceHistory(reset=True) + except Exception: + pass + + # Single control point for the bounded/resumable difficulty solve. Runs AFTER + # every per-solve setFromOptions (base and Stokes), so overrides here win. Gated + # on _difficulty_probe so it is active ONLY inside estimate_difficulty — a plain + # solve() is byte-for-byte unchanged and can never pick up a stale chain anchor. + # - iteration cap: cap nonlinear iterations for this probe; + # - resume anchor: once a chain is established (_resume_abs_target set from the + # first solve's ||F0||), terminate at the ORIGINAL ||F0|| via an absolute tol + # = tolerance*||F0||. This abs tol is LOOSER than the restart-relative + # rtol*||F_restart||, so it fires first — no need to zero rtol. So a chunked + # start/stop/restart terminates at the same point an uninterrupted solve would. + # Tolerances are saved and restored so nothing leaks into later solves. + _probe = bool(self._difficulty_probe) + _saved_tol = None + if _probe: + _saved_tol = self.snes.getTolerances() + if self._difficulty_max_it is not None: + self.snes.setTolerances(max_it=int(self._difficulty_max_it)) + if self._resume_abs_target is not None: + self.snes.setTolerances(atol=float(self._resume_abs_target)) + + try: + if self.consistent_jacobian == "continuation": + self._continuation_solve(gvec, verbose=verbose) + else: + self.snes.solve(None, gvec) + if divergence_retries > 0: + for _r in range(divergence_retries): + reason = self.snes.getConvergedReason() + if reason >= 0: + break + if verbose and uw.mpi.rank == 0: + print( + f"SNES DIVERGED (reason={reason}); " + f"warm-start retry {_r + 1}/{divergence_retries}", + flush=True, + ) + self.snes.solve(None, gvec) + finally: + if _saved_tol is not None: + self.snes.setTolerances(rtol=_saved_tol[0], atol=_saved_tol[1], + stol=_saved_tol[2], max_it=_saved_tol[3]) + + # Default-on difficulty report — single tail so every path is covered. + # bounded=True marks a report produced under an estimate_difficulty cap. + self._capture_solve_report(bounded=_probe) def _continuation_solve(self, gvec, verbose=False): """Picard -> Newton continuation via the constants[]-routed alpha. @@ -8048,6 +8270,9 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): self._rotated_freeslip_info = solve_rotated_freeslip_nonlinear( self, self._rotated_freeslip_bcs, verbose=verbose, zero_init_guess=zero_init_guess, picard=picard) + # 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. + self._capture_rotated_report(self._rotated_freeslip_info) return if time is not None: diff --git a/src/underworld3/systems/solve_report.py b/src/underworld3/systems/solve_report.py new file mode 100644 index 00000000..95b56efe --- /dev/null +++ b/src/underworld3/systems/solve_report.py @@ -0,0 +1,106 @@ +"""Solver difficulty / convergence reporting. + +Lives *with the solvers* (not ``utilities``): a :class:`SolveReport` is the record every +SNES-based solver leaves after a solve. It is populated by default on every solve (read via +``solver.solve_report``) and returned by ``solver.estimate_difficulty()``. + +Pure-Python (imports only the standard library) so it can be used without rebuilding the +Cython solver extension, and consumed by continuation / homotopy drivers. + +See ``underworld3.cython.petsc_generic_snes_solvers.SolverBaseClass``. +""" +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional, Tuple + +# PETSc SNESConvergedReason codes -> short names. Mirrors the compact map in +# SolverBaseClass._convergence_reasons; duplicated here so this module imports without the +# Cython solver extension present. +REASON_STRINGS = { + 1: "CONVERGED_FNORM_ABS", + 2: "CONVERGED_FNORM_RELATIVE", + 3: "CONVERGED_SNORM_RELATIVE", + 4: "CONVERGED_ITS", + 0: "ITERATING", + -1: "DIVERGED_FUNCTION_DOMAIN", + -2: "DIVERGED_FUNCTION_COUNT", + -3: "DIVERGED_LINEAR_SOLVE", + -4: "DIVERGED_FNORM_NAN", + -5: "DIVERGED_MAX_IT", + -6: "DIVERGED_LINE_SEARCH", + -7: "DIVERGED_INNER", + -8: "DIVERGED_LOCAL_MIN", + -9: "DIVERGED_DTOL", + -10: "DIVERGED_JACOBIAN_DOMAIN", + -11: "DIVERGED_TR_DELTA", +} + + +def reason_string(reason: int) -> str: + """Human-readable name for a PETSc SNES converged-reason code.""" + return REASON_STRINGS.get(int(reason), f"UNKNOWN_{reason}") + + +def contraction(history) -> Optional[float]: + """Geometric-mean per-iteration contraction factor ρ from a residual ladder. + + ρ = (‖F_last‖ / ‖F_first‖) ** (1 / (n-1)). Returns ``None`` when there are fewer than + two finite, positive residuals to compare (ρ is undefined for a single point). + """ + r = [h for h in history if h == h and h > 0.0] # finite & positive + if len(r) < 2: + return None + return (r[-1] / r[0]) ** (1.0 / (len(r) - 1)) + + +@dataclass(frozen=True) +class SolveReport: + """Difficulty / convergence record from a single solve. + + Populated by ``SolverBaseClass`` after every solve (default-on) and returned by + ``estimate_difficulty()``. Read-only; access the most recent via ``solver.solve_report`` + and the recent history via ``solver.solve_history``. + + Attributes + ---------- + reason, reason_str, converged + PETSc SNES converged-reason code, its name, and whether it indicates convergence. + nl_its, ksp_its + Nonlinear (SNES) iterations and cumulative linear (KSP) iterations of this solve. + fnorm, fnorm0, reduction + Final ‖F‖, initial ‖F‖ (first entry of the residual ladder, if captured), and their + ratio ``fnorm / fnorm0``. + rho + Geometric-mean per-iteration contraction of the residual ladder (``None`` if < 2 + points). A cheap difficulty indicator: ρ ≪ 1 is easy, ρ → 1 is a struggling solve. + fev + Function evaluations (includes line-search backtracks), if available. + history + The residual ladder ``getConvergenceHistory()`` recorded for this solve. + bounded + ``True`` when this report came from an iteration-capped (``estimate_difficulty``) + solve — i.e. work was truncated, not a genuine failure. + """ + + reason: int + reason_str: str + converged: bool + nl_its: int + ksp_its: int + fnorm: float + fnorm0: Optional[float] = None + reduction: Optional[float] = None + rho: Optional[float] = None + fev: Optional[int] = None + history: Tuple[float, ...] = () + bounded: bool = False + + def __str__(self) -> str: + rho = f"{self.rho:.3f}" if self.rho is not None else "n/a" + return ( + f"SolveReport({self.reason_str}, nl={self.nl_its}, ksp={self.ksp_its}, " + f"|F|={self.fnorm:.3e}, rho={rho}" + + (", bounded" if self.bounded else "") + + ")" + ) diff --git a/tests/test_1055_solve_report.py b/tests/test_1055_solve_report.py new file mode 100644 index 00000000..85c549ce --- /dev/null +++ b/tests/test_1055_solve_report.py @@ -0,0 +1,152 @@ +"""Default-on solver difficulty reporting + bounded/resumable difficulty solves. + +Covers the general base-solver capability (on ``SolverBaseClass``, so every SNES-based +solver gets it): + +1. ``test_solve_report_default_on`` — every ``solve()`` leaves a ``SolveReport`` (no opt-in), + with sensible iteration counts / residuals, and the property is read-only. +2. ``test_solve_report_scalar`` — the same report is populated for a scalar (Poisson) solver, + proving it lives on the base class, not Stokes only. +3. ``test_estimate_difficulty_bounded`` — an iteration-capped probe stops at the cap + (``DIVERGED_MAX_IT``), is flagged ``bounded``, and reports the effort spent. +4. ``test_resume_is_lossless_and_same_termination`` — a chunked start/stop/restart chain + reproduces the uninterrupted solution AND terminates at the same ‖F‖ (anchored to the + original ‖F0‖, not the tighter restart residual). +""" + +import numpy as np +import pytest +import sympy + +import underworld3 as uw +from underworld3.systems.solve_report import SolveReport, reason_string, contraction + + +def _linear_stokes(cell=0.25): + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0, 0), maxCoords=(1, 1), cellSize=cell + ) + v = uw.discretisation.MeshVariable("U", mesh, 2, degree=2) + p = uw.discretisation.MeshVariable("P", mesh, 1, degree=1, continuous=True) + st = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + st.constitutive_model = uw.constitutive_models.ViscousFlowModel + st.constitutive_model.Parameters.shear_viscosity_0 = 1.0 + x, y = mesh.X + st.bodyforce = sympy.Matrix([0, sympy.sin(np.pi * x) * sympy.sin(np.pi * y)]) + for w in ("Bottom", "Top"): + st.add_dirichlet_bc((sympy.oo, 0.0), w) + for w in ("Left", "Right"): + st.add_dirichlet_bc((0.0, sympy.oo), w) + return st, v + + +def _nonlinear_stokes(cell=0.2): + """Shear-thinning viscosity -> the default (Picard) tangent takes several iterations.""" + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0, 0), maxCoords=(1, 1), cellSize=cell + ) + v = uw.discretisation.MeshVariable("Un", mesh, 2, degree=2) + p = uw.discretisation.MeshVariable("Pn", mesh, 1, degree=1, continuous=True) + st = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + st.constitutive_model = uw.constitutive_models.ViscousFlowModel + edot = st.Unknowns.Einv2 + 1.0e-8 + st.constitutive_model.Parameters.shear_viscosity_0 = (1.0e-2 + edot) ** (-0.4) + st.tolerance = 1.0e-6 + x, y = mesh.X + st.bodyforce = sympy.Matrix([0, 2.0 * sympy.sin(np.pi * x) * sympy.sin(np.pi * y)]) + for w in ("Bottom", "Top"): + st.add_dirichlet_bc((sympy.oo, 0.0), w) + for w in ("Left", "Right"): + st.add_dirichlet_bc((0.0, sympy.oo), w) + return st, v + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_solve_report_default_on(): + st, _ = _linear_stokes() + assert st.solve_report is None # nothing before the first solve + st.solve() + r = st.solve_report + assert isinstance(r, SolveReport) + assert r.converged and r.reason > 0 + assert r.nl_its >= 1 and r.ksp_its >= 1 + assert r.fnorm == r.fnorm # finite (not NaN) + assert len(r.history) >= 1 + if r.fnorm0: + assert 0.0 < r.reduction <= 1.0 + 1e-12 + assert len(st.solve_history) == 1 # recorded in the trail + # read-only + with pytest.raises(AttributeError): + st.solve_report = 1 + with pytest.raises(AttributeError): + st.solve_history = 1 + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_solve_report_scalar(): + """Report is populated for a scalar solver too (it's on SolverBaseClass).""" + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0, 0), maxCoords=(1, 1), cellSize=0.25 + ) + T = uw.discretisation.MeshVariable("T", mesh, 1, degree=2) + poisson = uw.systems.Poisson(mesh, u_Field=T) + poisson.constitutive_model = uw.constitutive_models.DiffusionModel + poisson.constitutive_model.Parameters.diffusivity = 1.0 + poisson.f = 1.0 + poisson.add_dirichlet_bc(0.0, "Bottom") + poisson.add_dirichlet_bc(1.0, "Top") + poisson.solve() + r = poisson.solve_report + assert isinstance(r, SolveReport) and r.converged + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_estimate_difficulty_bounded(): + st, _ = _nonlinear_stokes() + rep = st.estimate_difficulty(max_nl_its=2, warm=False) + assert rep.bounded + assert rep.reason == -5 # DIVERGED_MAX_IT — the intended stop + assert rep.nl_its == 2 # exactly the cap + assert not rep.converged # truncated, not converged + assert rep is st.solve_report + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_resume_is_lossless_and_same_termination(): + # reference: one uninterrupted solve + ref, vref = _nonlinear_stokes() + ref.solve(zero_init_guess=True) + u_ref = vref.data.copy() + f_ref = ref.solve_report.fnorm + assert ref.solve_report.nl_its >= 3 # genuinely nonlinear + + # chunked: bounded probe then resume to completion (large cap) + st, v = _nonlinear_stokes() + rep1 = st.estimate_difficulty(max_nl_its=2, warm=False) + assert rep1.bounded and rep1.reason == -5 + rep2 = st.estimate_difficulty(max_nl_its=100, warm=True) + assert rep2.converged + + # lossless: warm-continued solution matches the uninterrupted reference + rel = np.linalg.norm(v.data - u_ref) / max(np.linalg.norm(u_ref), 1e-30) + assert rel < 1e-6, f"resume not lossless (rel-err {rel:.2e})" + + # same termination point (A5): the chunked chain converges at tolerance*||F0||, + # NOT the tighter tolerance*||F_restart|| + assert abs(rep2.fnorm - f_ref) / f_ref < 1e-3, ( + f"chunked termination {rep2.fnorm:.3e} != uninterrupted {f_ref:.3e}" + ) + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_solve_report_helpers(): + assert reason_string(2) == "CONVERGED_FNORM_RELATIVE" + assert reason_string(-5) == "DIVERGED_MAX_IT" + assert reason_string(999).startswith("UNKNOWN") + assert contraction([50.0, 1e-3, 1e-6, 1e-9]) is not None + assert contraction([50.0]) is None # < 2 points -> undefined