feat(projection): opt-in linear_solver() + post-processing how-to (#156, #158)

docs/developer/index.md

@@ -158,6 +158,7 @@ CHANGELOG
 subsystems/meshing
 subsystems/discretisation
 subsystems/solvers
+subsystems/boundary-stress-and-projection-postprocessing
 subsystems/petsc-jacobian-layout
 subsystems/constitutive-models
 subsystems/constitutive-models-theory

docs/developer/subsystems/boundary-stress-and-projection-postprocessing.md

@@ -0,0 +1,85 @@
+# Boundary stress recovery & projection post-processing
+
+Two related rules for getting accurate, affordable derived quantities (boundary
+tractions, radial stress, deviatoric-stress components) out of a solved Stokes
+field. Both came out of the spherical-benchmark post-processing work
+(issues [#156], [#157], [#158]).
+
+## 1. Use a *linear* solver for projections
+
+`uw.systems.Projection` (and the vector/tensor/multi-component variants) performs
+an **L2 projection** — a linear, symmetric-positive-definite mass-matrix solve
+(or a screened-Poisson smooth when `smoothing > 0`, still linear SPD). It
+inherits the general `SNES_Scalar` default stack:
+
+```text
+snes_type = newtonls    # Newton — but the problem is linear
+ksp_type  = gmres
+pc_type   = gamg        # AMG setup + repartition: the expensive part
+```
+
+That stack is robust for general nonlinear scalar PDEs, but for projection it is
+**unnecessarily heavy**: GAMG setup/repartition dominates cost and memory, and
+when you project many quantities in sequence (e.g. six deviatoric-stress
+components at large MPI scale) it becomes the bottleneck — runs have been
+OOM-killed mid-sequence on Gadi.
+
+Switch projectors used purely for output/post-processing to a lightweight linear
+solve with the opt-in helper:
+
+```python
+proj = uw.systems.Projection(mesh, target)
+proj.uw_function = expr
+proj.linear_solver()          # ksponly + CG + jacobi, drops the GAMG options
+proj.solve()
+```
+
+`linear_solver(pc="jacobi", rtol=1e-10)` sets `snes_type=ksponly`,
+`ksp_type=cg`, `pc_type=<pc>` and removes the now-unused GAMG options. It is
+**opt-in** — the global default is unchanged, so internal paths that rely on it
+(proxy variables, derivative evaluation, …) are unaffected. `jacobi` is fine for
+a well-conditioned mass matrix; use `bjacobi` or `icc` if CG iteration counts
+climb on distorted or high-degree meshes.
+
+## 2. Project components, compose analytically
+
+When recovering a **boundary stress** such as the radial normal stress
+`σ_rr = nᵀ σ n`, do **not** project the composed expression directly:
+
+```python
+# DON'T: project the fully-composed scalar
+sigma_rr = uw.systems.Projection(mesh, scalar)
+sigma_rr.uw_function = n.T * stokes.stress * n      # = nᵀ(τ − pI)n
+```
+
+The composite mixes quantities of **different FE character** — the deviatoric
+stress `τ ~ C:ε̇(u)` (built from velocity gradients, one order below the P2
+velocity) and the pressure `p` (often discontinuous P1) — so a single scalar
+projection cannot represent it as accurately as the parts. On the spherical
+Thieulot benchmark the direct path gave ≈2× the `σ_rr` error of the reference.
+
+Instead, **project the low-level components and compose analytically**:
+
+```python
+# DO: project the deviatoric-stress components, then form σ_rr symbolically
+tau_proj = <project the τ_ij components, e.g. with a tensor/multi-component projection>
+sigma_rr = n.T * tau_proj * n - p        # p used directly, not re-projected
+```
+
+This reproduces the `stokes.tau`-based benchmark values. The general principle:
+**project the smoothest available primitives and combine them in closed form** —
+projecting a derived composite throws away accuracy the primitives still have.
+
+Combine the two rules — project the τ components with `linear_solver()` for the
+cheap linear solve, then compose `σ_rr` analytically — for accurate boundary
+stress recovery that also scales.
+
+## See also
+
+- Issues [#156] (projection solver settings), [#157] (projection memory),
+  [#158] (direct vs τ-based `σ_rr`).
+- [`Projection`](../../api/) — the projection solver family.
+
+[#156]: https://github.com/underworldcode/underworld3/issues/156
+[#157]: https://github.com/underworldcode/underworld3/issues/157
+[#158]: https://github.com/underworldcode/underworld3/issues/158

src/underworld3/systems/solvers.py

@@ -2476,6 +2476,53 @@ def __init__(
     # Use SymbolicProperty for automatic unwrapping
     uw_function = SymbolicProperty(matrix_wrap=True, doc="Function to project onto mesh")
 
+    def linear_solver(self, pc="jacobi", rtol=1.0e-10):
+        """Switch this projector to a lightweight *linear* (SPD) solve.
+
+        An L2 projection (and the screened-Poisson smoother) is a **linear,
+        symmetric-positive-definite** problem, so the inherited
+        ``newtonls / gmres / gamg`` default is unnecessarily heavy — GAMG
+        setup/repartition dominates cost and memory at MPI scale, which is the
+        bottleneck for repeated post-processing projections (UW3 issue #156).
+        This replaces it with ``ksponly + CG + a cheap preconditioner`` — the
+        right tool for the mass/Helmholtz matrix — and removes the now-unused
+        GAMG options.
+
+        Opt-in: the default ``SNES_Scalar`` solver stack is unchanged for code
+        that relies on it. Call this on a projector used purely for output /
+        post-processing.
+
+        Parameters
+        ----------
+        pc : str, default "jacobi"
+            Preconditioner. ``"jacobi"`` is fine for a well-conditioned mass
+            matrix; use ``"bjacobi"`` or ``"icc"`` if CG iteration counts climb
+            on distorted or high-degree meshes.
+        rtol : float, default 1e-10
+            KSP relative tolerance.
+
+        Returns
+        -------
+        self (so the call can be chained).
+        """
+        self.petsc_options["snes_type"] = "ksponly"
+        self.petsc_options["ksp_type"] = "cg"
+        self.petsc_options["pc_type"] = pc
+        self.petsc_options["ksp_rtol"] = rtol
+        # GAMG-specific options are now unused; remove them to avoid PETSc
+        # "unused option" warnings (and any stale AMG configuration).
+        for _k in (
+            "pc_gamg_type",
+            "pc_gamg_repartition",
+            "pc_gamg_agg_nsmooths",
+            "pc_mg_type",
+        ):
+            try:
+                self.petsc_options.delValue(_k)
+            except Exception:
+                pass
+        return self
+
     @property
     def smoothing(self):
         r"""Smoothing coefficient :math:`\alpha` of the screened-Poisson form.

tests/test_0506_projection_linear_solver.py

@@ -0,0 +1,67 @@
+#!/usr/bin/env python3
+"""Regression: Projection.linear_solver() opt-in lightweight SPD solve (#156).
+
+An L2 projection is a linear SPD problem; the default newtonls/gmres/gamg stack
+is heavy (GAMG setup/repartition is a memory/communication bottleneck for
+repeated post-processing projections). ``linear_solver()`` switches a projector
+to ``ksponly + CG + cheap PC`` without changing the global default.
+"""
+import numpy as np
+import sympy
+import pytest
+
+import underworld3 as uw
+
+pytestmark = [pytest.mark.level_1, pytest.mark.tier_a]
+
+
+@pytest.fixture
+def mesh():
+    return uw.meshing.UnstructuredSimplexBox(
+        minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=0.1, qdegree=3
+    )
+
+
+def _project(mesh, expr, use_linear):
+    target = uw.discretisation.MeshVariable(
+        f"t_{'lin' if use_linear else 'def'}", mesh, 1, degree=2
+    )
+    proj = uw.systems.Projection(mesh, target)
+    proj.uw_function = expr
+    proj.smoothing = 0.0
+    if use_linear:
+        proj.linear_solver()
+    proj.solve()
+    return proj, np.array(target.data[:, 0]).copy()
+
+
+def test_linear_solver_sets_lightweight_options(mesh):
+    """linear_solver() sets the ksponly/CG/jacobi stack and is chainable."""
+    target = uw.discretisation.MeshVariable("t_opts", mesh, 1, degree=2)
+    proj = uw.systems.Projection(mesh, target)
+    ret = proj.linear_solver()
+    assert ret is proj, "linear_solver() should return self for chaining"
+    assert proj.petsc_options.getString("snes_type", "") == "ksponly"
+    assert proj.petsc_options.getString("ksp_type", "") == "cg"
+    assert proj.petsc_options.getString("pc_type", "") == "jacobi"
+    # GAMG options removed
+    assert not proj.petsc_options.hasName("pc_gamg_type")
+
+
+def test_linear_solver_matches_default_projection(mesh):
+    """The lightweight solve produces the same projection as the default."""
+    x, y = mesh.X
+    expr = sympy.sin(3 * x) * sympy.cos(2 * y) + x ** 2
+    _, T_default = _project(mesh, expr, use_linear=False)
+    _, T_linear = _project(mesh, expr, use_linear=True)
+    rel = np.linalg.norm(T_default - T_linear) / np.linalg.norm(T_default)
+    assert rel < 1.0e-4, f"linear_solver projection differs from default: rel L2 = {rel:.3e}"
+
+
+def test_linear_solver_is_opt_in(mesh):
+    """A projector that does NOT call linear_solver() keeps the default stack."""
+    target = uw.discretisation.MeshVariable("t_optin", mesh, 1, degree=2)
+    proj = uw.systems.Projection(mesh, target)
+    # default SNES_Scalar stack untouched
+    assert proj.petsc_options.getString("snes_type", "") == "newtonls"
+    assert proj.petsc_options.getString("pc_type", "") == "gamg"