diff --git a/docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md b/docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md new file mode 100644 index 000000000..4e99c7911 --- /dev/null +++ b/docs/developer/design/CONSTRAINED_FREESLIP_MULTIPLIER.md @@ -0,0 +1,249 @@ +# Constrained free-slip via a recoverable Lagrange multiplier (dynamic 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 + +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 Γ. 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 stabilisation `r` + +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,$$ + +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). + +> **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) + +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) +``` + +`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 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⁴·μ`. + +## Validation + +Two regression tests cover the shipped solver: + +- `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 +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, 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): + +| `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) 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: 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,λ]` fieldsplit — the shipped design + +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 +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`.) + +**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`, `_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.py` — box + annulus validation. +- `tests/test_1062_constrained_solcx.py` — SolCx analytic validation. diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index d005a3728..73257c4ec 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. 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 ## 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,20 @@ 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 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] + _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 +5610,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 +5674,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 +5757,14 @@ 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) 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) self.dm.copyDS(coarse_dm) @@ -5526,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): @@ -5557,6 +5933,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 +6078,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 +6161,18 @@ 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. + 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 +6476,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 +6509,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 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..5e55443a4 100644 --- a/src/underworld3/systems/solvers.py +++ b/src/underworld3/systems/solvers.py @@ -1936,6 +1936,240 @@ def delta_t(self): return self.constitutive_model.Parameters.dt_elastic +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", + "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 + # 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_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 + 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`. + + 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 : 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, + ) + + # 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 + + 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.0e4, degree=None): + r"""Register a multiplier-enforced normal-velocity constraint on ``boundary``. + + 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 + ---------- + 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, 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 + ------- + 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_Constrained 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) + + 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_1061_constrained_freeslip.py b/tests/test_1061_constrained_freeslip.py new file mode 100644 index 000000000..ea500ca93 --- /dev/null +++ b/tests/test_1061_constrained_freeslip.py @@ -0,0 +1,219 @@ +"""Constrained free-slip: a Lagrange multiplier INSIDE the saddle point. + +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. + +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_1061_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_Constrained(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_Constrained(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_Constrained(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 diff --git a/tests/test_1062_constrained_solcx.py b/tests/test_1062_constrained_solcx.py new file mode 100644 index 000000000..673ffacc6 --- /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}"