Per-iteration SNES update callbacks (pressure gauge + boundary-correct generic scatter)

[lmoresi]

Summary

Adds a user hook that runs at the start of every nonlinear (SNES) iteration via PETSc's SNESSetUpdate, for quantities that must track the current iterate rather than being recomputed once per timestep. Works on any solver.

solver.add_update_callback(fn)   # fn(solver, iteration)

Three parts:

1. Callback mechanism — add_update_callback(fn) + the SNESSetUpdate dispatcher (scatter iterate→fields, run callbacks, gather fields→iterate, refresh aux vec; also applied once to the converged iterate). With no callback registered the solve path is byte-for-byte unchanged (the hook is only installed when a callback exists).

2. Pressure-gauge helper — SNES_Stokes.set_pressure_gauge(boundary, reference=0.0) removes the surface-mean pressure each iteration, pinning a specific physical gauge on enclosed (pressure-null-space) problems. Converges to machine-zero surface mean.

3. Boundary-correct, solver-generic field scatter — the scatter that lets a callback read the iterate is now correct everywhere, including driven (non-zero Dirichlet) boundaries, on all solvers.

The scatter fix

A plain globalToLocal fills a field's interior and zero-Dirichlet DOFs but not its non-zero Dirichlet (driven) boundary DOFs — those aren't in the global vector; they're imposed on the local vector by DMPlexSNESComputeBoundaryFEM. Previously the dispatcher's scatter also assumed Stokes-only self.fields/_subdict, so a callback on a scalar/vector solver raised AttributeError.

• Base class _scatter_global_to_fields / _gather_fields_to_global rewritten to be single-field and boundary-correct (globalToLocal → DMPlexSNESComputeBoundaryFEM → self.u), serving scalar / vector / multi-component solvers and their subclasses.
• SNES_Stokes_SaddlePt overrides both to split its multi-field DM (velocity / pressure / block-constraint multipliers); the local index-set build is factored into _ensure_local_field_index_sets(), shared with the post-solve copy-back. Adds _multiplier_is cleanup to the rebuild reset.

This enables the gradient-plasticity / shear-band use case: an auxiliary Helmholtz/Projection smoother fired each iteration whose smoothed strain-rate field stays self-consistent with the velocity (measured error on a driven lid dropped from ~35% to 0).

Tests

tests/test_1016_snes_update_callbacks.py (tier_a / level_1, 7 passed):

• callback fires; pressure gauge → machine-zero surface mean; no-callback path unchanged;
• callback reads the driven-boundary velocity correctly (Stokes);
• callback works on scalar (Poisson) and vector (Vector_Projection) solvers, reading non-zero Dirichlet boundaries correctly (both raised AttributeError before generalisation);
• Helmholtz/Projection smoother self-consistent at convergence.

The boundary-scatter tests were confirmed to fail with the override disabled (genuine regression guards). test_1010, the Poisson and projection suites, and the constrained-Stokes tests still pass.

Docs

docs/advanced/solver-iteration-callbacks.md documents the generic, boundary-correct behaviour and both use cases.

Notes for review

• New public API (add_update_callback, set_pressure_gauge) is introduced here on a feature branch; nothing else depends on it yet, so it's folded into this one PR rather than pre-extracted to development.
• var.data (global-access path) still reads 0 on driven DOFs by design — the global vector never holds them; callbacks read through uw.function.evaluate/the aux-vec, which is now boundary-correct.

Underworld development team with AI support from Claude Code

[Copilot]

Pull request overview

Adds a solver-generic mechanism to run user-defined callbacks at the start of each PETSc SNES (Newton) iteration, enabling per-iterate updates such as pressure gauge fixing and auxiliary “inner” solves, and includes a boundary-correct iterate→field scatter so callbacks can reliably read driven Dirichlet boundary values.

Changes:

• Introduces add_update_callback() + SNESSetUpdate dispatcher, plus boundary-correct scatter/gather utilities (with Stokes saddle-point specialisation).
• Adds SNES_Stokes.set_pressure_gauge(boundary, reference=0.0) helper to enforce a surface-mean pressure gauge per iteration.
• Adds tests and documentation covering callback firing, boundary correctness, scalar/vector solver support, and a Helmholtz/Projection smoother use case.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
tests/test_1016_snes_update_callbacks.py	New regression tests for per-iteration callbacks, pressure gauge, and driven-boundary correctness across solvers.
src/underworld3/systems/solvers.py	Adds set_pressure_gauge() helper on SNES_Stokes.
src/underworld3/cython/petsc_generic_snes_solvers.pyx	Implements callback registry + SNES update hook, boundary-correct scatter/gather, and Stokes saddle-point overrides.
docs/advanced/solver-iteration-callbacks.md	New documentation for the callback API, boundary-correct semantics, and example use cases.
docs/advanced/index.md	Adds the new callbacks doc page to the advanced docs index.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[lmoresi]

Review — approved ✅

Reviewed the core .pyx change (+234/−53) closely since it touches the solver:

• The −53 lines are a clean extraction, not a behaviour change. The inline get_local_field_is block that built _pressure_is/_velocity_is/_multiplier_is is now the cached _ensure_local_field_index_sets helper with byte-identical logic, called from the post-solve copy-back (same site as before) and reused by the new mid-solve scatter. The existing Stokes copy-back path is preserved.
• No-op when unused is guaranteed by construction: _snes_update_callbacks is empty by default → _maybe_install_snes_update attaches no SNESSetUpdate hook → solve path byte-for-byte unchanged. Verified the docstring claim against the code.
• Boundary-correct scatter is right: _scatter_global_to_fields mirrors the post-solve copy-back (globalToLocal → DMPlexSNESComputeBoundaryFEM → copy into the field), so a callback reading v/p on a driven (non-zero Dirichlet) boundary sees the imposed value, not a stale one. Stokes correctly overrides both scatter/gather for its multi-field DM.
• set_pressure_gauge cleanly layers on add_update_callback (surface-mean removal each iteration) — the right altitude.
• CI green incl. the 2 new tier_a tests in test_1016 (which fail with the override off). Docs added under docs/advanced/.

Merging for the release.

Underworld development team with AI support from Claude Code

[lmoresi]

Thanks for the review — all three points addressed in b82acea:

1. Spurious rewire on add_update_callback — removed _needs_function_rewire = True. The SNESSetUpdate hook is already (re)attached at the start of every solve in _snes_solve_with_retries, so registering a callback no longer forces a JIT/pointwise-function rewire.

2. Final-iterate dispatch only on Stokes — centralised the final _dispatch_snes_update(self.snes, -1) in _snes_solve_with_retries, so every solver (scalar/vector/multi-component as well as Stokes) applies callbacks once to the converged iterate, matching the docs. Added a test asserting the iteration == -1 dispatch fires on a scalar (Poisson) solver.

3. Double-fire risk — removed the per-class dispatch from SNES_Stokes_SaddlePt.solve(); each solve calls the helper exactly once, so callbacks fire once on the final iterate.

Also renamed _maybe_install_snes_update → _attach_snes_update_hook to match the existing _attach_constant_nullspace / _attach_stokes_nullspace convention. test_1016 (8) and test_1010 (6) pass.