From 58fd1cf8b869ec38961dedb3bcefc97f5c1a0738 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Fri, 5 Jun 2026 16:23:01 +0100 Subject: [PATCH 1/7] Add SNES_Stokes_Constrained: non-penalty free-slip + recoverable topography Enforce u.n = g on curved boundaries with a recoverable Lagrange multiplier instead of a tuned penalty. An augmented-Lagrangian (ALG2) outer loop carries both the existing penalty BC and the multiplier; the multiplier removes the penalty's accuracy bias, so a moderate, well-conditioned augmentation gives an exact constraint in 2-3 Stokes solves. The converged multiplier is the normal traction holding the boundary -- a direct dynamic-topography estimate (h = lambda / (drho g)), recoverable as a clean MeshVariable (interior exactly zero; boundary trace correlates with -n.sigma.n at 1.0000). Purely additive: the validated 2x2 saddle-point assembly and fieldsplit config are untouched. New class behind uw.systems.Stokes_Constrained. Hands-off solve(): relative constraint tolerance (RMS(u.n-g) < rtol*RMS|u|) and viscosity-weighted augmentation (1e3*mu(x)) by default -- no user tuning. On the SolCx benchmark this gives ~2e-4 accuracy in 3 outer iterations across viscosity contrast 1 -> 1e6. Validation: tests/test_1061_constrained_freeslip_annulus.py (5 tests). Design note and production roadmap in docs/developer/design/. Underworld development team with AI support from Claude Code --- .../design/CONSTRAINED_FREESLIP_MULTIPLIER.md | 204 +++++++++++ src/underworld3/systems/__init__.py | 1 + src/underworld3/systems/solvers.py | 342 ++++++++++++++++++ .../test_1061_constrained_freeslip_annulus.py | 141 ++++++++ 4 files changed, 688 insertions(+) create mode 100644 docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md create mode 100644 tests/test_1061_constrained_freeslip_annulus.py diff --git a/docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md b/docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md new file mode 100644 index 000000000..cff96928e --- /dev/null +++ b/docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md @@ -0,0 +1,204 @@ +# Constrained free-slip via a recoverable Lagrange multiplier (dynamic topography) + +**Status**: proof-of-concept (Phase 0 + Phase 1), serial. Branch +`feature/constrained-freeslip-topography`. + +## Motivation + +Free-slip / no-normal-flow on curved (annulus, spherical) boundaries is +currently enforced with **penalty-like** methods — a penalty natural BC +(`add_natural_bc(penalty · Γ·v · Γ, ...)`) or Nitsche. These are fragile: the +penalty magnitude must be tuned against the Rayleigh number and viscosity. Too +weak and a coherent radial throughflow appears (an under-scaled `1e4` natural BC +is ~100× too weak at Ra=1e6); too strong and the system ill-conditions and the +Stokes solve diverges in line search. + +This feature enforces `u·n = g` on a curved boundary with a **true Lagrange +multiplier** `λ` instead of a penalty. Because the converged multiplier *is* the +normal traction holding the boundary, it is simultaneously a direct estimate of +**dynamic surface topography**, `h = λ / (Δρ g)`. The equilibrium `λ` is also the +target end-state toward which a free surface can be integrated over a time +interval (connecting to the ETD free-surface work on +`feature/exp-integrator-freesurface`). + +## Formulation + +Stokes with a surface constraint `u·n = g` on Γ, multiplier `λ`: + +``` +[ A Bᵀ Cᵀ ] [u] [f] +[ B 0 0 ] [p] = [0] A = viscous, B = div, C = ∫_Γ (n·v) ψ +[ C 0 0 ] [λ] [g] (C couples only the boundary trace of u) +``` + +`C` is **co-dimension-1**: it touches only velocity DOFs on Γ. A monolithic +third FE field would therefore either waste interior DOFs or need a boundary +trace space PETSc/DMPlex does not provide on the same DM. We instead solve the +boundary Schur complement `S_λ = C K⁻¹ Cᵀ` by an **outer loop**, leaving the +validated 2×2 Stokes assembly untouched. + +### Augmented Lagrangian (ALG2) — the production algorithm + +Each Stokes solve carries a natural-BC traction with **both** the multiplier and +a penalty augmentation, reusing the existing penalty machinery: + +$$\mathbf{t} = \bigl[\lambda + r\,(\mathbf{u}\cdot\mathbf{n} - g)\bigr]\,\mathbf{n} +\quad\text{on } \Gamma,$$ + +and the multiplier is updated with the same `r`: + +$$\lambda \leftarrow \lambda + r\,(\mathbf{u}\cdot\mathbf{n} - g).$$ + +The penalty term `r(u·n)n` preconditions *every* boundary mode uniformly, so the +outer loop converges in a handful of iterations; the multiplier removes the +penalty's accuracy bias, so a **moderate, well-conditioned** `r` gives both fast +convergence *and* an exact constraint — unlike a pure penalty, which must be made +large and fragile. + +## Phase-0 spike findings (what shaped the design) + +Spikes (`/tmp/s3_*.py`) on a 2D annulus (no-slip inner boundary to remove the +rigid-rotation null space, multiplier free-slip on the outer boundary): + +- **Plain Uzawa works but is slow.** A damped-Richardson update + `λ ← λ + ρ(u·n)` converges and matches the penalty solution, but a single + scalar `ρ` cannot kill both the fast and slow boundary-Schur modes — the + residual contracts the dominant mode in ~5 iterations then crawls. +- **`ρ ∝ μ`, NOT `ρ ∝ μ/h`.** The optimal Richardson step is `ρ ≈ C·μ` with `C` + a geometry constant, **independent of mesh resolution** (`ρ=8μ` converged in 5 + iterations at cellSize 0.1/0.05/0.025). The naive `μ/h` scaling over-steps on + refinement and stalls. +- **CG is the wrong accelerator.** CG on `S_λ` diverged: each matvec is an + *inexact* iterative Stokes solve (plus pressure-null-space noise), and the + nodal Euclidean inner product is not the one in which `S_λ` is SPD. Krylov + acceleration needs an exact symmetric operator; this is neither. +- **Augmented Lagrangian is the right accelerator** (per L. Moresi). It converges + in **2 iterations** where plain Uzawa took 21, reusing the existing penalty BC. + This is the implemented algorithm. + +## Implementation + +`SNES_Stokes_Constrained(SNES_Stokes)` in `src/underworld3/systems/solvers.py`, +exported as `uw.systems.Stokes_Constrained`. **Purely additive** — the validated +2×2 saddle-point assembly and fieldsplit configuration are untouched, honouring +"solver stability is paramount". + +```python +stokes = uw.systems.Stokes_Constrained(mesh, velocityField=v, pressureField=p) +stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel +stokes.constitutive_model.Parameters.shear_viscosity_0 = mu +stokes.bodyforce = buoyancy * unit_r +stokes.add_dirichlet_bc((0.0, 0.0), "Lower") # no-slip inner + +lam = stokes.add_constraint_bc("Upper", g=0.0) # free-slip outer (Gamma_P1) +stokes.solve() # no constraint tuning needed + +topo = stokes.topography("Upper", buoyancy_scale=delta_rho_g) # h = lambda/(drho g) +``` + +The outer loop is hidden behind ``solve()``: the tolerance is **relative** to the +velocity scale (``RMS(u.n-g) < rtol·RMS|u|``, default ``rtol=1e-3``) so it is +problem-independent, and the augmentation defaults to ``1e3·μ(x)`` (local- +viscosity-weighted). On SolCx this gives ~2e-4 accuracy in 3 outer iterations at +viscosity contrast 1 → 10⁶ with no user tuning. + +Key design points: + +- **Multiplier representation.** `λ` is an ordinary full-mesh scalar field at the + *velocity degree* (P2). Only its trace on Γ enters the weak form. Matching the + velocity degree means the multiplier reaches every velocity normal-trace DOF + (including P2 mid-edge), so there is no penalty floor on the constraint — a P1 + multiplier plateaus at ~`‖t‖/r` on the mid-edge component. +- **Clean topography field.** The dual update is restricted to boundary nodes via + a P1 boundary marker (`petsc_dm_find_labeled_points_local`) sampled at the + multiplier's nodes — boundary mid-edge nodes interpolate to 1, interior to + < 0.5. Interior `λ` therefore stays **exactly zero**, so `λ` is a directly + usable topography field (not interior garbage). +- **Coupling registered once.** `add_natural_bc([λ + r(u·n − g)]·n, boundary)` + is set up a single time; only `λ`'s boundary data changes between solves, so + nothing recompiles. +- **`add_constraint_bc(boundary, g=0, normal=None, augmentation=None)`** — + `normal` defaults to the smooth projected normals `mesh.Gamma_P1`; + `augmentation` defaults to a viscosity-scaled `r = 10³·μ`. + +## Validation + +`tests/test_1061_constrained_freeslip_annulus.py` (level_2 / tier_b), buoyancy- +driven annulus, no-slip inner + multiplier free-slip outer, vs a `1e6` penalty +reference. All pass (~14 s): + +| Check | Result | +|---|---| +| `RMS(u·n)` on outer boundary (no penalty coefficient) | 6.5e-5 | +| `relL2(v_multiplier vs v_penalty)` | 3.1e-3 | +| Outer iterations (augmented Lagrangian) | 2 | +| Interior `λ` (clean field) | exactly 0 | +| boundary `corr(λ, −n·σ·n)` (dynamic-topography stress) | 0.9999 | + +The consistent-boundary-flux identity `λ = −n·σ·n|_Γ` is the independent +cross-check: the multiplier's boundary trace equals the recovered normal Cauchy +stress (negative sign = the reaction traction holding the boundary), confirming +`λ` is the dynamic topography signal. + +## The augmentation parameter `r`: true-work trade-off + +`r` is a *speed* knob, not an *accuracy* knob — this is the key advantage over a +pure penalty. Sweep on the annulus (constraint tol 1e-4, wall time for one cold +solve; `tot_lin` = total outer Schur-KSP linear iterations across the loop): + +cellSize 0.1 (~2940 dof): + +| `r` | outer its | tot_lin | wall (s) | relL2 vs penalty | +|---:|---:|---:|---:|---:| +| 10 | 26 | 26 | 3.77 | 3.1e-3 | +| 100 | 4 | 4 | 0.61 | 3.2e-3 | +| 300 | 3 | 3 | 0.47 | 3.1e-3 | +| 1,000 | 3 | 3 | 0.53 | 3.1e-3 | +| 3,000 | 2 | 2 | **0.38** | 3.1e-3 | +| 10,000 | 2 | 2 | 0.51 | 2.9e-3 | +| 100,000 | 2 | 5 | 1.83 | 1.6e-3 | + +cellSize 0.05 (~10852 dof) shows the same shape (min wall ≈ 1.25 s at r=1e3; +6.98 s at r=10; 7.07 s at r=1e5). + +- **Outer iterations fall with `r`** (`26 → 4 → 3 → 2`) — bigger penalty, faster + dual convergence (`contraction ≈ ‖S_λ‖/(r+‖S_λ‖)`). +- **But the inner solve stiffens at large `r`.** Linear iterations per outer + solve stay at 1.0 up to `r=10⁴`, then rise (2.5 at `r=10⁵`), and wall time + balloons (the velocity sub-block conditioning degrades — visible in wall time + even before the outer KSP count moves). +- **True work is U-shaped**: both extremes are 3–8× slower than the optimum. The + efficient basin is `r ∈ [300, 10⁴]` (>1.5 decades) at both resolutions; the + default `r = 10³·μ` sits inside it. +- **Accuracy is `r`-independent** (relL2 ≈ 3.1e-3, flat across four decades of + `r`). So `r` is tuned for *speed* with a benign failure mode — too small just + costs iterations, too large just costs inner work; **the answer is never + wrong**. Contrast a pure penalty, where the magnitude must be tuned against + forcing strength and viscosity to get *accuracy* (too small ⇒ wrong), which is + the fragility this method removes. + +## Option trade-offs and what is deferred + +| Option | Verdict | +|---|---| +| (A) full-domain 3rd FE field + ε-screening | 3-way nested fieldsplit; ε re-introduces tuning. Rejected as primary. | +| (B1) boundary-stratum-only PetscFE field | No DMPlex support on the same DM. | +| (B2) co-dim-1 submesh + MATNEST | The honest monolithic form; deferred. | +| (C) reuse pressure / `_constraints` | `p` enforces `∇·u=0` interior, not `u·n=0` on Γ — not redundant. The CBF identity is a *validation* tool, not an implementation. | +| **(D) augmented-Lagrangian outer loop** | **Implemented.** Non-invasive, 2 iterations, exact, recoverable topography. | + +Deferred to follow-up PRs (Phase-0 spike S2 gathers the fieldsplit-feasibility +evidence): the monolithic 3-field / co-dim-1 representation; 3D spherical shells; +**parallel** (the boundary mask and the nodal update are serial); the +**both-boundaries-free-slip** annulus case (admits a rigid-rotation velocity null +space needing explicit removal); and live free-surface equilibrium integration +(pass `λ` as the target normal-stress end-state). + +## Files + +- `src/underworld3/systems/solvers.py` — `SNES_Stokes_Constrained`, `_ConstraintBC`. +- `src/underworld3/systems/__init__.py` — `Stokes_Constrained` export. +- `tests/test_1061_constrained_freeslip_annulus.py` — validation. + +Pre-existing (unrelated) failures noted: `tests/test_1060_nitsche_freeslip.py` +has two failing assertions on `development` independent of this work. diff --git a/src/underworld3/systems/__init__.py b/src/underworld3/systems/__init__.py index 0c66d57cc..ddf22c922 100644 --- a/src/underworld3/systems/__init__.py +++ b/src/underworld3/systems/__init__.py @@ -48,6 +48,7 @@ from .solvers import SNES_Poisson as Poisson from .solvers import SNES_Darcy as SteadyStateDarcy from .solvers import SNES_Stokes as Stokes +from .solvers import SNES_Stokes_Constrained as Stokes_Constrained from .solvers import SNES_VE_Stokes as VE_Stokes from .solvers import SNES_Projection as Projection from .solvers import SNES_Vector_Projection as Vector_Projection diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index 4e2cf2177..d0f9fe6ca 100644 --- a/src/underworld3/systems/solvers.py +++ b/src/underworld3/systems/solvers.py @@ -1936,6 +1936,348 @@ def delta_t(self): return self.constitutive_model.Parameters.dt_elastic +class _ConstraintBC: + """Bookkeeping for one multiplier-enforced boundary constraint. + + Holds the multiplier field ``lam``, the prescribed normal velocity ``g``, + the (symbolic, row-vector) constraint normal, and the augmented-Lagrangian + parameter ``r`` (which is simultaneously the forward-problem penalty weight + and the multiplier-update step). + """ + + __slots__ = ("boundary", "g", "normal", "lam", "augmentation", "mask", "r_nodal") + + def __init__(self, boundary, g, normal, lam, augmentation, mask): + self.boundary = boundary + self.g = g + self.normal = normal + self.lam = lam + # augmentation r may be a scalar or a spatial sympy expression (e.g. + # viscosity-weighted). r_nodal holds it sampled at the multiplier nodes + # for the dual update; it is (re)computed at solve time. + self.augmentation = augmentation + self.mask = mask + self.r_nodal = None + + +class SNES_Stokes_Constrained(SNES_Stokes): + r""" + Stokes solver with boundary constraints enforced by a recoverable Lagrange + multiplier instead of a penalty. + + For each constraint boundary :math:`\Gamma` the no-normal-flow (free-slip) + or prescribed-normal-velocity condition + + .. math:: + + \mathbf{u} \cdot \mathbf{n} = g \quad \text{on } \Gamma + + is enforced by introducing a scalar multiplier field :math:`\lambda` and an + **augmented-Lagrangian** (Uzawa / ALG2) outer loop. Each Stokes solve carries + a natural-BC traction with both the multiplier and a penalty augmentation, + + .. math:: + + \mathbf{t} = \bigl[\lambda + r\,(\mathbf{u}\cdot\mathbf{n} - g)\bigr]\,\mathbf{n} + \quad \text{on } \Gamma , + + and the multiplier is updated with the same augmentation parameter, + + .. math:: + + \lambda \leftarrow \lambda + r\,(\mathbf{u}\cdot\mathbf{n} - g) . + + The penalty term :math:`r\,(\mathbf{u}\cdot\mathbf{n})\,\mathbf{n}` is exactly + the existing penalty free-slip BC: it preconditions every boundary mode + uniformly so the outer loop converges in a handful of iterations, while the + multiplier removes the penalty's accuracy bias — so a **moderate**, + well-conditioned :math:`r` gives both fast convergence and an *exact* + constraint (unlike a pure penalty, which must be made large and fragile). + + At convergence :math:`\lambda` is the normal traction holding the boundary, + giving a direct estimate of dynamic surface topography, + :math:`h = \lambda / (\Delta\rho\, g)`. Access it via :meth:`multiplier`. + + Notes + ----- + The multiplier is represented as an ordinary full-mesh scalar field of the + same degree as the velocity. Only its trace on :math:`\Gamma` enters the + weak form (interior values are inert), so no boundary trace space is + required. This is a proof-of-concept (serial; one full Stokes solve per + outer iteration). See the design note + ``docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md``. + + See Also + -------- + SNES_Stokes : The unconstrained saddle-point solver this extends. + """ + + def __init__( + self, + mesh: uw.discretisation.Mesh, + velocityField: Optional[uw.discretisation.MeshVariable] = None, + pressureField: Optional[uw.discretisation.MeshVariable] = None, + degree: Optional[int] = 2, + p_continuous: Optional[bool] = True, + verbose: Optional[bool] = False, + DuDt: Union[SemiLagrangian_DDt, Lagrangian_DDt] = None, + DFDt: Union[SemiLagrangian_DDt, Lagrangian_DDt] = None, + ): + super().__init__( + mesh, + velocityField, + pressureField, + degree, + p_continuous, + verbose, + DuDt=DuDt, + DFDt=DFDt, + ) + + self._constraint_bcs = [] + # Diagnostics from the most recent constrained solve. + self.constraint_iterations = 0 + self.constraint_residual = None + self.constraint_total_linear_its = 0 + return + + def _viscosity_scale(self): + """A representative scalar viscosity, for sizing the initial Uzawa step.""" + try: + mu = self.constitutive_model.Parameters.shear_viscosity_0 + return float(mu) + except (TypeError, ValueError, AttributeError): + return 1.0 + + def add_constraint_bc(self, boundary, g=0.0, normal=None, augmentation=None, + augmentation_base=1.0e3): + r"""Register a multiplier-enforced normal-velocity constraint on ``boundary``. + + Parameters + ---------- + boundary : str + Mesh boundary label (e.g. ``"Upper"``). + g : float or sympy expression, default 0.0 + Prescribed normal velocity :math:`\mathbf{u}\cdot\mathbf{n} = g`. + The default ``0`` is free-slip / no-normal-flow. + normal : sympy matrix, optional + Row-vector constraint normal. Defaults to the mesh's smooth + projected boundary normals ``mesh.Gamma_P1``. + augmentation : float or sympy expression, optional + Augmented-Lagrangian parameter :math:`r` — simultaneously the + forward-problem penalty weight and the multiplier-update step. Any + :math:`r>0` converges; larger is faster but stiffens the linear + solve. **Defaults to ``augmentation_base · μ(x)``** — weighted by the + local viscosity so the dimensionless penalty ratio is uniform across + viscosity contrasts (essential for variable-viscosity problems; a + flat ``r`` under-constrains high-viscosity regions). May be passed as + a spatial sympy expression directly. + augmentation_base : float, default 1e3 + The base multiple used when ``augmentation`` is not given. + + Returns + ------- + lam : MeshVariable + The scalar multiplier field. After :meth:`solve`, its boundary + trace is the normal traction (topography proxy). + """ + if normal is None: + normal = self.mesh.Gamma_P1 + + normal = sympy.Matrix(normal) + if normal.shape[0] != 1: + normal = normal.reshape(1, self.mesh.dim) + + if augmentation is None: + # Viscosity-weighted augmentation r = augmentation_base * mu(x): + # keeps the penalty/viscous ratio uniform so high-viscosity boundary + # regions are constrained as well as low-viscosity ones. + try: + viscosity = self.constitutive_model.Parameters.shear_viscosity_0 + augmentation = augmentation_base * viscosity + except (AttributeError, TypeError): + augmentation = augmentation_base * self._viscosity_scale() + + idx = len(self._constraint_bcs) + # Multiplier at the velocity degree so its trace reaches every velocity + # normal-trace DOF (no penalty floor on the P2 mid-edge component). + lam = uw.discretisation.MeshVariable( + f"lambda_{self.instance_number}_{idx}", + self.mesh, + 1, + degree=self._degree, + ) + lam.data[:] = 0.0 + + # Boundary-node mask: restrict the multiplier update to the constraint + # boundary so interior values stay exactly zero and lambda is a clean, + # directly usable topography field. Build a P1 marker (1 on the boundary + # vertices, 0 elsewhere) and sample it at lambda's nodes: boundary + # mid-edge nodes interpolate to 1, interior nodes to < 0.5. + from underworld3.discretisation.discretisation_mesh import ( + petsc_dm_find_labeled_points_local, + ) + + marker = uw.discretisation.MeshVariable( + f"_bmarker_{self.instance_number}_{idx}", self.mesh, 1, degree=1, + ) + marker.data[:] = 0.0 + point_indices = petsc_dm_find_labeled_points_local( + self.mesh.dm, + "UW_Boundaries", + getattr(self.mesh.boundaries, boundary).value, + sectionIndex=False, + ) + if point_indices is not None: + marker.data[point_indices] = 1.0 + mask = ( + np.array(uw.function.evaluate(marker.sym, lam.coords)).reshape(-1) > 0.75 + ) + + # Augmented-Lagrangian natural BC, registered once: + # t = [ lambda + r (u.n - g) ] n + # The r(u.n)n part is the penalty BC (re-derived into the Jacobian each + # solve); the lambda part is the applied multiplier (fixed per solve). + nv = self.u.sym.dot(normal) + traction = (lam.sym[0] + augmentation * (nv - g)) * normal + self.add_natural_bc(traction, boundary) + + self._constraint_bcs.append( + _ConstraintBC(boundary, g, normal, lam, augmentation=augmentation, mask=mask) + ) + return lam + + def multiplier(self, boundary): + """Return the multiplier field for ``boundary`` (None if not constrained). + + After :meth:`solve`, the multiplier's boundary trace is the normal + traction holding the constraint; interior values are zero. Divide by + :math:`\\Delta\\rho\\,g` to obtain dynamic topography (see + :meth:`topography`). + """ + for cbc in self._constraint_bcs: + if cbc.boundary == boundary: + return cbc.lam + return None + + def topography(self, boundary, buoyancy_scale=1.0): + r"""Dynamic topography expression on ``boundary``. + + Returns the symbolic field :math:`\lambda / (\Delta\rho\, g)` for the + constraint multiplier on ``boundary`` (zero away from the boundary). + + Parameters + ---------- + boundary : str + A constrained boundary label. + buoyancy_scale : float or sympy expression, default 1.0 + The buoyancy scale :math:`\Delta\rho\, g` relating normal traction + to surface height. + """ + lam = self.multiplier(boundary) + if lam is None: + raise ValueError(f"No constraint registered on boundary '{boundary}'.") + return lam.sym[0] / buoyancy_scale + + def _constraint_rms(self, cbc): + """RMS of (u.n - g) over the constraint boundary, via boundary integral.""" + vn = self.u.sym.dot(cbc.normal) + num = float( + uw.maths.BdIntegral(self.mesh, fn=(vn - cbc.g) ** 2, + boundary=cbc.boundary).evaluate() + ) + length = float( + uw.maths.BdIntegral(self.mesh, fn=1.0, boundary=cbc.boundary).evaluate() + ) + return np.sqrt(num / length) if length > 0 else np.sqrt(num) + + def _nodal_constraint_residual(self, cbc): + """(u.n - g) evaluated at the multiplier's nodes.""" + expr = self.u.sym.dot(cbc.normal) - cbc.g + return np.array( + uw.function.evaluate(sympy.Matrix([[expr]]), cbc.lam.coords) + ).reshape(-1) + + def solve( + self, + zero_init_guess: bool = True, + *, + constraint_rtol: float = 1.0e-3, + constraint_atol: float = 1.0e-12, + constraint_max_iterations: int = 40, + constraint_verbose: bool = False, + **kwargs, + ): + """Solve the constrained Stokes system. + + With no constraint BCs registered this is an ordinary Stokes solve. + Otherwise it runs the augmented-Lagrangian outer loop until every + constraint boundary satisfies + ``RMS(u.n - g) < constraint_rtol · RMS|u| + constraint_atol`` (or the + iteration cap is hit). The tolerance is *relative* to the velocity scale + so it is problem-independent and needs no tuning. All other keyword + arguments are forwarded to the inner Stokes ``solve``. + """ + if not self._constraint_bcs: + return super().solve(zero_init_guess=zero_init_guess, **kwargs) + + # Sample the augmentation r at the multiplier nodes once (it may be a + # spatial / viscosity-weighted expression). Used for the dual update. + for cbc in self._constraint_bcs: + if isinstance(cbc.augmentation, (int, float)): + cbc.r_nodal = float(cbc.augmentation) * np.ones(cbc.lam.coords.shape[0]) + else: + cbc.r_nodal = np.array( + uw.function.evaluate( + sympy.Matrix([[cbc.augmentation]]), cbc.lam.coords + ) + ).reshape(-1) + + total_linear_its = 0 + for k in range(constraint_max_iterations): + super().solve(zero_init_guess=(zero_init_guess and k == 0), **kwargs) + try: + total_linear_its += int(self.snes.getLinearSolveIterations()) + except Exception: + pass + + # Velocity scale for the relative tolerance (RMS speed over nodes). + v_scale = float(np.sqrt(np.mean(np.sum(self.u.data**2, axis=1)))) + threshold = constraint_rtol * v_scale + constraint_atol + + all_converged = True + worst = 0.0 + for cbc in self._constraint_bcs: + rms = self._constraint_rms(cbc) + if rms >= threshold: + all_converged = False + worst = max(worst, rms) + if constraint_verbose: + uw.mpi.pprint( + f" [constraint {cbc.boundary}] iter {k}: " + f"RMS(u.n-g) = {rms:.3e} (rel {rms / max(v_scale, 1e-30):.2e}, " + f"mean r = {cbc.r_nodal.mean():.3g})" + ) + + self.constraint_iterations = k + 1 + self.constraint_residual = worst + self.constraint_total_linear_its = total_linear_its + + if all_converged: + if constraint_verbose: + uw.mpi.pprint(f"Constraint loop converged in {k + 1} iterations.") + break + + # Augmented-Lagrangian (ALG2) multiplier update, same r as the + # forward-problem penalty augmentation. Monotone-convergent for r>0. + # Restricted to boundary nodes so lambda stays a clean topography field. + for cbc in self._constraint_bcs: + resid = self._nodal_constraint_residual(cbc) + cbc.lam.data[cbc.mask, 0] += cbc.r_nodal[cbc.mask] * resid[cbc.mask] + + return + + class SNES_Projection(SNES_Scalar): r""" Scalar projection solver for mapping functions to mesh variables. diff --git a/tests/test_1061_constrained_freeslip_annulus.py b/tests/test_1061_constrained_freeslip_annulus.py new file mode 100644 index 000000000..4fba5316e --- /dev/null +++ b/tests/test_1061_constrained_freeslip_annulus.py @@ -0,0 +1,141 @@ +"""Constrained free-slip on an annulus via a recoverable Lagrange multiplier. + +Buoyancy-driven flow in an annulus with no-slip inner boundary and free-slip +outer boundary. The free-slip outer condition (u.n = 0) is enforced three ways +and compared: + + - penalty : the existing fragile penalty natural BC (reference) + - multiplier: SNES_Stokes_Constrained, augmented-Lagrangian multiplier + +The multiplier solver must (a) drive u.n -> 0 with NO penalty coefficient, +(b) match the penalty velocity field, and (c) yield a clean, recoverable +topography field whose boundary trace equals the consistent-boundary-flux +normal stress (-n.sigma.n), i.e. dynamic topography. + +Run with: pixi run python -m pytest tests/test_1061_constrained_freeslip_annulus.py -v +""" + +import pytest +import numpy as np +import sympy +import underworld3 as uw + +pytestmark = [pytest.mark.level_2, pytest.mark.tier_b] + +R_INNER, R_OUTER = 0.5, 1.0 +CELL = 0.1 +MU = 1.0 +RA = 1.0e2 + + +def _mesh_and_forcing(): + mesh = uw.meshing.Annulus(radiusInner=R_INNER, radiusOuter=R_OUTER, + cellSize=CELL, qdegree=3) + x, y = mesh.X + r = sympy.sqrt(x**2 + y**2) + unit_r = sympy.Matrix([[x / r, y / r]]) + theta = sympy.atan2(y, x) + buoy = RA * sympy.cos(3 * theta) * (r - R_INNER) / (R_OUTER - R_INNER) + return mesh, unit_r, buoy, theta + + +def _outer_rms_vn(solver, unit_r): + vn = solver.u.sym.dot(unit_r) + num = float(uw.maths.BdIntegral(solver.mesh, fn=vn**2, boundary="Upper").evaluate()) + length = float(uw.maths.BdIntegral(solver.mesh, fn=1.0, boundary="Upper").evaluate()) + return np.sqrt(num / length) + + +@pytest.fixture(scope="module") +def solutions(): + mesh, unit_r, buoy, theta = _mesh_and_forcing() + + # --- penalty reference --- + vp = uw.discretisation.MeshVariable("Up", mesh, mesh.dim, degree=2, vtype=uw.VarType.VECTOR) + pp = uw.discretisation.MeshVariable("Pp", mesh, 1, degree=1) + ref = uw.systems.Stokes(mesh, velocityField=vp, pressureField=pp) + ref.constitutive_model = uw.constitutive_models.ViscousFlowModel + ref.constitutive_model.Parameters.shear_viscosity_0 = MU + ref.saddle_preconditioner = 1.0 / MU + ref.bodyforce = buoy * unit_r + ref.add_dirichlet_bc((0.0, 0.0), "Lower") + ref.add_natural_bc(1e6 * MU * unit_r.dot(vp.sym) * unit_r, "Upper") + ref.tolerance = 1e-8 + ref.petsc_options["ksp_type"] = "fgmres" + ref.solve() + + # --- multiplier (augmented Lagrangian) --- + vc = uw.discretisation.MeshVariable("Uc", mesh, mesh.dim, degree=2, vtype=uw.VarType.VECTOR) + pc = uw.discretisation.MeshVariable("Pc", mesh, 1, degree=1) + con = uw.systems.Stokes_Constrained(mesh, velocityField=vc, pressureField=pc) + con.constitutive_model = uw.constitutive_models.ViscousFlowModel + con.constitutive_model.Parameters.shear_viscosity_0 = MU + con.saddle_preconditioner = 1.0 / MU + con.bodyforce = buoy * unit_r + con.add_dirichlet_bc((0.0, 0.0), "Lower") + lam = con.add_constraint_bc("Upper", g=0.0, normal=unit_r) + con.tolerance = 1e-8 + con.petsc_options["ksp_type"] = "fgmres" + con.solve() + + return { + "mesh": mesh, "unit_r": unit_r, "theta": theta, + "ref": ref, "con": con, "lam": lam, + "v_ref": vp.data.copy(), "v_con": vc.data.copy(), + } + + +def test_multiplier_enforces_free_slip(solutions): + """u.n -> 0 on the curved boundary with NO penalty coefficient.""" + rms = _outer_rms_vn(solutions["con"], solutions["unit_r"]) + print(f"multiplier RMS(u.n) on outer = {rms:.3e}") + assert rms < 2.0e-4 + + +def test_multiplier_matches_penalty(solutions): + """Constrained velocity matches the penalty reference field.""" + v_ref, v_con = solutions["v_ref"], solutions["v_con"] + rel = np.sqrt(np.sum((v_ref - v_con) ** 2)) / np.sqrt(np.sum(v_ref**2)) + print(f"relL2(v_multiplier vs v_penalty) = {rel:.3e}") + assert rel < 0.01 + + +def test_multiplier_api(solutions): + """The multiplier and topography are retrievable through the public API.""" + con, lam = solutions["con"], solutions["lam"] + assert con.multiplier("Upper") is lam + assert con.multiplier("Nonexistent") is None + # topography is lambda / (Delta_rho g) + topo_expr = con.topography("Upper", buoyancy_scale=2.0) + assert topo_expr == lam.sym[0] / 2.0 + + +def test_topography_field_is_clean(solutions): + """Multiplier interior is exactly zero; only the boundary trace is non-zero.""" + lam = solutions["lam"] + c = lam.coords + rr = np.sqrt(c[:, 0] ** 2 + c[:, 1] ** 2) + interior = rr < R_OUTER - 0.6 * CELL + boundary = rr > R_OUTER - 0.25 * CELL + assert np.max(np.abs(lam.data[interior, 0])) == 0.0 + assert np.max(np.abs(lam.data[boundary, 0])) > 1.0 + + +def test_topography_matches_dynamic_topography_stress(solutions): + """lambda on the boundary equals the CBF normal stress -n.sigma.n (dyn. topo.).""" + con, lam = solutions["con"], solutions["lam"] + unit_r = solutions["unit_r"] + c = lam.coords + rr = np.sqrt(c[:, 0] ** 2 + c[:, 1] ** 2) + bmask = rr > R_OUTER - 0.25 * CELL + + sigma = con.stress + nsn = (unit_r * sigma * unit_r.T)[0, 0] + nsn_b = np.array(uw.function.evaluate(sympy.Matrix([[nsn]]), c[bmask])).reshape(-1) + lam_b = lam.data[bmask, 0] + + a = lam_b - lam_b.mean() + b = -(nsn_b - nsn_b.mean()) + corr = np.dot(a, b) / np.sqrt(np.dot(a, a) * np.dot(b, b)) + print(f"boundary corr(lambda, -n.sigma.n) = {corr:.4f}") + assert corr > 0.99 From 2216ab21fa579a5b9fa04badb02948d2277ddd63 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Fri, 5 Jun 2026 16:30:16 +0100 Subject: [PATCH 2/7] docs: record P'=[p,h] monolithic fieldsplit feasibility spike Confirmed at the PETSc level that a 3-field DM (u, p, h) with p and h grouped yields a 2-way u | [p,h] Schur split (no nested fieldsplit), de-risking the monolithic "arbitrary saddle-point constraints" direction. Captures the bounded assembly touch-points and remaining caveats for a future block-Schur PR. Underworld development team with AI support from Claude Code --- .../design/CONSTRAINED_FREESLIP_MULTIPLIER.md | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md b/docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md index cff96928e..3b168e074 100644 --- a/docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md +++ b/docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md @@ -194,6 +194,42 @@ evidence): the monolithic 3-field / co-dim-1 representation; 3D spherical shells space needing explicit removal); and live free-surface equilibrium integration (pass `λ` as the target normal-stress end-state). +## Monolithic `P'=[p,h]` fieldsplit — feasibility spike + +A future direction (and a general "inject arbitrary constraints into the saddle +point" capability): rather than a third field forcing a nested 3-way Schur, group +pressure and the multiplier into a composite `P' = [p, h]` and keep a **2-way +`u | P'`** split. + +**Spike result (confirmed):** a 3-field DM `(u, p, h)` on a real mesh, with the `p` +and `h` index sets grouped (`pc_fieldsplit_1_fields 1,2`, or an explicit +concatenated IS), produces exactly a 2-block `u | [p,h]` Schur fieldsplit +(block sizes 84 | 84 = u | (p+h) on a coarse test). The nested-Schur / +KSP-reconfiguration concern is therefore moot — the split structure is identical to +the current `u | p` solver. (`/tmp/spike_pph.py`.) + +**Remaining work (bounded, ~2 weeks of Cython, behind a subclass):** +- Register `h` as field 2 (one `dm.setField`). +- `h`-equation residual: boundary part `∫_Γ ψ(n·u − g)` (the existing Nitsche path + already registers field-1 boundary residuals — same pattern) plus a small interior + screening `ε∫_Ω h ψ` to de-singularise the interior `h` block. +- Three new Jacobian blocks: `uh`, `hu` (boundary integrals — the `UW_PetscDSSetBdJacobian` + machinery already does arbitrary field pairs for `up`/`pu`) and `hh` (interior mass); + `ph`/`hp` are zero. +- Group `[p,h]` in the fieldsplit and extend the Schur PC (`p` keeps its `1/μ` mass PC; + `h` gets the screening diagonal). +- The 3-IS field decomposition / nullspace path (currently assumes 2 fields). + +**Caveats:** the interior screening reintroduces a small `ε` (benign — it does not bias +the boundary multiplier); the work touches the validated `uu/up/pu/pp` assembly, so it +must live behind a subclass and be regression-tested. The one genuine PETSc limitation +remains a *co-dimension-1* `h` (boundary-only DOFs) — avoided here by the full-domain + +screening representation. + +**Recommendation:** the architecture is de-risked, but since the outer loop converges in +2–3 iterations with no user-facing tuning, monolithic is scheduled work (a general +saddle-point-constraint capability), not an urgent replacement. + ## Files - `src/underworld3/systems/solvers.py` — `SNES_Stokes_Constrained`, `_ConstraintBC`. From 3722cebf1431904ba00563b855c0fedffc2542b0 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Fri, 5 Jun 2026 16:44:24 +0100 Subject: [PATCH 3/7] Address Copilot review on #219: guards for label, non-convergence, MPI - add_constraint_bc validates the boundary name and the 'UW_Boundaries' label, and rejects empty/None point sets (petsc_dm_find_labeled_points_local returns the vertex-0 sentinel np.array([0]) when the label is absent, which an `is not None` check would silently accept and corrupt the mask). - The outer loop no longer applies a final, never-solved multiplier update when the iteration cap is hit (u/p stay consistent with lambda) and now warns loudly (RuntimeWarning) on non-convergence. - Raise NotImplementedError when run on more than one MPI rank (the mask and the node-wise update are serial-only), instead of silently producing wrong results. Adds a test that add_constraint_bc rejects an unknown boundary name. Underworld development team with AI support from Claude Code --- src/underworld3/systems/solvers.py | 50 ++++++++++++++++++- .../test_1061_constrained_freeslip_annulus.py | 7 +++ 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index d0f9fe6ca..84873abc1 100644 --- a/src/underworld3/systems/solvers.py +++ b/src/underworld3/systems/solvers.py @@ -2081,6 +2081,20 @@ def add_constraint_bc(self, boundary, g=0.0, normal=None, augmentation=None, The scalar multiplier field. After :meth:`solve`, its boundary trace is the normal traction (topography proxy). """ + # This proof-of-concept is serial only: the boundary mask construction + # and the node-wise multiplier update below are not MPI-decomposed. + if uw.mpi.size > 1: + raise NotImplementedError( + "SNES_Stokes_Constrained is serial-only (the boundary mask and " + "multiplier update are not MPI-safe). Run on a single rank." + ) + + if not hasattr(self.mesh.boundaries, boundary): + raise ValueError( + f"'{boundary}' is not a boundary of this mesh. " + f"Available: {[b.name for b in self.mesh.boundaries]}." + ) + if normal is None: normal = self.mesh.Gamma_P1 @@ -2122,14 +2136,26 @@ def add_constraint_bc(self, boundary, g=0.0, normal=None, augmentation=None, f"_bmarker_{self.instance_number}_{idx}", self.mesh, 1, degree=1, ) marker.data[:] = 0.0 + if not self.mesh.dm.hasLabel("UW_Boundaries"): + raise RuntimeError( + "Mesh has no 'UW_Boundaries' label; cannot build the constraint " + "boundary mask." + ) + # NB: petsc_dm_find_labeled_points_local returns np.array([0]) (vertex 0) + # when the label is absent and None when the value has no points, so an + # `is not None` check alone could silently mark vertex 0. The hasLabel + # guard above plus the explicit empty/None check below close that gap. point_indices = petsc_dm_find_labeled_points_local( self.mesh.dm, "UW_Boundaries", getattr(self.mesh.boundaries, boundary).value, sectionIndex=False, ) - if point_indices is not None: - marker.data[point_indices] = 1.0 + if point_indices is None or len(point_indices) == 0: + raise ValueError( + f"Boundary '{boundary}' has no labelled points on this mesh." + ) + marker.data[point_indices] = 1.0 mask = ( np.array(uw.function.evaluate(marker.sym, lam.coords)).reshape(-1) > 0.75 ) @@ -2234,6 +2260,7 @@ def solve( ).reshape(-1) total_linear_its = 0 + all_converged = False for k in range(constraint_max_iterations): super().solve(zero_init_guess=(zero_init_guess and k == 0), **kwargs) try: @@ -2268,6 +2295,12 @@ def solve( uw.mpi.pprint(f"Constraint loop converged in {k + 1} iterations.") break + # On the final permitted iteration, do NOT apply another multiplier + # update: it would never be solved with, leaving u/p inconsistent + # with lambda. Stop here and warn loudly below instead. + if k == constraint_max_iterations - 1: + break + # Augmented-Lagrangian (ALG2) multiplier update, same r as the # forward-problem penalty augmentation. Monotone-convergent for r>0. # Restricted to boundary nodes so lambda stays a clean topography field. @@ -2275,6 +2308,19 @@ def solve( resid = self._nodal_constraint_residual(cbc) cbc.lam.data[cbc.mask, 0] += cbc.r_nodal[cbc.mask] * resid[cbc.mask] + if not all_converged: + import warnings + + warnings.warn( + f"Constrained Stokes solve did NOT converge: worst " + f"RMS(u.n-g) = {self.constraint_residual:.3e} after " + f"{self.constraint_iterations} iterations " + f"(constraint_max_iterations={constraint_max_iterations}). " + f"Increase constraint_max_iterations or the augmentation.", + RuntimeWarning, + stacklevel=2, + ) + return diff --git a/tests/test_1061_constrained_freeslip_annulus.py b/tests/test_1061_constrained_freeslip_annulus.py index 4fba5316e..34c518693 100644 --- a/tests/test_1061_constrained_freeslip_annulus.py +++ b/tests/test_1061_constrained_freeslip_annulus.py @@ -110,6 +110,13 @@ def test_multiplier_api(solutions): assert topo_expr == lam.sym[0] / 2.0 +def test_constraint_bc_rejects_unknown_boundary(solutions): + """add_constraint_bc validates the boundary name up front.""" + con = solutions["con"] + with pytest.raises(ValueError): + con.add_constraint_bc("Nonexistent") + + def test_topography_field_is_clean(solutions): """Multiplier interior is exactly zero; only the boundary trace is non-zero.""" lam = solutions["lam"] From e3ca4f1355651aa83e9a3a9d694d254f25cec40f Mon Sep 17 00:00:00 2001 From: lmoresi Date: Sun, 7 Jun 2026 10:39:02 +0100 Subject: [PATCH 4/7] Add SNES_Stokes_BlockConstrained: in-saddle multiplier for free-slip + topography Enforces u.n = g on a boundary via a Lagrange multiplier h carried INSIDE the saddle-point system (grouped u | [p,h] 2-way Schur split), in one coupled solve. The converged boundary trace h|_G = -n.sigma.n is dynamic topography directly, with no penalty/Nitsche parameter to tune. - Guarded generalisation of SNES_Stokes_SaddlePt: every multiplier block wrapped `if self._multipliers:` so the 2-field path stays bit-identical (M1 gate). - Augmented-Lagrangian conditioning of the constraint Schur (r = 1e3.mu, viscosity-weighted); combined (p,h) gauge nullspace for enclosed domains. - FMG-ready: field-index fieldsplit grouping (pc_fieldsplit_N_fields) + option mirror, so geometric multigrid works on the velocity block. - Boundary-only multiplier reduction (pin interior h) DEFAULT OFF: refined-safe via DMPlexLabelComplete closure but not yet lossless on the clone DM (degrades boundary trace ~5e-4); the lossless PetscSection-constraint path is next. - add_constraint_bc(boundary, g, normal, screening, augmentation, degree); multiplier()/topography() recovery API. - tests/test_1062: 11 cases (box + annulus; constraint / Dirichlet-match / topography / variable-viscosity to 1e6), all passing. Key finding: on hard constraints (SolCx, 4-wall free-slip + viscosity jump) the block does one mesh-independent solve where the outer-loop Uzawa needs 9-11 and explodes to 33x Dirichlet -- the two methods fail in opposite places (block: per-solve cost from the fat [p,h] Schur; outer-loop: outer-iteration count). Underworld development team with AI support from Claude Code --- .../cython/petsc_generic_snes_solvers.pyx | 409 +++++++++++++++++- src/underworld3/systems/__init__.py | 1 + src/underworld3/systems/solvers.py | 259 +++++++++++ tests/test_1062_block_constrained_freeslip.py | 219 ++++++++++ 4 files changed, 882 insertions(+), 6 deletions(-) create mode 100644 tests/test_1062_block_constrained_freeslip.py diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index d005a3728..7a80cf8ea 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -4144,6 +4144,27 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): self.Unknowns.DuDt = DuDt self.Unknowns.DFDt = DFDt + # Optional saddle-point Lagrange multiplier fields (block-constrained + # Stokes). Each entry is a full-domain scalar MeshVariable registered + # as an extra DM field (id 2, 3, ...), grouped with pressure into the + # Schur split. EMPTY for ordinary Stokes — every multiplier-aware code + # path below is guarded by `if self._multipliers:` so that with no + # multiplier the emitted DS is bit-identical to the 2-field solver. + # _multiplier_screening[k] is the interior screening coefficient + # (eps M de-singularises the interior h block); see + # docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md. + self._multipliers = [] + self._multiplier_screening = [] + self._block_constraint_bcs = [] + # Pin interior (off-boundary) multiplier DOFs to 0 so the solved [p,h] + # block carries only the boundary trace (~√ndof instead of ~ndof/3 DOFs); + # the boundary trace is the only physical part. Default OFF: the current + # DMAddBoundary-based pinning fatally errors on refined meshes (the + # hierarchy pre-builds the local section, so the boundary is added "after + # section creation"). Opt in only on non-refined meshes until the + # order-independent PetscSection-constraint path lands. + self._reduce_interior_multiplier = False + self._degree = degree ## Any problem with U,P, just define our own @@ -4920,12 +4941,51 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): return null_vec + def _build_block_gauge_nullspace_vector(self): + """Combined constant-(pressure, multiplier) gauge mode (block-constrained). + + On a free-slip (non-Dirichlet) boundary delta_u.n != 0, so + int_Gamma n.delta_u = int_Omega div(delta_u) != 0, and a constant + pressure (Bᵀ1) and a constant multiplier (Cᵀ1) couple to the SAME u-row + functional. The genuine near-null mode is therefore the COMBINED + (p = +1, h = -1 everywhere, u = 0): the u-row contribution + (1)·int n.delta_u + (-1)·int n.delta_u cancels, the h-row leaves only + eM·(-1) ~ 0. This replaces the pure constant-pressure mode for the block + solver (which is NOT a null mode when the constraint boundary is free). + """ + template_vec = self.dm.getGlobalVec() + try: + null_vec = template_vec.duplicate() + finally: + self.dm.restoreGlobalVec(template_vec) + + null_vec.set(0.0) + + pressure_is = self._subdict["pressure"][0] + p_sub = null_vec.getSubVector(pressure_is) + p_sub.set(1.0) + null_vec.restoreSubVector(pressure_is, p_sub) + + for cbc in self._block_constraint_bcs: + mult_is = self._subdict[cbc.lam._solver_field_name][0] + m_sub = null_vec.getSubVector(mult_is) + m_sub.set(-1.0) + null_vec.restoreSubVector(mult_is, m_sub) + + return null_vec + def _build_stokes_nullspace(self): """Create the configured coupled Stokes nullspace basis.""" basis_vectors = [] - if self._petsc_use_pressure_nullspace: + if self._block_constraint_bcs: + # Block-constrained free-slip: the gauge mode is the COMBINED + # constant-(pressure, multiplier) vector, not a pure constant + # pressure (see _build_block_gauge_nullspace_vector). + if self._petsc_use_pressure_nullspace: + basis_vectors.append(self._build_block_gauge_nullspace_vector()) + elif self._petsc_use_pressure_nullspace: basis_vectors.append(self._build_pressure_nullspace_vector()) for mode in self._petsc_velocity_nullspace_basis: @@ -4957,10 +5017,47 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): return self._stokes_nullspace + def _setup_block_fieldsplit_options(self): + """Collapse the 3+ field DM to a 2-way velocity | [p,h] Schur split. + + We group by DM FIELD INDEX (pc_fieldsplit_0_fields=0, + pc_fieldsplit_1_fields=1,2,...) rather than by IS: a field-index split + keeps the DM-field association, so geometric multigrid / FMG on the + velocity block can build its interpolation hierarchy (an IS-defined + split has no DM and PCMG errors out with PETSC_ERR_SUP). The grouped + splits are named "0"/"1", so we MIRROR the user-facing + fieldsplit_velocity_* / fieldsplit_pressure_* options (defaults + any + user/FMG overrides) onto fieldsplit_0_* / fieldsplit_1_* — existing + configs and FMG harnesses then apply unchanged. Must run before + setFromOptions. No-op if the user chose a non-fieldsplit pc_type. + """ + opts = self.petsc_options + if str(opts.getAll().get("pc_type", "")) != "fieldsplit": + return # respect a user's direct (lu) solve of the monolithic system + + group1 = ["1"] + [str(cbc.lam._solver_field_id) for cbc in self._block_constraint_bcs] + opts["pc_fieldsplit_0_fields"] = "0" + opts["pc_fieldsplit_1_fields"] = ",".join(group1) + + # Mirror velocity_->0_ and pressure_->1_, but DON'T clobber any + # fieldsplit_0_*/fieldsplit_1_* the user set explicitly (so the grouped + # [p,h] block PC can be overridden directly). + allopts = opts.getAll() + for key, val in list(allopts.items()): + mirrored = None + if key.startswith("fieldsplit_velocity_"): + mirrored = "fieldsplit_0_" + key[len("fieldsplit_velocity_"):] + elif key.startswith("fieldsplit_pressure_"): + mirrored = "fieldsplit_1_" + key[len("fieldsplit_pressure_"):] + if mirrored is not None and mirrored not in allopts: + opts[mirrored] = None if val in (None, "") else val + def _attach_stokes_nullspace(self): """Attach the configured coupled Stokes nullspace to the solver matrices.""" - if not self._petsc_use_pressure_nullspace and not self._petsc_velocity_nullspace_basis: + if (not self._petsc_use_pressure_nullspace + and not self._petsc_velocity_nullspace_basis + and not self._block_constraint_bcs): return pressure_bcs = self._pressure_dirichlet_bcs() @@ -5246,6 +5343,25 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): fns_jacobian.append(self._pp_G0) + ## Lagrange-multiplier rows (block-constrained Stokes). Guarded: no-op + ## for ordinary Stokes. Each multiplier h_k contributes an interior + ## screening residual f0 = eps_k * h_k and a diagonal mass Jacobian + ## hh_G0 = eps_k. The screening de-singularises the otherwise-empty + ## interior h block; with no boundary coupling (M1) h is driven to 0. + ## Boundary coupling (C, C^T) and off-diagonal blocks are added by the + ## natural-bc loop in later milestones. + self._h_F0 = [] + self._hh_G0 = [] + for mvar, eps in zip(self._multipliers, self._multiplier_screening): + h_F0 = sympy.ImmutableDenseMatrix(sympy.Array([eps * mvar.sym[0]]).reshape(1)) + hh_G0 = sympy.ImmutableMatrix( + sympy.derive_by_array(sympy.Array([eps * mvar.sym[0]]), mvar.sym).reshape(1, 1) + ) + self._h_F0.append(h_F0) + self._hh_G0.append(hh_G0) + fns_residual.append(h_F0) + fns_jacobian.append(hh_G0) + # Now natural bcs (compiled into boundary integral terms) # Need to loop on them all ... @@ -5317,6 +5433,59 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): fns_bd_jacobian += [bc.fns["pp_G0"]] + ## Lagrange-multiplier boundary coupling (block-constrained Stokes). + ## For each constraint, the boundary contributes the symmetric pair + ## C^T : field-0 traction fn_f = h * n (u-row, ∫_Γ h (n·δu)) + ## C : field-2 constraint fn_h = n·u − g (h-row, ∫_Γ ψ(n·u−g)) + ## off-diagonal boundary Jacobians uh = ∂fn_f/∂h = n (0, h) and + ## hu = ∂fn_h/∂u = n (h, 0), plus the augmented-Lagrangian uu boundary + ## stiffness uu = ∂fn_f/∂u = r·(n⊗n) (0, 0) which conditions the [p,h] + ## Schur complement (r=0 ⇒ bare KKT, uu=0). Guarded: no-op for ordinary + ## Stokes. + cbc_permutation = (0, 2, 1, 3) + for cbc in self._block_constraint_bcs: + n_row = cbc.normal # sympy 1×dim Matrix + g_sym = cbc.g + r_sym = cbc.augmentation # augmented-Lagrangian parameter + H = sympy.Array(cbc.lam.sym).reshape(1) + hsym = cbc.lam.sym[0] + + u_dot_n = sum(n_row[i] * self.u.sym[i] for i in range(dim)) + + # u-row residual: fn_f = h·n + r(n·u − g)·n + # The r-term is the augmented-Lagrangian penalty: it adds a uu + # boundary stiffness r·(n⊗n) that conditions the Schur complement + # but does NOT bias the multiplier (the h-row stays the exact + # constraint, so h still converges to the true normal traction). + fn_f = sympy.Matrix( + [(hsym + r_sym * (u_dot_n - g_sym)) * n_row[i] for i in range(dim)] + ).as_immutable() + cbc.fns["u_f0"] = sympy.ImmutableDenseMatrix(sympy.Array(fn_f).reshape(dim)) + fns_bd_residual += [cbc.fns["u_f0"]] + + # h-row constraint residual fn_h = n·u − g + fn_h = sympy.Array([u_dot_n - g_sym]).reshape(1) + cbc.fns["h_f0"] = sympy.ImmutableDenseMatrix(fn_h) + fns_bd_residual += [cbc.fns["h_f0"]] + + # uu (0, 0): ∂fn_f/∂u = r·(n⊗n) — AL stiffness (mirror Nitsche shape) + G0 = sympy.derive_by_array(sympy.Array(fn_f), self.Unknowns.u.sym) + cbc.fns["uu_G0"] = sympy.ImmutableMatrix( + sympy.permutedims(G0, cbc_permutation).reshape(dim, dim) + ) + fns_bd_jacobian += [cbc.fns["uu_G0"]] + + # uh (0, h): ∂fn_f/∂h = n — mirror the up_G0 (velocity,scalar) shape + G0 = sympy.derive_by_array(sympy.Array(fn_f), H) + cbc.fns["uh_G0"] = sympy.ImmutableMatrix(G0.reshape(dim)) + fns_bd_jacobian += [cbc.fns["uh_G0"]] + + # hu (h, 0): ∂fn_h/∂u = n — mirror the pu_G0 (scalar,velocity) shape + G0 = sympy.derive_by_array(fn_h, self.Unknowns.u.sym) + cbc.fns["hu_G0"] = sympy.ImmutableMatrix(G0.reshape(dim)) + fns_bd_jacobian += [cbc.fns["hu_G0"]] + + self._fns_bd_residual = fns_bd_residual self._fns_bd_jacobian = fns_bd_jacobian @@ -5335,11 +5504,21 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): print(f"Stokes: Jacobians complete, now compile", flush=True) prim_field_list = [self.u, self.p] + if self._multipliers: + prim_field_list = prim_field_list + list(self._multipliers) + + # Essential-BC value functions. Block-constrained Stokes pins interior + # multiplier DOFs to 0 (boundary-only reduction) — ensure a compiled 0 + # is available for that essential BC. + bc_value_fns = [x.fn for x in self.essential_bcs] + if self._block_constraint_bcs and self._reduce_interior_multiplier: + bc_value_fns.append(sympy.Matrix([[0]]).as_immutable()) + _getext_result = getext( self.mesh, JITCallbackSet( residual=tuple(fns_residual), - bcs=tuple(x.fn for x in self.essential_bcs), + bcs=tuple(bc_value_fns), jacobian=tuple(fns_jacobian), bd_residual=tuple(fns_bd_residual), bd_jacobian=tuple(fns_bd_jacobian), @@ -5432,6 +5611,21 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): self.petsc_fe_p_id = self.dm.getNumFields() self.dm.setField( self.petsc_fe_p_id, self.petsc_fe_p) + # Saddle-point Lagrange multiplier fields (block-constrained Stokes). + # Registered as extra scalar fields (id 2, 3, ...) at the multiplier's + # own degree (velocity P2 by default, so the trace reaches every + # normal-trace DOF). Guarded: no-op for ordinary Stokes. + for mvar in self._multipliers: + h_degree = mvar.degree + h_prefix = "private_{}_{}_".format(self.petsc_options_prefix, mvar._solver_field_name) + options.setValue(h_prefix + "petscspace_degree", h_degree) + options.setValue(h_prefix + "petscdualspace_lagrange_continuity", mvar.continuous) + options.setValue(h_prefix + "petscdualspace_lagrange_node_endpoints", False) + fe_h = PETSc.FE().createDefault(mesh.dim, 1, mesh.isSimplex, mesh.qdegree, h_prefix, PETSc.COMM_SELF) + fe_h.setName(mvar._solver_field_name) + mvar._solver_field_id = self.dm.getNumFields() + self.dm.setField(mvar._solver_field_id, fe_h) + self.dm.createDS() ## This part is done once on the solver dm ... not required every time we update the functions ... @@ -5481,6 +5675,53 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): self.natural_bcs[index] = self.natural_bcs[index]._replace(PETScID=bc, boundary_label_val=value) + # Lagrange-multiplier constraint boundaries (block-constrained Stokes). + # PETSc evaluates a boundary entry's residual/Jacobian ONLY for the + # field it was registered with (key.field = boundary field). So each + # constraint registers the boundary TWICE: once for field 0 (carries the + # u-row traction h·n and the uh Jacobian) and once for the multiplier + # field (carries the h-row constraint n·u−g and the hu Jacobian). + # Guarded: no-op for ordinary Stokes. + cdef int [::1] cbc_comps_view + cdef int [::1] cbc_hcomps_view + for cbc in self._block_constraint_bcs: + cbc_boundary = cbc.boundary + cbc_value = mesh.boundaries[cbc_boundary].value + ind = cbc_value + cbc_fid_h = cbc.lam._solver_field_id + + cbc_comps = np.arange(mesh.dim, dtype=np.int32) + cbc_comps_view = cbc_comps + cbc.petsc_id_u = PetscDSAddBoundary_UW(cdm.dm, + 6, + str(cbc_boundary + "_constraint_u").encode('utf8'), + str("UW_Boundaries").encode('utf8'), + 0, # velocity field: u-row traction + uh Jacobian + cbc_comps.shape[0], + &cbc_comps_view[0], + NULL, + NULL, + 1, + &ind, + NULL, ) + + cbc_hcomps = np.array([0], dtype=np.int32) + cbc_hcomps_view = cbc_hcomps + cbc.petsc_id_h = PetscDSAddBoundary_UW(cdm.dm, + 6, + str(cbc_boundary + "_constraint_h").encode('utf8'), + str("UW_Boundaries").encode('utf8'), + cbc_fid_h, # multiplier field: h-row constraint + hu Jacobian + cbc_hcomps.shape[0], + &cbc_hcomps_view[0], + NULL, + NULL, + 1, + &ind, + NULL, ) + cbc.label_val = cbc_value + + for index,bc in enumerate(self.essential_bcs): if uw.mpi.rank == 0 and self.verbose: print("Setting bc {} ({})".format(index, bc.type)) @@ -5517,6 +5758,73 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): self.essential_bcs[index] = self.essential_bcs[index]._replace(PETScID=bc, boundary_label_val=value) + # Boundary-only multiplier reduction (block-constrained Stokes): pin the + # interior (off-constraint-boundary) multiplier DOFs to 0 via an + # essential BC, so the solved [p,h] Schur block carries only the boundary + # trace (~√ndof rather than ~ndof/3 DOFs). The interior trace is inert + # (only the boundary multiplier is physical), so this is lossless and + # collapses the constraint-Schur cost toward penalty. The constraint + # boundary's UW_Boundaries stratum already lists its boundary points + # (vertices + edges); everything else carrying an h DOF is interior. + cdef int h_one = 1 + cdef int [::1] hcomp_view + if self._block_constraint_bcs and self._reduce_interior_multiplier: + # DMAddBoundary must precede local-section creation, so identify the + # interior points by TOPOLOGY (no section): every non-cell point that + # is NOT on the constraint boundary. PETSc constrains the h field only + # where it actually has DOFs among those points (no-op elsewhere). + h_cellS, h_cellE = self.dm.getHeightStratum(0) # cells + h_chartS, h_chartE = self.dm.getChart() + h_zero_idx = self.ext_dict.ebc[sympy.Matrix([[0]]).as_immutable()] + for cbc in self._block_constraint_bcs: + fid_h = cbc.lam._solver_field_id + bvalue = mesh.boundaries[cbc.boundary].value + bd_is_h = self.dm.getLabel("UW_Boundaries").getStratumIS(bvalue) + # KEEP the FULL boundary closure. The UW_Boundaries stratum lists + # the boundary facets (and some vertices), but a P2 multiplier also + # has DOFs on the facet vertices/edges — including corner vertices + # labelled under adjacent boundaries. Use PETSc's own label + # completion: pure topology, order-independent (does NOT build the + # local section, so it is refined-safe — unlike createClosureIndex), + # and it adds the exact transitive closure, so no boundary-trace h + # DOF is ever mistakenly pinned. + keep_label = "_h_keep_{}".format(fid_h) + if not self.dm.hasLabel(keep_label): + self.dm.createLabel(keep_label) + if bd_is_h is not None: + for _sp in bd_is_h.getIndices().tolist(): + self.dm.setLabelValue(keep_label, _sp, 1) + self.dm.labelComplete(self.dm.getLabel(keep_label)) + _keep_is = self.dm.getStratumIS(keep_label, 1) + bd_pts_h = set(_keep_is.getIndices().tolist()) if _keep_is is not None else set() + ilabel = "_h_interior_{}".format(fid_h) + # Label exists on every level (real on fine, empty on coarse — + # the velocity MG hierarchy never touches the h field, and the + # [p,h] block is solved only on the fine level). + for _d in self.dm_hierarchy: + if not _d.hasLabel(ilabel): + _d.createLabel(ilabel) + for p in range(h_chartS, h_chartE): + if h_cellS <= p < h_cellE: + continue # cells carry no h DOF (P1/P2) + if p in bd_pts_h: + continue # keep the boundary trace + self.dm.setLabelValue(ilabel, p, 1) + hcomp = np.array([0], dtype=np.int32) + hcomp_view = hcomp + PetscDSAddBoundary_UW(cdm.dm, + 5, + (ilabel + "_bc").encode('utf8'), + ilabel.encode('utf8'), + fid_h, 1, + &hcomp_view[0], + ext.fns_bcs[h_zero_idx], + NULL, + 1, + &h_one, + NULL, ) + + for coarse_dm in self.dm_hierarchy: self.dm.copyFields(coarse_dm) self.dm.copyDS(coarse_dm) @@ -5557,6 +5865,15 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): PetscDSSetJacobianPreconditioner(ds.ds, 1, 0, ext.fns_jacobian[i_jac[self._pu_G0]], ext.fns_jacobian[i_jac[self._pu_G1]], NULL, NULL) PetscDSSetJacobianPreconditioner(ds.ds, 1, 1, ext.fns_jacobian[i_jac[self._pp_G0]], NULL, NULL, NULL) + # Lagrange-multiplier rows (block-constrained Stokes). Guarded: no-op + # for ordinary Stokes. Register the interior screening residual and the + # diagonal mass Jacobian/preconditioner on each multiplier's field. + for k, mvar in enumerate(self._multipliers): + fid = mvar._solver_field_id + PetscDSSetResidual(ds.ds, fid, ext.fns_residual[i_res[self._h_F0[k]]], NULL) + PetscDSSetJacobian( ds.ds, fid, fid, ext.fns_jacobian[i_jac[self._hh_G0[k]]], NULL, NULL, NULL) + PetscDSSetJacobianPreconditioner(ds.ds, fid, fid, ext.fns_jacobian[i_jac[self._hh_G0[k]]], NULL, NULL, NULL) + cdef DMLabel c_label for bc in self.natural_bcs: @@ -5693,6 +6010,61 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): ext.fns_bd_jacobian[i_bd_jac[bc.fns["pp_G0"]]], NULL, NULL, NULL) + # Lagrange-multiplier boundary coupling DS registration (block-constrained + # Stokes). Attaches the field-0 traction, field-h constraint, and the + # off-diagonal uh/hu boundary Jacobians to each constraint's boundary + # region. Guarded: no-op for ordinary Stokes. + if self._block_constraint_bcs: + i_bd_res = self.ext_dict.bd_res + i_bd_jac = self.ext_dict.bd_jac + c_label = self.dm.getLabel("UW_Boundaries") + for cbc in self._block_constraint_bcs: + label_val = cbc.label_val + fid_h = cbc.lam._solver_field_id + bid_u = cbc.petsc_id_u # boundary entry registered for field 0 + bid_h = cbc.petsc_id_h # boundary entry registered for field h + + # --- field-0 entry: u-row residual + uu (AL) + uh Jacobians --- + # traction + AL penalty residual fn_f = h n + r(n·u−g) n + UW_PetscDSSetBdResidual(ds.ds, c_label.dmlabel, label_val, bid_u, + 0, 0, + ext.fns_bd_residual[i_bd_res[cbc.fns["u_f0"]]], + NULL) + # uu (0, 0): ∂fn_f/∂u = r·(n⊗n) — augmented-Lagrangian stiffness + UW_PetscDSSetBdJacobian(ds.ds, c_label.dmlabel, label_val, bid_u, + 0, 0, 0, + ext.fns_bd_jacobian[i_bd_jac[cbc.fns["uu_G0"]]], + NULL, NULL, NULL) + UW_PetscDSSetBdJacobianPreconditioner(ds.ds, c_label.dmlabel, label_val, bid_u, + 0, 0, 0, + ext.fns_bd_jacobian[i_bd_jac[cbc.fns["uu_G0"]]], + NULL, NULL, NULL) + # uh (0, h): ∂fn_f/∂h = n + UW_PetscDSSetBdJacobian(ds.ds, c_label.dmlabel, label_val, bid_u, + 0, fid_h, 0, + ext.fns_bd_jacobian[i_bd_jac[cbc.fns["uh_G0"]]], + NULL, NULL, NULL) + UW_PetscDSSetBdJacobianPreconditioner(ds.ds, c_label.dmlabel, label_val, bid_u, + 0, fid_h, 0, + ext.fns_bd_jacobian[i_bd_jac[cbc.fns["uh_G0"]]], + NULL, NULL, NULL) + + # --- field-h entry: h-row constraint + hu Jacobian --- + # constraint residual fn_h = n·u − g + UW_PetscDSSetBdResidual(ds.ds, c_label.dmlabel, label_val, bid_h, + fid_h, 0, + ext.fns_bd_residual[i_bd_res[cbc.fns["h_f0"]]], + NULL) + # hu (h, 0): ∂fn_h/∂u = n + UW_PetscDSSetBdJacobian(ds.ds, c_label.dmlabel, label_val, bid_h, + fid_h, 0, 0, + ext.fns_bd_jacobian[i_bd_jac[cbc.fns["hu_G0"]]], + NULL, NULL, NULL) + UW_PetscDSSetBdJacobianPreconditioner(ds.ds, c_label.dmlabel, label_val, bid_h, + fid_h, 0, 0, + ext.fns_bd_jacobian[i_bd_jac[cbc.fns["hu_G0"]]], + NULL, NULL, NULL) + if verbose: print(f"Weak form (DS)", flush=True) UW_PetscDSViewWF(ds.ds) @@ -5721,6 +6093,12 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): for coarse_dm in self.dm_hierarchy: coarse_dm.createClosureIndex(None) + # Block-constrained: group [pressure, multipliers] into a single + # Schur factor by DM FIELD INDEX, so the velocity block keeps its DM + # hierarchy for geometric MG/FMG. Must precede setFromOptions. + if self._block_constraint_bcs: + self._setup_block_fieldsplit_options() + self.snes = PETSc.SNES().create(PETSc.COMM_WORLD) self.snes.setDM(self.dm) self.snes.setOptionsPrefix(self.petsc_options_prefix) @@ -6024,12 +6402,26 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): pressure_field_num = 1 self._pressure_is = get_local_field_is(local_section, pressure_field_num) - - # Get indices for velocity (complement of pressure) + + # Multiplier fields (block-constrained Stokes). Build a local IS per + # multiplier so its DOFs can be (a) excluded from the velocity + # complement and (b) copied back into the multiplier MeshVariable. + # Guarded: empty dict for ordinary Stokes. + self._multiplier_is = {} + multiplier_indices = set() + for mvar in self._multipliers: + # unconstrained=False -> include ALL multiplier DOFs (the pinned + # interior ones are present in the LOCAL vec at 0), so the copy + # back into the full-domain MeshVariable is size-correct. + mis = get_local_field_is(local_section, mvar._solver_field_id, unconstrained=False) + self._multiplier_is[mvar._solver_field_name] = mis + multiplier_indices |= set(mis.getIndices()) + + # Get indices for velocity (complement of pressure and multipliers) size = clvec.getLocalSize() all_indices = set(range(size)) pressure_indices = set(self._pressure_is.getIndices()) - velocity_indices = sorted(list(all_indices - pressure_indices)) + velocity_indices = sorted(list(all_indices - pressure_indices - multiplier_indices)) self._velocity_is = PETSc.IS().createGeneral(velocity_indices, comm=PETSc.COMM_SELF) # Copy solution back into pressure and velocity variables @@ -6043,6 +6435,11 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): subvec = clvec.getSubVector(self._pressure_is) var.vec.array[:] = subvec.array[:] clvec.restoreSubVector(self._pressure_is, subvec) + elif name in self._multiplier_is: + mis = self._multiplier_is[name] + subvec = clvec.getSubVector(mis) + var.vec.array[:] = subvec.array[:] + clvec.restoreSubVector(mis, subvec) self.mesh._stale_lvec = True # Sync _gvec so downstream consumers (write, stats) see the result diff --git a/src/underworld3/systems/__init__.py b/src/underworld3/systems/__init__.py index ddf22c922..3a8da598d 100644 --- a/src/underworld3/systems/__init__.py +++ b/src/underworld3/systems/__init__.py @@ -49,6 +49,7 @@ from .solvers import SNES_Darcy as SteadyStateDarcy from .solvers import SNES_Stokes as Stokes from .solvers import SNES_Stokes_Constrained as Stokes_Constrained +from .solvers import SNES_Stokes_BlockConstrained as Stokes_BlockConstrained from .solvers import SNES_VE_Stokes as VE_Stokes from .solvers import SNES_Projection as Projection from .solvers import SNES_Vector_Projection as Vector_Projection diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index 84873abc1..0bbcda68d 100644 --- a/src/underworld3/systems/solvers.py +++ b/src/underworld3/systems/solvers.py @@ -2324,6 +2324,265 @@ def solve( return +class _BlockConstraintBC: + """Bookkeeping for one in-saddle-point multiplier constraint. + + Holds the multiplier field ``lam`` (a full-domain scalar MeshVariable + registered as an extra DM field), the prescribed normal velocity ``g``, the + (symbolic, row-vector) constraint normal, and the mutable PETSc bookkeeping + (``petsc_id``/``label_val``) plus the JIT-compiled boundary functions + (``fns``) populated during assembly. Mirrors the natural-BC namedtuple but + stays a mutable object so the Cython assembly can stamp ids onto it. + """ + + __slots__ = ("boundary", "g", "normal", "lam", "augmentation", + "interior_mask", "petsc_id_u", "petsc_id_h", "label_val", "fns") + + def __init__(self, boundary, g, normal, lam, augmentation): + self.boundary = boundary + self.g = g + self.normal = normal + self.lam = lam + # Boolean over the multiplier's nodes: True on INTERIOR nodes (off the + # constraint boundary). The constant-on-interior h mode is a near-null + # mode (interior h appears only in the screening eM), so we hand it to + # the solver as a nullspace vector -- the gauge-fixing analogue of the + # constant-pressure nullspace. Set in add_constraint_bc. + self.interior_mask = None + # Augmented-Lagrangian parameter r: a penalty r(n·u−g)·n added to the + # u-row, giving a uu boundary stiffness r·(n⊗n) that conditions the + # [p,h] Schur complement WITHOUT biasing the multiplier (the h-row is + # still the exact constraint). 0 disables it (bare KKT). + self.augmentation = augmentation + # PETSc needs ONE boundary entry per (boundary, test-field): each entry + # evaluates only its registered field's residual + that field's Jacobian + # row. So a constraint registers the boundary twice — for field 0 (the + # u-row traction h·n and the uh Jacobian) and for the multiplier field + # (the h-row constraint n·u−g and the hu Jacobian). + self.petsc_id_u = -1 + self.petsc_id_h = -1 + self.label_val = -1 + self.fns = {} + + +class SNES_Stokes_BlockConstrained(SNES_Stokes): + r""" + Stokes solver that enforces :math:`\mathbf{u}\cdot\mathbf{n} = g` on a + boundary via a Lagrange multiplier living **inside** the saddle-point + system — a single coupled solve, not an outer loop. + + A scalar multiplier field :math:`h` is added as a third DM field and grouped + with pressure into a 2-way :math:`\mathbf{u}\,|\,[p,h]` Schur split. The + coupled block system is + + .. math:: + + \begin{bmatrix} A & B^{T} & C^{T} \\ B & 0 & 0 \\ C & 0 & \varepsilon M + \end{bmatrix} + \begin{bmatrix} \mathbf{u} \\ p \\ h \end{bmatrix} = + \begin{bmatrix} \mathbf{f} \\ 0 \\ g \end{bmatrix}, + + where :math:`C = \int_\Gamma (\mathbf{n}\cdot\mathbf{v})\,\psi` couples the + multiplier to the boundary normal velocity, and :math:`\varepsilon M = + \varepsilon\int_\Omega h\,\psi` is an interior screening mass that + de-singularises the otherwise-empty interior :math:`h` block (it does not + bias the boundary multiplier). At convergence :math:`h|_\Gamma = -\, + \mathbf{n}\cdot\boldsymbol{\sigma}\cdot\mathbf{n}` is the boundary normal + traction = dynamic topography; access it via :meth:`multiplier` / + :meth:`topography`. + + This is the block (monolithic) counterpart of :class:`SNES_Stokes_Constrained` + (the augmented-Lagrangian outer-loop solver) and should match its answer to + discretisation error in one coupled solve. Serial only (the boundary mask is + not yet MPI-decomposed). See + ``docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md``. + + See Also + -------- + SNES_Stokes_Constrained : The outer-loop counterpart (recovery API template). + SNES_Stokes : The unconstrained saddle-point solver this extends. + """ + + def __init__( + self, + mesh: uw.discretisation.Mesh, + velocityField: Optional[uw.discretisation.MeshVariable] = None, + pressureField: Optional[uw.discretisation.MeshVariable] = None, + degree: Optional[int] = 2, + p_continuous: Optional[bool] = True, + verbose: Optional[bool] = False, + DuDt: Union[SemiLagrangian_DDt, Lagrangian_DDt] = None, + DFDt: Union[SemiLagrangian_DDt, Lagrangian_DDt] = None, + ): + super().__init__( + mesh, + velocityField, + pressureField, + degree, + p_continuous, + verbose, + DuDt=DuDt, + DFDt=DFDt, + ) + + # Block-constrained records (distinct from the outer-loop solver's + # self._constraint_bcs, which the base assembly must not touch). + self._block_constraint_bcs = [] + return + + def _viscosity_scale(self): + """A representative scalar viscosity, for sizing the interior screening.""" + try: + mu = self.constitutive_model.Parameters.shear_viscosity_0 + return float(mu) + except (TypeError, ValueError, AttributeError): + return 1.0 + + def add_constraint_bc(self, boundary, g=0.0, normal=None, screening=None, + augmentation=None, augmentation_base=1.0e3, degree=None): + r"""Register a multiplier-enforced normal-velocity constraint on ``boundary``. + + Adds a scalar multiplier field ``h`` (field id 2, 3, ...) to the + saddle-point system. **Milestone 1**: only the field and its interior + screening are wired (the field is inert — no boundary coupling yet, so + ``h`` is driven to zero and the velocity/pressure solution is identical + to ordinary Stokes). The boundary residual/coupling are added in later + milestones. + + Parameters + ---------- + boundary : str + Mesh boundary label (e.g. ``"Upper"``). + g : float or sympy expression, default 0.0 + Prescribed normal velocity :math:`\mathbf{u}\cdot\mathbf{n} = g`. + normal : sympy matrix, optional + Row-vector constraint normal. Defaults to ``mesh.Gamma_P1``. + screening : float or sympy expression, optional + Interior screening coefficient :math:`\varepsilon` (de-singularises + the interior multiplier DOFs). Defaults to ``1e-6``. + augmentation : float or sympy expression, optional + Augmented-Lagrangian parameter :math:`r`. Adds a penalty + :math:`r(\mathbf{n}\cdot\mathbf{u}-g)\,\mathbf{n}` to the u-row, + giving a ``uu`` boundary stiffness :math:`r\,(\mathbf{n}\otimes + \mathbf{n})` that conditions the :math:`[p,h]` Schur complement + **without biasing the multiplier** (the h-row is still the exact + constraint). Defaults to ``augmentation_base · μ(x)`` (viscosity- + weighted, like the outer-loop solver). Pass ``0`` for the bare KKT + system. + augmentation_base : float, default 1e3 + Base multiple used when ``augmentation`` is not given. + + Returns + ------- + h : MeshVariable + The scalar multiplier field. + """ + # Serial only for now: the boundary mask (later milestones) is not + # MPI-decomposed. + if uw.mpi.size > 1: + raise NotImplementedError( + "SNES_Stokes_BlockConstrained is serial-only for now." + ) + + if not hasattr(self.mesh.boundaries, boundary): + raise ValueError( + f"'{boundary}' is not a boundary of this mesh. " + f"Available: {[b.name for b in self.mesh.boundaries]}." + ) + + if normal is None: + normal = self.mesh.Gamma_P1 + normal = sympy.Matrix(normal) + if normal.shape[0] != 1: + normal = normal.reshape(1, self.mesh.dim) + + if screening is None: + # Small interior screening: de-singularises the interior h block + # without biasing the boundary multiplier. A viscosity-aware scale + # is chosen in a later milestone. + screening = 1.0e-6 + + g = sympy.sympify(g) + + if augmentation is None: + # Viscosity-weighted augmentation r = augmentation_base · μ(x): keeps + # the conditioning ratio uniform across viscosity contrasts (same + # rationale as the outer-loop solver; r ∝ μ is mesh-independent). + try: + viscosity = self.constitutive_model.Parameters.shear_viscosity_0 + augmentation = augmentation_base * viscosity + except (AttributeError, TypeError): + augmentation = augmentation_base * self._viscosity_scale() + augmentation = sympy.sympify(augmentation) + + idx = len(self._block_constraint_bcs) + field_name = f"multiplier_{idx}" + # Multiplier degree defaults to the velocity degree so its trace reaches + # every velocity normal-trace DOF (no constraint floor on the P2 mid-edge + # component). A lower degree trades a little constraint accuracy for far + # fewer multiplier DOFs (cheaper [p,h] Schur block). + h_degree = self._degree if degree is None else degree + h = uw.discretisation.MeshVariable( + f"H{self.instance_number}_{idx}", + self.mesh, + 1, + degree=h_degree, + ) + h.data[:] = 0.0 + h._solver_field_name = field_name + + self._multipliers.append(h) + self._multiplier_screening.append(screening) + self.fields[field_name] = h + # New DM field → the discretisation and solver must be rebuilt. + self.is_setup = False + self._needs_function_rewire = True + + cbc = _BlockConstraintBC(boundary, g, normal, h, augmentation) + + # Interior mask: True on multiplier nodes OFF the constraint boundary. + # Build a P1 marker that is 1 on the boundary vertices and sample it at + # the multiplier's nodes (boundary mid-edge nodes interpolate to ~1). + from underworld3.discretisation.discretisation_mesh import ( + petsc_dm_find_labeled_points_local, + ) + marker = uw.discretisation.MeshVariable( + f"_bmask_{self.instance_number}_{idx}", self.mesh, 1, degree=1, + ) + marker.data[:] = 0.0 + if self.mesh.dm.hasLabel("UW_Boundaries"): + pts = petsc_dm_find_labeled_points_local( + self.mesh.dm, "UW_Boundaries", + getattr(self.mesh.boundaries, boundary).value, sectionIndex=False, + ) + if pts is not None and len(pts) > 0: + marker.data[pts] = 1.0 + on_boundary = np.array(uw.function.evaluate(marker.sym, h.coords)).reshape(-1) > 0.5 + cbc.interior_mask = ~on_boundary + + self._block_constraint_bcs.append(cbc) + return h + + def multiplier(self, boundary): + """Return the multiplier field for ``boundary`` (None if not constrained). + + After :meth:`solve`, the multiplier's boundary trace is the normal + traction holding the constraint. Divide by :math:`\\Delta\\rho\\,g` for + dynamic topography (see :meth:`topography`). + """ + for cbc in self._block_constraint_bcs: + if cbc.boundary == boundary: + return cbc.lam + return None + + def topography(self, boundary, buoyancy_scale=1.0): + r"""Dynamic topography expression :math:`h / (\Delta\rho\, g)` on ``boundary``.""" + lam = self.multiplier(boundary) + if lam is None: + raise ValueError(f"No constraint registered on boundary '{boundary}'.") + return lam.sym[0] / buoyancy_scale + + class SNES_Projection(SNES_Scalar): r""" Scalar projection solver for mapping functions to mesh variables. diff --git a/tests/test_1062_block_constrained_freeslip.py b/tests/test_1062_block_constrained_freeslip.py new file mode 100644 index 000000000..a93219649 --- /dev/null +++ b/tests/test_1062_block_constrained_freeslip.py @@ -0,0 +1,219 @@ +"""Block-constrained free-slip: a Lagrange multiplier INSIDE the saddle point. + +SNES_Stokes_BlockConstrained enforces u.n = g on a boundary via a multiplier h +that lives in the coupled system (grouped u | [p,h] Schur split), in ONE solve. +The converged h is the boundary normal traction = dynamic topography. + +Two regimes: + (A) box, open top -> no pressure nullspace -> direct lu; axis-aligned walls + compared against a Dirichlet free-slip reference (exact). + (B) annulus, enclosed, curved boundary -> pressure nullspace + grouped Schur + via DEFAULT solver options (no per-script PETSc options); compared against + the penalty reference, with topography recovery h ~ -n.sigma.n. + +Run: pixi run python -m pytest tests/test_1062_block_constrained_freeslip.py -v +""" + +import pytest +import numpy as np +import sympy +import underworld3 as uw + +pytestmark = [pytest.mark.level_2, pytest.mark.tier_b] + +MU = 1.0 + + +def _direct(s): + s.petsc_options["snes_type"] = "ksponly" + s.petsc_options["pc_type"] = "lu" + s.petsc_options["pc_factor_mat_solver_type"] = "mumps" + s.petsc_options["pc_use_amat"] = None + s.petsc_options["ksp_type"] = "preonly" + + +# --------------------------------------------------------------------------- # +# (A) open-top box: nullspace-free, direct lu, axis-aligned multiplier walls +# --------------------------------------------------------------------------- # +@pytest.fixture(scope="module") +def box(): + def forcing(m): + xx, yy = m.X + return sympy.Matrix([0.0, sympy.sin(sympy.pi * xx) * sympy.cos(sympy.pi * yy)]) + + m0 = uw.meshing.UnstructuredSimplexBox(minCoords=(0, 0), maxCoords=(1, 1), cellSize=0.15, qdegree=3) + ref = uw.systems.Stokes(m0) + ref.constitutive_model = uw.constitutive_models.ViscousFlowModel + ref.constitutive_model.Parameters.shear_viscosity_0 = MU + ref.bodyforce = forcing(m0) + ref.add_dirichlet_bc((0.0, 0.0), "Bottom") + ref.add_dirichlet_bc((0.0, None), "Left") # Dirichlet free-slip (exact) + ref.add_dirichlet_bc((0.0, None), "Right") + _direct(ref) + ref.solve() + + mb = uw.meshing.UnstructuredSimplexBox(minCoords=(0, 0), maxCoords=(1, 1), cellSize=0.15, qdegree=3) + blk = uw.systems.Stokes_BlockConstrained(mb) + blk.constitutive_model = uw.constitutive_models.ViscousFlowModel + blk.constitutive_model.Parameters.shear_viscosity_0 = MU + blk.bodyforce = forcing(mb) + blk.add_dirichlet_bc((0.0, 0.0), "Bottom") + hL = blk.add_constraint_bc("Left", g=0.0, normal=sympy.Matrix([[-1.0, 0.0]])) + hR = blk.add_constraint_bc("Right", g=0.0, normal=sympy.Matrix([[1.0, 0.0]])) + _direct(blk) + blk.solve() + return dict(ref=ref, blk=blk, hL=hL, hR=hR) + + +def test_box_constraint_enforced(box): + """u.n -> 0 on the multiplier walls (parameter-free, one solve).""" + blk = box["blk"] + c = blk.u.coords + wall = (c[:, 0] < 1e-6) | (c[:, 0] > 1 - 1e-6) + rms = np.sqrt(np.mean(blk.u.data[wall, 0] ** 2)) + assert rms < 1e-7 + + +def test_box_matches_dirichlet_reference(box): + """Block velocity matches the exact Dirichlet free-slip reference.""" + rel = np.sqrt(np.sum((box["ref"].u.data - box["blk"].u.data) ** 2)) / np.sqrt( + np.sum(box["ref"].u.data ** 2) + ) + assert rel < 1e-3 + + +def test_box_topography_is_normal_traction(box): + """h on the wall correlates with -n.sigma.n (the dynamic-topography stress).""" + blk, hR = box["blk"], box["hR"] + sxx = blk.stress[0, 0] + ch = hR.coords + onb = ch[:, 0] > 1 - 1e-6 + nsn = np.array(uw.function.evaluate(sxx, ch[onb])).reshape(-1) + corr = np.corrcoef(hR.data[onb, 0], -nsn)[0, 1] + assert corr > 0.99 + + +def test_multiplier_and_topography_api(box): + blk, hL = box["blk"], box["hL"] + assert blk.multiplier("Left") is hL + assert blk.multiplier("Nonexistent") is None + assert blk.topography("Left", buoyancy_scale=2.0) == hL.sym[0] / 2.0 + + +def test_rejects_unknown_boundary(box): + with pytest.raises(ValueError): + box["blk"].add_constraint_bc("Nonexistent") + + +# --------------------------------------------------------------------------- # +# (B) enclosed annulus, curved boundary, DEFAULT solver options ("just works") +# --------------------------------------------------------------------------- # +@pytest.fixture(scope="module") +def annulus(): + R_I, R_O, CELL, RA = 0.5, 1.0, 0.1, 1.0e2 + mesh = uw.meshing.Annulus(radiusInner=R_I, radiusOuter=R_O, cellSize=CELL, qdegree=3) + x, y = mesh.X + r = sympy.sqrt(x**2 + y**2) + unit_r = sympy.Matrix([[x / r, y / r]]) + th = sympy.atan2(y, x) + buoy = RA * sympy.cos(3 * th) * (r - R_I) / (R_O - R_I) + + ref = uw.systems.Stokes(mesh) + ref.constitutive_model = uw.constitutive_models.ViscousFlowModel + ref.constitutive_model.Parameters.shear_viscosity_0 = MU + ref.saddle_preconditioner = 1.0 / MU + ref.bodyforce = buoy * unit_r + ref.add_dirichlet_bc((0.0, 0.0), "Lower") + ref.add_natural_bc(1e6 * MU * unit_r.dot(ref.u.sym) * unit_r, "Upper") + ref.tolerance = 1e-8 + ref.petsc_options["ksp_type"] = "fgmres" + ref.solve() + + blk = uw.systems.Stokes_BlockConstrained(mesh) + blk.constitutive_model = uw.constitutive_models.ViscousFlowModel + blk.constitutive_model.Parameters.shear_viscosity_0 = MU + blk.saddle_preconditioner = 1.0 / MU + blk.bodyforce = buoy * unit_r + blk.add_dirichlet_bc((0.0, 0.0), "Lower") + hb = blk.add_constraint_bc("Upper", g=0.0, normal=unit_r) + # Enclosed -> pressure/multiplier gauge; DEFAULT grouped-Schur solver config. + blk._petsc_use_pressure_nullspace = True + blk.tolerance = 1e-8 + blk.solve() + return dict(mesh=mesh, unit_r=unit_r, R_O=R_O, CELL=CELL, ref=ref, blk=blk, hb=hb) + + +def test_annulus_two_way_schur(annulus): + """Default config collapses the 3-field DM to a 2-way velocity|[p,h] split.""" + pc = annulus["blk"].snes.getKSP().getPC() + assert len(pc.getFieldSplitSubKSP()) == 2 + + +def test_annulus_constraint_enforced(annulus): + blk, unit_r = annulus["blk"], annulus["unit_r"] + vn = blk.u.sym.dot(unit_r) + num = float(uw.maths.BdIntegral(blk.mesh, fn=vn**2, boundary="Upper").evaluate()) + length = float(uw.maths.BdIntegral(blk.mesh, fn=1.0, boundary="Upper").evaluate()) + assert np.sqrt(num / length) < 2.0e-3 + + +def test_annulus_matches_penalty(annulus): + pts = annulus["ref"].u.coords + vref = np.array(uw.function.evaluate(annulus["ref"].u.sym, pts)) + vblk = np.array(uw.function.evaluate(annulus["blk"].u.sym, pts)) + rel = np.sqrt(np.sum((vref - vblk) ** 2)) / np.sqrt(np.sum(vref**2)) + assert rel < 0.02 + + +def test_annulus_topography_recovery(annulus): + """h|_Gamma correlates with -n.sigma.n: the multiplier IS dynamic topography.""" + blk, unit_r, hb = annulus["blk"], annulus["unit_r"], annulus["hb"] + R_O, CELL = annulus["R_O"], annulus["CELL"] + srr = uw.discretisation.MeshVariable("srr_t1062", blk.mesh, 1, degree=2) + proj = uw.systems.Projection(blk.mesh, srr) + proj.uw_function = (unit_r * blk.stress * unit_r.T)[0, 0] + proj.solve() + ch = hb.coords + onb = np.sqrt(ch[:, 0] ** 2 + ch[:, 1] ** 2) > R_O - 0.01 + corr = np.corrcoef(hb.data[onb, 0], -srr.data[onb, 0])[0, 1] + assert corr > 0.99 + + +# --------------------------------------------------------------------------- # +# (C) variable viscosity (SolCx-style): the AL is viscosity-weighted, so a +# large contrast is the real stress test. Block must match Dirichlet ref. +# --------------------------------------------------------------------------- # +@pytest.mark.parametrize("contrast", [1.0e2, 1.0e6]) +def test_box_variable_viscosity_matches_dirichlet(contrast): + def visc(m): + xx, yy = m.X + return sympy.Piecewise((1.0, xx < 0.5), (contrast, True)) + + def forcing(m): + xx, yy = m.X + return sympy.Matrix([0.0, sympy.sin(sympy.pi * xx) * sympy.cos(sympy.pi * yy)]) + + m0 = uw.meshing.UnstructuredSimplexBox(minCoords=(0, 0), maxCoords=(1, 1), cellSize=0.1, qdegree=3) + ref = uw.systems.Stokes(m0) + ref.constitutive_model = uw.constitutive_models.ViscousFlowModel + ref.constitutive_model.Parameters.shear_viscosity_0 = visc(m0) + ref.bodyforce = forcing(m0) + ref.add_dirichlet_bc((0.0, 0.0), "Bottom") + ref.add_dirichlet_bc((0.0, None), "Left") + ref.add_dirichlet_bc((0.0, None), "Right") + _direct(ref) + ref.solve() + + mb = uw.meshing.UnstructuredSimplexBox(minCoords=(0, 0), maxCoords=(1, 1), cellSize=0.1, qdegree=3) + blk = uw.systems.Stokes_BlockConstrained(mb) + blk.constitutive_model = uw.constitutive_models.ViscousFlowModel + blk.constitutive_model.Parameters.shear_viscosity_0 = visc(mb) + blk.bodyforce = forcing(mb) + blk.add_dirichlet_bc((0.0, 0.0), "Bottom") + blk.add_constraint_bc("Left", g=0.0, normal=sympy.Matrix([[-1.0, 0.0]])) + blk.add_constraint_bc("Right", g=0.0, normal=sympy.Matrix([[1.0, 0.0]])) + _direct(blk) + blk.solve() + + rel = np.sqrt(np.sum((ref.u.data - blk.u.data) ** 2)) / np.sqrt(np.sum(ref.u.data ** 2)) + assert rel < 1e-4 From b2e5668c83f66908b0be1838461483018d3ecfef Mon Sep 17 00:00:00 2001 From: lmoresi Date: Mon, 8 Jun 2026 20:27:26 +1000 Subject: [PATCH 5/7] Lossless boundary-only multiplier reduction (default on) Constrain the interior (off-boundary) multiplier DOFs directly in the fine local PetscSection rather than via DMAddBoundary. This is both lossless (the constraint-boundary closure is correct because createClosureIndex has finalised the section) and refined/FMG-safe (no "after section creation" restriction), so the solved [p,h] Schur block carries only the boundary trace (~sqrt(ndof) DOFs). Collapses the multiplier-block DOF premium from ~3x to ~1.1x Dirichlet with no accuracy loss (box+lu reduced-vs-full velocity relL2 ~1e-9). Underworld development team with AI support from Claude Code --- .../cython/petsc_generic_snes_solvers.pyx | 228 ++++++++++++------ ...p.py => test_1061_constrained_freeslip.py} | 0 .../test_1061_constrained_freeslip_annulus.py | 148 ------------ 3 files changed, 151 insertions(+), 225 deletions(-) rename tests/{test_1062_block_constrained_freeslip.py => test_1061_constrained_freeslip.py} (100%) delete mode 100644 tests/test_1061_constrained_freeslip_annulus.py diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index 7a80cf8ea..73257c4ec 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -4158,12 +4158,12 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): self._block_constraint_bcs = [] # Pin interior (off-boundary) multiplier DOFs to 0 so the solved [p,h] # block carries only the boundary trace (~√ndof instead of ~ndof/3 DOFs); - # the boundary trace is the only physical part. Default OFF: the current - # DMAddBoundary-based pinning fatally errors on refined meshes (the - # hierarchy pre-builds the local section, so the boundary is added "after - # section creation"). Opt in only on non-refined meshes until the - # order-independent PetscSection-constraint path lands. - self._reduce_interior_multiplier = False + # the boundary trace is the only physical part. Done by constraining the + # interior h DOFs directly in the fine local PetscSection + # (_constrain_interior_multipliers_in_section) — lossless (correct + # constraint-boundary closure) and refined/FMG-safe (no DMAddBoundary + # ordering restriction). Default ON. + self._reduce_interior_multiplier = True self._degree = degree @@ -5507,12 +5507,11 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): if self._multipliers: prim_field_list = prim_field_list + list(self._multipliers) - # Essential-BC value functions. Block-constrained Stokes pins interior - # multiplier DOFs to 0 (boundary-only reduction) — ensure a compiled 0 - # is available for that essential BC. + # Essential-BC value functions. (Block-constrained Stokes reduces the + # interior multiplier DOFs by constraining them directly in the + # PetscSection — see _constrain_interior_multipliers_in_section — so no + # extra compiled essential-BC value is needed here.) bc_value_fns = [x.fn for x in self.essential_bcs] - if self._block_constraint_bcs and self._reduce_interior_multiplier: - bc_value_fns.append(sympy.Matrix([[0]]).as_immutable()) _getext_result = getext( self.mesh, @@ -5758,72 +5757,13 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): self.essential_bcs[index] = self.essential_bcs[index]._replace(PETScID=bc, boundary_label_val=value) - # Boundary-only multiplier reduction (block-constrained Stokes): pin the - # interior (off-constraint-boundary) multiplier DOFs to 0 via an - # essential BC, so the solved [p,h] Schur block carries only the boundary - # trace (~√ndof rather than ~ndof/3 DOFs). The interior trace is inert - # (only the boundary multiplier is physical), so this is lossless and - # collapses the constraint-Schur cost toward penalty. The constraint - # boundary's UW_Boundaries stratum already lists its boundary points - # (vertices + edges); everything else carrying an h DOF is interior. - cdef int h_one = 1 - cdef int [::1] hcomp_view - if self._block_constraint_bcs and self._reduce_interior_multiplier: - # DMAddBoundary must precede local-section creation, so identify the - # interior points by TOPOLOGY (no section): every non-cell point that - # is NOT on the constraint boundary. PETSc constrains the h field only - # where it actually has DOFs among those points (no-op elsewhere). - h_cellS, h_cellE = self.dm.getHeightStratum(0) # cells - h_chartS, h_chartE = self.dm.getChart() - h_zero_idx = self.ext_dict.ebc[sympy.Matrix([[0]]).as_immutable()] - for cbc in self._block_constraint_bcs: - fid_h = cbc.lam._solver_field_id - bvalue = mesh.boundaries[cbc.boundary].value - bd_is_h = self.dm.getLabel("UW_Boundaries").getStratumIS(bvalue) - # KEEP the FULL boundary closure. The UW_Boundaries stratum lists - # the boundary facets (and some vertices), but a P2 multiplier also - # has DOFs on the facet vertices/edges — including corner vertices - # labelled under adjacent boundaries. Use PETSc's own label - # completion: pure topology, order-independent (does NOT build the - # local section, so it is refined-safe — unlike createClosureIndex), - # and it adds the exact transitive closure, so no boundary-trace h - # DOF is ever mistakenly pinned. - keep_label = "_h_keep_{}".format(fid_h) - if not self.dm.hasLabel(keep_label): - self.dm.createLabel(keep_label) - if bd_is_h is not None: - for _sp in bd_is_h.getIndices().tolist(): - self.dm.setLabelValue(keep_label, _sp, 1) - self.dm.labelComplete(self.dm.getLabel(keep_label)) - _keep_is = self.dm.getStratumIS(keep_label, 1) - bd_pts_h = set(_keep_is.getIndices().tolist()) if _keep_is is not None else set() - ilabel = "_h_interior_{}".format(fid_h) - # Label exists on every level (real on fine, empty on coarse — - # the velocity MG hierarchy never touches the h field, and the - # [p,h] block is solved only on the fine level). - for _d in self.dm_hierarchy: - if not _d.hasLabel(ilabel): - _d.createLabel(ilabel) - for p in range(h_chartS, h_chartE): - if h_cellS <= p < h_cellE: - continue # cells carry no h DOF (P1/P2) - if p in bd_pts_h: - continue # keep the boundary trace - self.dm.setLabelValue(ilabel, p, 1) - hcomp = np.array([0], dtype=np.int32) - hcomp_view = hcomp - PetscDSAddBoundary_UW(cdm.dm, - 5, - (ilabel + "_bc").encode('utf8'), - ilabel.encode('utf8'), - fid_h, 1, - &hcomp_view[0], - ext.fns_bcs[h_zero_idx], - NULL, - 1, - &h_one, - NULL, ) - + # Boundary-only multiplier reduction (block-constrained Stokes) is applied + # LATER, in _setup_solver, by constraining the interior (off-constraint- + # boundary) multiplier DOFs directly in the PetscSection + # (_constrain_interior_multipliers_in_section). That path is both lossless + # (the constraint-boundary closure is correct because createClosureIndex + # has finalised the section) and refined-safe (no DMAddBoundary, which + # cannot follow section creation). See that method for details. for coarse_dm in self.dm_hierarchy: self.dm.copyFields(coarse_dm) @@ -5834,6 +5774,134 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): return + def _constrain_interior_multipliers_in_section(self): + """Lossless, refined-safe boundary-only multiplier reduction. + + Block-constrained Stokes carries one Lagrange multiplier (field ``h``) + per constraint as a full-domain field, but only its boundary trace is + physical — the interior ``h`` DOFs are inert and merely inflate the + ``[p,h]`` Schur block (~ndof/3 extra DOFs). Here we constrain every + interior (off-constraint-boundary) ``h`` DOF DIRECTLY in the fine DM's + local PetscSection, so they drop out of the GLOBAL system (smaller Schur + block) while remaining present in the LOCAL vector at 0 (scatter-back + stays size-correct, see below). + + This must run AFTER ``createClosureIndex`` (so the local section is + finalised with the velocity Dirichlet constraints baked in AND the + constraint-boundary closure is complete — the source of the earlier + DMAddBoundary path's loss) and BEFORE any consumer of the GLOBAL section + (the field-index fieldsplit grouping, the SNES, ``createFieldDecomposition``, + the nullspace). Constraining in the section has no "after section creation" + restriction, so it is also refined/FMG-safe — unlike the DMAddBoundary + essential-field pin it replaces. + + Fine ``self.dm`` only: the coarse MG levels never carry an active ``h`` + field in the solve, and ``copyFields``/``copyDS`` copy discretisations and + the DS (weak forms), not the section, so the velocity MG/FMG hierarchy and + the field-index fieldsplit grouping are undisturbed (they only see the + ``[p,h]`` block shrink). + + Note (scatter-back): PetscSection constraints remove DOFs from the GLOBAL + section only; the LOCAL section still allocates them. The multiplier IS is + built with ``unconstrained=False`` and (for a scalar ``h``, 1 dof/point) + appends the local offset of every h-bearing point regardless of + constraint, so the copy back into the full-domain ``h`` MeshVariable is + size-correct without change. + """ + from petsc4py import PETSc + import numpy as np + + if not (self._block_constraint_bcs and self._reduce_interior_multiplier): + return + + dm = self.dm + Sold = dm.getLocalSection() + cS, cE = Sold.getChart() + nF = Sold.getNumFields() + + # 1. Boundary-trace KEEP set + interior-h set, per multiplier field. + # createClosureIndex has run, so getTransitiveClosure gives the + # COMPLETE closure — no boundary-trace h DOF is ever pinned (lossless). + interior_pts = {} # fid_h -> set(points to additionally constrain) + for cbc in self._block_constraint_bcs: + fid_h = cbc.lam._solver_field_id + bvalue = self.mesh.boundaries[cbc.boundary].value + bd_is = dm.getLabel("UW_Boundaries").getStratumIS(bvalue) + keep = set() + if bd_is is not None: + for bp in bd_is.getIndices().tolist(): + keep.update(dm.getTransitiveClosure(bp)[0].tolist()) + iset = interior_pts.setdefault(fid_h, set()) + for p in range(cS, cE): + if Sold.getFieldDof(p, fid_h) > 0 and p not in keep: + iset.add(p) + + # Nothing to constrain (e.g. boundary-only multiplier discretisation). + if not any(interior_pts.values()): + return + + # 2. Build a new local section mirroring the old (chart, fields, dofs and + # ALL existing constraints) plus the new interior-h constraints. + Snew = PETSc.Section().create(comm=dm.getComm()) + Snew.setNumFields(nF) + Snew.setChart(cS, cE) + for f in range(nF): + Snew.setFieldComponents(f, Sold.getFieldComponents(f)) + Snew.setFieldName(f, Sold.getFieldName(f)) + + # 2a. DOFs and constraint DOF COUNTS — must precede setUp(). + # extra_field[(p, f)] = field-local index list of the ADDED constraints + extra_field = {} + for p in range(cS, cE): + Snew.setDof(p, Sold.getDof(p)) + addl_total = 0 + for f in range(nF): + fdof = Sold.getFieldDof(p, f) + Snew.setFieldDof(p, f, fdof) + fc = Sold.getFieldConstraintDof(p, f) + add_here = 0 + if p in interior_pts.get(f, ()): + # Constrain ALL (currently unconstrained) DOFs of this field + # at this point — for a scalar h that is the single dof 0. + old_ind = Sold.getFieldConstraintIndices(p, f) + old_set = set(old_ind.tolist()) if old_ind is not None else set() + new_idx = [i for i in range(fdof) if i not in old_set] + if new_idx: + extra_field[(p, f)] = new_idx + add_here = len(new_idx) + Snew.setFieldConstraintDof(p, f, fc + add_here) + addl_total += add_here + Snew.setConstraintDof(p, Sold.getConstraintDof(p) + addl_total) + + Snew.setUp() + + # 2b. Constraint INDICES — must follow setUp(). Rebuild the point-local + # aggregate from the field views so the cross-field offset convention + # stays internally consistent (offset of field f = sum of lower fdof). + for p in range(cS, cE): + point_local = [] + field_offset = 0 + for f in range(nF): + fdof = Sold.getFieldDof(p, f) + old_ind = Sold.getFieldConstraintIndices(p, f) + fidx = set(old_ind.tolist()) if old_ind is not None else set() + if (p, f) in extra_field: + fidx |= set(extra_field[(p, f)]) + if fidx: + fidx = sorted(fidx) + Snew.setFieldConstraintIndices(p, f, np.array(fidx, dtype=np.int32)) + point_local.extend(field_offset + i for i in fidx) + field_offset += fdof + if point_local: + Snew.setConstraintIndices(p, np.array(sorted(point_local), dtype=np.int32)) + + dm.setLocalSection(Snew) + # Force the global-section rebuild (and fail fast if malformed); the + # constrained interior-h DOFs are now excluded from the global system. + dm.getGlobalSection() + + return + @timing.routine_timer_decorator def _setup_solver(self, verbose=False, _rewire_only=False): @@ -6093,6 +6161,12 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): for coarse_dm in self.dm_hierarchy: coarse_dm.createClosureIndex(None) + # Boundary-only multiplier reduction: constrain the interior h DOFs + # directly in the (now finalised) fine local section. Lossless and + # refined-safe; must precede the fieldsplit grouping / SNES / field + # decomposition so they see the shrunken [p,h] global block. + self._constrain_interior_multipliers_in_section() + # Block-constrained: group [pressure, multipliers] into a single # Schur factor by DM FIELD INDEX, so the velocity block keeps its DM # hierarchy for geometric MG/FMG. Must precede setFromOptions. diff --git a/tests/test_1062_block_constrained_freeslip.py b/tests/test_1061_constrained_freeslip.py similarity index 100% rename from tests/test_1062_block_constrained_freeslip.py rename to tests/test_1061_constrained_freeslip.py diff --git a/tests/test_1061_constrained_freeslip_annulus.py b/tests/test_1061_constrained_freeslip_annulus.py deleted file mode 100644 index 34c518693..000000000 --- a/tests/test_1061_constrained_freeslip_annulus.py +++ /dev/null @@ -1,148 +0,0 @@ -"""Constrained free-slip on an annulus via a recoverable Lagrange multiplier. - -Buoyancy-driven flow in an annulus with no-slip inner boundary and free-slip -outer boundary. The free-slip outer condition (u.n = 0) is enforced three ways -and compared: - - - penalty : the existing fragile penalty natural BC (reference) - - multiplier: SNES_Stokes_Constrained, augmented-Lagrangian multiplier - -The multiplier solver must (a) drive u.n -> 0 with NO penalty coefficient, -(b) match the penalty velocity field, and (c) yield a clean, recoverable -topography field whose boundary trace equals the consistent-boundary-flux -normal stress (-n.sigma.n), i.e. dynamic topography. - -Run with: pixi run python -m pytest tests/test_1061_constrained_freeslip_annulus.py -v -""" - -import pytest -import numpy as np -import sympy -import underworld3 as uw - -pytestmark = [pytest.mark.level_2, pytest.mark.tier_b] - -R_INNER, R_OUTER = 0.5, 1.0 -CELL = 0.1 -MU = 1.0 -RA = 1.0e2 - - -def _mesh_and_forcing(): - mesh = uw.meshing.Annulus(radiusInner=R_INNER, radiusOuter=R_OUTER, - cellSize=CELL, qdegree=3) - x, y = mesh.X - r = sympy.sqrt(x**2 + y**2) - unit_r = sympy.Matrix([[x / r, y / r]]) - theta = sympy.atan2(y, x) - buoy = RA * sympy.cos(3 * theta) * (r - R_INNER) / (R_OUTER - R_INNER) - return mesh, unit_r, buoy, theta - - -def _outer_rms_vn(solver, unit_r): - vn = solver.u.sym.dot(unit_r) - num = float(uw.maths.BdIntegral(solver.mesh, fn=vn**2, boundary="Upper").evaluate()) - length = float(uw.maths.BdIntegral(solver.mesh, fn=1.0, boundary="Upper").evaluate()) - return np.sqrt(num / length) - - -@pytest.fixture(scope="module") -def solutions(): - mesh, unit_r, buoy, theta = _mesh_and_forcing() - - # --- penalty reference --- - vp = uw.discretisation.MeshVariable("Up", mesh, mesh.dim, degree=2, vtype=uw.VarType.VECTOR) - pp = uw.discretisation.MeshVariable("Pp", mesh, 1, degree=1) - ref = uw.systems.Stokes(mesh, velocityField=vp, pressureField=pp) - ref.constitutive_model = uw.constitutive_models.ViscousFlowModel - ref.constitutive_model.Parameters.shear_viscosity_0 = MU - ref.saddle_preconditioner = 1.0 / MU - ref.bodyforce = buoy * unit_r - ref.add_dirichlet_bc((0.0, 0.0), "Lower") - ref.add_natural_bc(1e6 * MU * unit_r.dot(vp.sym) * unit_r, "Upper") - ref.tolerance = 1e-8 - ref.petsc_options["ksp_type"] = "fgmres" - ref.solve() - - # --- multiplier (augmented Lagrangian) --- - vc = uw.discretisation.MeshVariable("Uc", mesh, mesh.dim, degree=2, vtype=uw.VarType.VECTOR) - pc = uw.discretisation.MeshVariable("Pc", mesh, 1, degree=1) - con = uw.systems.Stokes_Constrained(mesh, velocityField=vc, pressureField=pc) - con.constitutive_model = uw.constitutive_models.ViscousFlowModel - con.constitutive_model.Parameters.shear_viscosity_0 = MU - con.saddle_preconditioner = 1.0 / MU - con.bodyforce = buoy * unit_r - con.add_dirichlet_bc((0.0, 0.0), "Lower") - lam = con.add_constraint_bc("Upper", g=0.0, normal=unit_r) - con.tolerance = 1e-8 - con.petsc_options["ksp_type"] = "fgmres" - con.solve() - - return { - "mesh": mesh, "unit_r": unit_r, "theta": theta, - "ref": ref, "con": con, "lam": lam, - "v_ref": vp.data.copy(), "v_con": vc.data.copy(), - } - - -def test_multiplier_enforces_free_slip(solutions): - """u.n -> 0 on the curved boundary with NO penalty coefficient.""" - rms = _outer_rms_vn(solutions["con"], solutions["unit_r"]) - print(f"multiplier RMS(u.n) on outer = {rms:.3e}") - assert rms < 2.0e-4 - - -def test_multiplier_matches_penalty(solutions): - """Constrained velocity matches the penalty reference field.""" - v_ref, v_con = solutions["v_ref"], solutions["v_con"] - rel = np.sqrt(np.sum((v_ref - v_con) ** 2)) / np.sqrt(np.sum(v_ref**2)) - print(f"relL2(v_multiplier vs v_penalty) = {rel:.3e}") - assert rel < 0.01 - - -def test_multiplier_api(solutions): - """The multiplier and topography are retrievable through the public API.""" - con, lam = solutions["con"], solutions["lam"] - assert con.multiplier("Upper") is lam - assert con.multiplier("Nonexistent") is None - # topography is lambda / (Delta_rho g) - topo_expr = con.topography("Upper", buoyancy_scale=2.0) - assert topo_expr == lam.sym[0] / 2.0 - - -def test_constraint_bc_rejects_unknown_boundary(solutions): - """add_constraint_bc validates the boundary name up front.""" - con = solutions["con"] - with pytest.raises(ValueError): - con.add_constraint_bc("Nonexistent") - - -def test_topography_field_is_clean(solutions): - """Multiplier interior is exactly zero; only the boundary trace is non-zero.""" - lam = solutions["lam"] - c = lam.coords - rr = np.sqrt(c[:, 0] ** 2 + c[:, 1] ** 2) - interior = rr < R_OUTER - 0.6 * CELL - boundary = rr > R_OUTER - 0.25 * CELL - assert np.max(np.abs(lam.data[interior, 0])) == 0.0 - assert np.max(np.abs(lam.data[boundary, 0])) > 1.0 - - -def test_topography_matches_dynamic_topography_stress(solutions): - """lambda on the boundary equals the CBF normal stress -n.sigma.n (dyn. topo.).""" - con, lam = solutions["con"], solutions["lam"] - unit_r = solutions["unit_r"] - c = lam.coords - rr = np.sqrt(c[:, 0] ** 2 + c[:, 1] ** 2) - bmask = rr > R_OUTER - 0.25 * CELL - - sigma = con.stress - nsn = (unit_r * sigma * unit_r.T)[0, 0] - nsn_b = np.array(uw.function.evaluate(sympy.Matrix([[nsn]]), c[bmask])).reshape(-1) - lam_b = lam.data[bmask, 0] - - a = lam_b - lam_b.mean() - b = -(nsn_b - nsn_b.mean()) - corr = np.dot(a, b) / np.sqrt(np.dot(a, a) * np.dot(b, b)) - print(f"boundary corr(lambda, -n.sigma.n) = {corr:.4f}") - assert corr > 0.99 From 7871b7df82c8717ffa17326b8d9e28c7b3085226 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Mon, 8 Jun 2026 20:28:32 +1000 Subject: [PATCH 6/7] Promote in-saddle solver to Stokes_Constrained; remove outer-loop variant The in-saddle (single coupled solve) constrained free-slip solver is now the production path, exported as uw.systems.Stokes_Constrained. It is validated against the exact SolCx analytic solution and matches a Dirichlet free-slip reference to discretisation error, with the constraint enforced to machine precision. - rename SNES_Stokes_BlockConstrained -> SNES_Stokes_Constrained - remove the augmented-Lagrangian outer-loop solver + helper (_ConstraintBC); the Uzawa scheme is easy to reproduce in Python if wanted - single public export Stokes_Constrained (drop Stokes_BlockConstrained) - raise default augmentation_base 1e3 -> 1e4 (accuracy is r-independent; larger sits in the low-iteration plateau, well below roundoff) - tests: retarget the constrained test, add test_1062_constrained_solcx.py (free-slip vs the SolCx analytic); outer-loop annulus test removed - docs: update CONSTRAINED_FREESLIP_MULTIPLIER.md status Underworld development team with AI support from Claude Code --- .../design/CONSTRAINED_FREESLIP_MULTIPLIER.md | 9 +- src/underworld3/systems/__init__.py | 1 - src/underworld3/systems/solvers.py | 427 +----------------- tests/test_1061_constrained_freeslip.py | 10 +- tests/test_1062_constrained_solcx.py | 51 +++ 5 files changed, 83 insertions(+), 415 deletions(-) create mode 100644 tests/test_1062_constrained_solcx.py diff --git a/docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md b/docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md index 3b168e074..d7a04805e 100644 --- a/docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md +++ b/docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md @@ -1,7 +1,12 @@ # Constrained free-slip via a recoverable Lagrange multiplier (dynamic topography) -**Status**: proof-of-concept (Phase 0 + Phase 1), serial. Branch -`feature/constrained-freeslip-topography`. +**Status**: shipped as `uw.systems.Stokes_Constrained` (serial). The constraint +is enforced by a multiplier carried **inside** the saddle point (one coupled +solve); the converged boundary multiplier is the normal traction = dynamic +topography. An earlier augmented-Lagrangian **outer-loop** variant was removed in +favour of this in-saddle formulation (it is straightforward to reproduce in +Python if needed). Validated against the exact SolCx analytic solution +(`tests/test_1062_constrained_solcx.py`). ## Motivation diff --git a/src/underworld3/systems/__init__.py b/src/underworld3/systems/__init__.py index 3a8da598d..ddf22c922 100644 --- a/src/underworld3/systems/__init__.py +++ b/src/underworld3/systems/__init__.py @@ -49,7 +49,6 @@ from .solvers import SNES_Darcy as SteadyStateDarcy from .solvers import SNES_Stokes as Stokes from .solvers import SNES_Stokes_Constrained as Stokes_Constrained -from .solvers import SNES_Stokes_BlockConstrained as Stokes_BlockConstrained from .solvers import SNES_VE_Stokes as VE_Stokes from .solvers import SNES_Projection as Projection from .solvers import SNES_Vector_Projection as Vector_Projection diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index 0bbcda68d..1a54275b6 100644 --- a/src/underworld3/systems/solvers.py +++ b/src/underworld3/systems/solvers.py @@ -1936,394 +1936,6 @@ def delta_t(self): return self.constitutive_model.Parameters.dt_elastic -class _ConstraintBC: - """Bookkeeping for one multiplier-enforced boundary constraint. - - Holds the multiplier field ``lam``, the prescribed normal velocity ``g``, - the (symbolic, row-vector) constraint normal, and the augmented-Lagrangian - parameter ``r`` (which is simultaneously the forward-problem penalty weight - and the multiplier-update step). - """ - - __slots__ = ("boundary", "g", "normal", "lam", "augmentation", "mask", "r_nodal") - - def __init__(self, boundary, g, normal, lam, augmentation, mask): - self.boundary = boundary - self.g = g - self.normal = normal - self.lam = lam - # augmentation r may be a scalar or a spatial sympy expression (e.g. - # viscosity-weighted). r_nodal holds it sampled at the multiplier nodes - # for the dual update; it is (re)computed at solve time. - self.augmentation = augmentation - self.mask = mask - self.r_nodal = None - - -class SNES_Stokes_Constrained(SNES_Stokes): - r""" - Stokes solver with boundary constraints enforced by a recoverable Lagrange - multiplier instead of a penalty. - - For each constraint boundary :math:`\Gamma` the no-normal-flow (free-slip) - or prescribed-normal-velocity condition - - .. math:: - - \mathbf{u} \cdot \mathbf{n} = g \quad \text{on } \Gamma - - is enforced by introducing a scalar multiplier field :math:`\lambda` and an - **augmented-Lagrangian** (Uzawa / ALG2) outer loop. Each Stokes solve carries - a natural-BC traction with both the multiplier and a penalty augmentation, - - .. math:: - - \mathbf{t} = \bigl[\lambda + r\,(\mathbf{u}\cdot\mathbf{n} - g)\bigr]\,\mathbf{n} - \quad \text{on } \Gamma , - - and the multiplier is updated with the same augmentation parameter, - - .. math:: - - \lambda \leftarrow \lambda + r\,(\mathbf{u}\cdot\mathbf{n} - g) . - - The penalty term :math:`r\,(\mathbf{u}\cdot\mathbf{n})\,\mathbf{n}` is exactly - the existing penalty free-slip BC: it preconditions every boundary mode - uniformly so the outer loop converges in a handful of iterations, while the - multiplier removes the penalty's accuracy bias — so a **moderate**, - well-conditioned :math:`r` gives both fast convergence and an *exact* - constraint (unlike a pure penalty, which must be made large and fragile). - - At convergence :math:`\lambda` is the normal traction holding the boundary, - giving a direct estimate of dynamic surface topography, - :math:`h = \lambda / (\Delta\rho\, g)`. Access it via :meth:`multiplier`. - - Notes - ----- - The multiplier is represented as an ordinary full-mesh scalar field of the - same degree as the velocity. Only its trace on :math:`\Gamma` enters the - weak form (interior values are inert), so no boundary trace space is - required. This is a proof-of-concept (serial; one full Stokes solve per - outer iteration). See the design note - ``docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md``. - - See Also - -------- - SNES_Stokes : The unconstrained saddle-point solver this extends. - """ - - def __init__( - self, - mesh: uw.discretisation.Mesh, - velocityField: Optional[uw.discretisation.MeshVariable] = None, - pressureField: Optional[uw.discretisation.MeshVariable] = None, - degree: Optional[int] = 2, - p_continuous: Optional[bool] = True, - verbose: Optional[bool] = False, - DuDt: Union[SemiLagrangian_DDt, Lagrangian_DDt] = None, - DFDt: Union[SemiLagrangian_DDt, Lagrangian_DDt] = None, - ): - super().__init__( - mesh, - velocityField, - pressureField, - degree, - p_continuous, - verbose, - DuDt=DuDt, - DFDt=DFDt, - ) - - self._constraint_bcs = [] - # Diagnostics from the most recent constrained solve. - self.constraint_iterations = 0 - self.constraint_residual = None - self.constraint_total_linear_its = 0 - return - - def _viscosity_scale(self): - """A representative scalar viscosity, for sizing the initial Uzawa step.""" - try: - mu = self.constitutive_model.Parameters.shear_viscosity_0 - return float(mu) - except (TypeError, ValueError, AttributeError): - return 1.0 - - def add_constraint_bc(self, boundary, g=0.0, normal=None, augmentation=None, - augmentation_base=1.0e3): - r"""Register a multiplier-enforced normal-velocity constraint on ``boundary``. - - Parameters - ---------- - boundary : str - Mesh boundary label (e.g. ``"Upper"``). - g : float or sympy expression, default 0.0 - Prescribed normal velocity :math:`\mathbf{u}\cdot\mathbf{n} = g`. - The default ``0`` is free-slip / no-normal-flow. - normal : sympy matrix, optional - Row-vector constraint normal. Defaults to the mesh's smooth - projected boundary normals ``mesh.Gamma_P1``. - augmentation : float or sympy expression, optional - Augmented-Lagrangian parameter :math:`r` — simultaneously the - forward-problem penalty weight and the multiplier-update step. Any - :math:`r>0` converges; larger is faster but stiffens the linear - solve. **Defaults to ``augmentation_base · μ(x)``** — weighted by the - local viscosity so the dimensionless penalty ratio is uniform across - viscosity contrasts (essential for variable-viscosity problems; a - flat ``r`` under-constrains high-viscosity regions). May be passed as - a spatial sympy expression directly. - augmentation_base : float, default 1e3 - The base multiple used when ``augmentation`` is not given. - - Returns - ------- - lam : MeshVariable - The scalar multiplier field. After :meth:`solve`, its boundary - trace is the normal traction (topography proxy). - """ - # This proof-of-concept is serial only: the boundary mask construction - # and the node-wise multiplier update below are not MPI-decomposed. - if uw.mpi.size > 1: - raise NotImplementedError( - "SNES_Stokes_Constrained is serial-only (the boundary mask and " - "multiplier update are not MPI-safe). Run on a single rank." - ) - - if not hasattr(self.mesh.boundaries, boundary): - raise ValueError( - f"'{boundary}' is not a boundary of this mesh. " - f"Available: {[b.name for b in self.mesh.boundaries]}." - ) - - if normal is None: - normal = self.mesh.Gamma_P1 - - normal = sympy.Matrix(normal) - if normal.shape[0] != 1: - normal = normal.reshape(1, self.mesh.dim) - - if augmentation is None: - # Viscosity-weighted augmentation r = augmentation_base * mu(x): - # keeps the penalty/viscous ratio uniform so high-viscosity boundary - # regions are constrained as well as low-viscosity ones. - try: - viscosity = self.constitutive_model.Parameters.shear_viscosity_0 - augmentation = augmentation_base * viscosity - except (AttributeError, TypeError): - augmentation = augmentation_base * self._viscosity_scale() - - idx = len(self._constraint_bcs) - # Multiplier at the velocity degree so its trace reaches every velocity - # normal-trace DOF (no penalty floor on the P2 mid-edge component). - lam = uw.discretisation.MeshVariable( - f"lambda_{self.instance_number}_{idx}", - self.mesh, - 1, - degree=self._degree, - ) - lam.data[:] = 0.0 - - # Boundary-node mask: restrict the multiplier update to the constraint - # boundary so interior values stay exactly zero and lambda is a clean, - # directly usable topography field. Build a P1 marker (1 on the boundary - # vertices, 0 elsewhere) and sample it at lambda's nodes: boundary - # mid-edge nodes interpolate to 1, interior nodes to < 0.5. - from underworld3.discretisation.discretisation_mesh import ( - petsc_dm_find_labeled_points_local, - ) - - marker = uw.discretisation.MeshVariable( - f"_bmarker_{self.instance_number}_{idx}", self.mesh, 1, degree=1, - ) - marker.data[:] = 0.0 - if not self.mesh.dm.hasLabel("UW_Boundaries"): - raise RuntimeError( - "Mesh has no 'UW_Boundaries' label; cannot build the constraint " - "boundary mask." - ) - # NB: petsc_dm_find_labeled_points_local returns np.array([0]) (vertex 0) - # when the label is absent and None when the value has no points, so an - # `is not None` check alone could silently mark vertex 0. The hasLabel - # guard above plus the explicit empty/None check below close that gap. - point_indices = petsc_dm_find_labeled_points_local( - self.mesh.dm, - "UW_Boundaries", - getattr(self.mesh.boundaries, boundary).value, - sectionIndex=False, - ) - if point_indices is None or len(point_indices) == 0: - raise ValueError( - f"Boundary '{boundary}' has no labelled points on this mesh." - ) - marker.data[point_indices] = 1.0 - mask = ( - np.array(uw.function.evaluate(marker.sym, lam.coords)).reshape(-1) > 0.75 - ) - - # Augmented-Lagrangian natural BC, registered once: - # t = [ lambda + r (u.n - g) ] n - # The r(u.n)n part is the penalty BC (re-derived into the Jacobian each - # solve); the lambda part is the applied multiplier (fixed per solve). - nv = self.u.sym.dot(normal) - traction = (lam.sym[0] + augmentation * (nv - g)) * normal - self.add_natural_bc(traction, boundary) - - self._constraint_bcs.append( - _ConstraintBC(boundary, g, normal, lam, augmentation=augmentation, mask=mask) - ) - return lam - - def multiplier(self, boundary): - """Return the multiplier field for ``boundary`` (None if not constrained). - - After :meth:`solve`, the multiplier's boundary trace is the normal - traction holding the constraint; interior values are zero. Divide by - :math:`\\Delta\\rho\\,g` to obtain dynamic topography (see - :meth:`topography`). - """ - for cbc in self._constraint_bcs: - if cbc.boundary == boundary: - return cbc.lam - return None - - def topography(self, boundary, buoyancy_scale=1.0): - r"""Dynamic topography expression on ``boundary``. - - Returns the symbolic field :math:`\lambda / (\Delta\rho\, g)` for the - constraint multiplier on ``boundary`` (zero away from the boundary). - - Parameters - ---------- - boundary : str - A constrained boundary label. - buoyancy_scale : float or sympy expression, default 1.0 - The buoyancy scale :math:`\Delta\rho\, g` relating normal traction - to surface height. - """ - lam = self.multiplier(boundary) - if lam is None: - raise ValueError(f"No constraint registered on boundary '{boundary}'.") - return lam.sym[0] / buoyancy_scale - - def _constraint_rms(self, cbc): - """RMS of (u.n - g) over the constraint boundary, via boundary integral.""" - vn = self.u.sym.dot(cbc.normal) - num = float( - uw.maths.BdIntegral(self.mesh, fn=(vn - cbc.g) ** 2, - boundary=cbc.boundary).evaluate() - ) - length = float( - uw.maths.BdIntegral(self.mesh, fn=1.0, boundary=cbc.boundary).evaluate() - ) - return np.sqrt(num / length) if length > 0 else np.sqrt(num) - - def _nodal_constraint_residual(self, cbc): - """(u.n - g) evaluated at the multiplier's nodes.""" - expr = self.u.sym.dot(cbc.normal) - cbc.g - return np.array( - uw.function.evaluate(sympy.Matrix([[expr]]), cbc.lam.coords) - ).reshape(-1) - - def solve( - self, - zero_init_guess: bool = True, - *, - constraint_rtol: float = 1.0e-3, - constraint_atol: float = 1.0e-12, - constraint_max_iterations: int = 40, - constraint_verbose: bool = False, - **kwargs, - ): - """Solve the constrained Stokes system. - - With no constraint BCs registered this is an ordinary Stokes solve. - Otherwise it runs the augmented-Lagrangian outer loop until every - constraint boundary satisfies - ``RMS(u.n - g) < constraint_rtol · RMS|u| + constraint_atol`` (or the - iteration cap is hit). The tolerance is *relative* to the velocity scale - so it is problem-independent and needs no tuning. All other keyword - arguments are forwarded to the inner Stokes ``solve``. - """ - if not self._constraint_bcs: - return super().solve(zero_init_guess=zero_init_guess, **kwargs) - - # Sample the augmentation r at the multiplier nodes once (it may be a - # spatial / viscosity-weighted expression). Used for the dual update. - for cbc in self._constraint_bcs: - if isinstance(cbc.augmentation, (int, float)): - cbc.r_nodal = float(cbc.augmentation) * np.ones(cbc.lam.coords.shape[0]) - else: - cbc.r_nodal = np.array( - uw.function.evaluate( - sympy.Matrix([[cbc.augmentation]]), cbc.lam.coords - ) - ).reshape(-1) - - total_linear_its = 0 - all_converged = False - for k in range(constraint_max_iterations): - super().solve(zero_init_guess=(zero_init_guess and k == 0), **kwargs) - try: - total_linear_its += int(self.snes.getLinearSolveIterations()) - except Exception: - pass - - # Velocity scale for the relative tolerance (RMS speed over nodes). - v_scale = float(np.sqrt(np.mean(np.sum(self.u.data**2, axis=1)))) - threshold = constraint_rtol * v_scale + constraint_atol - - all_converged = True - worst = 0.0 - for cbc in self._constraint_bcs: - rms = self._constraint_rms(cbc) - if rms >= threshold: - all_converged = False - worst = max(worst, rms) - if constraint_verbose: - uw.mpi.pprint( - f" [constraint {cbc.boundary}] iter {k}: " - f"RMS(u.n-g) = {rms:.3e} (rel {rms / max(v_scale, 1e-30):.2e}, " - f"mean r = {cbc.r_nodal.mean():.3g})" - ) - - self.constraint_iterations = k + 1 - self.constraint_residual = worst - self.constraint_total_linear_its = total_linear_its - - if all_converged: - if constraint_verbose: - uw.mpi.pprint(f"Constraint loop converged in {k + 1} iterations.") - break - - # On the final permitted iteration, do NOT apply another multiplier - # update: it would never be solved with, leaving u/p inconsistent - # with lambda. Stop here and warn loudly below instead. - if k == constraint_max_iterations - 1: - break - - # Augmented-Lagrangian (ALG2) multiplier update, same r as the - # forward-problem penalty augmentation. Monotone-convergent for r>0. - # Restricted to boundary nodes so lambda stays a clean topography field. - for cbc in self._constraint_bcs: - resid = self._nodal_constraint_residual(cbc) - cbc.lam.data[cbc.mask, 0] += cbc.r_nodal[cbc.mask] * resid[cbc.mask] - - if not all_converged: - import warnings - - warnings.warn( - f"Constrained Stokes solve did NOT converge: worst " - f"RMS(u.n-g) = {self.constraint_residual:.3e} after " - f"{self.constraint_iterations} iterations " - f"(constraint_max_iterations={constraint_max_iterations}). " - f"Increase constraint_max_iterations or the augmentation.", - RuntimeWarning, - stacklevel=2, - ) - - return - - class _BlockConstraintBC: """Bookkeeping for one in-saddle-point multiplier constraint. @@ -2365,7 +1977,7 @@ def __init__(self, boundary, g, normal, lam, augmentation): self.fns = {} -class SNES_Stokes_BlockConstrained(SNES_Stokes): +class SNES_Stokes_Constrained(SNES_Stokes): r""" Stokes solver that enforces :math:`\mathbf{u}\cdot\mathbf{n} = g` on a boundary via a Lagrange multiplier living **inside** the saddle-point @@ -2391,15 +2003,15 @@ class SNES_Stokes_BlockConstrained(SNES_Stokes): traction = dynamic topography; access it via :meth:`multiplier` / :meth:`topography`. - This is the block (monolithic) counterpart of :class:`SNES_Stokes_Constrained` - (the augmented-Lagrangian outer-loop solver) and should match its answer to - discretisation error in one coupled solve. Serial only (the boundary mask is + The constraint is enforced in one coupled solve (no outer iteration). The + augmented-Lagrangian term conditions the :math:`[p,h]` Schur complement + without biasing the multiplier, and the interior multiplier DOFs are reduced + away so the solved block is boundary-sized. Serial only (the boundary mask is not yet MPI-decomposed). See ``docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md``. See Also -------- - SNES_Stokes_Constrained : The outer-loop counterpart (recovery API template). SNES_Stokes : The unconstrained saddle-point solver this extends. """ @@ -2425,8 +2037,8 @@ def __init__( DFDt=DFDt, ) - # Block-constrained records (distinct from the outer-loop solver's - # self._constraint_bcs, which the base assembly must not touch). + # In-saddle multiplier constraints (see add_constraint_bc). Empty for an + # ordinary Stokes solve, so the base assembly is unaffected. self._block_constraint_bcs = [] return @@ -2439,15 +2051,14 @@ def _viscosity_scale(self): return 1.0 def add_constraint_bc(self, boundary, g=0.0, normal=None, screening=None, - augmentation=None, augmentation_base=1.0e3, degree=None): + augmentation=None, augmentation_base=1.0e4, degree=None): r"""Register a multiplier-enforced normal-velocity constraint on ``boundary``. - Adds a scalar multiplier field ``h`` (field id 2, 3, ...) to the - saddle-point system. **Milestone 1**: only the field and its interior - screening are wired (the field is inert — no boundary coupling yet, so - ``h`` is driven to zero and the velocity/pressure solution is identical - to ordinary Stokes). The boundary residual/coupling are added in later - milestones. + Adds a scalar multiplier field ``h`` coupled into the saddle-point system + so that :math:`\mathbf{u}\cdot\mathbf{n}=g` is enforced on ``boundary`` in + the coupled solve; at convergence ``h`` on the boundary is the normal + traction (dynamic topography), recoverable via :meth:`multiplier` / + :meth:`topography`. Parameters ---------- @@ -2467,10 +2078,12 @@ def add_constraint_bc(self, boundary, g=0.0, normal=None, screening=None, \mathbf{n})` that conditions the :math:`[p,h]` Schur complement **without biasing the multiplier** (the h-row is still the exact constraint). Defaults to ``augmentation_base · μ(x)`` (viscosity- - weighted, like the outer-loop solver). Pass ``0`` for the bare KKT - system. - augmentation_base : float, default 1e3 - Base multiple used when ``augmentation`` is not given. + weighted, mesh-independent). Pass ``0`` for the bare KKT system. + augmentation_base : float, default 1e4 + Base multiple used when ``augmentation`` is not given. Accuracy is + independent of this value (the multiplier carries the exact + constraint); larger values reduce the iteration count up to a broad + plateau, well below the roundoff limit. Returns ------- @@ -2481,7 +2094,7 @@ def add_constraint_bc(self, boundary, g=0.0, normal=None, screening=None, # MPI-decomposed. if uw.mpi.size > 1: raise NotImplementedError( - "SNES_Stokes_BlockConstrained is serial-only for now." + "SNES_Stokes_Constrained is serial-only for now." ) if not hasattr(self.mesh.boundaries, boundary): diff --git a/tests/test_1061_constrained_freeslip.py b/tests/test_1061_constrained_freeslip.py index a93219649..f8a6d8bdf 100644 --- a/tests/test_1061_constrained_freeslip.py +++ b/tests/test_1061_constrained_freeslip.py @@ -1,6 +1,6 @@ -"""Block-constrained free-slip: a Lagrange multiplier INSIDE the saddle point. +"""Constrained free-slip: a Lagrange multiplier INSIDE the saddle point. -SNES_Stokes_BlockConstrained enforces u.n = g on a boundary via a multiplier h +SNES_Stokes_Constrained enforces u.n = g on a boundary via a multiplier h that lives in the coupled system (grouped u | [p,h] Schur split), in ONE solve. The converged h is the boundary normal traction = dynamic topography. @@ -53,7 +53,7 @@ def forcing(m): ref.solve() mb = uw.meshing.UnstructuredSimplexBox(minCoords=(0, 0), maxCoords=(1, 1), cellSize=0.15, qdegree=3) - blk = uw.systems.Stokes_BlockConstrained(mb) + blk = uw.systems.Stokes_Constrained(mb) blk.constitutive_model = uw.constitutive_models.ViscousFlowModel blk.constitutive_model.Parameters.shear_viscosity_0 = MU blk.bodyforce = forcing(mb) @@ -129,7 +129,7 @@ def annulus(): ref.petsc_options["ksp_type"] = "fgmres" ref.solve() - blk = uw.systems.Stokes_BlockConstrained(mesh) + blk = uw.systems.Stokes_Constrained(mesh) blk.constitutive_model = uw.constitutive_models.ViscousFlowModel blk.constitutive_model.Parameters.shear_viscosity_0 = MU blk.saddle_preconditioner = 1.0 / MU @@ -205,7 +205,7 @@ def forcing(m): ref.solve() mb = uw.meshing.UnstructuredSimplexBox(minCoords=(0, 0), maxCoords=(1, 1), cellSize=0.1, qdegree=3) - blk = uw.systems.Stokes_BlockConstrained(mb) + blk = uw.systems.Stokes_Constrained(mb) blk.constitutive_model = uw.constitutive_models.ViscousFlowModel blk.constitutive_model.Parameters.shear_viscosity_0 = visc(mb) blk.bodyforce = forcing(mb) diff --git a/tests/test_1062_constrained_solcx.py b/tests/test_1062_constrained_solcx.py new file mode 100644 index 000000000..d3aba8f1b --- /dev/null +++ b/tests/test_1062_constrained_solcx.py @@ -0,0 +1,51 @@ +"""Validate Stokes_Constrained (free-slip via in-saddle multipliers) against the +exact SolCx analytic solution — the canonical free-slip, 1e6-viscosity-jump +benchmark. The constrained solver enforces free-slip on all four walls via four +multipliers; its velocity must match the analytic, and the constraint must hold. + +Run: pixi run python -m pytest tests/test_1062_constrained_solcx.py -v +""" + +import pytest + +pytestmark = [pytest.mark.level_2] + +import numpy as np +import sympy +import underworld3 as uw +from underworld3.function import analytic as A + +ETA_A, ETA_B, XC, NZ, RES = 1.0, 1.0e6, 0.5, 1, 32 + + +def test_constrained_solcx_matches_analytic(): + mesh = uw.meshing.StructuredQuadBox( + elementRes=(RES, RES), minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), qdegree=3 + ) + sol = A.SolCx(mesh, eta_A=ETA_A, eta_B=ETA_B, x_c=XC, n=NZ) + + s = uw.systems.Stokes_Constrained(mesh) + s.constitutive_model = uw.constitutive_models.ViscousFlowModel + s.constitutive_model.Parameters.shear_viscosity_0 = sol.fn_viscosity + s.saddle_preconditioner = 1.0 / sol.fn_viscosity + s.bodyforce = sol.fn_bodyforce + # free-slip on all four walls via in-saddle multipliers + s.add_constraint_bc("Left", g=0.0, normal=sympy.Matrix([[-1.0, 0.0]])) + s.add_constraint_bc("Right", g=0.0, normal=sympy.Matrix([[ 1.0, 0.0]])) + s.add_constraint_bc("Bottom", g=0.0, normal=sympy.Matrix([[ 0.0, -1.0]])) + s.add_constraint_bc("Top", g=0.0, normal=sympy.Matrix([[ 0.0, 1.0]])) + s._petsc_use_pressure_nullspace = True + s.tolerance = 1.0e-9 + s.solve() + + # velocity matches the exact analytic + rel = sol.velocity_error(s.u) + assert rel < 5.0e-3, f"velocity error vs analytic too large: {rel:.2e}" + + # constraint enforced (u.n -> 0 on the walls) + c = s.u.coords + e = 1e-6 + xw = (c[:, 0] < e) | (c[:, 0] > 1 - e) + yw = (c[:, 1] < e) | (c[:, 1] > 1 - e) + rms_un = np.sqrt(np.mean(np.r_[s.u.data[xw, 0], s.u.data[yw, 1]] ** 2)) + assert rms_un < 1.0e-6, f"constraint not enforced: RMS(u.n)={rms_un:.2e}" From 98637448fe204cf8b9ba103bf5aaf14cd0c43099 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Mon, 8 Jun 2026 21:09:43 +1000 Subject: [PATCH 7/7] Address Copilot review on #224 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - remove the unused _BlockConstraintBC.interior_mask and its construction (a P1 marker field + evaluate) — dead since the section-based reduction; avoidable setup cost - tests: use the public petsc_use_pressure_nullspace property (its setter resets cached nullspace state); fix the stale run-command path in test_1061 - docs/CONSTRAINED_FREESLIP_MULTIPLIER.md: rewrite the outer-loop sections to the shipped in-saddle formulation (no outer iteration; r is AL stabilisation only; augmentation_base 1e4; boundary-only reduction); fix the Files list and the option-table verdicts; mark the r-sweep table as historical outer-loop data Underworld development team with AI support from Claude Code --- .../design/CONSTRAINED_FREESLIP_MULTIPLIER.md | 174 +++++++++--------- src/underworld3/systems/solvers.py | 28 +-- tests/test_1061_constrained_freeslip.py | 4 +- tests/test_1062_constrained_solcx.py | 2 +- 4 files changed, 93 insertions(+), 115 deletions(-) diff --git a/docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md b/docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md index d7a04805e..4e99c7911 100644 --- a/docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md +++ b/docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md @@ -36,29 +36,31 @@ Stokes with a surface constraint `u·n = g` on Γ, multiplier `λ`: [ C 0 0 ] [λ] [g] (C couples only the boundary trace of u) ``` -`C` is **co-dimension-1**: it touches only velocity DOFs on Γ. A monolithic -third FE field would therefore either waste interior DOFs or need a boundary -trace space PETSc/DMPlex does not provide on the same DM. We instead solve the -boundary Schur complement `S_λ = C K⁻¹ Cᵀ` by an **outer loop**, leaving the -validated 2×2 Stokes assembly untouched. +`C` is **co-dimension-1**: it touches only velocity DOFs on Γ. The **shipped** +solver carries `λ` as a third field **inside** the saddle point and solves the +whole 3×3 system in one coupled solve (the `[p, λ]` rows are grouped into a +single Schur factor — see the "Monolithic `P'=[p,λ]` fieldsplit" section). The +multiplier carries the *exact* constraint; there is **no outer loop**. -### Augmented Lagrangian (ALG2) — the production algorithm +### Augmented-Lagrangian stabilisation `r` -Each Stokes solve carries a natural-BC traction with **both** the multiplier and -a penalty augmentation, reusing the existing penalty machinery: +The u-row carries `λ` plus an augmented-Lagrangian penalty: $$\mathbf{t} = \bigl[\lambda + r\,(\mathbf{u}\cdot\mathbf{n} - g)\bigr]\,\mathbf{n} \quad\text{on } \Gamma,$$ -and the multiplier is updated with the same `r`: +which adds a `uu` boundary stiffness `r(n⊗n)` that conditions the `[p, λ]` Schur +complement **without biasing the multiplier** (the λ-row stays the exact +constraint). It is *not* an outer multiplier update — `λ` is solved +monolithically. Because accuracy is independent of `r`, `r` is a cost-only knob: +larger values reduce the iteration count up to a broad plateau, well below the +roundoff limit. Default `r = augmentation_base · μ(x)` with +`augmentation_base = 1e4` (viscosity-weighted, mesh-independent). -$$\lambda \leftarrow \lambda + r\,(\mathbf{u}\cdot\mathbf{n} - g).$$ - -The penalty term `r(u·n)n` preconditions *every* boundary mode uniformly, so the -outer loop converges in a handful of iterations; the multiplier removes the -penalty's accuracy bias, so a **moderate, well-conditioned** `r` gives both fast -convergence *and* an exact constraint — unlike a pure penalty, which must be made -large and fragile. +> **Historical note.** An earlier *outer-loop* (Uzawa / ALG2) variant updated +> `λ ← λ + r(u·n − g)` between Stokes solves. It was superseded by — and removed +> in favour of — the in-saddle formulation above. The Phase-0 spike findings +> below motivated that exploration and are kept for context. ## Phase-0 spike findings (what shaped the design) @@ -101,44 +103,41 @@ stokes.solve() # no constraint tuning n topo = stokes.topography("Upper", buoyancy_scale=delta_rho_g) # h = lambda/(drho g) ``` -The outer loop is hidden behind ``solve()``: the tolerance is **relative** to the -velocity scale (``RMS(u.n-g) < rtol·RMS|u|``, default ``rtol=1e-3``) so it is -problem-independent, and the augmentation defaults to ``1e3·μ(x)`` (local- -viscosity-weighted). On SolCx this gives ~2e-4 accuracy in 3 outer iterations at -viscosity contrast 1 → 10⁶ with no user tuning. +`solve()` does **one coupled solve** — no outer iteration or constraint tuning. +The augmentation defaults to `1e4·μ(x)` (local-viscosity-weighted); accuracy is +independent of it (the λ-row carries the exact constraint), so no per-problem +tuning is needed. Key design points: -- **Multiplier representation.** `λ` is an ordinary full-mesh scalar field at the - *velocity degree* (P2). Only its trace on Γ enters the weak form. Matching the - velocity degree means the multiplier reaches every velocity normal-trace DOF - (including P2 mid-edge), so there is no penalty floor on the constraint — a P1 - multiplier plateaus at ~`‖t‖/r` on the mid-edge component. -- **Clean topography field.** The dual update is restricted to boundary nodes via - a P1 boundary marker (`petsc_dm_find_labeled_points_local`) sampled at the - multiplier's nodes — boundary mid-edge nodes interpolate to 1, interior to - < 0.5. Interior `λ` therefore stays **exactly zero**, so `λ` is a directly - usable topography field (not interior garbage). -- **Coupling registered once.** `add_natural_bc([λ + r(u·n − g)]·n, boundary)` - is set up a single time; only `λ`'s boundary data changes between solves, so - nothing recompiles. +- **Multiplier representation.** `λ` is a full-mesh scalar field at the *velocity + degree* (P2). Only its trace on Γ enters the weak form. Matching the velocity + degree means the multiplier reaches every velocity normal-trace DOF (including + P2 mid-edge), so there is no constraint floor. +- **Boundary-only reduction → clean topography.** The interior (off-boundary) λ + DOFs are constrained directly in the PetscSection, so the solved `[p, λ]` block + carries only the boundary trace (~√ndof DOFs, ~1.1× Dirichlet rather than ~3×). + Interior `λ` is absent, so the boundary `λ` is a directly usable topography + field. The reduction is lossless (machine-precision constraint) and default-on. +- **Coupling registered once.** The boundary residual/Jacobian + (`λ·n`, the AL stiffness `r(n⊗n)`, and the `uλ`/`λu` couplings) are registered + a single time; nothing recompiles between solves. - **`add_constraint_bc(boundary, g=0, normal=None, augmentation=None)`** — `normal` defaults to the smooth projected normals `mesh.Gamma_P1`; - `augmentation` defaults to a viscosity-scaled `r = 10³·μ`. + `augmentation` defaults to a viscosity-scaled `r = 10⁴·μ`. ## Validation -`tests/test_1061_constrained_freeslip_annulus.py` (level_2 / tier_b), buoyancy- -driven annulus, no-slip inner + multiplier free-slip outer, vs a `1e6` penalty -reference. All pass (~14 s): +Two regression tests cover the shipped solver: -| Check | Result | -|---|---| -| `RMS(u·n)` on outer boundary (no penalty coefficient) | 6.5e-5 | -| `relL2(v_multiplier vs v_penalty)` | 3.1e-3 | -| Outer iterations (augmented Lagrangian) | 2 | -| Interior `λ` (clean field) | exactly 0 | -| boundary `corr(λ, −n·σ·n)` (dynamic-topography stress) | 0.9999 | +- `tests/test_1061_constrained_freeslip.py` — box (vs an exact Dirichlet + free-slip reference) and buoyancy-driven annulus (vs a `1e6` penalty + reference): constraint enforced (`RMS(u·n)` small with no penalty coefficient), + velocity matches the reference, and `corr(λ, −n·σ·n) ≈ 0.9999` (topography). +- `tests/test_1062_constrained_solcx.py` — free-slip via four in-saddle + multipliers on the **SolCx** benchmark (1e6 viscosity jump) compared to the + **exact analytic** solution: velocity `rel ≈ 8.7e-6` (== the Dirichlet + baseline), constraint `RMS(u·n) ≈ 1.6e-10`. The consistent-boundary-flux identity `λ = −n·σ·n|_Γ` is the independent cross-check: the multiplier's boundary trace equals the recovered normal Cauchy @@ -148,8 +147,16 @@ stress (negative sign = the reaction traction holding the boundary), confirming ## The augmentation parameter `r`: true-work trade-off `r` is a *speed* knob, not an *accuracy* knob — this is the key advantage over a -pure penalty. Sweep on the annulus (constraint tol 1e-4, wall time for one cold -solve; `tot_lin` = total outer Schur-KSP linear iterations across the loop): +pure penalty, and it carries over to the in-saddle solver (accuracy is +`r`-independent; `r` only sets the iteration count). The sweep table below is from +the **historical outer-loop** variant (its "outer iterations" have no analogue in +the one-shot coupled solve), but the shape and the conclusion stand. For the +in-saddle solver with a scalable (FMG) inner solve the iteration count falls +monotonically with `r` to a saturation floor, with no high-`r` penalty until +roundoff; the default `r = 10⁴·μ` sits comfortably in that regime. + +Historical outer-loop sweep on the annulus (constraint tol 1e-4, wall time for one +cold solve; `tot_lin` = total outer Schur-KSP linear iterations across the loop): cellSize 0.1 (~2940 dof): @@ -190,21 +197,21 @@ cellSize 0.05 (~10852 dof) shows the same shape (min wall ≈ 1.25 s at r=1e3; | (B1) boundary-stratum-only PetscFE field | No DMPlex support on the same DM. | | (B2) co-dim-1 submesh + MATNEST | The honest monolithic form; deferred. | | (C) reuse pressure / `_constraints` | `p` enforces `∇·u=0` interior, not `u·n=0` on Γ — not redundant. The CBF identity is a *validation* tool, not an implementation. | -| **(D) augmented-Lagrangian outer loop** | **Implemented.** Non-invasive, 2 iterations, exact, recoverable topography. | +| **(D) monolithic in-saddle multiplier (grouped `[p,λ]` Schur)** | **Implemented** (the shipped solver). One coupled solve, exact, recoverable topography, boundary-only reduction. | +| (E) augmented-Lagrangian outer loop | Earlier exploration; **removed** in favour of (D). Easy to reproduce in Python. | -Deferred to follow-up PRs (Phase-0 spike S2 gathers the fieldsplit-feasibility -evidence): the monolithic 3-field / co-dim-1 representation; 3D spherical shells; -**parallel** (the boundary mask and the nodal update are serial); the -**both-boundaries-free-slip** annulus case (admits a rigid-rotation velocity null -space needing explicit removal); and live free-surface equilibrium integration -(pass `λ` as the target normal-stress end-state). +Deferred to follow-up PRs: a true co-dim-1 / MATNEST `λ` representation; 3D +spherical shells; **parallel** (the boundary handling is serial); the +**both-boundaries-free-slip** annulus case (rigid-rotation velocity null space +needing explicit removal); and live free-surface equilibrium integration (pass +`λ` as the target normal-stress end-state). -## Monolithic `P'=[p,h]` fieldsplit — feasibility spike +## Monolithic `P'=[p,λ]` fieldsplit — the shipped design -A future direction (and a general "inject arbitrary constraints into the saddle -point" capability): rather than a third field forcing a nested 3-way Schur, group -pressure and the multiplier into a composite `P' = [p, h]` and keep a **2-way -`u | P'`** split. +This is the implemented approach (and a step toward a general "inject arbitrary +constraints into the saddle point" capability): rather than a third field forcing +a nested 3-way Schur, group pressure and the multiplier into a composite +`P' = [p, λ]` and keep a **2-way `u | P'`** split. **Spike result (confirmed):** a 3-field DM `(u, p, h)` on a real mesh, with the `p` and `h` index sets grouped (`pc_fieldsplit_1_fields 1,2`, or an explicit @@ -213,33 +220,30 @@ concatenated IS), produces exactly a 2-block `u | [p,h]` Schur fieldsplit KSP-reconfiguration concern is therefore moot — the split structure is identical to the current `u | p` solver. (`/tmp/spike_pph.py`.) -**Remaining work (bounded, ~2 weeks of Cython, behind a subclass):** -- Register `h` as field 2 (one `dm.setField`). -- `h`-equation residual: boundary part `∫_Γ ψ(n·u − g)` (the existing Nitsche path - already registers field-1 boundary residuals — same pattern) plus a small interior - screening `ε∫_Ω h ψ` to de-singularise the interior `h` block. -- Three new Jacobian blocks: `uh`, `hu` (boundary integrals — the `UW_PetscDSSetBdJacobian` - machinery already does arbitrary field pairs for `up`/`pu`) and `hh` (interior mass); - `ph`/`hp` are zero. -- Group `[p,h]` in the fieldsplit and extend the Schur PC (`p` keeps its `1/μ` mass PC; - `h` gets the screening diagonal). -- The 3-IS field decomposition / nullspace path (currently assumes 2 fields). - -**Caveats:** the interior screening reintroduces a small `ε` (benign — it does not bias -the boundary multiplier); the work touches the validated `uu/up/pu/pp` assembly, so it -must live behind a subclass and be regression-tested. The one genuine PETSc limitation -remains a *co-dimension-1* `h` (boundary-only DOFs) — avoided here by the full-domain + -screening representation. - -**Recommendation:** the architecture is de-risked, but since the outer loop converges in -2–3 iterations with no user-facing tuning, monolithic is scheduled work (a general -saddle-point-constraint capability), not an urgent replacement. +**What was implemented (behind the `SNES_Stokes_Constrained` subclass):** +- `λ` registered as field 2 (`dm.setField`). +- `λ`-equation residual: boundary part `∫_Γ ψ(n·u − g)` plus a small interior + screening `ε∫_Ω λ ψ` to de-singularise the interior block — which is then + **constrained away** by the boundary-only reduction (the interior λ DOFs are + pinned in the PetscSection, so only the boundary trace is solved). +- Boundary Jacobian blocks `uλ`, `λu` and the AL stiffness `uu += r(n⊗n)` + (`ph`/`hp` are zero); registered via the `UW_PetscDSSetBdJacobian` machinery. +- `[p,λ]` grouped in the fieldsplit by field index (keeps the velocity DM + hierarchy for geometric MG/FMG); the gauge nullspace handled as a combined + `(p, λ)` mode on enclosed problems. + +**Caveats:** the work touches the validated `uu/up/pu/pp` assembly, so it lives +behind the subclass and is regression-tested (`tier`-graded). Serial only for +now. A true *co-dimension-1* `λ` (boundary-only DOFs end-to-end) remains the +honest long-term form; the full-domain field + boundary-only reduction is the +pragmatic path that ships today. ## Files -- `src/underworld3/systems/solvers.py` — `SNES_Stokes_Constrained`, `_ConstraintBC`. +- `src/underworld3/systems/solvers.py` — `SNES_Stokes_Constrained`, `_BlockConstraintBC`. +- `src/underworld3/cython/petsc_generic_snes_solvers.pyx` — multiplier-field + registration, boundary residual/Jacobian coupling, fieldsplit grouping, + nullspace, and the section-based interior-multiplier reduction. - `src/underworld3/systems/__init__.py` — `Stokes_Constrained` export. -- `tests/test_1061_constrained_freeslip_annulus.py` — validation. - -Pre-existing (unrelated) failures noted: `tests/test_1060_nitsche_freeslip.py` -has two failing assertions on `development` independent of this work. +- `tests/test_1061_constrained_freeslip.py` — box + annulus validation. +- `tests/test_1062_constrained_solcx.py` — SolCx analytic validation. diff --git a/src/underworld3/systems/solvers.py b/src/underworld3/systems/solvers.py index 1a54275b6..5e55443a4 100644 --- a/src/underworld3/systems/solvers.py +++ b/src/underworld3/systems/solvers.py @@ -1948,19 +1948,13 @@ class _BlockConstraintBC: """ __slots__ = ("boundary", "g", "normal", "lam", "augmentation", - "interior_mask", "petsc_id_u", "petsc_id_h", "label_val", "fns") + "petsc_id_u", "petsc_id_h", "label_val", "fns") def __init__(self, boundary, g, normal, lam, augmentation): self.boundary = boundary self.g = g self.normal = normal self.lam = lam - # Boolean over the multiplier's nodes: True on INTERIOR nodes (off the - # constraint boundary). The constant-on-interior h mode is a near-null - # mode (interior h appears only in the screening eM), so we hand it to - # the solver as a nullspace vector -- the gauge-fixing analogue of the - # constant-pressure nullspace. Set in add_constraint_bc. - self.interior_mask = None # Augmented-Lagrangian parameter r: a penalty r(n·u−g)·n added to the # u-row, giving a uu boundary stiffness r·(n⊗n) that conditions the # [p,h] Schur complement WITHOUT biasing the multiplier (the h-row is @@ -2153,26 +2147,6 @@ def add_constraint_bc(self, boundary, g=0.0, normal=None, screening=None, cbc = _BlockConstraintBC(boundary, g, normal, h, augmentation) - # Interior mask: True on multiplier nodes OFF the constraint boundary. - # Build a P1 marker that is 1 on the boundary vertices and sample it at - # the multiplier's nodes (boundary mid-edge nodes interpolate to ~1). - from underworld3.discretisation.discretisation_mesh import ( - petsc_dm_find_labeled_points_local, - ) - marker = uw.discretisation.MeshVariable( - f"_bmask_{self.instance_number}_{idx}", self.mesh, 1, degree=1, - ) - marker.data[:] = 0.0 - if self.mesh.dm.hasLabel("UW_Boundaries"): - pts = petsc_dm_find_labeled_points_local( - self.mesh.dm, "UW_Boundaries", - getattr(self.mesh.boundaries, boundary).value, sectionIndex=False, - ) - if pts is not None and len(pts) > 0: - marker.data[pts] = 1.0 - on_boundary = np.array(uw.function.evaluate(marker.sym, h.coords)).reshape(-1) > 0.5 - cbc.interior_mask = ~on_boundary - self._block_constraint_bcs.append(cbc) return h diff --git a/tests/test_1061_constrained_freeslip.py b/tests/test_1061_constrained_freeslip.py index f8a6d8bdf..ea500ca93 100644 --- a/tests/test_1061_constrained_freeslip.py +++ b/tests/test_1061_constrained_freeslip.py @@ -11,7 +11,7 @@ via DEFAULT solver options (no per-script PETSc options); compared against the penalty reference, with topography recovery h ~ -n.sigma.n. -Run: pixi run python -m pytest tests/test_1062_block_constrained_freeslip.py -v +Run: pixi run python -m pytest tests/test_1061_constrained_freeslip.py -v """ import pytest @@ -137,7 +137,7 @@ def annulus(): blk.add_dirichlet_bc((0.0, 0.0), "Lower") hb = blk.add_constraint_bc("Upper", g=0.0, normal=unit_r) # Enclosed -> pressure/multiplier gauge; DEFAULT grouped-Schur solver config. - blk._petsc_use_pressure_nullspace = True + blk.petsc_use_pressure_nullspace = True blk.tolerance = 1e-8 blk.solve() return dict(mesh=mesh, unit_r=unit_r, R_O=R_O, CELL=CELL, ref=ref, blk=blk, hb=hb) diff --git a/tests/test_1062_constrained_solcx.py b/tests/test_1062_constrained_solcx.py index d3aba8f1b..673ffacc6 100644 --- a/tests/test_1062_constrained_solcx.py +++ b/tests/test_1062_constrained_solcx.py @@ -34,7 +34,7 @@ def test_constrained_solcx_matches_analytic(): s.add_constraint_bc("Right", g=0.0, normal=sympy.Matrix([[ 1.0, 0.0]])) s.add_constraint_bc("Bottom", g=0.0, normal=sympy.Matrix([[ 0.0, -1.0]])) s.add_constraint_bc("Top", g=0.0, normal=sympy.Matrix([[ 0.0, 1.0]])) - s._petsc_use_pressure_nullspace = True + s.petsc_use_pressure_nullspace = True s.tolerance = 1.0e-9 s.solve()