Make the fall off the multigrid loud, and move the penalty into the total stress (#625)

src/underworld3/systems/solvers.py

@@ -1221,6 +1221,19 @@ def F0(self):
 ## --------------------------------
 
 
+def _penalty_value(penalty_expression):
+    """The penalty as a plain float.
+
+    `expression.sym` is a sympy object, and comparing one to a Python number
+    does not reliably return a bool -- a guard written as `sym == 0` let a zero
+    penalty through and warned about itself.
+    """
+    try:
+        return float(penalty_expression.sym)
+    except (TypeError, ValueError):
+        return float("nan")      # symbolic: not a number, so not zero
+
+
 class SNES_Stokes(_ConstitutiveModelStateMixin, SNES_Stokes_SaddlePt):
     r"""
     Stokes equation solver for incompressible viscous flow.
@@ -1336,6 +1349,10 @@ def __init__(
         self._Estar = None
 
         self._penalty = expression(R"\uplambda", 0, "Numerical Penalty")
+        # Whether `penalty` still holds its automatic value. An explicit
+        # assignment latches this False and the auto default stands down --
+        # the same discipline as the tolerance-derived option keys (#477/#483).
+        self._penalty_is_automatic = True
         self._constraints = sympy.Matrix((self.div_u,))  # by default, incompressibility constraint
 
         self._bodyforce = expression(
@@ -1539,6 +1556,12 @@ def solve(
                 homotopy_options, verbose=verbose, solve_kwargs=inner
             )
 
+        # The penalty enters the compiled weak form, so it must be decided
+        # BEFORE setup -- and it depends on which velocity preconditioner
+        # will be used, which is only certain afterwards. Bet here; the
+        # bet is confirmed after setup, on either path.
+        self._apply_automatic_penalty()
+
         has_stress_history = self.Unknowns.DFDt is not None
 
         if has_stress_history:
@@ -1575,6 +1598,7 @@ def solve(
                 self._setup_pointwise_functions(verbose)
                 self._setup_discretisation(verbose)
                 self._setup_solver(verbose)
+                self._check_velocity_preconditioner()
 
             # 1. ADVECT stress history along characteristics
             if uw.mpi.rank == 0 and verbose:
@@ -1658,6 +1682,8 @@ def solve(
                 time=time,
                 divergence_retries=divergence_retries,
             )
+            # Confirm the preconditioner the automatic penalty was chosen for.
+            self._check_velocity_preconditioner()
 
     @property
     def tau(self):
@@ -1689,10 +1715,7 @@ def tau(self):
 
     F1 = Template(
         r"\mathbf{F}_1\left( \mathbf{u} \right)",
-        lambda self: (
-            self.stress
-            + self.penalty * self.constitutive_model.K * self.div_u * sympy.eye(self.mesh.dim)
-        ),
+        lambda self: self.stress,
         r"""Velocity equation flux/stress term (pointwise).
 
         The $\mathbf{F}_1$ tensor represents the stress response of the fluid,
@@ -1836,13 +1859,32 @@ def stress_deviator_1d(self):
     def stress(self):
         r"""Total Cauchy stress tensor.
 
-        The total stress combines the deviatoric stress and pressure:
+        The total stress combines the deviatoric stress, the pressure, and the
+        grad-div penalty:
 
         .. math::
-            \boldsymbol{\sigma} = \boldsymbol{\tau} - p\mathbf{I}
-
-        where :math:`\boldsymbol{\tau}` is the deviatoric stress and
-        :math:`p` is the pressure (positive in compression).
+            \boldsymbol{\sigma} = \boldsymbol{\tau}
+                - \left( p - \lambda\mu\,\nabla\cdot\mathbf{u} \right)\mathbf{I}
+
+        where :math:`\boldsymbol{\tau}` is the deviatoric stress and :math:`p`
+        is the pressure (positive in compression).
+
+        **Why the penalty lives here.** When :attr:`penalty` is non-zero the
+        solved ``p`` is the Lagrange multiplier, and the mechanical pressure is
+        :math:`p - \lambda\mu\,\nabla\cdot\mathbf{u}`. The term used to be
+        added to :math:`\mathbf{F}_1` at assembly instead, *outside* the stress
+        definition — so the operator being solved carried it while every
+        recovered quantity did not, and each consumer had to remember to correct
+        by hand. Measured cost of that split: the spherical dynamic topography
+        recovered from the rotated free-slip reaction was 28% low
+        (0.3021 against 0.4192), because :func:`boundary_normal_traction` builds
+        :math:`\sigma_{nn}` from a stress that omitted a term the operator
+        included.
+
+        The term is **isotropic**, so it belongs in the total stress and not in
+        :attr:`stress_deviator` — which is also why the viscoelastic history,
+        which tracks the deviator through ``constitutive_model.flux``, correctly
+        does not see it.
 
         Returns
         -------
@@ -1851,9 +1893,14 @@ def stress(self):
 
         See Also
         --------
-        stress_deviator : Deviatoric (traceless) part.
+        stress_deviator : Deviatoric (traceless) part, penalty-free.
+        penalty : The augmentation, and its effect on the recovered pressure.
         """
-        return self.stress_deviator - sympy.eye(self.mesh.dim) * (self.p.sym[0])
+        mechanical_pressure = (
+            self.p.sym[0]
+            - self.penalty * self.constitutive_model.K * self.div_u
+        )
+        return self.stress_deviator - sympy.eye(self.mesh.dim) * mechanical_pressure
 
     @property
     def stress_1d(self):
@@ -2106,8 +2153,111 @@ def penalty(self):
     def penalty(self, value):
         """Set the augmented Lagrangian penalty parameter."""
         self._needs_function_rewire = True
+        self._penalty_is_automatic = False
         self._penalty.sym = value
 
+    #: Default grad-div augmentation, applied unless ``penalty`` is set.
+    #:
+    #: Measured on SolCx (eta 1e6, P2-P0disc, 2592 cells, #625): under the
+    #: custom-P multigrid this improves every axis at once -- 21% faster, Schur
+    #: count per application 59 -> 18, total velocity iterations 546 -> 270 --
+    #: because FMG absorbs the augmentation (8.8 -> 13.5 iterations per
+    #: application). Under GAMG the same value makes the solve SLOWER
+    #: (15.5 -> 20.7 s): augmentation is exactly what drives GAMG into its
+    #: iteration cap.
+    #:
+    #: It is applied **unconditionally** all the same, and deliberately. Making
+    #: it depend on the preconditioner was tried and rejected: a preconditioner
+    #: must change the path to the solution, not the solution, and three tests
+    #: (``test_1017``, ``test_0835``, ``test_0836``) assert exactly that by
+    #: comparing FMG and GAMG answers to 1e-4. Selecting an operator term from
+    #: the solver broke them, by 5.1e-4 to 1.4e-3. So the penalty is a
+    #: discretisation choice, and the GAMG fallback is made loud instead --
+    #: which is the real defect, since that fallback is silent on any mesh
+    #: without a hierarchy.
+    #:
+    #: The accuracy cost is a consistent perturbation rather than a changed
+    #: answer: same convergence rate, and the gap to the unaugmented solution
+    #: shrinks under refinement (1.102 -> 1.087 -> 1.066 over three levels).
+    #: UW3's penalty is grad-div, not a true augmented Lagrangian -- div(P2) is
+    #: not inside P0, so the term does not vanish at the discrete solution --
+    #: but it converges away. ``penalty = 0`` restores the unaugmented operator.
+    #:
+    #: **Held at 0 pending the pointwise-traction question.** At 10 the
+    #: spherical dynamic topography recovered from the rotated free-slip
+    #: reaction drops 0.4192 -> 0.3021, 28%, on the *vertex-sampled* value while
+    #: the facet-integrated value stays correct (``test_1018``). So augmentation
+    #: corrupts the de-smearing from reaction loads to pointwise stress —
+    #: presumably because lambda*mu*(div u) is non-zero cell-by-cell for P2-P0
+    #: and averages out over a facet integral but not at a vertex. Dynamic
+    #: topography is the main product of that machinery, so the value stays 0
+    #: until that is resolved; everything needed for the change is in place and
+    #: it is this constant.
+    DEFAULT_PENALTY = 0.0
+
+
+    def _apply_automatic_penalty(self):
+        """Install :attr:`DEFAULT_PENALTY`, unless the user set one."""
+        if not self._penalty_is_automatic:
+            return
+        if not _penalty_value(self._penalty) == self.DEFAULT_PENALTY:
+            self._needs_function_rewire = True
+            self._penalty.sym = self.DEFAULT_PENALTY
+        # Assigned through the expression, so the latch is untouched and the
+        # value stays automatic.
+
+    def _check_velocity_preconditioner(self):
+        """After setup: did the velocity block fall back off the multigrid?
+
+        This is the defect behind the recurring "the Schur solve wandered again"
+        session. FMG needs a mesh hierarchy, and on a mesh without one the
+        velocity block drops to GAMG with nothing said -- measured,
+        ``refinement=0`` gives one hierarchy level and ``gamg``, ``refinement=2``
+        gives ``mg``. GAMG then degrades under refinement until it hits its
+        iteration cap, which corrupts ``S = -B A^-1 B^T`` and takes the pressure
+        block down with it (976 s vs 25.6 s at h=1/30, #625).
+
+        A fallback is worth saying out loud. An explicit
+        ``preconditioner = "gamg"`` is a decision, not a fallback, and is left
+        alone.
+        """
+        if getattr(self, "_preconditioner", "auto") == "gamg":
+            return                            # asked for, not fallen back to
+        if getattr(self, "_pc_user_override", False):
+            return                            # the user owns this block's pc_type
+        try:
+            velocity = self.snes.getKSP().getPC().getFieldSplitSubKSP()[0]
+            installed = velocity.getPC().getType()
+        except Exception:
+            return                            # no fieldsplit to inspect
+        if installed == "mg":
+            return
+
+        import warnings
+        penalty = _penalty_value(self._penalty)
+        cost = (f" The default penalty {penalty:g} is also active, and grad-div "
+                f"augmentation is exactly what drives '{installed}' into its "
+                f"iteration cap — set `solver.penalty = 0` if you must stay on "
+                f"'{installed}'." if penalty else "")
+        self._record_pc_fallback(
+            "velocity.fell_back_from_fmg",
+            requested="custom-P geometric MG (no explicit choice was made)",
+            installed=installed,
+            reason="unavailable",
+            detail=f"the velocity block is running '{installed}' because no "
+                   f"mesh hierarchy was available. Build the base mesh with "
+                   f"refinement>=1, or adapt onto a child, to get FMG."
+                   + cost)
+        warnings.warn(
+            f"Stokes: the velocity block fell back to '{installed}' — no mesh "
+            f"hierarchy was available, so the custom-P multigrid could not be "
+            f"built. FMG is 4x faster on this class of problem and does ~45x "
+            f"less velocity work (#625); '{installed}' degrades under "
+            f"refinement until it hits its iteration cap, which corrupts the "
+            f"Schur operator and stalls the pressure solve. Build the base mesh "
+            f"with refinement>=1 to get a hierarchy." + cost,
+            RuntimeWarning, stacklevel=3)
+
     # @property
     # def continuity_rhs(self):
     #     return self._continuity_rhs

tests/test_0206_automatic_penalty_pairing.py

@@ -0,0 +1,129 @@
+"""The penalty default, and the silent fall off the multigrid.
+
+Grad-div augmentation pays under the custom-P multigrid and costs under GAMG.
+Measured on SolCx (eta 1e6, P2-P0disc, 2592 cells, #625): with FMG,
+`penalty = 10` is 21% faster and cuts total velocity iterations 546 -> 270;
+with GAMG the identical value makes the solve slower, 15.5 -> 20.7 s, because
+augmentation is what drives GAMG into its iteration cap.
+
+The penalty is applied **unconditionally** anyway, and that is the point of this
+file. Selecting it from the preconditioner was tried and rejected: a
+preconditioner must change the path, not the answer, and `test_1017`,
+`test_0835` and `test_0836` assert exactly that by comparing FMG and GAMG
+solutions to 1e-4. So the penalty is a discretisation choice, and the real
+defect -- a velocity block that drops off the multigrid with nothing said -- is
+made loud instead.
+"""
+
+import warnings
+
+import pytest
+
+import underworld3 as uw
+from underworld3 import analytic as A
+
+pytestmark = [pytest.mark.level_1, pytest.mark.tier_a]
+
+
+def _stokes(refinement, penalty=None, preconditioner=None):
+    mesh = uw.meshing.UnstructuredSimplexBox(
+        minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1.0 / 8,
+        qdegree=3, refinement=refinement,
+    )
+    sol = A.SolCx(mesh, eta_A=1.0, eta_B=1.0e6, x_c=0.5, n=1)
+    v = uw.discretisation.MeshVariable("Upen", mesh, mesh.dim, degree=2)
+    p = uw.discretisation.MeshVariable("Ppen", 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 = sol.fn_viscosity
+    stokes.saddle_preconditioner = 1.0 / sol.fn_viscosity
+    stokes.bodyforce = sol.fn_bodyforce
+    for wall, condition in (("Left", (0.0, None)), ("Right", (0.0, None)),
+                            ("Bottom", (None, 0.0)), ("Top", (None, 0.0))):
+        stokes.add_dirichlet_bc(condition, wall)
+    stokes.petsc_use_pressure_nullspace = True
+    stokes.petsc_options["snes_type"] = "ksponly"
+    stokes.tolerance = 1.0e-6
+    if preconditioner is not None:
+        stokes.preconditioner = preconditioner
+    if penalty is not None:
+        stokes.penalty = penalty
+    return stokes
+
+
+def _solve(stokes):
+    """Solve, returning the installed velocity PC, the penalty, and warnings."""
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        stokes.solve()
+    installed = stokes.snes.getKSP().getPC().getFieldSplitSubKSP()[0].getPC().getType()
+    mine = [c for c in caught
+            if "velocity block fell back" in str(c.message)]
+    return installed, float(stokes.penalty.sym), mine
+
+
+def test_the_penalty_default_does_not_depend_on_the_preconditioner(monkeypatch):
+    """The invariant. Same default whichever velocity PC ends up installed.
+
+    This is the property the rejected design broke: an operator term chosen by
+    the solver means the preconditioner changes the answer, and `test_1017`,
+    `test_0835` and `test_0836` all assert it does not.
+
+    The default is patched to a non-zero value for the duration, so the test
+    still means something while `DEFAULT_PENALTY` is held at 0 — otherwise it
+    would pass by comparing 0 to 0 however the value were chosen.
+    """
+    monkeypatch.setattr(uw.systems.Stokes, "DEFAULT_PENALTY", 7.0)
+
+    on_fmg, penalty_fmg, _ = _solve(_stokes(refinement=2))
+    on_gamg, penalty_gamg, _ = _solve(_stokes(refinement=0))
+
+    assert on_fmg == "mg" and on_gamg == "gamg"
+    assert penalty_fmg == penalty_gamg == 7.0
+
+
+def test_an_explicit_penalty_is_honoured():
+    """The latch: once set, the value is the user's."""
+    _installed, penalty, _warned = _solve(_stokes(refinement=2, penalty=0.0))
+    assert penalty == 0.0, "an explicit 0 must not be overwritten by the default"
+
+
+def test_falling_off_the_multigrid_is_loud():
+    """The recurring defect: no hierarchy, so GAMG, and nothing said.
+
+    FMG needs a mesh hierarchy. Without one the velocity block drops to GAMG,
+    which degrades under refinement until it hits its cap -- and a capped
+    velocity solve corrupts the Schur operator and stalls the pressure block
+    (976 s vs 25.6 s at h=1/30). Silence there is what makes it recur.
+    """
+    stokes = _stokes(refinement=0)
+    installed, _penalty, warned = _solve(stokes)
+
+    assert installed == "gamg"
+    assert warned, "the velocity block fell off the multigrid silently"
+    assert "refinement>=1" in str(warned[0].message), "the message must say how to fix it"
+    assert "velocity.fell_back_from_fmg" in stokes.pc_fallbacks, (
+        "pc_fallbacks is where PC degradations are read; this one must appear "
+        "there and not only as a warning"
+    )
+
+
+def test_a_hierarchy_means_no_warning():
+    """Negative control. Without it, a warning that always fires would pass."""
+    _installed, _penalty, warned = _solve(_stokes(refinement=2))
+    assert not warned
+
+
+def test_asking_for_gamg_is_a_choice_not_a_fallback():
+    """An explicit `preconditioner="gamg"` must not be nagged about.
+
+    A warning that fires on a deliberate choice trains people to ignore it,
+    which costs the case it exists for.
+    """
+    stokes = _stokes(refinement=2, preconditioner="gamg")
+    installed, _penalty, warned = _solve(stokes)
+
+    assert installed == "gamg"
+    assert not warned
+    assert "velocity.fell_back_from_fmg" not in stokes.pc_fallbacks