From 6c3a9765ca605484c8a68ea2358bd9aaabb37857 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Tue, 15 Sep 2026 07:17:16 +1000 Subject: [PATCH 01/23] =?UTF-8?q?feat:=20the=20discrete=20adjoint=20?= =?UTF-8?q?=E2=80=94=20in=20the=20solvers,=20and=20over=20the=20run=20tran?= =?UTF-8?q?script?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The exact content of the split: adjoint_solve / dual_of / sensitivity on every solver (composite on Stokes; the consistent tangent assembled for the adjoint whichever tangent the forward used); uw.adjoint.TranscriptAdjoint, a reverse driver over the transcript with no problem-specific wiring, with fields read through their value and their gradient assembled as FEM loads; uw.adjoint.inner over owned dofs; an adjoint verdict on every recorded operator and transcript_adjoint_segments; the consistent Newton tangent as the default with Picard opt-in; the sinking-blob example with its Taylor test through the library (1.00000); tests 0018-0020, serial and np=2. The history of how each piece was built and reviewed is on the docs/timestepping-pattern branch before its split commit (baae0817, 258ee6e8, fa83b19d, 790a97c5, 567c8c93, fba44975, dd928797). Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- CLAUDE.md | 9 +- .../design/run-plan-and-transcript.md | 55 ++ .../guides/HOW-TO-WRITE-UW3-SCRIPTS.md | 123 ++++ docs/developer/guides/adversarial-review.md | 6 + .../adjoint/sinker_transcript/README.md | 228 ++++++ .../forward_sinker_transcript.py | 272 ++++++++ .../generate_target_transcript.py | 22 + .../inverse_sinker_transcript.py | 307 ++++++++ .../sinker_transcript/taylor_test_library.py | 52 ++ .../taylor_test_transcript.py | 54 ++ src/underworld3/__init__.py | 2 + src/underworld3/adjoint.py | 486 +++++++++++++ .../cython/petsc_generic_snes_solvers.pyx | 657 +++++++++++++++++- src/underworld3/model.py | 44 ++ src/underworld3/swarm.py | 28 +- src/underworld3/systems/ddt.py | 24 + .../utilities/transcript_report.py | 73 +- tests/test_0018_adjoint_support_record.py | 360 ++++++++++ tests/test_0019_adjoint_solve.py | 290 ++++++++ tests/test_0020_transcript_adjoint.py | 211 ++++++ tests/test_0641_wave_c_api_shims.py | 8 +- tests/test_1057_yield_homotopy_solve.py | 7 +- 22 files changed, 3299 insertions(+), 19 deletions(-) create mode 100644 docs/examples/adjoint/sinker_transcript/README.md create mode 100644 docs/examples/adjoint/sinker_transcript/forward_sinker_transcript.py create mode 100644 docs/examples/adjoint/sinker_transcript/generate_target_transcript.py create mode 100644 docs/examples/adjoint/sinker_transcript/inverse_sinker_transcript.py create mode 100644 docs/examples/adjoint/sinker_transcript/taylor_test_library.py create mode 100644 docs/examples/adjoint/sinker_transcript/taylor_test_transcript.py create mode 100644 src/underworld3/adjoint.py create mode 100644 tests/test_0018_adjoint_support_record.py create mode 100644 tests/test_0019_adjoint_solve.py create mode 100644 tests/test_0020_transcript_adjoint.py diff --git a/CLAUDE.md b/CLAUDE.md index 82745c019..f2648f501 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -344,9 +344,12 @@ to impose `v·n̂ = 0`: a consistency error against the faceted assembly. See `docs/developer/subsystems/rotated-freeslip.md` ("Which normal to use"). - Works **inside the nonlinear SNES** and with **geometric FMG**. It honours - `solver.consistent_jacobian`: use `True` (consistent Newton) for smooth nonlinear - rheologies; `"continuation"` (staged Picard→Newton) for robustness far from the - solution. The rotated constraint is transparent to the tangent. + `solver.consistent_jacobian`: `True` (consistent Newton, the default — the + residual is symbolic, so the tangent is exact); `"continuation"` (staged + Picard→Newton) for robustness far from the solution; `False` (Picard) only where + a hard-yield viscoplastic solve needs it as an entry requirement. The rotated + constraint is transparent to the tangent, and the rotated path takes no warm-up + sweep before Newton. - The constraint **reaction** is the boundary normal traction σ_nn (`solver.boundary_normal_traction(boundary)` / `solver.dynamic_topography(...)`) — no augmented-Lagrangian splitting. diff --git a/docs/developer/design/run-plan-and-transcript.md b/docs/developer/design/run-plan-and-transcript.md index 58ad60264..66a7ec41a 100644 --- a/docs/developer/design/run-plan-and-transcript.md +++ b/docs/developer/design/run-plan-and-transcript.md @@ -146,6 +146,61 @@ insists that it is.** The signature requires a `dt`, so there is no container for "the next task". If the event clock is the general thing, the timestep is the common case rather than the definition. +## Where the adjoint lives, and where it stops + +The transcript now supplies two of the three things a discrete adjoint needs: +the ordered operator list, and the state each operator was linearised about +(a snapshot before the operator, bit-exact on restore). The third — the +linearisation itself — is a contract on each operator, not a pass over the +record: an operator provides it or declines with a reason. + +The declining is recorded first. Every `solve`, `history_shift` and +`swarm_advect` event carries `adjoint: {supported, reason}`, written when it +ran. The verdicts are structural: an implicit step is a residual (Jacobian +transpose for the state, symbolic derivative for a parameter); a rotated +constraint solves inside its own Krylov loop with no transpose path; an +unconverged solve is linearised about a state it never reached; a +semi-Lagrangian trace is differentiable in the velocity but its interpolation +at the departure points is not materialised; a particle step is adjointable +exactly when the particle set is fixed across it, which `swarm.advection` +checks by counting. + +Read back, the verdicts partition the run (`transcript_adjoint_segments`). +That partition is what data assimilation needs rather than perfect +invertibility: strong-constraint adjoint within a segment where every operator +is smooth, and across a refusal a control variable with an error covariance — +weak-constraint 4D-Var, with the joins chosen by the run. The optimiser needs +a descent direction that is the same inexact direction each iteration, not an +exact gradient; the exact discrete adjoint is the verification anchor where +the operators admit it, and the segments say where that anchor holds. + +Two things follow from the residual being symbolic. First, every first +derivative is always available: ∂R/∂u and ∂R/∂m are differentiated, not +approximated, so the gradient is never in question — and the tangent the +forward *iteration* used is irrelevant to it. Picard iterations spoil +nothing; the converged state is the same, and the adjoint assembles ∂R/∂u at +that state itself. Second, the same is not automatically true at second +order. A Hessian — for posterior covariance, or a Newton step on the outer +optimisation — needs ∂²R/∂u², ∂²R/∂u∂m, and a yield law written with `Min` +or a softmin has a second derivative that is a distribution at the yield +surface. Those terms exist symbolically, but they have to be handled with +care rather than differentiated and trusted. + +What follows from it, in order: `adjoint_solve`, `dual_of` and +`sensitivity` on the solvers — landed, checked against finite differences +on Poisson, on a non-symmetric SUPG step, and on Stokes with a linear and a +strain-rate-dependent viscosity, with the consistent tangent assembled for +the adjoint whichever tangent the forward iteration used; the reverse driver +(`uw.adjoint.TranscriptAdjoint`) — landed: it walks the transcript backwards, +restores each step's snapshot, replays each solve to its own input state, and +reads what each solve depends on from its residual, checked to 1e-7 against +finite differences on a two-solver run, including a field read through its +gradient (the Crank–Nicolson old flux), assembled as a FEM load rather than +by parts; the two transport operators materialised — interpolation at +departure points and ∂X_dep/∂v, which lift the semi-Lagrangian refusal; and +a Taylor test in the library (`test_0020`, and the sinker example through +the library at 1.00000). + ## Inferred plan, then declared plan The plan is **inferred** today — the figure takes the most common step as the diff --git a/docs/developer/guides/HOW-TO-WRITE-UW3-SCRIPTS.md b/docs/developer/guides/HOW-TO-WRITE-UW3-SCRIPTS.md index c5fcd6629..5dfeb9a4f 100644 --- a/docs/developer/guides/HOW-TO-WRITE-UW3-SCRIPTS.md +++ b/docs/developer/guides/HOW-TO-WRITE-UW3-SCRIPTS.md @@ -642,6 +642,129 @@ shown, and your own filters still apply. "The velocity block fell back to gamg" changes what the numbers mean, and a record that kept the residual norms but not that line would be an account of the run with the explanation removed. +**Every operator says whether it admits a discrete adjoint.** Each `solve`, +`history_shift` and `swarm_advect` event carries a verdict, written when the +operator ran: + +```json +{"kind": "solve", "name": "SNES_Stokes(v)", + "adjoint": {"supported": true, + "reason": "implicit residual: Jacobian transpose for the state, symbolic derivative of the residual for a parameter"}} +``` + +The verdict is structural — about the operator as configured, not about +whether a driver exists yet — so a run says where its adjoint breaks *while it +runs*. What refuses, and why: + +- a rotated constraint (free-slip or fault contact): the solve runs on a + rotated operator inside its own Krylov loop, with no transpose path; +- a solve that did not converge: a linearisation about a state the solve + never reached is not the adjoint of anything — the outcome overrides the + structural verdict after the fact; +- a semi-Lagrangian history: the departure-point trace is differentiable in + the velocity, but the interpolation at the departure points is not + materialised as an operator; +- a swarm step whose particle set changed — `swarm.advection` records the + count before and after, and a particle removed on leaving the domain + changes the dimension of the state. The rule is one line: a particle step + is adjointable exactly when the particle set is fixed across it. + +An Eulerian or SUPG history is supported — an implicit step is a residual, +and the SUPG adjoint that passed its Taylor test at 1.00000 is exactly that +case. The text transcript notes where the adjoint breaks, once per change +rather than on every step. + +`uw.transcript_adjoint_segments(source)` reads the verdicts back as the +partition they imply — maximal runs of steps whose every operator admits an +adjoint, separated by the steps where one refused. That partition is the +assimilation window's structure: strong-constraint adjoint within a segment; +across a refusal, a control variable and an error covariance, which is +weak-constraint 4D-Var with the joins chosen by the run rather than by hand. +Nothing is approximated silently — the refusal says what the model was +allowed to be wrong about. + +**The adjoint of one solve is built in.** For a solver whose verdict is +"supported", the discrete adjoint is two calls, with no hand algebra: + +```python +b = -solver.dual_of(T.sym[0] - T_target.sym[0]) # -dJ/dT for J = 1/2 int (T - T*)^2 +mu, reason = solver.adjoint_solve(b, target=mu_var) # K^T mu = b, K the SNES Jacobian +dJ_dkappa = solver.sensitivity(mu_var, kappa) # int (dF/dkappa) . mu, symbolic dF/dkappa +``` + +`dual_of` assembles the right-hand side on the solver's own space, so the +Dirichlet nodes are excluded and the multiplier comes back zero there — the +homogenised adjoint conditions, without stating them. `sensitivity` follows +the parameter through the constitutive model's own symbol (the residual holds +`\upkappa`, whose value is your `kappa`), so the chain rule reaches it. + +One thing to get right, because `solve()` moves it: a time step's residual is +`F(u_new; u_old, v, dt)`, and the history manager shifts `u_old` out of its +slot in the post-solve hook. Put the step's input back before linearising — +`solver.DuDt.psi_star[0].array[...] = u_old` — or the sensitivity is a few +per cent wrong on a SUPG step (measured). + +Stokes takes the same transpose on its composite (u, p) system, with +`target=(u_adj, p_adj)` and `dual_of` taking a velocity-space expression. +With a linear viscosity the operator is symmetric, and this reproduces the +second-solver construction in `docs/examples/adjoint`. With a strain-rate- or +pressure-dependent viscosity the adjoint is the transpose of the **consistent +tangent** ∂R/∂u, which that construction cannot build. Picard iterations in +the forward solve spoil nothing — the converged state is the same, and ∂R/∂u +is a function of that state alone — but they leave the SNES holding the +frozen-viscosity Jacobian *kernel*. So when the forward ran Picard on a +nonlinear residual, `adjoint_solve` switches the kernel to the consistent +tangent for its assembly (a JIT rebuild; the DM and KSP are kept), transposes +that, and puts the Picard kernel back for the next forward solve. The +verdict says so. + +**The whole run, backwards.** With `model.record_every = 1` the transcript is +a forward tape — the operators per step, and the state each step started +from — and `uw.adjoint.TranscriptAdjoint` walks it in reverse with no +problem-specific wiring: + +```python +final = model.save_state() # the N+1th level +back = uw.adjoint.TranscriptAdjoint(model, final) +result = back.gradient(misfit_integrand, parameters=[eta0], fields=[beta]) +result["parameters"][eta0] # dJ/d eta0 +result["fields"][beta] # dJ/d beta_0, as a dual field +``` + +For each solve, in reverse order of the record, it restores the step's +snapshot, replays the solves before it, replays it, and puts each history's +input back where the post-solve hook shifted it — so the residual is +linearised at the solve's own input state without anyone touching +`psi_star`. The residual then says what the solve read: every field in +`F0`/`F1` other than the unknown gets the dual `(dR/df)^T mu`, a history +slot's dual goes to the field it tracks at the previous level, and every +parameter gets `mu^T dR/dm`. A field read through its *gradient* — a +Crank–Nicolson step (θ = 0.5, the `AdvDiffusion` default) reads the old +level as `κ∇T_old` — gets the gradient part of the load too: the dual is +assembled as the FEM load `∫ g₀ φⱼ + g₁·∇φⱼ` by a generic solver's residual +at zero, so there is no integration by parts and no boundary term to drop. +A dual is held as a field (one coefficient per node); pair it with a +direction using `uw.adjoint.inner(field, dual, direction)`, which sums over +the owned degrees of freedom (a NumPy dot on `.array` counts a partition's +ghost nodes twice), so a control `c` with `f_0 = f_0(c)` finishes with +`inner(f, dual, d f_0 / d c)`. + +Two things the tape has to contain. Every solve must be inside a step — a +Stokes solve taken before the loop to make `v_0` is invisible to the walk, +and its dependence on the parameters with it. And a driver that runs the +forward model more than once must reset the Eulerian history each time it +sets the initial condition (`adv.DuDt.initialise_history()`), or the second +run reads the first run's history. Checked in `tests/test_0020` on a +two-solver, two-step sinking blob: the viscosity gradient and the dual on +the initial level set both match central finite differences to 1e-4 +(measured 1e-7), at θ = 1 and at the default θ = 0.5, serially and on two +ranks. + +All of it is checked against central finite differences in `tests/test_0019`: +Poisson; one SUPG step, where the Jacobian is not symmetric and a transpose +taken the wrong way round would show; Stokes with a constant viscosity; and +Stokes with η(ε̇) under the consistent tangent. + The figure marks the same three states per solve — converged, converged with a fieldsplit block that hit its iteration cap, and diverged. The middle one is worth the separate mark: a capped block did not solve, so the Schur operator diff --git a/docs/developer/guides/adversarial-review.md b/docs/developer/guides/adversarial-review.md index c9e45368b..08f312945 100644 --- a/docs/developer/guides/adversarial-review.md +++ b/docs/developer/guides/adversarial-review.md @@ -82,6 +82,12 @@ silently understates the run. See [HOW-TO-WRITE-UW3-SCRIPTS](HOW-TO-WRITE-UW3-SCRIPTS.md) and `docs/developer/design/run-plan-and-transcript.md`. +**Every operator gives an adjoint verdict.** A `solve`, `history_shift` or +`swarm_advect` event carries `adjoint: {supported, reason}`, written when it +ran. A new history scheme without `_adjoint_support()` fails +`tests/test_0018_adjoint_support_record.py`; a new operation on model state +that records no verdict lets a run claim invertibility it does not have. + **Named quantities keep their names.** A coefficient written as `uw.expression(r"\rho_0 \alpha g", ...)` appears in the description under that name. An anonymous float collapses into the assembled product and the diff --git a/docs/examples/adjoint/sinker_transcript/README.md b/docs/examples/adjoint/sinker_transcript/README.md new file mode 100644 index 000000000..80384778a --- /dev/null +++ b/docs/examples/adjoint/sinker_transcript/README.md @@ -0,0 +1,228 @@ +# The sinking-blob adjoint, written in the timestepping pattern + +Same problem and same mathematics as `../supg/`. What changes is the +scaffolding — and the point of the exercise is that the scaffolding is now +library machinery rather than something this project invented for itself. + +Run it: + +```bash +python generate_target_transcript.py # twin experiment: v(T) at the true centre +python forward_sinker_transcript.py # one forward run, printing its transcript +python taylor_test_transcript.py # the gate, with the hand-rolled adjoint below +python taylor_test_library.py # the same gate, with uw.adjoint.TranscriptAdjoint +``` + +The second Taylor test replaces everything in `inverse_sinker_transcript.py` +with one library call: the transcript is the tape, each solve's residual says +what it reads, and the backward pass needs no wiring from this script beyond +the misfit and the control. The initial Stokes solve goes on the tape as a +zero-length step (`solve_forward(..., initial_on_tape=True)`) so the walk +sees `beta_0 -> v_0`; the hand-rolled version accounts for that solve itself. + +--- + +## What the forward model looks like now + +```python +uw.reset_default_model() +uwmodel = uw.get_default_model() +uwmodel.set_reference_quantities( + domain_depth=uw.quantity(500, "km"), + material_viscosity=uw.quantity(1e21, "Pa*s"), + lithostatic_pressure=REF_DENSITY * GRAVITY * DOMAIN_DEPTH, +) + +# ... mesh, variables, solvers ... + +uwmodel.clear_transcript() +uwmodel.tracker.time = uw.quantity(0.0, "Myr") +uwmodel.tracker.step = 0 +uwmodel.record_every = 1 + +for _ in range(nsteps): + with uwmodel.step(dt, label="sink"): + adv.solve(timestep=dt, zero_init_guess=False) + stokes.solve(zero_init_guess=False) +``` + +and it prints its own account of itself: + +``` + history_shift:EulerianSUPG(beta) -> solve:SNES_Stokes(v)> + +... +restorable: 5 of 5 +``` + +### Three scales, not ten constants + +`../supg/forward_sinker_supg.py` opens with + +```python +REF_LENGTH = 500e3 +REF_DENSITY = 3300.0 +REF_GRAVITY = 9.81 +REF_VISCOSITY = 1e21 +REF_PRESSURE = REF_DENSITY * REF_GRAVITY * REF_LENGTH +REF_TIME = REF_VISCOSITY / REF_PRESSURE +REF_VELOCITY = REF_LENGTH / REF_TIME +DENSITY_BACKGROUND = 3200.0 / REF_DENSITY +... +DT_FIXED = 570.0 +``` + +and every number after that is a ratio you have to keep in your head. Here the +three scales that actually fix this problem are declared once — a length, a +viscosity, and the lithostatic stress `rho g L` — and everything downstream is +written in the units it is quoted in: `50 km`, `3300 kg/m^3`, `1.1159 Myr`. + +Density is deliberately **not** one of the reference quantities. The Stokes +sinker uses it only as a ratio, which is exactly why `REF_DENSITY` cancelled +out of every nondimensional number the old script produced. + +The nondimensional problem that reaches the solver is bit-for-bit the one the +old script assembled by hand — the interface at `t = T` sits at 271.2 km and +371.2 km on the centreline either way, against 0.5425 and 0.7425 of the box +before. The difference is that the scaling is now stated and checked rather +than spread across a module header. + +### There is no checkpoint dictionary + +The old forward model carried its own recording: + +```python +ck = {"B": [b0.copy()], "V": [...], "P": [...], "dt": []} +for k in range(nsteps): + adv.solve(timestep=dt, zero_init_guess=False) + stokes.solve(zero_init_guess=False) + ck["B"].append(np.asarray(beta.array)[:, 0, 0].copy()) + ck["V"].append(np.asarray(v.array)[:, 0, :].copy()) + ck["P"].append(np.asarray(p.array)[:, 0, 0].copy()) + ck["dt"].append(dt) +``` + +Four lists holding precisely the arrays the adjoint turned out to want. That is +a recording you can only write once you already know the adjoint — which is the +wrong way round, and is a large part of why an adjoint is normally a rewrite of +the forward model rather than an addition to it. + +It is also silently incomplete. It does not hold the transport history; it +works only because backward Euler happens to make `psi_star` recoverable from +`B[k]`. Change `theta`, add a second history level, put the level set on a +swarm, and the dictionary is quietly wrong in a way nothing detects. + +`model.record_every = 1` replaces all of it with a request. The snapshot each +step keeps is the *whole* model state — every mesh variable, every swarm, every +registered state-bearer including the DDt history, and the clock — captured +before the step's operators ran, which is the only correct point. + +--- + +## What the adjoint looks like now + +The three blocks and the backward recursion are untouched. What changed is +where the states come from. + +| | `../supg/` | here | +|---|---|---| +| the record | `ck` dict built by the forward loop | `model.transcript`, built by the library | +| a state | `ck["B"][k]`, `ck["V"][k]`, `ck["P"][k]` | `model.load_state(transcript[k].snapshot)` | +| the history | reconstructed from `ck["B"][k]` | restored with everything else | +| the clock | not recorded | restored with everything else | +| what ran | assumed | `transcript[k].events`, in order | + +```python +def linearisation_state(self, k): + """Restore the point step k's transport residual was linearised at.""" + entry = self.transcript[k] + self.state_at(k) # beta_k, V_k, history + beta_in = np.asarray(beta.array)[:, 0, 0].copy() + adv.solve(timestep=entry.dt, zero_init_guess=False) # -> beta_{k+1} + self.psi_star.array[:, 0, 0] = beta_in # where the residual reads it +``` + +Two things are worth stating plainly about that replay. + +**It is exact.** Restoring a snapshot and re-solving reproduces the step to the +last bit (measured: `maxdiff 0.00e+00` on every step). Re-*running* the script +does not — warm starts and preconditioner reuse are solver history, not model +state, so two independent runs of the same problem diverge at the 1e-13 level +from the first step. If you need to look at a step twice, restore it. + +**It costs one extra transport solve per step**, and that is what buys the +recording being generic: the forward run does not have to know an adjoint is +coming. Trading a solve for not having to write a bespoke tape is the right +trade at this size; on a long run you would raise `record_every` and recompute +between restore points, which is the standard checkpointing schedule and is +what `record_every` / `record_limit` are for. + +**One line is still scheme-specific.** Putting `beta_in` back into `psi_star` +is a statement about backward Euler, not about the record: the residual of step +`k` reads the step's input from the history slot, and the solve's post-hook has +already shifted it forward. A `solver.adjoint(...)` method would own that line; +today the driver does. + +### The one thing the transcript cannot hold + +An N-step run has N+1 time levels, and the transcript records *steps*. So +`solve_forward` returns `(transcript, final_state)`, and `state_at(N)` reads the +final state rather than a transcript entry. That asymmetry is real, not an +oversight — the last level is the run's output, not the input to anything. + +--- + +## Two library changes this exercise produced + +**`model.clear_transcript()` (new).** A driver that runs the same model many times +— an inversion, a parameter sweep — needs each run to have its own account. +Without it the transcript is the concatenation of every run the process has done, +and `rewind()` walks back into the previous one. The Taylor test runs the +forward model thirteen times; it found this immediately. + +**A snapshot no longer rescales the mesh (fixed).** `mesh.X.coords` is the +unit-aware view and returns metres once a model declares a length scale; the +DM coordinate vector that restore writes back into holds model units. Capture +took the first and restore wrote the second, so **every restore multiplied the +mesh by the length scale** — 500 km became 250,000,000 km. Nothing raised: +shapes matched, fields came back correctly, only the geometry was wrong. The +symptom is that `uw.function.evaluate` starts returning the value at one corner +for every sample point, because every sample point is now outside the domain. + +`model.rewind()` goes straight through that path, which is how it surfaced. +Covered now by `tests/test_0012_snapshot_units_coords.py`. + +--- + +## The gate + +`taylor_test_transcript.py`, control at (300 km, 350 km), true centre +(250 km, 375 km), five steps, contrast 1000: + +``` +J = 5.473279e-10 adjoint dJ/dcx0 = 2.022362e-14 /m dJ/dcy0 = -9.883878e-16 /m +h=5.0 km: FD dJ/dcx0 2.021281e-14 ratio 0.99947 | FD dJ/dcy0 -9.842465e-16 ratio 0.99581 +h=0.5 km: FD dJ/dcx0 2.022355e-14 ratio 1.00000 | FD dJ/dcy0 -9.883048e-16 ratio 0.99992 +h=0.05 km: FD dJ/dcx0 2.022356e-14 ratio 1.00000 | FD dJ/dcy0 -9.884208e-16 ratio 1.00003 +``` + +Same quality as `../supg/` (1.00000 / 0.99993). The gradients are small +because the control is now a length in metres rather than a fraction of the +box; multiply by 5e5 to compare with the old numbers. + +--- + +## What is unchanged, and still load-bearing + +The three physics choices from `../supg/` that make this adjoint exact: + +- **Eulerian SUPG transport**, so one timestep is a residual and every + sensitivity is a SymPy derivative of it. +- **Backward Euler (`theta = 1`)**, so `beta_old` appears in the residual + undifferentiated and the history coupling is a pointwise expression rather + than a weak form. +- **A fixed timestep**, so the objective does not depend on the control through + the schedule. This was the dominant defect in the original adjoint. + +See `../supg/README.md` for the derivation and the ablation that established +each of them. diff --git a/docs/examples/adjoint/sinker_transcript/forward_sinker_transcript.py b/docs/examples/adjoint/sinker_transcript/forward_sinker_transcript.py new file mode 100644 index 000000000..f6cad0e46 --- /dev/null +++ b/docs/examples/adjoint/sinker_transcript/forward_sinker_transcript.py @@ -0,0 +1,272 @@ +# %% [markdown] +""" +# Sinking Blob — Forward Model, written in the timestepping pattern + +Same physics as `../supg/forward_sinker_supg.py` and the same Eulerian SUPG +transport. What changes is the *scaffolding*, and only the scaffolding: + +**The model and its reference quantities come first.** Not a wall of `REF_*` +constants and hand-divided ratios — a declaration of the three scales this +problem actually uses (a length, a viscosity, and the lithostatic stress +`rho g L`), after which every number in the script is written in the units it +is quoted in. `500 km`, `1e21 Pa s`, `1.116 Myr`. The nondimensional problem +that reaches the solver is bit-for-bit the one the old script assembled by +hand; the difference is that here the scaling is stated once and checked, +rather than spread over ten module constants. + +**The timestep is a `model.step(dt)` block.** Which makes it a transaction: +the clock reads the end of the interval for the whole block (where an implicit +scheme centres its residual), the advance commits only on clean exit, and +everything the block did lands in `model.transcript` in order. + +**There is no checkpoint dictionary.** The old script carried its own +`ck = {"B": [...], "V": [...], "P": [...], "dt": [...]}` — a hand-rolled +recording of exactly the arrays the adjoint happened to need, which is a thing +you can only write once you already know what the adjoint is. Here +`model.record_every = 1` asks each step to keep the state it started from, and +that snapshot is the *whole* model state: fields, transport history, clock. +The adjoint in `inverse_sinker_transcript.py` reads the transcript instead. + +The three physics choices from the SUPG version are unchanged and still +load-bearing for the adjoint: Eulerian SUPG transport, backward Euler +(`theta = 1`), and a fixed timestep. +""" + +# %% +import os +import numpy as np +import sympy +import underworld3 as uw + +# --- the scales this problem is written in ---------------------------------- +# Three quantities fix the scaling completely: a length, a viscosity, and a +# stress. Density is NOT one of them — it enters the Stokes sinker only as a +# ratio, which is why the old script's REF_DENSITY cancelled out of every +# nondimensional number it produced. +DOMAIN_DEPTH = uw.quantity(500, "km") +REF_VISCOSITY = uw.quantity(1e21, "Pa*s") +REF_DENSITY = uw.quantity(3300, "kg/m**3") +GRAVITY = uw.quantity(9.81, "m/s**2") +LITHOSTATIC_PRESSURE = REF_DENSITY * GRAVITY * DOMAIN_DEPTH + +# --- geometry and materials, quoted in their own units ---------------------- +RESOLUTION = 16 +NSTEPS = 5 + +DENSITY_BACKGROUND = uw.quantity(3200, "kg/m**3") +DENSITY_BLOCK = uw.quantity(3300, "kg/m**3") + +BLOB_CENTER = (uw.quantity(250, "km"), uw.quantity(375, "km")) # 0.5, 0.75 of the box +BLOB_RADIUS = uw.quantity(50, "km") +SMOOTHING_WIDTH = 1.5 * DOMAIN_DEPTH / RESOLUTION + +# Fixed timestep. `stokes.estimate_dt()` at t=0 for the true configuration +# returns about this; fixing it keeps the objective from depending on the +# control through the schedule, which was the dominant defect in the original +# adjoint. 570 dimensionless units of eta/(rho g L) is 1.1159 Myr. +DT = uw.quantity(1.1158811388500671, "Myr") + +OUTPUT_DIR = os.path.join(os.path.dirname(__file__), "output", "forward") + + +# %% +def build_model(viscosity_contrast, resolution=RESOLUTION, stokes_tolerance=1.0e-10, + theta=1.0): + """Declare the model, then the mesh, then the solvers — in that order. + + Reference quantities must precede mesh creation, so the model is the first + line of the script rather than something the mesh conjures for you. + + Returned as a dict so `inverse_sinker_transcript.py` can reuse the SAME + objects the forward run used — the adjoint reads the transport solver's + residual (`adv.F0`, `adv.F1`) and its assembled Jacobian, so it must be the + very solver that produced the run, not a rebuilt copy. + """ + uw.reset_default_model() + uwmodel = uw.get_default_model() + uwmodel.set_reference_quantities( + domain_depth=DOMAIN_DEPTH, + material_viscosity=REF_VISCOSITY, + lithostatic_pressure=LITHOSTATIC_PRESSURE, + ) + + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=1.0 / resolution, + regular=False, + qdegree=3, + ) + x, y = mesh.X + + v = uw.discretisation.MeshVariable("v", mesh, 2, degree=2) + p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1) + beta = uw.discretisation.MeshVariable( + "beta", mesh, vtype=uw.VarType.SCALAR, degree=3, continuous=True + ) + + stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + stokes.add_dirichlet_bc((sympy.oo, 0.0), "Top") + stokes.add_dirichlet_bc((sympy.oo, 0.0), "Bottom") + stokes.add_dirichlet_bc((0.0, sympy.oo), "Left") + stokes.add_dirichlet_bc((0.0, sympy.oo), "Right") + + if viscosity_contrast > 1e3: + penalty = 100.0 + elif viscosity_contrast > 1e1: + penalty = 10.0 + else: + penalty = 1.0 + stokes.penalty = penalty + stokes.tolerance = stokes_tolerance + stokes.petsc_options.delValue("ksp_monitor") + + # The level set carries a LENGTH (it is a signed distance), so the tanh + # smoothing width is a length too. With the model declared, that is just + # what the expression says; without it, both were nondimensional numbers + # whose relationship to the mesh you had to keep in your head. + smoothing_nd = _nd(SMOOTHING_WIDTH / DOMAIN_DEPTH) + indicator = 0.5 * (1.0 - sympy.tanh(beta.sym[0] / smoothing_nd)) + eta = sympy.exp(indicator * sympy.log(viscosity_contrast)) + density_ratio = _nd(DENSITY_BACKGROUND / REF_DENSITY) + indicator * _nd( + (DENSITY_BLOCK - DENSITY_BACKGROUND) / REF_DENSITY) + + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = eta + stokes.bodyforce = sympy.Matrix([0, -density_ratio]) + + # Eulerian SUPG transport of the level set. AdvDiffusion's default DDt + # plugin is EulerianSUPG; theta=1 is backward Euler (see module docstring). + adv = uw.systems.AdvDiffusion(mesh, u_Field=beta, V_fn=v.sym) + adv.constitutive_model = uw.constitutive_models.DiffusionModel + adv.constitutive_model.Parameters.diffusivity = 0.0 + adv.Unknowns.DuDt.theta = theta + adv.tolerance = stokes_tolerance + adv.petsc_options.delValue("ksp_monitor") + + return dict(uwmodel=uwmodel, mesh=mesh, x=x, y=y, v=v, p=p, beta=beta, + stokes=stokes, adv=adv, eta=eta, density=density_ratio, + indicator=indicator, penalty=penalty, contrast=viscosity_contrast) + + +def _nd(q): + """The plain number a dimensionless quantity stands for.""" + try: + return float(q.to("dimensionless").magnitude) + except AttributeError: + return float(q) + + +# %% +def beta0_nodal(model, centre): + """Initial level set at the beta nodes, and its derivative w.r.t. the centre. + + `centre` is a pair of lengths. The mesh is unit-square in model units, so + the control is converted once, here, and the gradient this function + returns is therefore d(beta_0)/d(centre) in the SAME units — which is what + makes the adjoint's final dot product dimensionally honest. + """ + scale = _length_scale() + cx, cy = (_mag(c) / scale for c in centre) + R = _mag(BLOB_RADIUS) / scale + + X = np.asarray(model["beta"].coords)[:, :2] / scale + r = np.sqrt((X[:, 0] - cx) ** 2 + (X[:, 1] - cy) ** 2) + dbeta_dc = np.stack([-(X[:, 0] - cx) / r, -(X[:, 1] - cy) / r], axis=1) / scale + return r - R, dbeta_dc + + +def _length_scale(): + """Metres per model length unit.""" + return float(uw.get_default_model().get_fundamental_scales()["length"].to("m").magnitude) + + +def _mag(q): + return float(q.to("m").magnitude) + + +# %% +def solve_forward(model, centre, nsteps=NSTEPS, dt=DT, initial_on_tape=False): + """Run the forward model from `centre`. + + Returns `(transcript, final_state)`. The transcript is the record of the run: + one entry per step, holding the interval it covered, the operators it + applied in order, and the state it started from. `final_state` is the one + state the transcript cannot hold — an N-step run has N+1 time levels, and the + transcript records steps. + """ + uwmodel = model["uwmodel"] + beta, stokes, adv = model["beta"], model["stokes"], model["adv"] + + b0, _ = beta0_nodal(model, centre) + beta.array[:, 0, 0] = b0 + # The Eulerian history initialises itself only on its FIRST solve, so a + # solver reused for a second independent run silently carries the previous + # run's psi_star. An inversion driver runs the forward model many times; + # reset the history explicitly every time the initial condition is set. + adv.Unknowns.DuDt.initialise_history() + + # A new run gets a new transcript and a clock at zero. Without the clear, the + # transcript would be the concatenation of every run this process has done and + # rewind() would walk back into the previous one. + uwmodel.clear_transcript() + uwmodel.tracker.time = uw.quantity(0.0, "Myr") + uwmodel.tracker.step = 0 + uwmodel.tracker.dt = None + uwmodel.record_every = 1 # keep the state every step started from + uwmodel.record_limit = None # this run is short; keep all of them + + if initial_on_tape: + # v_0 = Stokes(beta_0) as a zero-length step, so the library's backward + # pass (uw.adjoint.TranscriptAdjoint) sees beta_0 -> v_0. The + # hand-rolled adjoint in inverse_sinker_transcript.py accounts for + # this solve itself and expects one transport per entry, so it keeps + # the solve off the tape. + with uwmodel.step(0 * dt, label="initial"): + stokes.solve(zero_init_guess=True) + else: + stokes.solve(zero_init_guess=True) + for _ in range(nsteps): + with uwmodel.step(dt, label="sink"): + adv.solve(timestep=dt, zero_init_guess=False) + stokes.solve(zero_init_guess=False) + + return uwmodel.transcript, uwmodel.save_state() + + +# %% +if __name__ == "__main__": + import sys + + contrast = float(sys.argv[1]) if len(sys.argv) > 1 else 1000.0 + model = build_model(contrast) + transcript, _final = solve_forward(model, BLOB_CENTER) + uwmodel, mesh, beta = model["uwmodel"], model["mesh"], model["beta"] + + uw.pprint(f"contrast {contrast:g}, dt {DT}, {NSTEPS} steps") + uw.pprint(f"clock now {uwmodel.tracker.time.to('Myr')}, " + f"step {uwmodel.tracker.step}") + uw.pprint("") + uw.pprint("the transcript:") + for entry in transcript: + uw.pprint(f" {entry}") + uw.pprint(f" restorable: {len(uwmodel.restore_points)} of {len(transcript)}") + uw.pprint("") + + # Where is the interface? Sampled on the vertical centreline. Coordinates + # are dimensional now, so the sample line is quoted in km. + scale = _length_scale() + line = np.column_stack([np.full(400, 250e3), np.linspace(175e3, 475e3, 400)]) / scale + + def crossings(): + vals = np.asarray(uw.function.evaluate(beta.sym[0], line)).ravel() + sgn = np.where(np.diff(np.sign(vals)) != 0)[0] + return [f"{line[i, 1] * scale / 1e3:.1f} km" for i in sgn] + + uw.pprint(f"interface at t=T on x=250 km: {crossings()}") + + # And the point of recording it: put the run back one step and look again. + uwmodel.rewind() + uw.pprint(f"after rewind(): clock {uwmodel.tracker.time.to('Myr')}, " + f"step {uwmodel.tracker.step}, transcript {len(uwmodel.transcript)} steps") + uw.pprint(f"interface one step earlier : {crossings()}") diff --git a/docs/examples/adjoint/sinker_transcript/generate_target_transcript.py b/docs/examples/adjoint/sinker_transcript/generate_target_transcript.py new file mode 100644 index 000000000..fd93a9e8f --- /dev/null +++ b/docs/examples/adjoint/sinker_transcript/generate_target_transcript.py @@ -0,0 +1,22 @@ +"""Twin-experiment target for the transcript forward model: run at the true centre, +save v(T). Kept on disk (not in memory) so a candidate run on a freshly built +mesh reads it back through the coordinate-remapping `read_timestep`, exactly as +the original project does.""" +import os +import underworld3 as uw +from forward_sinker_transcript import build_model, solve_forward, BLOB_CENTER + +TRUE_VISCOSITY_CONTRAST = 1000.0 +TARGET_FILENAME = "sinker_target_transcript" +TARGET_INDEX = 0 +OUTPUT_DIR = os.path.join(os.path.dirname(__file__), "output", "target") + +if __name__ == "__main__": + if uw.mpi.rank == 0: + os.makedirs(OUTPUT_DIR, exist_ok=True) + model = build_model(TRUE_VISCOSITY_CONTRAST) + transcript, _ = solve_forward(model, BLOB_CENTER) + model["mesh"].write_timestep(TARGET_FILENAME, index=TARGET_INDEX, + outputPath=OUTPUT_DIR, meshVars=[model["v"]]) + uw.pprint(f"true centre {BLOB_CENTER}; {len(transcript)} steps; " + f"saved v(T) to {OUTPUT_DIR}/{TARGET_FILENAME}") diff --git a/docs/examples/adjoint/sinker_transcript/inverse_sinker_transcript.py b/docs/examples/adjoint/sinker_transcript/inverse_sinker_transcript.py new file mode 100644 index 000000000..25ce8bc8f --- /dev/null +++ b/docs/examples/adjoint/sinker_transcript/inverse_sinker_transcript.py @@ -0,0 +1,307 @@ +# %% [markdown] +""" +# Sinking Blob — Adjoint driven from the model transcript + +The mathematics is identical to `../supg/inverse_sinker_supg.py`: one implicit +transport step is a residual + + F(beta_new; beta_old, v, dt) = 0, R_j = integral[ F0 phi_j + F1 . grad(phi_j) ] + +so its adjoint is a transpose solve against the SAME Jacobian the SNES already +assembles, and the couplings to the previous level and to the velocity are +SymPy derivatives of `F0` and `F1`. + +What changes is where the backward pass gets its states from. + +## The old version carried its own recording + + ck = {"B": [...], "V": [...], "P": [...], "dt": [...]} + +Four lists of numpy arrays, appended inside the forward loop, holding exactly +the quantities this adjoint turned out to need. That is a recording you can +only write once you already know the adjoint — which is the wrong way round, +and is why an adjoint is normally a rewrite of the forward model rather than an +addition to it. It is also silently incomplete: it does not hold the transport +history, so it only works because backward Euler happens to make `psi_star` +recoverable from `B[k]`. + +## This version reads the run's own record + +`model.record_every = 1` asks each step to keep the state it started from — +the whole state, fields and transport history and clock together — and +`model.transcript` is the ordered account of what each step did. The backward pass +walks that list: + + for k in reversed(range(len(transcript))) + model.load_state(transcript[k].snapshot) # the state step k started from + adv.solve(timestep=transcript[k].dt) # replay it: bit-for-bit + ... + +Two things are worth stating plainly about that replay. It is exact — +restoring a snapshot and re-solving reproduces the step to the last bit, where +re-running the script does not, because warm starts and preconditioner reuse +are solver history rather than model state. And it costs one extra transport +solve per step, which is what buys the recording being generic: the run does +not have to know an adjoint is coming. + +## The three blocks (unchanged) + +For a multiplier `mu` living on the step's unknown, with `mu_h` the finite +element function whose nodal values are `mu`: + + K = dF/d(beta_new) the assembled SNES Jacobian + L^T mu = dual of (dF0/d beta_old) mu_h + (dF1/d beta_old) . grad(mu_h) + G^T mu = dual of (dF0/d v_i) mu_h + (dF1/d v_i) . grad(mu_h) + +pointwise because `theta = 1` (backward Euler) leaves `beta_old` in the +residual undifferentiated, and because the velocity never enters as `grad(v)`. + +## The backward recursion + + RHS_N = S_N + K_l^T mu_{l+1} = -RHS_{l+1} l = N-1 .. 0 + RHS_l = S_l + L_l^T mu_{l+1} + dL/d(beta_0) = RHS_0 + +then `dJ/dc = dL/d(beta_0) . d(beta_0)/dc` as a plain dot product, because +`dL/d(beta_0)` is already a dual. +""" + +# %% +import numpy as np +import sympy +import underworld3 as uw +from scipy import sparse +from scipy.sparse.linalg import splu + +from forward_sinker_transcript import ( + build_model, beta0_nodal, solve_forward, BLOB_CENTER, DT, NSTEPS, +) + + +# %% +def _jacobian_csr(solver): + """Assembled Jacobian of a solver's residual, as scipy CSR. + + `uw.systems.Projection`'s SNES Jacobian IS the mass matrix of its space, so + the same helper gives both the transport Jacobian and the mass matrix used + to build duals. PETSc leaves the matrix zeroed until a Jacobian evaluation + is forced, hence `computeJacobian`. + """ + jac = solver.snes.getJacobian() + A, P = jac[0], jac[1] + solver.mesh.update_lvec() + xv = solver.snes.getSolution().duplicate() + xv.set(0.0) + solver.snes.computeJacobian(xv, A, P) + M = sparse.csr_matrix(A.getValuesCSR()[::-1]) + if M.nnz == 0 or abs(M).max() == 0.0: + M = sparse.csr_matrix(P.getValuesCSR()[::-1]) + return M + + +# %% +class AdjointMachinery: + """Everything the backward pass needs, built once for a model.""" + + def __init__(self, model, target_v, stokes_tolerance=1.0e-10): + mesh = model["mesh"] + self.model = model + self.uwmodel = model["uwmodel"] + self.mesh = mesh + beta, v = model["beta"], model["v"] + adv, stokes = model["adv"], model["stokes"] + + # --- multiplier field, and the projection used for every dual --- + self.mu = uw.discretisation.MeshVariable( + "mu", mesh, vtype=uw.VarType.SCALAR, degree=3, continuous=True) + self.scratch = uw.discretisation.MeshVariable( + "scratch", mesh, vtype=uw.VarType.SCALAR, degree=3, continuous=True) + self.proj = uw.systems.Projection(mesh, self.scratch) + self.proj.smoothing = 0.0 + self.proj.tolerance = stokes_tolerance + self.proj.uw_function = sympy.sympify(1.0) + self.proj.solve() + self.M3 = _jacobian_csr(self.proj) + + # --- adjoint Stokes: same operator, homogenised free-slip BCs --- + self.u_adj = uw.discretisation.MeshVariable("u_adj", mesh, 2, degree=2) + self.q_adj = uw.discretisation.MeshVariable("q_adj", mesh, 1, degree=1) + self.f_adj = uw.discretisation.MeshVariable("f_adj", mesh, 2, degree=2) + sa = uw.systems.Stokes(mesh, velocityField=self.u_adj, pressureField=self.q_adj) + sa.constitutive_model = uw.constitutive_models.ViscousFlowModel + sa.constitutive_model.Parameters.shear_viscosity_0 = model["eta"] + sa.penalty = model["penalty"] + sa.tolerance = stokes_tolerance + sa.petsc_options.delValue("ksp_monitor") + sa.add_essential_bc((sympy.oo, 0.0), "Top") + sa.add_essential_bc((sympy.oo, 0.0), "Bottom") + sa.add_essential_bc((0.0, sympy.oo), "Left") + sa.add_essential_bc((0.0, sympy.oo), "Right") + sa.bodyforce = self.f_adj.sym + self.stokes_adj = sa + + # The velocity-block integrand contains grad(mu) of a P3 field, so it + # does NOT live in the P2 velocity space. The adjoint body force must be + # its L2 PROJECTION onto that space, not its nodal interpolant. + self.f_proj_var = uw.discretisation.MeshVariable("f_proj", mesh, 2, degree=2) + self.f_proj = uw.systems.Vector_Projection(mesh, self.f_proj_var) + self.f_proj.smoothing = 0.0 + self.f_proj.tolerance = stokes_tolerance + + # --- Stokes sensitivity --- + dSigma_dbeta = uw.function.derivative(stokes.F1, beta.sym[0]) + d_density_dbeta = uw.function.derivative(model["density"], beta.sym[0]) + self.stokes_sensitivity = ( + uw.maths.tensor.rank2_inner_product(sa.Unknowns.E, dSigma_dbeta) + + d_density_dbeta * self.u_adj.sym[1] + ) + + # --- transport residual derivatives (built by refresh, after the run) --- + self.psi_star = adv.Unknowns.DuDt.psi_star[0] + self.hist_integrand = None + self.vel_integrand = None + + # --- misfit --- + self.v_target = target_v + self.misfit_integrand = sympy.Rational(1, 2) * ( + (v.sym - target_v.sym).dot(v.sym - target_v.sym)) + + # --- the run being differentiated (set by `attach`) --- + self.transcript = [] + self.final_state = None + + # ------------------------------------------------------------------ + def attach(self, transcript, final_state): + """Point the backward pass at a completed forward run. + + `transcript` is `model.transcript` — one entry per step, each holding the + interval and the state the step started from. `final_state` is the + snapshot taken after the last step, which is the one state the transcript + does not hold: the transcript records STEPS, and there are N+1 levels to + an N-step run. + + Also (re)builds the transport-residual derivatives. `adv.F0` / `adv.F1` + are LIVE templates that re-evaluate when the solver's parameters + change, so reading `.sym` before the solver has been configured by a + solve snapshots a residual with the WRONG timestep. + """ + self.transcript = list(transcript) + self.final_state = final_state + + mesh, adv, v = self.mesh, self.model["adv"], self.model["v"] + F0 = adv.F0.sym[0, 0] + F1 = adv.F1.sym + mu_s = self.mu.sym[0] + grad_mu = mesh.vector.gradient(mu_s) + + def contract(wrt): + """(dF0/dwrt) mu + (dF1/dwrt) . grad(mu): the integrand whose dual + is the transposed block applied to mu.""" + d0 = uw.function.derivative(F0, wrt) + d1 = uw.function.derivative(F1, wrt) + out = d0 * mu_s + for i in range(mesh.dim): + out = out + d1[i] * grad_mu[i] + return out + + self.hist_integrand = contract(self.psi_star.sym[0]) + self.vel_integrand = [contract(v.sym[i]) for i in range(mesh.dim)] + + # ------------------------------------------------------------------ + def dual_of(self, expr): + """integral[ expr * phi_j ] for every P3 basis function phi_j.""" + self.proj.uw_function = expr + self.proj.solve() + return self.M3 @ np.asarray(self.scratch.array)[:, 0, 0] + + def state_at(self, level): + """Restore the model to time level `level` of the recorded run. + + Level `l` for `l < N` is the state step `l` started from, which the + transcript holds. Level `N` is the state the run finished in. + + The clock comes back with the fields, so `model.tracker.time` follows + the backward pass — which is a small thing, but it means a diagnostic + written during the adjoint is labelled with the time it belongs to. + """ + if level < len(self.transcript): + self.uwmodel.load_state(self.transcript[level].snapshot) + else: + self.uwmodel.load_state(self.final_state) + + def linearisation_state(self, k): + """Restore the point step `k`'s transport residual was linearised at. + + The residual of step `k` is `F(beta_{k+1}; beta_k, V_k, dt)`, so the + model must hold the step's OUTPUT in the unknown and the step's INPUT + in the history slot. The transcript snapshot gives the input; replaying + the transport solve gives the output, bit for bit. The solve's + post-hook then shifts the history forward, so the input is put back + where the residual reads it. + """ + entry = self.transcript[k] + beta, adv = self.model["beta"], self.model["adv"] + + self.state_at(k) + beta_in = np.asarray(beta.array)[:, 0, 0].copy() + adv.solve(timestep=entry.dt, zero_init_guess=False) + self.psi_star.array[:, 0, 0] = beta_in + + def transport_jacobian(self, k): + """K_k = dF/d(beta_{k+1}) for transport step k, at that step's state.""" + self.linearisation_state(k) + return _jacobian_csr(self.model["adv"]) + + def transport_duals(self, k, mu_dual): + """The adjoint of ONE transport step, applied to multiplier `mu_dual`. + + Returns (L^T mu, G^T mu): the dual on the previous level `beta_k`, and + the dual on the velocity `V_k` as NODAL VALUES of a body-force field, + which is what the adjoint Stokes solve wants. This is the operation a + `solver.adjoint(...)` method would provide. + """ + self.linearisation_state(k) + self.mu.array[:, 0, 0] = mu_dual + hist_dual = self.dual_of(self.hist_integrand) + self.f_proj.uw_function = sympy.Matrix([self.vel_integrand]) + self.f_proj.solve() + vel_field = np.asarray(self.f_proj_var.array)[:, 0, :].copy() + return hist_dual, vel_field + + def stokes_dual(self, level, bodyforce_field, cold=False): + """Solve the adjoint Stokes problem at `level` with the given body-force + NODAL VALUES, and return the dual of the resulting beta-sensitivity.""" + self.state_at(level) + self.f_adj.array[:, 0, :] = bodyforce_field + self.stokes_adj.solve(zero_init_guess=cold) + return self.dual_of(self.stokes_sensitivity) + + def misfit(self): + """J at the run's final state.""" + self.state_at(len(self.transcript)) + return float(uw.maths.Integral(self.mesh, self.misfit_integrand).evaluate()) + + +# %% +def compute_adjoint_gradient(machinery, centre): + """Backward pass over the recorded run. Returns (dJ/dcx0, dJ/dcy0, J).""" + m = machinery.model + v = m["v"] + N = len(machinery.transcript) + + J = machinery.misfit() + + # dJ/dv_N is M_v (v_N - v_target), so the body-force FIELD is -(v_N - v_target) + V_N = np.asarray(v.array)[:, 0, :].copy() + VT = np.asarray(machinery.v_target.array)[:, 0, :] + rhs = machinery.stokes_dual(N, -(V_N - VT), cold=True) + + for level in range(N - 1, -1, -1): + K = machinery.transport_jacobian(level) + mu = splu(K.T.tocsc()).solve(-rhs) + hist_dual, vel_field = machinery.transport_duals(level, mu) + rhs = machinery.stokes_dual(level, -vel_field) + hist_dual + + _, dbeta0_dc = beta0_nodal(m, centre) + return float(rhs @ dbeta0_dc[:, 0]), float(rhs @ dbeta0_dc[:, 1]), J diff --git a/docs/examples/adjoint/sinker_transcript/taylor_test_library.py b/docs/examples/adjoint/sinker_transcript/taylor_test_library.py new file mode 100644 index 000000000..321b3047e --- /dev/null +++ b/docs/examples/adjoint/sinker_transcript/taylor_test_library.py @@ -0,0 +1,52 @@ +"""The same Taylor test, with the library's backward pass in place of the +hand-rolled AdjointMachinery: `uw.adjoint.TranscriptAdjoint` walks the run's +own transcript. The control is still the blob centre, reached through the +dual on beta_0 and the chain rule d(beta_0)/d(centre). +""" +import sys +import numpy as np +import sympy +import underworld3 as uw + +from forward_sinker_transcript import build_model, solve_forward, beta0_nodal +from generate_target_transcript import (TRUE_VISCOSITY_CONTRAST, TARGET_FILENAME, + TARGET_INDEX, OUTPUT_DIR as TARGET_DIR) + +CANDIDATE = (uw.quantity(300, "km"), uw.quantity(350, "km")) # true is (250, 375) +H_VALUES = [uw.quantity(5.0, "km"), uw.quantity(0.5, "km"), uw.quantity(0.05, "km")] + + +if __name__ == "__main__": + hs = [uw.quantity(float(a), "km") for a in sys.argv[1:]] or H_VALUES + + model = build_model(TRUE_VISCOSITY_CONTRAST) + mesh, beta, v = model["mesh"], model["beta"], model["v"] + v_target = uw.discretisation.MeshVariable("v_target", mesh, 2, degree=2) + v_target.read_timestep(data_filename=TARGET_FILENAME, data_name="v", + index=TARGET_INDEX, outputPath=TARGET_DIR) + misfit = sympy.Rational(1, 2) * (v.sym - v_target.sym).dot(v.sym - v_target.sym) + + transcript, final_state = solve_forward(model, CANDIDATE, initial_on_tape=True) + uw.pprint("the run being differentiated:") + for entry in transcript: + uw.pprint(f" {entry}") + uw.pprint("") + + back = uw.adjoint.TranscriptAdjoint(model["uwmodel"], final_state) + result = back.gradient(misfit, fields=[beta]) + dual = result["fields"][beta][:, 0, 0] + _, dbeta0_dc = beta0_nodal(model, CANDIDATE) + gx, gy, J = float(dual @ dbeta0_dc[:, 0]), float(dual @ dbeta0_dc[:, 1]), result["J"] + uw.pprint(f"J = {J:.6e} adjoint dJ/dcx0 = {gx:.6e} /m dJ/dcy0 = {gy:.6e} /m") + + def Jof(c): + solve_forward(model, c, initial_on_tape=True) + return float(uw.maths.Integral(mesh, misfit).evaluate()) + + for h in hs: + cx, cy = CANDIDATE + two_h = 2.0 * float(h.to("m").magnitude) + fx = (Jof((cx + h, cy)) - Jof((cx - h, cy))) / two_h + fy = (Jof((cx, cy + h)) - Jof((cx, cy - h))) / two_h + uw.pprint(f"h={h}: FD dJ/dcx0 {fx:.6e} ratio {fx/gx:8.5f} | " + f"FD dJ/dcy0 {fy:.6e} ratio {fy/gy:8.5f}") diff --git a/docs/examples/adjoint/sinker_transcript/taylor_test_transcript.py b/docs/examples/adjoint/sinker_transcript/taylor_test_transcript.py new file mode 100644 index 000000000..da5dfa2a2 --- /dev/null +++ b/docs/examples/adjoint/sinker_transcript/taylor_test_transcript.py @@ -0,0 +1,54 @@ +"""Taylor test for the transcript-driven adjoint. Central finite difference on both +control components; the ratio FD/adjoint should sit at 1 and stay there as h +shrinks. + +The control is the blob centre, a pair of LENGTHS, so both the adjoint gradient +and the finite difference are dJ/d(centre) per metre. +""" +import sys +import numpy as np +import underworld3 as uw + +from forward_sinker_transcript import build_model, solve_forward +from generate_target_transcript import (TRUE_VISCOSITY_CONTRAST, TARGET_FILENAME, + TARGET_INDEX, OUTPUT_DIR as TARGET_DIR) +from inverse_sinker_transcript import AdjointMachinery, compute_adjoint_gradient + +CANDIDATE = (uw.quantity(300, "km"), uw.quantity(350, "km")) # true is (250, 375) +H_VALUES = [uw.quantity(5.0, "km"), uw.quantity(0.5, "km"), uw.quantity(0.05, "km")] + + +if __name__ == "__main__": + hs = [uw.quantity(float(a), "km") for a in sys.argv[1:]] or H_VALUES + + model = build_model(TRUE_VISCOSITY_CONTRAST) + mesh = model["mesh"] + v_target = uw.discretisation.MeshVariable("v_target", mesh, 2, degree=2) + v_target.read_timestep(data_filename=TARGET_FILENAME, data_name="v", + index=TARGET_INDEX, outputPath=TARGET_DIR) + mach = AdjointMachinery(model, v_target) + + transcript, final_state = solve_forward(model, CANDIDATE) + mach.attach(transcript, final_state) + + uw.pprint("the run being differentiated:") + for entry in transcript: + uw.pprint(f" {entry}") + uw.pprint("") + + gx, gy, J = compute_adjoint_gradient(mach, CANDIDATE) + uw.pprint(f"J = {J:.6e} adjoint dJ/dcx0 = {gx:.6e} /m " + f"dJ/dcy0 = {gy:.6e} /m") + + def Jof(c): + """Rerun the forward model from centre `c` and evaluate the misfit.""" + solve_forward(model, c) + return float(uw.maths.Integral(mesh, mach.misfit_integrand).evaluate()) + + for h in hs: + cx, cy = CANDIDATE + two_h = 2.0 * float(h.to("m").magnitude) + fx = (Jof((cx + h, cy)) - Jof((cx - h, cy))) / two_h + fy = (Jof((cx, cy + h)) - Jof((cx, cy - h))) / two_h + uw.pprint(f"h={h}: FD dJ/dcx0 {fx:.6e} ratio {fx/gx:8.5f} | " + f"FD dJ/dcy0 {fy:.6e} ratio {fy/gy:8.5f}") diff --git a/src/underworld3/__init__.py b/src/underworld3/__init__.py index bf3faf592..03ee873b2 100644 --- a/src/underworld3/__init__.py +++ b/src/underworld3/__init__.py @@ -222,12 +222,14 @@ def view(): ThermalConvectionConfig, create_thermal_convection_model, ) +from . import adjoint from .utilities.transcript_report import ( transcript_diagram, transcript_flowchart, transcript_table, transcript_figure, transcript_key, + transcript_adjoint_segments, ) from .parameters import ParameterRegistry, ParameterType from .materials import MaterialRegistry, MaterialProperty diff --git a/src/underworld3/adjoint.py b/src/underworld3/adjoint.py new file mode 100644 index 000000000..8f1b11ba7 --- /dev/null +++ b/src/underworld3/adjoint.py @@ -0,0 +1,486 @@ +"""Reverse-mode over a recorded run. + +A transcript with snapshots is the forward tape: the ordered operators per +step, and the state each step started from. Every operator's linearisation +comes from the residual itself — the residual is SymPy, so the coupling of +one solve to the fields it reads is a symbolic derivative, and the adjoint of +one solve is a transpose against the Jacobian the SNES assembles +(``solver.adjoint_solve``). This module chains those backwards. + +The chain rule, per solve, in reverse order of the record. A solve +:math:`R(u; f_1, f_2, \\dots, m) = 0` reads fields :math:`f_i` and parameters +:math:`m`. Given the accumulated dual :math:`\\bar u = \\partial J/\\partial u` +on its unknown, + +.. math:: + + K^T \\mu = -\\bar u, \\qquad + \\bar f_i \\mathrel{+}= (\\partial R/\\partial f_i)^T \\mu, \\qquad + \\bar m \\mathrel{+}= \\mu^T \\partial R/\\partial m. + +A history slot read by the solve (``psi_star[0]``) is the tracked field at +the step's input, so its dual is the dual on that field at the previous +level — which is where the walk goes next. + +Each solve is linearised at ITS OWN input state: the step's snapshot is +restored, the solves before it in the step are replayed, the histories it +reads are captured, it is replayed, and the captured inputs are put back in +the history slots the post-solve hook shifted. That is what +``docs/examples/adjoint`` did by hand for one solver; here the record says +which operators ran and the residuals say what they read. +""" +from __future__ import annotations + +import re +from typing import Dict, Iterable, Optional + +import numpy as np +import sympy + +import underworld3 as uw +from underworld3.cython.generic_solvers import SNES_Scalar as _SNES_Scalar +from underworld3.cython.generic_solvers import SNES_Vector as _SNES_Vector +from underworld3.utilities._api_tools import Template + + +class _ScalarLoad(_SNES_Scalar): + r"""A generic scalar solver used as an ASSEMBLER: its residual at zero is + :math:`\int g_0\,\phi_j + \mathbf g_1\cdot\nabla\phi_j`, which is the + dual of a field read through its value (``g0``) and its gradient (``g1``) + in one assembly — no integration by parts, no boundary term to get wrong. + """ + + _solver_terms = (("_g0", "value part of the load"), ("_g1", "gradient part")) + F0 = Template(r"g_0", lambda self: sympy.Matrix([[self._g0]]), + "value part of a dual load") + F1 = Template(r"\mathbf{g}_1", lambda self: self._g1, + "gradient part of a dual load (1 x cdim)") + + +class _VectorLoad(_SNES_Vector): + """The vector-field counterpart of :class:`_ScalarLoad`: ``g0`` a row of + ``dim`` components, ``g1`` a ``dim x cdim`` matrix.""" + + _solver_terms = (("_g0", "value part of the load"), ("_g1", "gradient part")) + F0 = Template(r"\mathbf{g}_0", lambda self: self._g0, "value part of a dual load") + F1 = Template(r"\mathbf{G}_1", lambda self: self._g1, + "gradient part of a dual load (dim x cdim)") + + +class _Scratch: + """Fields and assemblers on a space, reused rather than re-created. + + A fresh MeshVariable per dual leaked sixteen registered variables per + ``gradient()`` call and slowed the sixth call tenfold (found in review). + Variables are handed out and taken back; an assembler is built once per + space and re-pointed at each load. + """ + + def __init__(self): + self.free = {} + self.assemblers = {} + + @staticmethod + def space(variable): + return (variable.mesh, getattr(variable, "num_components", 1), + str(variable.vtype), int(variable.degree), + bool(getattr(variable, "continuous", True))) + + def take(self, like): + key = self.space(like) + pool = self.free.setdefault(key, []) + if pool: + var = pool.pop() + var.array[...] = 0.0 + return var + mesh, n, _, degree, continuous = key + return uw.discretisation.MeshVariable( + f"_adj_scratch_{_counter()}", mesh, num_components=n, + vtype=like.vtype, degree=degree, continuous=continuous) + + def give(self, var): + self.free.setdefault(self.space(var), []).append(var) + + def assembler(self, like): + key = self.space(like) + asm = self.assemblers.get(key) + if asm is None: + target = self.take(like) + mesh, n = key[0], key[1] + asm = (_ScalarLoad(mesh, u_Field=target) if n == 1 + else _VectorLoad(mesh, u_Field=target)) + # Every solver family's setup reads constitutive_model._solver_is_setup + # unguarded, so the assembler carries an inert model: its F0/F1 + # templates override the model's flux entirely. + if n == 1: + asm.constitutive_model = uw.constitutive_models.DiffusionModel + asm.constitutive_model.Parameters.diffusivity = 1.0 + else: + asm.constitutive_model = uw.constitutive_models.ViscousFlowModel + asm.constitutive_model.Parameters.shear_viscosity_0 = 1.0 + asm.consistent_jacobian = False # a residual evaluation only + asm.petsc_options.delValue("ksp_monitor") + self.assemblers[key] = asm + return asm + + +_shared_scratch = _Scratch() + + +def dual_on(variable, value, grad=None, scratch=None): + r"""The dual of a load on ``variable``'s space, held as a field. + + :math:`b_j = \int v\,\phi_j + \mathbf g\cdot\nabla\phi_j` for every basis + function of ``variable``: ``value`` is the part read through the field's + value, ``grad`` (optional; ``1 x cdim`` for a scalar field, ``dim x cdim`` + for a vector one) the part read through its gradient — a Crank–Nicolson + step reads the old level's flux this way. Assembled as the residual of a + generic solver at zero, so it is exactly the FEM load, with no linear + solve and no integration by parts. The returned field comes from + ``scratch`` (a :class:`_Scratch` pool; the module's shared one by + default) — give it back with ``scratch.give(field)`` when done. + """ + scratch = _shared_scratch if scratch is None else scratch + mesh = variable.mesh + n = getattr(variable, "num_components", 1) + dim, cdim = mesh.dim, mesh.cdim + asm = scratch.assembler(variable) + if grad is None: + grad = sympy.zeros(1, cdim) if n == 1 else sympy.zeros(dim, cdim) + asm._g0 = value + asm._g1 = sympy.Matrix(grad) + asm._needs_function_rewire = True # the templates re-evaluate + asm._build(False, False, None) + out_var = scratch.take(variable) + gvec = asm.dm.getGlobalVec() + gvec.set(0.0) + mesh.update_lvec() + asm.dm.setAuxiliaryVec(mesh.lvec, None) + F = gvec.duplicate() + asm.snes.computeFunction(gvec, F) + lvec = asm.dm.getLocalVec() + lvec.set(0.0) + asm.dm.globalToLocal(F, lvec) + out_var.vec.array[:] = lvec.array[:] + asm.dm.restoreLocalVec(lvec) + F.destroy() + asm.dm.restoreGlobalVec(gvec) + mesh._stale_lvec = True + try: + out_var._sync_lvec_to_gvec() + except AttributeError: + pass + return out_var + + +def inner(variable, a, b): + r"""``a . b`` over the OWNED degrees of freedom of ``variable``'s space, + reduced across ranks. + + A dual is a covector on the basis, so ``dJ = sum_j dual_j * delta_j`` + is the right pairing — but ``.array`` on a rank holds ghost nodes as + well, so a plain NumPy dot counts shared nodes twice (found in review: + a different number on each rank, neither the finite difference). This + routes both through the field's global vector, which holds each degree + of freedom once. + """ + mesh = variable.mesh + dm = mesh.dm + field = variable.field_id + _is, subdm = dm.createSubDM(field) + ga = subdm.getGlobalVec() + gb = subdm.getGlobalVec() + la = subdm.getLocalVec() + lb = subdm.getLocalVec() + la.array[:] = np.asarray(a).ravel() + lb.array[:] = np.asarray(b).ravel() + subdm.localToGlobal(la, ga) + subdm.localToGlobal(lb, gb) + value = float(ga.dot(gb)) + subdm.restoreLocalVec(la); subdm.restoreLocalVec(lb) + subdm.restoreGlobalVec(ga); subdm.restoreGlobalVec(gb) + return value + + +_n = [0] + + +def _counter(): + _n[0] += 1 + return _n[0] + + +class TranscriptAdjoint: + """The backward pass over a model's recorded run. + + Parameters + ---------- + model : uw.Model + With ``model.transcript`` holding one entry per step, each with the + snapshot it started from (``model.record_every = 1``). + final_state + ``model.save_state()`` taken after the last step — the one level the + transcript does not hold, since it records steps and an N-step run has + N+1 levels. + """ + + def __init__(self, model, final_state): + self.model = model + self.final_state = final_state + self.steps = list(model.transcript) + self._scratch = _Scratch() + missing = [s.index for s in self.steps if not s.restorable] + if missing: + raise RuntimeError( + f"steps {missing} kept no snapshot; set model.record_every = 1 " + f"before the run so every step keeps the state it started from") + + # ------------------------------------------------------------------ + def gradient(self, misfit, parameters: Iterable = (), fields: Iterable = ()): + r"""``dJ/dm`` for each parameter, and the dual on each field at level 0. + + Parameters + ---------- + misfit : sympy expression + :math:`J` as an integrand over the mesh in the fields at the final + level, e.g. ``(v.sym - v_target.sym).dot(v.sym - v_target.sym) / 2``. + parameters + Named expressions (``uw.expression``) the residuals depend on. + fields + MeshVariables whose INITIAL values are controls; the result holds + :math:`\partial J/\partial f_0` as a dual field, so a control + :math:`c` with :math:`f_0 = f_0(c)` finishes with a dot product + against :math:`\partial f_0/\partial c`. + + Returns + ------- + dict + ``{"J": float, "parameters": {expr: float}, "fields": {var: dual}}`` + """ + parameters = list(parameters) + fields = list(fields) + model = self.model + scratch = self._scratch + + # J and its dual on every field it touches, at the final level — and + # the explicit dJ/dm, for a misfit that names a parameter directly. + model.load_state(self.final_state) + J = float(uw.maths.Integral(self._mesh(), misfit).evaluate()) + acc: Dict[str, object] = {} + peeled = _peel(misfit) + for var, symbols in self._fields_in(misfit): + dJ = [sympy.diff(peeled, s) for s in symbols] + self._accumulate(acc, var, dual_on(var, _as_expression(dJ), None, scratch)) + + grad = {p: 0.0 for p in parameters} + for p in parameters: + explicit = sympy.diff(_peel_except(misfit, p), p) + if explicit != 0: + grad[p] += float(uw.maths.Integral(self._mesh(), explicit).evaluate()) + + for k in range(len(self.steps) - 1, -1, -1): + step = self.steps[k] + solves = [e for e in step.events if e["kind"] == "solve"] + for j in range(len(solves) - 1, -1, -1): + solver = model.part_object(solves[j]["part"]) + if solver is None: + raise RuntimeError( + f"step {step.index}: no live object for part " + f"{solves[j]['part']!r} — the backward pass needs the " + f"solvers of the run in this process") + u = solver.u + rhs = acc.get(u.name) + # Decided on every rank together: the gate guards a collective + # solve, and a misfit supported on one rank's cells deadlocked + # here when each rank looked only at its own values. + if not self._nonzero(rhs): + continue + inputs = self._linearise_at(step, solves, j) + mu = self._adjoint(solver, rhs) + scratch.give(acc.pop(u.name)) # consumed: this level's output + + for p in parameters: + grad[p] += solver.sensitivity(mu, p) + + for var, symbols, derivatives in self._reads(solver, u): + value = [solver.adjoint_integrand(mu, s) for s in symbols] + g1 = None + if derivatives: + # g1[i, k] = the contraction with respect to d f_i / d x_k + cdim = self._mesh().cdim + g1 = sympy.zeros(len(symbols), cdim) + for (i, k), atom in derivatives.items(): + g1[i, k] = solver.adjoint_integrand(mu, atom) + target = inputs.get(var.name, var) # a history -> its field + self._accumulate(acc, target, + dual_on(target, _as_expression(value), g1, scratch)) + scratch.give(mu) + + out_fields = {} + for var in fields: + held = acc.get(var.name) + out_fields[var] = (np.zeros_like(np.asarray(var.array)) if held is None + else np.array(held.array, copy=True)) + for held in acc.values(): + scratch.give(held) + return {"J": J, "parameters": grad, "fields": out_fields} + + # ------------------------------------------------------------------ + def _mesh(self): + return next(iter(self.model._variables.values())).mesh + + def _tokens(self): + """``{token: variable}`` — how each variable PRINTS inside a residual. + + A variable prints as its symbol, which need not be its name and can + carry nested braces (a history slot is ``{\\psi^{*}_{...}}``), so the + token is taken from the symbol's own text: everything before the + coordinate arguments and, for a vector, before the component index. + Matching on the name found the user's fields and silently missed + every history, which cut the chain at the first step.""" + out = {} + for var in self.model._variables.values(): + if not hasattr(var, "sym"): + continue + text = str(_symbols_of(var)[0]).rsplit("(", 1)[0] # drop (N.x, N.y) + if getattr(var, "num_components", 1) > 1: + text = text.rsplit("_{", 1)[0] # drop _{ i } + out[text] = var + return out + + def _fields_in(self, expression): + text = str(_peel(expression)) + return [(v, _symbols_of(v)) for token, v in self._tokens().items() + if token in text] + + def _reads(self, solver, unknown): + """What a solver's residual reads, other than its unknown. + + ``(variable, value symbols, {(component, direction): derivative atom})`` + per variable. A component prints as ``{v}_{ 0 }``; a derivative + carries a comma — ``{v}_{ 0,1}`` for a vector, ``{T}_{,1}`` for a + scalar — and is read through the gradient part of the load. + """ + f0 = _peel(solver.F0.sym) + f1 = _peel(solver.F1.sym) + text = str(f0) + str(f1) + atoms = set(f0.atoms(sympy.Function)) | set(f1.atoms(sympy.Function)) + found = [] + for token, var in self._tokens().items(): + if var is unknown or token not in text: + continue + derivatives = {} + pattern = re.compile(re.escape(token) + r"_\{ ?(\d*),(\d+)\}\(") + for atom in atoms: + m = pattern.match(str(atom)) + if m: + i = int(m.group(1)) if m.group(1) else 0 + derivatives[(i, int(m.group(2)))] = atom + found.append((var, _symbols_of(var), derivatives)) + return found + + def _linearise_at(self, step, solves, j): + """Restore the step's snapshot, replay solves 0..j, and put each + history's input back where the post-solve hook shifted it. Returns + ``{history-slot name: tracked field}`` for the histories solve j read.""" + model = self.model + model.load_state(step.snapshot) + for e in solves[:j]: + self._replay(model.part_object(e["part"]), step) + solver = model.part_object(solves[j]["part"]) + histories = [h for h in (getattr(solver, "DuDt", None), getattr(solver, "DFDt", None)) + if h is not None and getattr(h, "psi_star", None)] + captured = [] + for h in histories: + tracked = self._tracked_field(h) + if tracked is None: + continue + captured.append((h, tracked, np.array(tracked.array, copy=True))) + self._replay(solver, step) + inputs = {} + for h, tracked, before in captured: + h.psi_star[0].array[...] = before + inputs[h.psi_star[0].name] = tracked + return inputs + + def _replay(self, solver, step): + if getattr(solver, "DuDt", None) is not None: + solver.solve(timestep=step.dt, zero_init_guess=False) + else: + solver.solve(zero_init_guess=False) + + def _tracked_field(self, history): + text = str(history.psi_fn) + for token, var in self._tokens().items(): + if token in text: + return var + return None + + def _adjoint(self, solver, rhs): + u = solver.u + scratch = self._scratch + mu = scratch.take(u) + neg = scratch.take(u) + neg.array[...] = -np.asarray(rhs.array) + if getattr(solver, "p", None) is not None and hasattr(solver, "_subdict"): + p_adj = scratch.take(solver.p) + _, reason = solver.adjoint_solve((neg, None), target=(mu, p_adj)) + scratch.give(p_adj) + else: + _, reason = solver.adjoint_solve(neg, target=mu) + scratch.give(neg) + if reason <= 0: + raise RuntimeError(f"adjoint of {type(solver).__name__}({u.name}) did not converge ({reason})") + return mu + + def _accumulate(self, acc, var, dual): + held = acc.get(var.name) + if held is None: + acc[var.name] = dual + return + held.array[...] = np.asarray(held.array) + np.asarray(dual.array) + self._scratch.give(dual) + + @staticmethod + def _nonzero(dual): + """Whether the dual is nonzero ANYWHERE — reduced across ranks, and + safe on a rank that holds no degrees of freedom of the space.""" + if dual is None: + local = 0.0 + else: + values = np.asarray(dual.array) + local = float(np.abs(values).max()) if values.size else 0.0 + return uw.mpi.comm.allreduce(local, op=uw.MPI.MAX) > 0.0 + + +def _symbols_of(var): + n = getattr(var, "num_components", 1) + return [var.sym[i] for i in range(n)] if n > 1 else [var.sym[0]] + + +def _as_expression(components): + """A scalar for a scalar field, a row vector for a vector one — the shape + a projection onto that field's space expects.""" + return components[0] if len(components) == 1 else sympy.Matrix([components]) + + +def _peel_except(expression, wrt, depth=8): + """Expand every named expression except ``wrt`` (see the solver's + ``_peel_except``): ``_peel`` would substitute the parameter's value and + the derivative of a number is zero.""" + for _ in range(depth): + named = [e for e in uw.function.fn_extract_expressions(expression) + if e is not wrt and e != wrt] + if not named: + break + expression = expression.subs({e: e.sym for e in named}) + return expression + + +def _peel(expression, depth=8): + for _ in range(depth): + named = uw.function.fn_extract_expressions(expression) + if not named: + break + expression = expression.subs({e: e.sym for e in named}) + return expression diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index 2647afb16..38dc9275c 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -131,7 +131,7 @@ class SolverBaseClass(uw_object): # Jacobian tangent selection — validated property, see the # consistent_jacobian docstring below for the mode semantics. - self.consistent_jacobian = False + self.consistent_jacobian = True # Picard->Newton continuation parameter (constants[]-routed so it can be # ramped at solve time without a JIT recompile). 0 = Picard, 1 = Newton. # Created LAZILY (see _get_newton_alpha) only when continuation is used, @@ -394,15 +394,23 @@ class SolverBaseClass(uw_object): dispatch; the residual is never affected, so the converged solution always satisfies the exact constitutive law. - ``False`` (default) + ``True`` (default) + Unwrap the flux before differentiation so the tangent captures + :math:`\partial\eta/\partial(\nabla v)` (full Newton). The + residual is symbolic, so this tangent is exact and cheap, and it + is the tangent the adjoint transposes. On the saddle-point + solvers' standard path only (not the scalar/vector solvers, not + the rotated or fault-contact path), a cold start first takes one + ``nrichardson`` residual sweep (``solve(picard=-1)`` turns it + off); see the note at that line for what it is and is not. + ``False`` Differentiate the residual flux *as wrapped* — the effective viscosity is frozen, giving a Picard / defect-correction tangent. - Bit-identical to the long-standing behaviour. Globally robust; - load-bearing for the tuned hard-yield viscoplastic paths. - ``True`` - Unwrap the flux before differentiation so the tangent captures - :math:`\partial\eta/\partial(\nabla v)` (full Newton). Fast near - the solution; its yield kink can stall the line search far from it. + Linearly convergent, and the SNES then holds a Jacobian that is + not :math:`\partial R/\partial u`. Opt in where the hard-yield + viscoplastic solves need it (the notch class), where Picard is an + entry requirement rather than an accelerator. Was the default + until 2026-09. ``"continuation"`` Picard :math:`\rightarrow` Newton. Blend :math:`J(\alpha) = J_{\mathrm{picard}} + \alpha\,(J_{\mathrm{newton}} @@ -1130,6 +1138,7 @@ class SolverBaseClass(uw_object): self._needs_dm_rebuild = False self._needs_bc_reregister = False self._needs_function_rewire = False + self._adjoint_kernel_installed = False else: self._needs_dm_rebuild = True self._needs_bc_reregister = True @@ -1392,6 +1401,431 @@ class SolverBaseClass(uw_object): except Exception: pass + def _adjoint_support(self): + """Whether this solve, as configured, admits a discrete adjoint. + + ``(supported, reason)``. The verdict is STRUCTURAL — about the + operator, not about whether a driver exists yet — and it is written + into the transcript when the solve is recorded, so a run says where + its adjoint breaks while it runs, not three hours into an inversion. + + An implicit step is a residual, so its adjoint is the Jacobian + transpose for the state and the symbolic derivative of the residual + for a parameter. What removes that: + + * a rotated constraint — free-slip or fault contact — solves on a + rotated operator inside its own Krylov loop (``rotated_bc.py``), + and there is no transpose path through it; + * an unconverged solve, which is caught after the fact by + :meth:`_record_solve_outcome`: a linearisation about a state the + solve never reached is not the adjoint of anything. + + A subclass whose operator is not a residual overrides this and says + why. The contract is enforced by + ``tests/test_0018_adjoint_support_record.py``. + """ + mechanisms = self._constraint_mechanisms() + rotated = mechanisms["rotated_freeslip"] + mechanisms["fault_contact"] + if rotated: + return (False, + f"{len(rotated)} rotated constraint(s): the solve runs on a " + f"rotated operator with its own Krylov loop, and there is " + f"no transpose path through it") + if not self.consistent_jacobian and not self._residual_is_linear_in_unknown(): + return (True, + "implicit residual: Jacobian transpose for the state, symbolic " + "derivative of the residual for a parameter. The forward solve " + "used the Picard tangent, so the consistent tangent is assembled " + "for the adjoint (a rebuild of the Jacobian kernel)") + return (True, + "implicit residual: Jacobian transpose for the state, symbolic " + "derivative of the residual for a parameter") + + # ------------------------------------------------------------------ + # The discrete adjoint of one solve + # ------------------------------------------------------------------ + + def adjoint_support(self): + """``(supported, reason)``: whether this solve admits a discrete adjoint. + + The same verdict the transcript records on the solve event. See + :meth:`_adjoint_support` for what refuses and why. + """ + return self._adjoint_support() + + def adjoint_solve(self, rhs, target=None): + r"""Solve the adjoint of the LAST solve: :math:`K^T \mu = b`. + + An implicit step is a residual :math:`R(u; m) = 0`, and the SNES + already assembles its Jacobian :math:`K = \partial R / \partial u`. + The adjoint state is the transpose solve against that same matrix, + taken at the state the forward solve ended in — so call this after + ``solve()``, on the solver that did the solving, before anything + moves the fields. + + The essential boundary conditions come out homogenised for free: + PETSc's global vector holds only the unconstrained degrees of freedom, + so :math:`K` is the operator on those, and the multiplier written back + to ``target`` is zero on every Dirichlet node. + + **The state matters, and ``solve()`` moves it.** A time step's + residual is :math:`F(u_{n+1}; u_n, v, \\Delta t)`: the step's OUTPUT in + the unknown and its INPUT in the history slot. The history manager + shifts that slot forward in its post-solve hook, so straight after + ``solve()`` the slot holds :math:`u_{n+1}`, and a sensitivity read + there is 5% wrong on a SUPG step (measured, ``test_0019``). Put the + step's input back — ``solver.DuDt.psi_star[0].array[...] = u_n`` — + before calling this and :meth:`sensitivity`. A driver that walks the + transcript restores the step's snapshot, replays the solve, and does + exactly that; see ``docs/examples/adjoint``. + + Parameters + ---------- + rhs : numpy.ndarray or petsc4py.PETSc.Vec + The right-hand side :math:`b`, in the global ordering of this + solver's DM — a DUAL vector, an integral against the basis, not a + field. :meth:`dual_of` builds one from an expression. + target : MeshVariable, optional + A variable on the same space as the unknown to receive + :math:`\mu` as a field. Constrained nodes are set to zero. + + Returns + ------- + (numpy.ndarray, int) + :math:`\mu` in the global ordering, and the KSP converged reason + (positive means converged). + + Raises + ------ + RuntimeError + If this solve refuses an adjoint (:meth:`adjoint_support` says + why), or no solve has run. + """ + supported, why = self._adjoint_support() + if not supported: + raise RuntimeError(f"adjoint_solve: this solve refuses an adjoint — {why}") + if self.snes is None or (not self.is_setup + and not getattr(self, "_adjoint_kernel_installed", False)): + raise RuntimeError( + "adjoint_solve: no forward solve to take the adjoint of. Call " + "solve() first; the adjoint is taken about the state it ended in.") + import numpy as np + + cdef DM dm + cdef Vec cmesh_lvec + + # The kernel first: if the forward ran Picard, the rebuild below may + # replace the DS the SNES assembles with, so every handle taken from + # the DM must be taken AFTER it. + tangent = self._consistent_tangent_for_adjoint() + dm = self.dm + + # The Jacobian at the state the forward solve ended in. + gvec = self.dm.getGlobalVec() + self.dm.localToGlobal(self.u.vec, gvec) + self.mesh.update_lvec() + cmesh_lvec = self.mesh.lvec + ierr = DMSetAuxiliaryVec_UW(dm.dm, NULL, 0, 0, cmesh_lvec.vec); CHKERRQ(ierr) + jacobian = self.snes.getJacobian() # (J, P, callback, args) + J, P = jacobian[0], jacobian[1] + alpha_before = self._newton_alpha_for_adjoint() + self.snes.computeJacobian(gvec, J, P) + self._restore_newton_alpha(alpha_before) + + b = gvec.duplicate() + if isinstance(rhs, PETSc.Vec): + rhs.copy(b) + elif hasattr(rhs, "vec") and hasattr(rhs, "array"): + # A dual held as a FIELD on the unknown's space: one coefficient + # per node. localToGlobal keeps the unconstrained ones, which is + # the restriction to this solver's rows. + self.dm.localToGlobal(rhs.vec, b) + else: + values = np.asarray(rhs, dtype=float).ravel() + bad = uw.mpi.comm.allreduce(int(values.size != b.getLocalSize()), op=uw.MPI.MAX) + if bad: + raise ValueError( + f"adjoint_solve: rhs has {values.size} entries; this solver's " + f"global vector has {b.getLocalSize()} on rank {uw.mpi.rank}") + b.array[:] = values + x = gvec.duplicate() + x.set(0.0) + + ksp = self.snes.getKSP() + ksp.setOperators(J, P) + ksp.solveTranspose(b, x) + reason = int(ksp.getConvergedReason()) + self._restore_tangent(tangent) + + if target is not None: + # Homogeneous constraints: the local vector is zeroed before the + # scatter, so every constrained node reads zero rather than the + # forward Dirichlet value that ComputeBoundaryFEM would insert. + lvec = self.dm.getLocalVec() + lvec.set(0.0) + self.dm.globalToLocal(x, lvec) + target.vec.array[:] = lvec.array[:] + self.mesh._stale_lvec = True + try: + target._sync_lvec_to_gvec() + except AttributeError: + pass + self.dm.restoreLocalVec(lvec) + + out = np.array(x.array, copy=True) + self.dm.restoreGlobalVec(gvec) + b.destroy(); x.destroy() + + try: + model = uw.get_default_model() + part, label = self._transcript_identity() + model._record_step_event( + "adjoint_solve", label, part=part, + converged=reason > 0, ksp_reason=reason) + except Exception: + pass + return out, reason + + def dual_of(self, expression): + r"""The dual of an expression on this solver's unknown space. + + :math:`b_j = \int e \, \phi_j` for every basis function of the + unknown — the form :meth:`adjoint_solve` wants its right-hand side in. + A misfit :math:`J = \tfrac12 \int (u - u^*)^2` has + :math:`\partial J/\partial u` dual ``dual_of(u - u_target)``. + + Assembled as the residual of a projection at zero: the projection's + residual is :math:`\int (0 - e)\phi_j`, so no linear solve is taken. + """ + import numpy as np + + proj = self._dual_projection() + proj.uw_function = expression + proj._build(False, False, None) + gvec = proj.dm.getGlobalVec() + gvec.set(0.0) + proj.mesh.update_lvec() + cdef DM dm = proj.dm + cdef Vec cmesh_lvec = proj.mesh.lvec + ierr = DMSetAuxiliaryVec_UW(dm.dm, NULL, 0, 0, cmesh_lvec.vec); CHKERRQ(ierr) + F = gvec.duplicate() + proj.snes.computeFunction(gvec, F) + out = -np.array(F.array, copy=True) + expected = self.dm.getGlobalVec() + n_solver = expected.getLocalSize() + self.dm.restoreGlobalVec(expected) + bad = uw.mpi.comm.allreduce(int(out.size != n_solver), op=uw.MPI.MAX) + if bad: + raise RuntimeError( + f"dual_of: the dual has {out.size} entries and the solver's " + f"global vector {n_solver} on rank {uw.mpi.rank}; the two spaces " + f"constrain different nodes. Essential conditions added through " + f"a route other than add_dirichlet_bc are not mirrored onto the " + f"dual space.") + F.destroy() + proj.dm.restoreGlobalVec(gvec) + return out + + def _dual_projection(self): + """One projection onto the unknown's space, built on first use. + + It carries the solver's essential conditions, homogenised: PETSc's + global vector holds only unconstrained degrees of freedom, so the + projection's global ordering agrees with the solver's only if the two + constrain the same nodes. Rebuilt if the solver's conditions change. + """ + import sympy + + u = self.u + n = getattr(u, "num_components", 1) + signature = tuple( + (str(bc.boundary), tuple(int(c) for c in bc.components)) + for bc in self.essential_bcs) + proj = getattr(self, "_dual_projection_solver", None) + if proj is None or getattr(self, "_dual_projection_signature", None) != signature: + scratch = uw.discretisation.MeshVariable( + f"_dual_{type(self).__name__}_{self.instance_number}_" + f"{abs(hash(signature)) % 100000}", + self.mesh, num_components=n, vtype=u.vtype, degree=u.degree, + continuous=getattr(u, "continuous", True)) + if n == 1: + proj = uw.systems.Projection(self.mesh, scratch) + else: + proj = uw.systems.Vector_Projection(self.mesh, scratch) + proj.smoothing = 0.0 + proj.petsc_options.delValue("ksp_monitor") + for boundary, components in signature: + if n == 1: + proj.add_dirichlet_bc(0.0, boundary) + else: + conds = [0.0 if i in components else sympy.oo for i in range(n)] + proj.add_dirichlet_bc(tuple(conds), boundary) + self._dual_projection_solver = proj + self._dual_projection_signature = signature + return proj + + def adjoint_integrand(self, mu, wrt): + r"""The integrand whose integral is :math:`\mu^T \partial R/\partial m`. + + :math:`(\partial F_0/\partial m)\,\mu + (\partial F_1/\partial m)\cdot\nabla\mu`, + with :math:`F_0`, :math:`F_1` the residual templates AS IMPLEMENTED — + read after the solve, because they are live — and the derivatives + taken symbolically. For a scalar parameter its integral is the + sensitivity; for a field, project it to get the dual on that field. + """ + F0 = self._peel_except(self.F0.sym, wrt) + F1 = self._peel_except(self.F1.sym, wrt) + import sympy + + d0 = sympy.diff(F0, wrt) + d1 = sympy.diff(F1, wrt) + mu_sym = mu.sym + if getattr(mu, "num_components", 1) == 1: + grad_mu = self.mesh.vector.gradient(mu_sym[0]) + out = d0[0] * mu_sym[0] if hasattr(d0, "shape") else d0 * mu_sym[0] + for i in range(self.mesh.dim): + out = out + d1[i] * grad_mu[i] + return out + grad_mu = self.mesh.vector.jacobian(mu_sym) + out = 0 + for i in range(self.mesh.dim): + out = out + d0[i] * mu_sym[i] + return out + uw.maths.tensor.rank2_inner_product(d1, grad_mu) + + def _consistent_tangent_for_adjoint(self): + """Make sure the Jacobian kernel the adjoint assembles is dR/du. + + The converged state is the same whichever tangent the forward + iteration used, and dR/du is a function of that state alone — so + Picard iterations spoil nothing. What they leave behind is a SNES + whose Jacobian KERNEL is the frozen-coefficient one, and assembling + that at the converged state gives the wrong matrix to transpose. If + the forward ran Picard on a nonlinear residual, switch the kernel to + the consistent tangent here (a JIT rebuild of the pointwise + functions; the DM, SNES and KSP are kept) and return a token so + :meth:`_restore_tangent` can put the Picard kernel back for the next + forward solve. Returns None when nothing had to change. + """ + if self.consistent_jacobian is not False: + return None + if self._residual_is_linear_in_unknown(): + return None # the two tangents coincide + if getattr(self, "_adjoint_kernel_installed", False): + return "picard" # still there from the last adjoint + self._consistent_jacobian = True + self._needs_function_rewire = True + self._build(False, False, None) + # The rewire hands back a new DM and SNES on the saddle-point class. + # snes.solve() would set them up itself; a direct computeJacobian + # will not, and segfaults on the unset-up SNES (measured). + self.snes.setUp() + return "picard" + + def _restore_tangent(self, token): + """Put the Picard setting back, but LEAVE the consistent kernel + installed: a second adjoint (a second misfit on the same forward) + reuses it, and the next forward solve's own build rewires to Picard. + Tearing it down here made a second adjoint_solve refuse with "no + forward solve" (found in review).""" + if token is None: + return + self._consistent_jacobian = False + self._adjoint_kernel_installed = True + self._needs_function_rewire = True # the next forward solve rewires + + def _newton_alpha_for_adjoint(self): + """Under ``"continuation"``, put the tangent at full Newton for the + adjoint's Jacobian assembly; return what alpha was so it can be put back. + + The forward solve ramps alpha from 0 towards 1 as the residual drops + and may converge before it arrives, leaving the SNES holding a blend. + The blend is a fine tangent to converge on and the wrong matrix to + transpose: the adjoint wants dR/du, which is alpha = 1. Returns None + when continuation is not in use, and nothing is touched. + """ + if self.consistent_jacobian != "continuation": + return None + import sympy + + before = self._get_newton_alpha().sym + self._set_newton_alpha(1.0) + return before + + def _restore_newton_alpha(self, before): + if before is None: + return + self._get_newton_alpha().sym = before + try: + self._update_constants(record=False) + except Exception: + pass + + def _residual_is_linear_in_unknown(self): + """Whether the residual templates are linear in the unknown. + + Scale every occurrence of the unknown (and its derivatives) by ``s`` + and ask whether the second derivative in ``s`` vanishes. A viscosity + that depends on the strain rate fails this; a constant one passes. The + adjoint cares because a nonlinear residual solved with the Picard + tangent leaves the SNES holding a Jacobian that is NOT + :math:`\\partial R/\\partial u`. + """ + import sympy + + try: + name = self.u.name + s = sympy.Symbol("s_scale_adjoint") + for template in ("F0", "F1", "PF0"): + form = getattr(self, template, None) + if form is None: + continue + expression = self._peel_except(form.sym, None) + atoms = [a for a in expression.atoms(sympy.Function) + if str(a).startswith("{" + name)] + if not atoms: + continue + scaled = expression.subs({a: s * a for a in atoms}) + second = sympy.diff(scaled, s, 2) + if isinstance(second, sympy.MatrixBase): + if any(x != 0 for x in second): + return False + elif second != 0: + return False + return True + except Exception: + return False + + @staticmethod + def _peel_except(expression, wrt, depth=8): + """Expand every named expression in ``expression`` except ``wrt``. + + A parameter reaches the residual through the constitutive model's own + named symbol — ``Parameters.diffusivity = kappa`` puts ``\\upkappa`` in + ``F1`` with ``kappa`` as its value — so differentiating the residual + as written with respect to ``kappa`` gives zero. ``fn_unwrap`` will not + do either: it substitutes every constant's VALUE, and the derivative + of a number is zero too. This substitutes each named expression by its + definition, one level at a time, and stops at ``wrt`` so the chain + rule has something to hold on to. + """ + for _ in range(depth): + named = [e for e in uw.function.fn_extract_expressions(expression) + if e is not wrt and e != wrt] + if not named: + break + expression = expression.subs({e: e.sym for e in named}) + return expression + + def sensitivity(self, mu, wrt): + r"""``d J / d m`` for a scalar parameter ``wrt``, given the adjoint state. + + :math:`\int` of :meth:`adjoint_integrand` — with :math:`\mu` the + solution of :math:`K^T \mu = -\partial J/\partial u`, this is the + implicit part of the gradient; add :math:`\partial J/\partial m` if + the misfit depends on the parameter directly. + """ + return float(uw.maths.Integral(self.mesh, self.adjoint_integrand(mu, wrt)).evaluate()) + def _constraint_mechanisms(self): """Every way a constraint can have been put on this solver. @@ -2570,7 +3004,10 @@ class SolverBaseClass(uw_object): # SymPy, so the weak form can be written into the transcript # exactly as implemented. model._describe_part(self, part, label) - model._record_step_event("solve", label, part=part) + supported, why = self._adjoint_support() + model._record_step_event( + "solve", label, part=part, + adjoint={"supported": bool(supported), "reason": why}) except Exception: pass @@ -9632,6 +10069,190 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): self.dm.restoreLocalVec(xlocal) self.dm.restoreGlobalVec(gvec) + def _adjoint_support(self): + """The base verdict, plus the tangent the forward solve used. + + A nonlinear rheology solved with the default Picard tangent leaves + the SNES holding the frozen-viscosity operator, not + :math:`\\partial R/\\partial u`. The forward solve converges either way + (defect correction), so nothing complains — and the transposed adjoint + would be silently wrong. Refused here, with the fix in the reason. + """ + supported, why = SolverBaseClass._adjoint_support(self) + if not supported: + return supported, why + # The verdict is written on EVERY solve inside a step, so it must be + # cheap: the symbolic test, cached until the residual can change. The + # numerical probe (two assemblies) belongs in adjoint_solve, where an + # adjoint is actually being taken and the cost is paid once. + key = (self.consistent_jacobian, id(getattr(self, "_constitutive_model", None)), + self.is_setup, self._needs_function_rewire) + cached = getattr(self, "_adjoint_linearity_cache", None) + if cached is None or cached[0] != key: + cached = (key, self._residual_is_linear_in_unknown()) + self._adjoint_linearity_cache = cached + nonlinear = not cached[1] + if not self.consistent_jacobian and nonlinear: + return (True, + why + ". The forward solve used the Picard tangent, so the " + "consistent tangent is assembled for the adjoint (a rebuild " + "of the Jacobian kernel)") + return supported, why + + def adjoint_solve(self, rhs, target=None): + r"""Solve :math:`K^T (\\mu, \\lambda) = b` on the composite (u, p) system. + + The same transpose the scalar solvers take, on the two-field DM. + With a linear viscosity :math:`K` is symmetric and this reproduces + the second-solver construction of ``docs/examples/adjoint``; with a + strain-rate- or pressure-dependent viscosity it is the transpose of + the consistent tangent, which that construction cannot build. + + Parameters + ---------- + rhs : numpy.ndarray or petsc4py.PETSc.Vec + The dual on the composite global vector; :meth:`dual_of` builds + one from a velocity-space expression. + target : (MeshVariable, MeshVariable), optional + ``(u_adj, p_adj)`` on the velocity and pressure spaces, to receive + the multipliers as fields. Constrained nodes are set to zero. + + Returns + ------- + (numpy.ndarray, int) + """ + supported, why = self._adjoint_support() + if not supported: + raise RuntimeError(f"adjoint_solve: this solve refuses an adjoint — {why}") + if self.snes is None or (not self.is_setup + and not getattr(self, "_adjoint_kernel_installed", False)): + raise RuntimeError( + "adjoint_solve: no forward solve to take the adjoint of. Call " + "solve() first; the adjoint is taken about the state it ended in.") + + import numpy as np + + tangent = self._consistent_tangent_for_adjoint() # before any DM vector + gvec = self.dm.getGlobalVec() + gvec.setArray(0.0) + self._gather_fields_to_global(gvec) + self.mesh.update_lvec() + self.dm.setAuxiliaryVec(self.mesh.lvec, None) + jacobian = self.snes.getJacobian() + J, P = jacobian[0], jacobian[1] + alpha_before = self._newton_alpha_for_adjoint() + self.snes.computeJacobian(gvec, J, P) + self._restore_newton_alpha(alpha_before) + + b = gvec.duplicate() + if isinstance(rhs, PETSc.Vec): + rhs.copy(b) + elif isinstance(rhs, (tuple, list)) and len(rhs) == 2 \ + and hasattr(rhs[0], "vec"): + # (u_dual, p_dual) held as fields: restrict each to its block. + b.setArray(0.0) + for name, var in zip(("velocity", "pressure"), rhs): + if var is None or name not in self._subdict: + continue + gis, subdm = self._subdict[name] + sub = b.getSubVector(gis) + subdm.localToGlobal(var.vec, sub) + b.restoreSubVector(gis, sub) + else: + values = np.asarray(rhs, dtype=float).ravel() + bad = uw.mpi.comm.allreduce(int(values.size != b.getLocalSize()), op=uw.MPI.MAX) + if bad: + raise ValueError( + f"adjoint_solve: rhs has {values.size} entries; this solver's " + f"composite global vector has {b.getLocalSize()} on rank " + f"{uw.mpi.rank}") + b.array[:] = values + x = gvec.duplicate() + x.set(0.0) + + ksp = self.snes.getKSP() + ksp.setOperators(J, P) + ksp.solveTranspose(b, x) + reason = int(ksp.getConvergedReason()) + self._restore_tangent(tangent) + + if target is not None: + u_adj, p_adj = target + lvec = self.dm.getLocalVec() + lvec.set(0.0) + self.dm.globalToLocal(x, lvec) + self._ensure_local_field_index_sets(lvec, self.dm.getLocalSection()) + sub = lvec.getSubVector(self._velocity_is) + u_adj.vec.array[:] = sub.array[:] + lvec.restoreSubVector(self._velocity_is, sub) + sub = lvec.getSubVector(self._pressure_is) + p_adj.vec.array[:] = sub.array[:] + lvec.restoreSubVector(self._pressure_is, sub) + self.dm.restoreLocalVec(lvec) + self.mesh._stale_lvec = True + for var in (u_adj, p_adj): + try: + var._sync_lvec_to_gvec() + except AttributeError: + pass + + out = np.array(x.array, copy=True) + self.dm.restoreGlobalVec(gvec) + b.destroy(); x.destroy() + + try: + model = uw.get_default_model() + part, label = self._transcript_identity() + model._record_step_event( + "adjoint_solve", label, part=part, + converged=reason > 0, ksp_reason=reason) + except Exception: + pass + return out, reason + + def dual_of(self, expression): + r"""The dual of a VELOCITY-space expression on the composite vector. + + :math:`\\int \\mathbf e \\cdot \\boldsymbol\\phi_j` on the velocity + degrees of freedom, zero on the pressure ones — the form + :meth:`adjoint_solve` wants for a misfit in the velocity. + """ + import numpy as np + + proj = self._dual_projection() + proj.uw_function = expression + proj._build(False, False, None) + pg = proj.dm.getGlobalVec() + pg.set(0.0) + proj.mesh.update_lvec() + cdef DM pdm = proj.dm + cdef Vec pmesh_lvec = proj.mesh.lvec + ierr = DMSetAuxiliaryVec_UW(pdm.dm, NULL, 0, 0, pmesh_lvec.vec); CHKERRQ(ierr) + F = pg.duplicate() + proj.snes.computeFunction(pg, F) + velocity_dual = -np.array(F.array, copy=True) + F.destroy() + proj.dm.restoreGlobalVec(pg) + + gvec = self.dm.getGlobalVec() + gvec.setArray(0.0) + gis, _subdm = self._subdict["velocity"] + sub = gvec.getSubVector(gis) + n = sub.getLocalSize() + bad = uw.mpi.comm.allreduce(int(n != velocity_dual.size), op=uw.MPI.MAX) + if bad: + gvec.restoreSubVector(gis, sub) + self.dm.restoreGlobalVec(gvec) + raise RuntimeError( + f"dual_of: the velocity dual has {velocity_dual.size} entries and " + f"the composite velocity block {n}; the two spaces constrain " + f"different nodes.") + sub.array[:] = velocity_dual + gvec.restoreSubVector(gis, sub) + out = np.array(gvec.array, copy=True) + self.dm.restoreGlobalVec(gvec) + return out + def _ensure_local_field_index_sets(self, clvec, local_section): """Build (once) and cache the LOCAL index sets that decompose a parent-DM local vector into the per-field MeshVariable storage: velocity, pressure @@ -10028,9 +10649,27 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): # still evaluates the Newton branch pointwise, and IEEE 0*NaN = NaN); # with the guard in place both tangents are finite everywhere. See # docs/developer/design/nonlinear-solver-homotopy-warmstart.md (Layer 1). + # Not for a linear solve. ``ksponly`` is the user's declaration that + # the problem is linear, and a sweep before it is not merely + # redundant: under Eisenstat-Walker the real solve then starts from a + # reduced residual and is handed a loose tolerance — measured 12% + # error against 2% on the spherical-shell Nitsche response + # (test_1064) when this ran before ksponly. Read the DECLARATION: + # ``snes.getType()`` is whatever the previous solve's setFromOptions + # left, so a type set between solves was missed (found in review). + # + # What this sweep IS: one ``nrichardson`` step, x <- x - lambda F(x), + # with no linear solve and no frozen tangent. It is not a Picard step + # and it is nearly inert (1-12% residual reduction on a linear Stokes, + # measured); whether it earns its place on a nonlinear cold start is + # a benchmark item. ``picard=-1`` switches it off explicitly. + declared = self.petsc_options.getString("snes_type", snes_type) or snes_type if (picard == 0 and self.consistent_jacobian is True + and declared != "ksponly" and (zero_init_guess or self._solution_is_trivially_zero())): picard = 1 + if picard < 0: + picard = 0 if verbose and uw.mpi.rank == 0: print(f"SNES solve - picard = {picard}", flush=True) diff --git a/src/underworld3/model.py b/src/underworld3/model.py index 8d0845db1..1098be38f 100644 --- a/src/underworld3/model.py +++ b/src/underworld3/model.py @@ -449,6 +449,7 @@ class Model(PintNativeModelMixin, BaseModel): _transcript_dir: Any = PrivateAttr(default=None) _transcript_fh: Any = PrivateAttr(default=None) _transcript_format: Any = PrivateAttr(default=None) + _transcript_last_refusals: Any = PrivateAttr(default=None) _transcript_columns: Any = PrivateAttr(default=None) _announced_transcript: Any = PrivateAttr(default=None) # The automatic run directory carries BOTH renderings: the text one is for @@ -1304,6 +1305,7 @@ def _render_transcript_text(self, payload): # Column names are written lazily, with the first step, because the # time unit is not known until a step carries one. self._transcript_columns = None + self._transcript_last_refusals = None return "\n".join(lines) if kind == "step": @@ -1382,6 +1384,29 @@ def _render_transcript_text(self, payload): notes.append(f" ~~ ... and {len(order) - 8} more distinct warning(s) " f"in the record") + # Where the adjoint breaks. Written when the set of refusals + # CHANGES from the previous step, not on every step — a + # semi-Lagrangian run refuses identically three hundred times, and + # a note that repeats is a note nobody reads. + refusals = tuple(sorted({ + (e.get("name", "?"), e["adjoint"].get("reason", "")) + for e in events + if isinstance(e.get("adjoint"), dict) + and e["adjoint"].get("supported") is False + })) + previous = self._transcript_last_refusals or () + if refusals != previous: + self._transcript_last_refusals = refusals + if refusals: + for name, why in refusals: + notes.append(f" -- no adjoint through {name}: {why}") + else: + # Only after a refusal has cleared. A run whose every step + # admits an adjoint says nothing about it — one aligned + # line per step is the format's promise. + notes.append(" -- adjoint: every operator in this step " + "admits one again") + return "\n".join([ f"{prefix}" f" {payload['index']:>5d} {t1:>14.6g} {dt:>14.6g} " @@ -1679,6 +1704,25 @@ def _record_solve_outcome(self, part: str, report) -> None: event["deadline_expired"] = True if getattr(report, "bounded", False): event["bounded"] = True + + # The structural verdict was written before the solve. A solve + # that did not converge is linearised about a state it never + # reached, and that is not the adjoint of anything — so the + # outcome overrides it, and says why. + if not event["converged"]: + event["adjoint"] = { + "supported": False, + "reason": f"the solve did not converge ({event['reason']}); " + f"a linearisation about an unreached state is " + f"not an adjoint", + } + elif event.get("capped") and event.get("adjoint", {}).get("supported"): + event["adjoint"] = { + "supported": True, + "reason": event["adjoint"]["reason"] + + "; the forward solve was inexact (a block hit " + "its cap) and the adjoint inherits that", + } return def _record_warning(self, message, category, filename, lineno) -> None: diff --git a/src/underworld3/swarm.py b/src/underworld3/swarm.py index fe9ae544a..fd20c04e3 100644 --- a/src/underworld3/swarm.py +++ b/src/underworld3/swarm.py @@ -5190,21 +5190,43 @@ def advection( self._note_advection(delta_t_model, substeps, order, n_before) return + def _adjoint_support(self, n_before, n_after): + """Whether this advection admits a discrete adjoint. + + The Runge-Kutta step in position is an explicit ODE step and is + differentiable; migration is a permutation. What is not is a change + in the NUMBER of particles — a particle removed on leaving the domain + changes the dimension of the state, and there is no linear map to + transpose. So the rule is one line: the step is adjointable exactly + when the particle set is fixed across it. That is checkable, and this + checks it. + """ + if n_after != n_before: + return (False, + f"the particle set changed: {n_before} -> {n_after} " + f"({n_before - n_after:+d} removed on leaving the domain, " + f"or repopulated); the state changed dimension") + return (True, + "explicit Runge-Kutta step in position on a fixed particle set; " + "migration is a permutation") + def _note_advection(self, dt, substeps, order, n_before): """Tell the model's open step that this swarm moved. - Recorded with the particle count before and after: a swarm that - quietly lost forty particles to the boundary is the kind of thing a - run should say. + Recorded with the particle count before and after, because that count + is the adjoint verdict — and because a swarm that quietly lost forty + particles to the boundary is the kind of thing a run should say. A no-op outside a ``model.step`` block. """ try: n_after = uw.mpi.comm.allreduce(max(self.local_size, 0), op=uw.MPI.SUM) + supported, why = self._adjoint_support(n_before, n_after) uw.get_default_model()._record_step_event( "swarm_advect", f"{type(self).__name__}#{self.instance_number}", part=f"{type(self).__name__}#{self.instance_number}", dt=float(dt), substeps=int(substeps), order=int(order), n_before=int(n_before), n_after=int(n_after), + adjoint={"supported": bool(supported), "reason": why}, ) except Exception: pass diff --git a/src/underworld3/systems/ddt.py b/src/underworld3/systems/ddt.py index 5e3c13ed5..888d08aca 100644 --- a/src/underworld3/systems/ddt.py +++ b/src/underworld3/systems/ddt.py @@ -629,6 +629,16 @@ def _init_coefficient_expressions(self, order, theta, with_exp): if with_exp: _update_exp_values(self._exp_coeffs, None, None) + def _adjoint_support(self): + """Whether this history's shift admits a discrete adjoint. + + ``(supported, reason)``, written onto the ``history_shift`` event. + The default is a refusal that names the class, so a scheme added + without a verdict shows up in the transcript as undeclared rather + than passing as either. + """ + return (False, f"{type(self).__name__} declares no linearisation") + def _note_history_shift(self, dt, **detail): """Tell the model's open step that this history advanced. @@ -649,12 +659,14 @@ def _note_history_shift(self, dt, **detail): try: import underworld3 as uw + supported, why = self._adjoint_support() part = f"{type(self).__name__}#{self.instance_number}" uw.get_default_model()._part_objects[part] = self uw.get_default_model()._record_step_event( "history_shift", self._history_label(), dt=float(dt), part=part, tracks=self._tracked_expression(), + adjoint={"supported": bool(supported), "reason": why}, **detail, ) except Exception: @@ -1053,6 +1065,8 @@ class Symbolic(_DDtBase): Lagrangian : Swarm-based material tracking. """ + def _adjoint_support(self): + return (True, "implicit residual: the linearisation is the owning solver's Jacobian") @timing.routine_timer_decorator def __init__( @@ -1317,6 +1331,8 @@ class Eulerian(_DDtBase): Symbolic : For purely symbolic history (no mesh storage). """ + def _adjoint_support(self): + return (True, "implicit residual: the linearisation is the owning solver's Jacobian") @timing.routine_timer_decorator def __init__( @@ -2316,6 +2332,8 @@ class SemiLagrangian(_DDtBase): Lagrangian : For full particle-following Lagrangian tracking. """ + def _adjoint_support(self): + return (False, "the departure-point trace is differentiable in the velocity, but the interpolation at the departure points is not materialised as an operator") @timing.routine_timer_decorator def __init__( @@ -3739,6 +3757,8 @@ class Lagrangian(_DDtBase): Lagrangian_Swarm : For user-provided swarms. """ + def _adjoint_support(self): + return (True, "a copy on the particle set; valid while that set is fixed across the step — the swarm's advection record says whether it was") instances = ( 0 # count how many of these there are in order to create unique private mesh variable ids @@ -4057,6 +4077,8 @@ class Lagrangian_Swarm(_DDtBase): Eulerian : Pure mesh-based history (no particle tracking). """ + def _adjoint_support(self): + return (True, "a copy on the particle set; valid while that set is fixed across the step — the swarm's advection record says whether it was") instances = ( 0 # count how many of these there are in order to create unique private mesh variable ids @@ -4368,6 +4390,8 @@ class IntegrationPointSemiLagrangian(_DDtBase): velocity history caches it by evaluation at each time level. """ + def _adjoint_support(self): + return (False, "the departure-point trace is differentiable in the velocity, but the interpolation at the departure points is not materialised as an operator") def __init__( self, diff --git a/src/underworld3/utilities/transcript_report.py b/src/underworld3/utilities/transcript_report.py index 612b519da..daeb6e730 100644 --- a/src/underworld3/utilities/transcript_report.py +++ b/src/underworld3/utilities/transcript_report.py @@ -29,7 +29,8 @@ import zlib __all__ = ["transcript_diagram", "transcript_flowchart", - "transcript_table", "transcript_figure", "transcript_key"] + "transcript_table", "transcript_figure", "transcript_key", + "transcript_adjoint_segments"] # --- palette --------------------------------------------------------------- @@ -413,6 +414,76 @@ def transcript_key(source, run=-1, out=None, format="markdown"): handle.write(text) return text +def transcript_adjoint_segments(source, run=-1): + """Where a run can be inverted, and where it cannot. + + Each recorded operator carries a verdict — ``adjoint: {supported, + reason}`` — written when it ran. This reads them back as the partition + they imply: maximal runs of consecutive steps whose every operator admits + a discrete adjoint, separated by the steps where one refused. + + That partition is the assimilation window's structure. Strong-constraint + adjoint within a segment; across a refusal, a control variable and an + error covariance — weak-constraint 4D-Var, with the joins chosen by the + run rather than by hand. Nothing is approximated silently: the refusal + says what the model was allowed to be wrong about. + + Returns + ------- + list of dict + ``{"first", "last", "steps", "supported", "refusals"}`` per segment, + in order. ``first``/``last`` are step indices as recorded; + ``refusals`` is a sorted list of ``(operator name, reason)`` for an + unsupported segment, empty for a supported one. Steps that were + abandoned are left out — they are not part of the run's state + history. + """ + runs = _as_runs(source) + entry = _pick_run(runs, run) + steps = [s for s in entry["steps"] if s.get("completed")] + + def verdict(step): + refusals = set() + undeclared = set() + for event in step.get("events", []): + if event.get("kind") not in ("solve", "history_shift", "swarm_advect"): + continue + verdict = event.get("adjoint") + if not isinstance(verdict, dict) or "supported" not in verdict: + # Recorded before verdicts existed. Not a refusal, not a + # pass: say so rather than read absence as either. + undeclared.add((event.get("name", "?"), + "recorded without an adjoint verdict")) + elif verdict.get("supported") is not True: + # Anything but a literal True is a refusal — a None, a 0 or a + # string is not a verdict this reader may take as support. + refusals.add((event.get("name", "?"), verdict.get("reason", ""))) + return tuple(sorted(refusals | undeclared)) + + segments = [] + for position, step in enumerate(steps): + refusals = verdict(step) + supported = not refusals + index = step.get("index") + # A rewind replays an index, so consecutive records can carry the same + # or a smaller index. Segments follow the RECORD's order (positions); + # an index that goes backwards ends the segment rather than folding a + # replayed step into the one it replaced. + if segments and segments[-1]["supported"] == supported \ + and tuple(segments[-1]["refusals"]) == refusals \ + and index is not None and segments[-1]["last"] is not None \ + and index > segments[-1]["last"]: + segments[-1]["last"] = index + segments[-1]["last_position"] = position + segments[-1]["steps"] += 1 + continue + segments.append({ + "first": index, "last": index, + "first_position": position, "last_position": position, "steps": 1, + "supported": supported, "refusals": list(refusals), + }) + return segments + def _pick_run(runs, index): populated = [r for r in runs if r.get("steps")] diff --git a/tests/test_0018_adjoint_support_record.py b/tests/test_0018_adjoint_support_record.py new file mode 100644 index 000000000..0bbfda380 --- /dev/null +++ b/tests/test_0018_adjoint_support_record.py @@ -0,0 +1,360 @@ +"""Every recorded operator says whether it admits a discrete adjoint. + +The verdict is written when the operator runs, not when someone asks for a +gradient, so a run says where its adjoint breaks while it runs — instead of +that being discovered three hours into an inversion. + +The verdicts are STRUCTURAL: about the operator as configured, not about +whether a driver exists yet. What they encode: + + * an implicit step is a residual: Jacobian transpose for the state, + symbolic derivative for a parameter — supported; + * a rotated constraint solves on a rotated operator in its own Krylov loop + with no transpose path — refused; + * a solve that did not converge is linearised about a state it never + reached — refused, after the fact; + * a semi-Lagrangian history's interpolation at the departure points is not + materialised as an operator — refused, naming what is missing; + * a swarm step is adjointable exactly when the particle set is fixed across + it — checked by counting. + +``transcript_adjoint_segments`` reads the verdicts back as the partition they +imply: strong-constraint within a segment, weak-constraint across a refusal. +""" + +import warnings + +import numpy as np +import pytest +import sympy + +pytestmark = [pytest.mark.level_1, pytest.mark.tier_a] + + +def _fresh_model(): + import underworld3 as uw + + uw.reset_default_model() + return uw, uw.get_default_model() + + +@pytest.fixture(scope="module") +def mesh(): + import underworld3 as uw + + return uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1.0 / 8, qdegree=3 + ) + + +def _poisson(uw, mesh, name): + T = uw.discretisation.MeshVariable(name, mesh, 1, degree=2) + solver = uw.systems.Poisson(mesh, u_Field=T) + solver.constitutive_model = uw.constitutive_models.DiffusionModel + solver.constitutive_model.Parameters.diffusivity = 1.0 + solver.f = 1.0 + solver.add_dirichlet_bc(0.0, "Top") + solver.add_dirichlet_bc(0.0, "Bottom") + solver.petsc_options.delValue("ksp_monitor") + return solver + + +def _stokes(uw, mesh, tag): + V = uw.discretisation.MeshVariable(f"V_{tag}", mesh, 2, degree=2) + P = uw.discretisation.MeshVariable(f"P_{tag}", mesh, 1, degree=1) + stokes = uw.systems.Stokes(mesh, velocityField=V, pressureField=P) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = 1.0 + stokes.bodyforce = sympy.Matrix([0.0, -1.0]) + stokes.petsc_options.delValue("ksp_monitor") + return stokes, V + + +def _adjoint_of(model, kind): + events = [e for e in model.transcript[0].events if e["kind"] == kind] + assert events, f"no {kind} event was recorded" + return events[-1]["adjoint"] + + +# --------------------------------------------------------------------------- +# solves +# --------------------------------------------------------------------------- + + +def test_a_plain_implicit_solve_is_supported(mesh): + uw, model = _fresh_model() + solver = _poisson(uw, mesh, "T_adj_ok") + model.tracker.time, model.tracker.step = 0.0, 0 + with model.step(0.1): + solver.solve() + verdict = _adjoint_of(model, "solve") + assert verdict["supported"] is True + assert "Jacobian transpose" in verdict["reason"] + + +def test_a_rotated_constraint_refuses_and_says_why(mesh): + """The rotated solve runs its own Krylov loop on a rotated operator; there + is no transpose path through it. The transcript must say so rather than + let the solve pass as an ordinary residual.""" + uw, model = _fresh_model() + stokes, _ = _stokes(uw, mesh, "rot") + stokes.add_dirichlet_bc((0.0, 0.0), "Bottom") + stokes.add_dirichlet_bc((0.0, sympy.oo), "Left") + stokes.add_dirichlet_bc((0.0, sympy.oo), "Right") + stokes.add_rotated_freeslip_bc(0.0, "Top") + model.tracker.time, model.tracker.step = 0.0, 0 + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + with model.step(0.1): + stokes.solve() + verdict = _adjoint_of(model, "solve") + assert verdict["supported"] is False + assert "rotated" in verdict["reason"] + + +def test_an_unconverged_solve_is_refused_after_the_fact(mesh): + """The structural verdict is written before the solve. A solve that then + diverged is linearised about a state it never reached — the outcome must + override the verdict.""" + uw, model = _fresh_model() + solver = _poisson(uw, mesh, "T_adj_div") + solver.petsc_options["ksp_max_it"] = 1 + solver.petsc_options["ksp_rtol"] = 1.0e-30 + solver.petsc_options["snes_max_it"] = 1 + model.tracker.time, model.tracker.step = 0.0, 0 + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + with model.step(0.1): + solver.solve() + verdict = _adjoint_of(model, "solve") + assert verdict["supported"] is False + assert "did not converge" in verdict["reason"] + + +# --------------------------------------------------------------------------- +# histories +# --------------------------------------------------------------------------- + + +def _advdiff(uw, mesh, tag, order=1): + T = uw.discretisation.MeshVariable(f"T_{tag}", mesh, 1, degree=2) + V = uw.discretisation.MeshVariable(f"V_{tag}", mesh, 2, degree=2) + x, y = mesh.X + V.array[:, 0, :] = np.asarray( + uw.function.evaluate(sympy.Matrix([[-(y - 0.5), (x - 0.5)]]), V.coords) + ).reshape(-1, 2) + solver = uw.systems.AdvDiffusion(mesh, u_Field=T, V_fn=V.sym, order=order) + solver.constitutive_model = uw.constitutive_models.DiffusionModel + solver.constitutive_model.Parameters.diffusivity = 1.0e-3 + solver.petsc_options.delValue("ksp_monitor") + return solver + + +def test_an_eulerian_history_is_supported(mesh): + """An implicit step IS a residual; the SUPG adjoint that passed its Taylor + test at 1.00000 is exactly this case.""" + uw, model = _fresh_model() + solver = _advdiff(uw, mesh, "eul") + model.tracker.time, model.tracker.step = 0.0, 0 + with model.step(0.01): + solver.solve(timestep=0.01) + verdict = _adjoint_of(model, "history_shift") + assert verdict["supported"] is True + assert "owning solver" in verdict["reason"] + + +def test_a_semi_lagrangian_history_refuses_and_names_what_is_missing(): + uw, model = _fresh_model() + # Its own mesh: the semi-Lagrangian trace-back fails point location on + # the module's shared mesh after the Eulerian test has run on it under + # pytest, though the same sequence passes as a script. Not chased here. + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1.0 / 8, qdegree=3 + ) + T = uw.discretisation.MeshVariable("T_sl", mesh, 1, degree=2) + V = uw.discretisation.MeshVariable("V_sl", mesh, 2, degree=2) + x, y = mesh.X + V.array[:, 0, :] = np.asarray( + uw.function.evaluate(sympy.Matrix([[-(y - 0.5), (x - 0.5)]]), V.coords) + ).reshape(-1, 2) + solver = uw.systems.AdvDiffusionSLCN(mesh, u_Field=T, V_fn=V.sym) + solver.constitutive_model = uw.constitutive_models.DiffusionModel + solver.constitutive_model.Parameters.diffusivity = 1.0e-3 + solver.petsc_options.delValue("ksp_monitor") + model.tracker.time, model.tracker.step = 0.0, 0 + with model.step(0.01): + solver.solve(timestep=0.01) + verdict = _adjoint_of(model, "history_shift") + assert verdict["supported"] is False + assert "departure" in verdict["reason"] + assert "not materialised" in verdict["reason"] + + +def test_every_history_scheme_declares_a_verdict(): + """The base refuses by naming the class, so a scheme added without a + verdict shows up as undeclared rather than passing as either.""" + from underworld3.systems import ddt + + base = ddt._DDtBase._adjoint_support + silent = [] + for name in dir(ddt): + cls = getattr(ddt, name) + if (isinstance(cls, type) and issubclass(cls, ddt._DDtBase) + and cls is not ddt._DDtBase + and cls._adjoint_support is base): + silent.append(name) + assert silent == [], f"these history schemes declare no adjoint verdict: {silent}" + + +# --------------------------------------------------------------------------- +# swarms +# --------------------------------------------------------------------------- + + +def _swarm_in_flow(uw, mesh, V_fn_matrix): + swarm = uw.swarm.Swarm(mesh) + swarm.populate(fill_param=2) + return swarm + + +def test_a_swarm_step_on_a_fixed_particle_set_is_supported(mesh): + """A rotating flow keeps every particle inside the box.""" + uw, model = _fresh_model() + x, y = mesh.X + V = sympy.Matrix([[-(y - 0.5), (x - 0.5)]]) + swarm = _swarm_in_flow(uw, mesh, V) + model.tracker.time, model.tracker.step = 0.0, 0 + with model.step(0.05): + swarm.advection(V, 0.05) + event = [e for e in model.transcript[0].events if e["kind"] == "swarm_advect"][-1] + assert event["n_before"] == event["n_after"] > 0 + assert event["adjoint"]["supported"] is True + assert "fixed particle set" in event["adjoint"]["reason"] + + +def test_a_swarm_step_that_loses_particles_refuses_with_the_count(): + """Particles leave the domain, the box's own return-to-bounds is switched + off, so the migrate deletes them: the state changed dimension, and no + linear map can be transposed across that. + + On its own mesh: switching the return-to-bounds off is a change to the + mesh, and the module's shared one is used by the semi-Lagrangian test.""" + uw, model = _fresh_model() + own = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1.0 / 8, qdegree=3 + ) + V = sympy.Matrix([[10.0, 0.0]]) # everything exits to the right + swarm = _swarm_in_flow(uw, own, V) + own.return_coords_to_bounds = None + model.tracker.time, model.tracker.step = 0.0, 0 + with model.step(0.5): + swarm.advection(V, 0.5) + event = [e for e in model.transcript[0].events if e["kind"] == "swarm_advect"][-1] + assert event["n_after"] < event["n_before"], event + assert event["adjoint"]["supported"] is False + assert f"{event['n_before']} -> {event['n_after']}" in event["adjoint"]["reason"] + + +# --------------------------------------------------------------------------- +# the partition +# --------------------------------------------------------------------------- + + +def _step(index, events, completed=True): + return {"kind": "step", "index": index, "label": None, "t0": index * 0.1, + "t1": (index + 1) * 0.1, "dt": 0.1, "completed": completed, + "restorable": True, "wall": 0.1, "events": events} + + +def _solve(name, supported, reason="r"): + return {"kind": "solve", "name": name, "part": f"{name}#1", + "adjoint": {"supported": supported, "reason": reason}} + + +def test_segments_partition_the_window_at_the_refusals(): + import underworld3 as uw + + steps = [_step(i, [_solve("Stokes(v)", True)]) for i in range(6)] + steps[3]["events"] = [_solve("Stokes(v)", False, "the particle set changed")] + runs = [{"run": {"kind": "run", "model": "t"}, "steps": steps, "notes": []}] + + segments = uw.transcript_adjoint_segments(runs) + assert [(s["first"], s["last"], s["supported"]) for s in segments] == [ + (0, 2, True), (3, 3, False), (4, 5, True) + ] + assert segments[1]["refusals"] == [("Stokes(v)", "the particle set changed")] + + +def test_segments_leave_abandoned_steps_out_and_flag_undeclared_verdicts(): + import underworld3 as uw + + steps = [ + _step(0, [_solve("Stokes(v)", True)]), + _step(1, [_solve("Stokes(v)", True)], completed=False), + _step(1, [{"kind": "solve", "name": "Old(v)", "part": "Old(v)#1"}]), + ] + runs = [{"run": {"kind": "run", "model": "t"}, "steps": steps, "notes": []}] + segments = uw.transcript_adjoint_segments(runs) + assert [s["steps"] for s in segments] == [1, 1] + assert segments[1]["supported"] is False + assert "without an adjoint verdict" in segments[1]["refusals"][0][1] + + +# --------------------------------------------------------------------------- +# the text transcript +# --------------------------------------------------------------------------- + + +def test_the_text_transcript_notes_a_refusal_once_per_change(tmp_path): + """A semi-Lagrangian run refuses identically every step. The note is + written when the set of refusals CHANGES — on the first refusing step, and + again when it clears — not three hundred times.""" + uw, model = _fresh_model() + model.transcript_file = tmp_path / "t.log" + model.transcript_format = "text" + model.tracker.time, model.tracker.step = 0.0, 0 + + def refusing(): + model._record_step_event( + "history_shift", "SemiLagrangian(T)", dt=0.1, + part="SemiLagrangian#9", + adjoint={"supported": False, "reason": "not materialised"}) + + for n in range(6): + with model.step(0.1): + if n in (1, 2, 4): + refusing() + + lines = (tmp_path / "t.log").read_text().splitlines() + refused = [l for l in lines if "no adjoint through" in l] + cleared = [l for l in lines if "admits one again" in l] + assert len(refused) == 2, refused # steps 1 and 4, not 1, 2 and 4 + assert len(cleared) == 2, cleared # steps 3 and 5 + # a clean run says nothing: step 0 is one line, with no note under it + assert "admits one" not in lines[lines.index(next(l for l in lines if l.strip().startswith("0 "))) + 1] + + +def test_segments_take_only_a_literal_true_as_support(): + """A verdict of None, 0 or "false" is not support (found in review).""" + import underworld3 as uw + + steps = [] + for i, value in enumerate([True, None, 0, "false", True]): + steps.append(_step(i, [{"kind": "solve", "name": "S", "part": "S#1", + "adjoint": {"supported": value, "reason": "r"}}])) + segments = uw.transcript_adjoint_segments( + [{"run": {"kind": "run", "model": "t"}, "steps": steps, "notes": []}]) + assert [s["supported"] for s in segments] == [True, False, True] + assert segments[1]["steps"] == 3 + + +def test_segments_do_not_fold_a_replayed_step_into_the_one_it_replaced(): + """After a rewind the record carries index 3, then 3 again. They are two + records, and the segment boundary must sit between them.""" + import underworld3 as uw + + steps = [_step(i, [_solve("S", True)]) for i in (0, 1, 2, 3)] + steps.append(_step(3, [_solve("S", True)])) # the replay + segments = uw.transcript_adjoint_segments( + [{"run": {"kind": "run", "model": "t"}, "steps": steps, "notes": []}]) + assert [(s["first_position"], s["last_position"]) for s in segments] == [(0, 3), (4, 4)] diff --git a/tests/test_0019_adjoint_solve.py b/tests/test_0019_adjoint_solve.py new file mode 100644 index 000000000..31d5e267d --- /dev/null +++ b/tests/test_0019_adjoint_solve.py @@ -0,0 +1,290 @@ +"""The discrete adjoint of one solve, checked against finite differences. + +``solver.adjoint_solve(b)`` solves :math:`K^T \\mu = b` against the Jacobian +the SNES already assembled; ``solver.sensitivity(mu, m)`` integrates the +symbolic :math:`\\partial R/\\partial m` against it. Together they are the +gradient of a misfit through one implicit solve, with no hand algebra. + +The check is the only one that counts: the adjoint gradient against a +central finite difference in the parameter. A symmetric operator (Poisson) +cannot tell a transpose from the operator itself, so the second case is one +SUPG advection–diffusion step, whose Jacobian is not symmetric — a wrong +transpose fails there. +""" + +import numpy as np +import pytest +import sympy + +pytestmark = [pytest.mark.level_1, pytest.mark.tier_a] + + +def _fresh(): + import underworld3 as uw + + uw.reset_default_model() + return uw, uw.get_default_model() + + +def _mesh(uw): + return uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1.0 / 8, qdegree=3 + ) + + +def _misfit(uw, mesh, T): + return float(uw.maths.Integral(mesh, sympy.Rational(1, 2) * T.sym[0] ** 2).evaluate()) + + +def test_poisson_gradient_in_the_diffusivity_matches_finite_differences(): + """J = 1/2 int T^2 for the Poisson solve with source 1 and diffusivity + kappa. K^T mu = -dJ/dT, then dJ/dkappa = int (dF1/dkappa) . grad(mu).""" + uw, model = _fresh() + mesh = _mesh(uw) + T = uw.discretisation.MeshVariable("T_adj", mesh, 1, degree=2) + mu = uw.discretisation.MeshVariable("mu_adj", mesh, 1, degree=2) + kappa = uw.expression(r"\kappa", 1.0, "diffusivity") + + solver = uw.systems.Poisson(mesh, u_Field=T) + solver.constitutive_model = uw.constitutive_models.DiffusionModel + solver.constitutive_model.Parameters.diffusivity = kappa + solver.f = 1.0 + solver.add_dirichlet_bc(0.0, "Top") + solver.add_dirichlet_bc(0.0, "Bottom") + solver.petsc_options.delValue("ksp_monitor") + solver.tolerance = 1.0e-12 + + def J_at(value): + kappa.sym = sympy.Float(value) + solver.solve(zero_init_guess=True) + return _misfit(uw, mesh, T) + + J0 = J_at(1.0) + b = -solver.dual_of(T.sym[0]) # -dJ/dT as a dual + _, reason = solver.adjoint_solve(b, target=mu) + assert reason > 0, reason + adjoint = solver.sensitivity(mu, kappa) + + h = 1.0e-4 + fd = (J_at(1.0 + h) - J_at(1.0 - h)) / (2 * h) + kappa.sym = sympy.Float(1.0) + assert adjoint == pytest.approx(fd, rel=1.0e-4), (adjoint, fd) + # the multiplier honours the homogenised Dirichlet conditions + top = np.abs(np.asarray(mu.coords)[:, 1] - 1.0) < 1.0e-10 + assert np.abs(np.asarray(mu.array)[top, 0, 0]).max() < 1.0e-12 + + +def test_one_supg_step_gradient_matches_finite_differences_where_the_jacobian_is_not_symmetric(): + """One implicit advection–diffusion step from a fixed initial state, + restored by snapshot before every evaluation so the finite difference + and the adjoint see the same step. SUPG makes K non-symmetric, so a + transpose taken the wrong way round fails here and not on Poisson.""" + uw, model = _fresh() + mesh = _mesh(uw) + x, y = mesh.X + T = uw.discretisation.MeshVariable("T_sup", mesh, 1, degree=2) + V = uw.discretisation.MeshVariable("V_sup", mesh, 2, degree=2) + mu = uw.discretisation.MeshVariable("mu_sup", mesh, 1, degree=2) + V.array[:, 0, :] = np.asarray( + uw.function.evaluate(sympy.Matrix([[-(y - 0.5), (x - 0.5)]]), V.coords) + ).reshape(-1, 2) + T.array[:, 0, 0] = np.asarray( + uw.function.evaluate(sympy.exp(-(((x - 0.3) ** 2 + (y - 0.5) ** 2) / 0.02)), T.coords) + ).ravel() + kappa = uw.expression(r"\kappa", 1.0e-2, "diffusivity") + + solver = uw.systems.AdvDiffusion(mesh, u_Field=T, V_fn=V.sym) + solver.constitutive_model = uw.constitutive_models.DiffusionModel + solver.constitutive_model.Parameters.diffusivity = kappa + solver.petsc_options.delValue("ksp_monitor") + solver.tolerance = 1.0e-12 + dt = 0.05 + + start = model.save_state() + + def J_at(value): + model.load_state(start) + kappa.sym = sympy.Float(value) + solver.solve(timestep=dt, zero_init_guess=True) + return _misfit(uw, mesh, T) + + T_old = np.array(T.array, copy=True) + J0 = J_at(1.0e-2) + # The residual of the step is F(T_new; T_old, v, dt). solve() shifted the + # history forward in its post-hook, so the slot now holds T_new; put the + # step's INPUT back where the residual reads it before linearising. + solver.DuDt.psi_star[0].array[...] = T_old + b = -solver.dual_of(T.sym[0]) + _, reason = solver.adjoint_solve(b, target=mu) + assert reason > 0, reason + adjoint = solver.sensitivity(mu, kappa) + + h = 1.0e-5 + fd = (J_at(1.0e-2 + h) - J_at(1.0e-2 - h)) / (2 * h) + assert adjoint == pytest.approx(fd, rel=1.0e-3), (adjoint, fd) + + +def test_a_refusing_solve_raises_with_its_reason(): + uw, model = _fresh() + mesh = _mesh(uw) + V = uw.discretisation.MeshVariable("V_ref", mesh, 2, degree=2) + P = uw.discretisation.MeshVariable("P_ref", mesh, 1, degree=1) + stokes = uw.systems.Stokes(mesh, velocityField=V, pressureField=P) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = 1.0 + stokes.add_rotated_freeslip_bc(0.0, "Top") + with pytest.raises(RuntimeError, match="rotated"): + stokes.adjoint_solve(np.zeros(1)) + + +def test_adjoint_before_any_solve_says_to_solve_first(): + uw, model = _fresh() + mesh = _mesh(uw) + T = uw.discretisation.MeshVariable("T_none", mesh, 1, degree=2) + solver = uw.systems.Poisson(mesh, u_Field=T) + with pytest.raises(RuntimeError, match="solve\\(\\) first"): + solver.adjoint_solve(np.zeros(1)) + + +def _stokes(uw, mesh, tag, viscosity): + V = uw.discretisation.MeshVariable(f"V_{tag}", mesh, 2, degree=2) + P = uw.discretisation.MeshVariable(f"P_{tag}", mesh, 1, degree=1) + stokes = uw.systems.Stokes(mesh, velocityField=V, pressureField=P) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = viscosity(stokes) + x, y = mesh.X + stokes.bodyforce = sympy.Matrix([0.0, -sympy.sin(sympy.pi * x) * sympy.sin(sympy.pi * y)]) + for b in ("Top", "Bottom"): + stokes.add_dirichlet_bc((0.0, 0.0), b) + for b in ("Left", "Right"): + stokes.add_dirichlet_bc((0.0, sympy.oo), b) + stokes.petsc_options.delValue("ksp_monitor") + stokes.tolerance = 1.0e-12 + return stokes, V, P + + +def _kinetic(uw, mesh, V): + return float(uw.maths.Integral(mesh, sympy.Rational(1, 2) * V.sym.dot(V.sym)).evaluate()) + + +def test_stokes_gradient_in_the_viscosity_matches_finite_differences(): + """Linear viscosity: K is symmetric, and the composite transpose must + reproduce what the example builds as a second Stokes solver.""" + uw, model = _fresh() + mesh = _mesh(uw) + eta0 = uw.expression(r"\eta_0", 1.0, "viscosity") + stokes, V, P = _stokes(uw, mesh, "lin", lambda s: eta0) + u_adj = uw.discretisation.MeshVariable("u_adj_lin", mesh, 2, degree=2) + p_adj = uw.discretisation.MeshVariable("p_adj_lin", mesh, 1, degree=1) + + def J_at(value): + eta0.sym = sympy.Float(value) + stokes.solve(zero_init_guess=True) + return _kinetic(uw, mesh, V) + + J_at(1.0) + b = -stokes.dual_of(V.sym) + _, reason = stokes.adjoint_solve(b, target=(u_adj, p_adj)) + assert reason > 0, reason + adjoint = stokes.sensitivity(u_adj, eta0) + + h = 1.0e-4 + fd = (J_at(1.0 + h) - J_at(1.0 - h)) / (2 * h) + assert adjoint == pytest.approx(fd, rel=1.0e-4), (adjoint, fd) + # dJ/deta for a viscous flow driven by a fixed body force is negative + assert adjoint < 0 + + +def test_a_nonlinear_rheology_solved_with_picard_still_gives_the_right_gradient(): + """Picard iterations spoil nothing: the converged state is the same, and + dR/du is a function of that state alone. What Picard leaves behind is a + Jacobian KERNEL that is the frozen-viscosity one — so the adjoint + assembles the consistent tangent itself, and the gradient matches finite + differences exactly as it does under Newton. The Picard kernel is put + back for the next forward solve.""" + uw, model = _fresh() + mesh = _mesh(uw) + eta0 = uw.expression(r"\eta_0", 1.0, "prefactor") + stokes, V, P = _stokes(uw, mesh, "pic", lambda s: eta0 / (1 + 4 * s.Unknowns.Einv2)) + stokes.consistent_jacobian = False # Picard, explicitly + u_adj = uw.discretisation.MeshVariable("u_adj_pic", mesh, 2, degree=2) + p_adj = uw.discretisation.MeshVariable("p_adj_pic", mesh, 1, degree=1) + + def J_at(value): + eta0.sym = sympy.Float(value) + stokes.solve(zero_init_guess=True) + assert stokes.solve_report.converged, stokes.solve_report + return _kinetic(uw, mesh, V) + + J_at(1.0) + supported, why = stokes.adjoint_support() + assert supported is True and "Picard" in why, why + b = -stokes.dual_of(V.sym) + _, reason = stokes.adjoint_solve(b, target=(u_adj, p_adj)) + assert reason > 0, reason + adjoint = stokes.sensitivity(u_adj, eta0) + assert stokes.consistent_jacobian is False # put back + + h = 1.0e-4 + fd = (J_at(1.0 + h) - J_at(1.0 - h)) / (2 * h) + assert adjoint == pytest.approx(fd, rel=1.0e-3), (adjoint, fd) + + +def test_a_nonlinear_rheology_with_the_consistent_tangent_matches_finite_differences(): + """The case the composite transpose exists for: eta(strain rate), where + the forward operator is not the adjoint operator.""" + uw, model = _fresh() + mesh = _mesh(uw) + eta0 = uw.expression(r"\eta_0", 1.0, "prefactor") + stokes, V, P = _stokes(uw, mesh, "nl", lambda s: eta0 / (1 + 4 * s.Unknowns.Einv2)) + stokes.consistent_jacobian = True + u_adj = uw.discretisation.MeshVariable("u_adj_nl", mesh, 2, degree=2) + p_adj = uw.discretisation.MeshVariable("p_adj_nl", mesh, 1, degree=1) + + def J_at(value): + eta0.sym = sympy.Float(value) + stokes.solve(zero_init_guess=True) + assert stokes.solve_report.converged, stokes.solve_report + return _kinetic(uw, mesh, V) + + J_at(1.0) + b = -stokes.dual_of(V.sym) + _, reason = stokes.adjoint_solve(b, target=(u_adj, p_adj)) + assert reason > 0, reason + adjoint = stokes.sensitivity(u_adj, eta0) + + h = 1.0e-4 + fd = (J_at(1.0 + h) - J_at(1.0 - h)) / (2 * h) + assert adjoint == pytest.approx(fd, rel=1.0e-3), (adjoint, fd) + + +def test_a_nonlinear_rheology_under_continuation_matches_finite_differences(): + """"continuation" solves Picard to a loose tolerance, then Newton, then + puts alpha back to 0. A fresh Jacobian assembly for the adjoint would + therefore be the PICARD tangent unless alpha is set to 1 for it — the + solve is right, and the matrix left behind is the wrong one to transpose.""" + uw, model = _fresh() + mesh = _mesh(uw) + eta0 = uw.expression(r"\eta_0", 1.0, "prefactor") + stokes, V, P = _stokes(uw, mesh, "cont", lambda s: eta0 / (1 + 4 * s.Unknowns.Einv2)) + stokes.consistent_jacobian = "continuation" + u_adj = uw.discretisation.MeshVariable("u_adj_cont", mesh, 2, degree=2) + p_adj = uw.discretisation.MeshVariable("p_adj_cont", mesh, 1, degree=1) + + def J_at(value): + eta0.sym = sympy.Float(value) + stokes.solve(zero_init_guess=True) + assert stokes.solve_report.converged, stokes.solve_report + return _kinetic(uw, mesh, V) + + J_at(1.0) + assert float(stokes._get_newton_alpha().sym) == 0.0 # what the solve leaves behind + b = -stokes.dual_of(V.sym) + _, reason = stokes.adjoint_solve(b, target=(u_adj, p_adj)) + assert reason > 0, reason + adjoint = stokes.sensitivity(u_adj, eta0) + assert float(stokes._get_newton_alpha().sym) == 0.0 # and is put back + + h = 1.0e-4 + fd = (J_at(1.0 + h) - J_at(1.0 - h)) / (2 * h) + assert adjoint == pytest.approx(fd, rel=1.0e-3), (adjoint, fd) diff --git a/tests/test_0020_transcript_adjoint.py b/tests/test_0020_transcript_adjoint.py new file mode 100644 index 000000000..b102ca114 --- /dev/null +++ b/tests/test_0020_transcript_adjoint.py @@ -0,0 +1,211 @@ +"""The backward pass over a recorded run, against finite differences. + +A two-solver, multi-step problem in the shape of the sinking blob: a level +set ``beta`` carried by an SUPG advection–diffusion solve in a velocity +``v``, and a Stokes solve whose body force is ``-Ra * beta``. The misfit is +on the final velocity. ``uw.adjoint.TranscriptAdjoint`` walks the transcript +backwards with no problem-specific wiring: the residuals say what each solve +reads, the transcript says what ran and holds the state each step started +from. + +Two controls, two checks: a scalar parameter (the viscosity) against a +central finite difference, and the initial level set as a FIELD control — +the dual on beta_0 dotted with a perturbation direction against the finite +difference of J along that direction. +""" + +import numpy as np +import pytest +import sympy + +pytestmark = [pytest.mark.level_2, pytest.mark.tier_a] + + +def _build(): + import underworld3 as uw + + uw.reset_default_model() + model = uw.get_default_model() + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=1.0 / 8, qdegree=3 + ) + x, y = mesh.X + beta = uw.discretisation.MeshVariable("beta", mesh, 1, degree=2) + v = uw.discretisation.MeshVariable("v", mesh, 2, degree=2) + p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1) + eta0 = uw.expression(r"\eta_0", 1.0, "viscosity") + Ra = uw.expression(r"Ra", 50.0, "buoyancy number") + + stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = eta0 + stokes.bodyforce = sympy.Matrix([0.0, -Ra * beta.sym[0]]) + for b in ("Top", "Bottom"): + stokes.add_dirichlet_bc((0.0, 0.0), b) + for b in ("Left", "Right"): + stokes.add_dirichlet_bc((0.0, sympy.oo), b) + stokes.petsc_options.delValue("ksp_monitor") + stokes.tolerance = 1.0e-12 + + adv = uw.systems.AdvDiffusion(mesh, u_Field=beta, V_fn=v.sym, theta=1.0) + adv.constitutive_model = uw.constitutive_models.DiffusionModel + adv.constitutive_model.Parameters.diffusivity = 1.0e-3 + adv.petsc_options.delValue("ksp_monitor") + adv.tolerance = 1.0e-12 + + def beta0(centre=(0.5, 0.6)): + X = np.asarray(beta.coords) + r = np.sqrt((X[:, 0] - centre[0]) ** 2 + (X[:, 1] - centre[1]) ** 2) + return r - 0.2 + + return dict(uw=uw, model=model, mesh=mesh, beta=beta, v=v, p=p, eta0=eta0, + Ra=Ra, stokes=stokes, adv=adv, beta0=beta0) + + +DT, NSTEPS = 0.02, 2 + + +def _forward(m, beta_initial): + """Run from ``beta_initial`` and return (final_state, J).""" + uw, model = m["uw"], m["model"] + model.clear_transcript() + model.tracker.time, model.tracker.step = 0.0, 0 + model.record_every = 1 + m["beta"].array[:, 0, 0] = beta_initial + m["v"].array[...] = 0.0 + m["p"].array[...] = 0.0 + # The Eulerian history initialises itself only on its FIRST solve; a + # driver that runs the forward model more than once must reset it each + # time the initial condition is set, or the second run reads the first + # run's history. + m["adv"].DuDt.initialise_history() + # Every solve inside a step, so every solve is on the tape: a Stokes + # solve taken before the first step would carry beta_0 -> v_0 with no + # record, and the walk could not see its dependence on the viscosity. + for _ in range(NSTEPS): + with model.step(DT): + m["stokes"].solve(zero_init_guess=True) # v_k from beta_k + m["adv"].solve(timestep=DT, zero_init_guess=False) + return model.save_state(), _misfit(m) + + +def _misfit_expr(m): + v = m["v"] + return sympy.Rational(1, 2) * v.sym.dot(v.sym) + + +def _misfit(m): + return float(m["uw"].maths.Integral(m["mesh"], _misfit_expr(m)).evaluate()) + + +def test_parameter_gradient_matches_finite_differences(): + m = _build() + uw, eta0 = m["uw"], m["eta0"] + b0 = m["beta0"]() + + final, J = _forward(m, b0) + result = uw.adjoint.TranscriptAdjoint(m["model"], final).gradient( + _misfit_expr(m), parameters=[eta0]) + assert result["J"] == pytest.approx(J, rel=1e-10) + adjoint = result["parameters"][eta0] + + h = 1.0e-4 + eta0.sym = sympy.Float(1.0 + h) + _, Jp = _forward(m, b0) + eta0.sym = sympy.Float(1.0 - h) + _, Jm = _forward(m, b0) + eta0.sym = sympy.Float(1.0) + fd = (Jp - Jm) / (2 * h) + assert adjoint == pytest.approx(fd, rel=1.0e-4), (adjoint, fd) + + +def test_initial_field_gradient_matches_a_directional_finite_difference(): + m = _build() + uw, beta = m["uw"], m["beta"] + b0 = m["beta0"]() + + final, J = _forward(m, b0) + result = uw.adjoint.TranscriptAdjoint(m["model"], final).gradient( + _misfit_expr(m), fields=[beta]) + dual = result["fields"][beta][:, 0, 0] + + # a smooth direction in beta_0, and J along it + X = np.asarray(beta.coords) + direction = np.sin(np.pi * X[:, 0]) * np.sin(np.pi * X[:, 1]) + h = 1.0e-3 + _, Jp = _forward(m, b0 + h * direction) + _, Jm = _forward(m, b0 - h * direction) + fd = (Jp - Jm) / (2 * h) + # over the OWNED degrees of freedom: a NumPy dot on .array counts the + # ghost nodes of a partition twice (found in review at np=2) + adjoint = uw.adjoint.inner(beta, dual, direction) + assert adjoint == pytest.approx(fd, rel=1.0e-4), (adjoint, fd) + + +def test_a_misfit_that_names_the_parameter_gets_its_explicit_term(): + """dJ/dm = the implicit part through the solves plus dJ/dm at the final + level for a misfit written in terms of m (found in review: omitted).""" + m = _build() + uw, eta0 = m["uw"], m["eta0"] + b0 = m["beta0"]() + misfit = eta0 * _misfit_expr(m) + + final, J = _forward(m, b0) + adjoint = uw.adjoint.TranscriptAdjoint(m["model"], final).gradient( + misfit, parameters=[eta0])["parameters"][eta0] + + def J_at(value): + eta0.sym = sympy.Float(value) + _forward(m, b0) + return float(uw.maths.Integral(m["mesh"], misfit).evaluate()) + + h = 1.0e-4 + fd = (J_at(1.0 + h) - J_at(1.0 - h)) / (2 * h) + eta0.sym = sympy.Float(1.0) + assert adjoint == pytest.approx(fd, rel=1.0e-4), (adjoint, fd) + + +def test_gradient_reuses_its_scratch_fields(): + """Sixteen registered variables leaked per call and the sixth call took + ten times the first (found in review).""" + m = _build() + uw, eta0 = m["uw"], m["eta0"] + final, _ = _forward(m, m["beta0"]()) + back = uw.adjoint.TranscriptAdjoint(m["model"], final) + back.gradient(_misfit_expr(m), parameters=[eta0]) + n_after_first = len(m["model"]._variables) + for _ in range(3): + back.gradient(_misfit_expr(m), parameters=[eta0]) + assert len(m["model"]._variables) == n_after_first + + +def test_a_crank_nicolson_step_reads_the_old_flux_and_the_gradient_still_matches(): + """theta = 0.5 (the AdvDiffusion default) reads the previous level through + its GRADIENT — kappa grad(T_old) in the residual — so the dual on that + level has a gradient part. This was refused by name; now it is assembled + as the FEM load int g1 . grad(phi_j), and the field gradient matches + finite differences as it does at theta = 1.""" + m = _build() + uw, beta = m["uw"], m["beta"] + adv = uw.systems.AdvDiffusion(m["mesh"], u_Field=beta, V_fn=m["v"].sym) # default theta + adv.constitutive_model = uw.constitutive_models.DiffusionModel + adv.constitutive_model.Parameters.diffusivity = 1.0e-2 # a visible flux term + adv.petsc_options.delValue("ksp_monitor") + adv.tolerance = 1.0e-12 + assert adv.theta == 0.5 + m["adv"] = adv + b0 = m["beta0"]() + + final, J = _forward(m, b0) + result = uw.adjoint.TranscriptAdjoint(m["model"], final).gradient( + _misfit_expr(m), fields=[beta], parameters=[m["eta0"]]) + dual = result["fields"][beta][:, 0, 0] + + X = np.asarray(beta.coords) + direction = np.sin(np.pi * X[:, 0]) * np.sin(np.pi * X[:, 1]) + h = 1.0e-3 + _, Jp = _forward(m, b0 + h * direction) + _, Jm = _forward(m, b0 - h * direction) + fd = (Jp - Jm) / (2 * h) + adjoint = uw.adjoint.inner(beta, dual, direction) + assert adjoint == pytest.approx(fd, rel=1.0e-4), (adjoint, fd) diff --git a/tests/test_0641_wave_c_api_shims.py b/tests/test_0641_wave_c_api_shims.py index b143ef893..87fdb2bd4 100644 --- a/tests/test_0641_wave_c_api_shims.py +++ b/tests/test_0641_wave_c_api_shims.py @@ -250,9 +250,13 @@ def test_invalid_values_raise(self, mesh, stokes): with pytest.raises(ValueError, match="consistent_jacobian"): stokes.consistent_jacobian = value - def test_default_is_false(self, mesh): + def test_default_is_the_consistent_tangent(self, mesh): + """Newton by default: the residual is symbolic, so the tangent is + exact and cheap, and it is the matrix the adjoint transposes. Picard + (``False``) is the opt-in for the hard-yield solves that need it as + an entry requirement (flipped 2026-09).""" solver = uw.systems.Poisson(mesh) - assert solver.consistent_jacobian is False + assert solver.consistent_jacobian is True # --------------------------------------------------------------------------- diff --git a/tests/test_1057_yield_homotopy_solve.py b/tests/test_1057_yield_homotopy_solve.py index 8bd74fd64..6606ce8ec 100644 --- a/tests/test_1057_yield_homotopy_solve.py +++ b/tests/test_1057_yield_homotopy_solve.py @@ -272,9 +272,14 @@ def test_homotopy_rescues_a_solve_the_cold_start_cannot_do(): """ import numpy as np - # (a) the direct cold solve of the sharp law FAILS + # (a) the direct cold solve of the sharp law FAILS — under the Picard + # tangent, which is what the hard-yield entry problem is about. With the + # consistent tangent (the default since 2026-09) this cold solve + # CONVERGES on its own, so Picard is set explicitly here: the rescue + # being demonstrated is of the Picard solve that homotopy was built for. mesh, cold, cm_cold = _yielding_box("c", 0.30) cm_cold.yield_mode = "min" + cold.consistent_jacobian = False cold.solve() cold_reason = int(cold.snes.getConvergedReason()) From e3054850174f14096217531fcc7b02492448c29e Mon Sep 17 00:00:00 2001 From: lmoresi Date: Wed, 16 Sep 2026 06:23:15 +1000 Subject: [PATCH 02/23] feat: a misfit read through a gradient, and the fault-segments example uw.adjoint.misfit_duals assembles dJ/df for each field a misfit reads, through the field's value and its gradient, as one load. A misfit on a stress or a strain rate reads the velocity through its gradient, and the driver's misfit dual carried only the value part. The driver now routes through it. docs/examples/adjoint/fault_segments: a dipping fault as a weak plane in a transversely isotropic viscosity, three segments of unknown weak-plane viscosity, observed through the surface uplift rate and the shear stress near four points. The adjoint gradient matches central differences to 1.00000 on every segment, and L-BFGS recovers (0.0101, 0.1001, 0.0300) from a uniform 0.1 in thirteen evaluations. test_0021 is the same check in miniature. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- .../adjoint/fault_segments/fault_segments.py | 184 ++++++++++++++++++ .../fault_segments/plot_fault_segments.py | 40 ++++ src/underworld3/adjoint.py | 72 +++++-- ...st_0021_adjoint_misfit_through_gradient.py | 89 +++++++++ 4 files changed, 365 insertions(+), 20 deletions(-) create mode 100644 docs/examples/adjoint/fault_segments/fault_segments.py create mode 100644 docs/examples/adjoint/fault_segments/plot_fault_segments.py create mode 100644 tests/test_0021_adjoint_misfit_through_gradient.py diff --git a/docs/examples/adjoint/fault_segments/fault_segments.py b/docs/examples/adjoint/fault_segments/fault_segments.py new file mode 100644 index 000000000..3d7db883a --- /dev/null +++ b/docs/examples/adjoint/fault_segments/fault_segments.py @@ -0,0 +1,184 @@ +"""Strength of a fault, segment by segment, from surface uplift and stress. + +A dipping fault under horizontal shortening, represented as a weak plane in a +transversely isotropic viscosity (no cut in the mesh). Its weak-plane +viscosity is different on each of a few segments down dip, and those are the +unknowns. The observations are the uplift rate along the top surface and the +shear stress in the bulk near a handful of points, taken from a run at the +true strengths. + +The gradient of the misfit with respect to each segment's strength comes +from the solver's own discrete adjoint: a transpose against the Jacobian it +assembled, and the symbolic derivative of its residual with respect to the +named strength. Nothing here is differenced. + +Run it: + + python fault_segments.py # twin experiment: gradient check, then the inversion + python fault_segments.py -uw_check_only 1 +""" +import math + +import numpy as np +import sympy + +import underworld3 as uw +from underworld3.adjoint import misfit_duals, inner + +params = uw.Params( + cell_size=uw.Param(1 / 24, "mesh cell size (box is 2 x 1)"), + dip=uw.Param(45.0, "fault dip, degrees"), + n_segments=uw.Param(3, "segments of independent strength down dip"), + band=uw.Param(0.08, "half-width of the weak band, in box units"), + true_strengths=uw.Param("0.01,0.1,0.03", "weak-plane viscosity per segment, true run"), + initial_strength=uw.Param(0.1, "starting guess, every segment"), + n_stress_points=uw.Param(4, "shear-stress observation points in the bulk"), + check_only=uw.Param(0, "1: gradient check against finite differences, no inversion"), +) + +# --- the model --------------------------------------------------------------- +mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0.0, 0.0), maxCoords=(2.0, 1.0), + cellSize=params.cell_size, qdegree=3) +x, y = mesh.X + +v = uw.discretisation.MeshVariable("v", mesh, mesh.dim, degree=2) +p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1, continuous=True) +v_obs = uw.discretisation.MeshVariable("v_obs", mesh, mesh.dim, degree=2) + +# The fault: a plane through (0.6, 0) at the given dip, a weak band around it, +# and segments of equal length along it. +theta = math.radians(params.dip) +t_hat = sympy.Matrix([[math.cos(theta), math.sin(theta)]]) # along dip +n_hat = sympy.Matrix([[-math.sin(theta), math.cos(theta)]]) # the director +d = (x - 0.6) * n_hat[0] + y * n_hat[1] # distance from the plane +s = (x - 0.6) * t_hat[0] + y * t_hat[1] # position along it +length = 1.0 / math.sin(theta) +band = sympy.exp(-(d / params.band) ** 2) + +n_seg = int(params.n_segments) +strengths = [uw.expression(rf"\eta_{k + 1}", params.initial_strength, + f"weak-plane viscosity of segment {k + 1}") + for k in range(n_seg)] + +def segment(k): + """A smooth indicator for segment k along the fault, in [0, 1].""" + edge = params.band + lo, hi = k * length / n_seg, (k + 1) * length / n_seg + on = 1 if k == 0 else (1 + sympy.tanh((s - lo) / edge)) / 2 + off = 1 if k == n_seg - 1 else (1 - sympy.tanh((s - hi) / edge)) / 2 + return on * off + +eta_0 = 1 +eta_1 = eta_0 - band * sum((eta_0 - strengths[k]) * segment(k) for k in range(n_seg)) + +stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) +stokes.constitutive_model = uw.constitutive_models.TransverseIsotropicFlowModel +stokes.constitutive_model.Parameters.shear_viscosity_0 = eta_0 +stokes.constitutive_model.Parameters.shear_viscosity_1 = eta_1 +stokes.constitutive_model.Parameters.director = n_hat +stokes.tolerance = 1e-8 + +# Shortening from both sides, a no-slip base, and a free top: the surface +# velocity is the uplift rate. +stokes.add_essential_bc((0.0, 0.0), "Bottom") +stokes.add_essential_bc((0.5, None), "Left") +stokes.add_essential_bc((-0.5, None), "Right") + +# --- the observations ------------------------------------------------------ +w_top = sympy.exp(-((1 - y) / params.band) ** 2) +rng = np.random.default_rng(7) +points = [(0.3 + 1.4 * i / max(int(params.n_stress_points) - 1, 1), 0.3 + 0.3 * (i % 2)) + for i in range(int(params.n_stress_points))] +w_points = sum(sympy.exp(-((x - px) ** 2 + (y - py) ** 2) / (2 * params.band) ** 2) + for px, py in points) + +def shear_stress(field): + e = mesh.vector.strain_tensor(field.sym) + return 2 * eta_0 * e[0, 1] + +misfit = (w_top * (v.sym[1] - v_obs.sym[1]) ** 2 + + w_points * (shear_stress(v) - shear_stress(v_obs)) ** 2) / 2 + +def set_strengths(values): + for expr, value in zip(strengths, values): + expr.sym = float(value) + +def J_and_gradient(): + """The misfit and dJ/d(log strength) for each segment, by the adjoint.""" + stokes.solve(zero_init_guess=True) + J = float(uw.maths.Integral(mesh, misfit).evaluate()) + dJ_dv = misfit_duals(misfit, [v])[v] + mu = uw.discretisation.MeshVariable(f"mu_{uw.adjoint._counter()}", mesh, mesh.dim, degree=2) + lam = uw.discretisation.MeshVariable(f"lam_{uw.adjoint._counter()}", mesh, 1, degree=1) + dJ_dv.array[...] = -np.asarray(dJ_dv.array) + _, reason = stokes.adjoint_solve((dJ_dv, None), target=(mu, lam)) + assert reason > 0, reason + grad = np.array([stokes.sensitivity(mu, expr) * float(expr.sym) for expr in strengths]) + return J, grad + +# --- the truth, and the twin ------------------------------------------------- +true_values = [float(t) for t in str(params.true_strengths).split(",")][:n_seg] +set_strengths(true_values) +stokes.solve(zero_init_guess=True) +v_obs.array[...] = np.asarray(v.array) +uw.pprint(f"true strengths {true_values}") + +set_strengths([params.initial_strength] * n_seg) +J0, g0 = J_and_gradient() +uw.pprint(f"initial J = {J0:.6e} dJ/dlog eta = {g0}") + +# Gradient check: central differences in each log-strength. +h = 1e-3 +for k in range(n_seg): + base = math.log(params.initial_strength) + fd = [] + for sign in (+1, -1): + vals = [params.initial_strength] * n_seg + vals[k] = math.exp(base + sign * h) + set_strengths(vals) + stokes.solve(zero_init_guess=True) + fd.append(float(uw.maths.Integral(mesh, misfit).evaluate())) + fd = (fd[0] - fd[1]) / (2 * h) + uw.pprint(f"segment {k + 1}: adjoint {g0[k]: .6e} finite difference {fd: .6e} " + f"ratio {fd / g0[k]:.5f}") +set_strengths([params.initial_strength] * n_seg) + +if int(params.check_only): + raise SystemExit + +# --- the inversion --------------------------------------------------------------- +from scipy.optimize import minimize + +history = [] + +def objective(log_eta): + set_strengths(np.exp(log_eta)) + J, grad = J_and_gradient() + history.append((J, np.exp(log_eta).copy())) + uw.pprint(f" J = {J:.6e} strengths = {np.exp(log_eta)}") + return J, grad + +result = minimize(objective, np.log([params.initial_strength] * n_seg), jac=True, + method="L-BFGS-B", options={"maxiter": 40, "gtol": 1e-10}) +uw.pprint(f"recovered {np.exp(result.x)} true {true_values} " + f"after {len(history)} evaluations") + +# --- what the figure needs ----------------------------------------------------- +# Uplift-rate profiles along the top at the truth, the start and the answer, the +# weak-plane viscosity on a grid, and the path the strengths took. +xs = np.linspace(0.0, 2.0, 161) +top = np.column_stack([xs, np.full_like(xs, 1.0 - 1e-6)]) +profiles = {} +for label, values in (("true", true_values), ("initial", [params.initial_strength] * n_seg), + ("recovered", list(np.exp(result.x)))): + set_strengths(values) + stokes.solve(zero_init_guess=True) + profiles[label] = np.asarray(uw.function.evaluate(v.sym[1], top)).ravel() +set_strengths(true_values) +gx, gy = np.meshgrid(np.linspace(0, 2, 201), np.linspace(0, 1, 101)) +grid = np.column_stack([gx.ravel(), gy.ravel()]) +eta_1_grid = np.asarray(uw.function.evaluate(eta_1, grid)).reshape(gx.shape) +np.savez("fault_segments_data.npz", xs=xs, gx=gx, gy=gy, eta_1=eta_1_grid, + points=np.array(points), true=np.array(true_values), + history=np.array([[J, *vals] for J, vals in history]), + **{f"uplift_{k}": val for k, val in profiles.items()}) diff --git a/docs/examples/adjoint/fault_segments/plot_fault_segments.py b/docs/examples/adjoint/fault_segments/plot_fault_segments.py new file mode 100644 index 000000000..a379293fc --- /dev/null +++ b/docs/examples/adjoint/fault_segments/plot_fault_segments.py @@ -0,0 +1,40 @@ +"""The figure for the fault-segments example, from fault_segments_data.npz.""" +import numpy as np +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +d = np.load("fault_segments_data.npz") +history, true = d["history"], d["true"] +n_seg = len(true) + +fig, axes = plt.subplots(1, 3, figsize=(11, 3.2), gridspec_kw={"width_ratios": [2.2, 1.6, 1.4]}) + +ax = axes[0] +ax.contourf(d["gx"], d["gy"], np.log10(d["eta_1"]), levels=np.linspace(-2.2, 0, 12), cmap="viridis") +ax.plot(d["points"][:, 0], d["points"][:, 1], "wx", ms=7, mew=1.5) +ax.set_aspect("equal") +ax.set_xlim(0, 2); ax.set_ylim(0, 1) +ax.set_xlabel("$x$"); ax.set_ylabel("$y$") +ax.set_title(r"$\log_{10}\eta_1$ at the true strengths; $\times$ stress points", fontsize=10) + +ax = axes[1] +for label, style in (("true", "k-"), ("initial", "C3--"), ("recovered", "C0:")): + ax.plot(d["xs"], d[f"uplift_{label}"], style, lw=1.6, label=label) +ax.set_xlabel("$x$ along the surface"); ax.set_ylabel("uplift rate $v_y$") +ax.set_title("surface uplift rate", fontsize=10) +ax.legend(fontsize=8, frameon=False) + +ax = axes[2] +its = np.arange(len(history)) +for k in range(n_seg): + ax.semilogy(its, history[:, 1 + k], "o-", ms=3, lw=1, color=f"C{k}", label=f"segment {k + 1}") + ax.axhline(true[k], color=f"C{k}", lw=0.8, ls="--") +ax.set_xlabel("misfit evaluation"); ax.set_ylabel("weak-plane viscosity") +ax.set_title("strengths (dashed: true)", fontsize=10) +ax.legend(fontsize=8, frameon=False) + +fig.tight_layout() +fig.savefig("fault_segments.png", dpi=180) +fig.savefig("fault_segments.pdf") +print("wrote fault_segments.png / .pdf") diff --git a/src/underworld3/adjoint.py b/src/underworld3/adjoint.py index 8f1b11ba7..8e90e9e46 100644 --- a/src/underworld3/adjoint.py +++ b/src/underworld3/adjoint.py @@ -202,6 +202,53 @@ def inner(variable, a, b): return value +def _token_of(var): + """How ``var`` prints inside an expression: everything before the + coordinate arguments and, for a vector, before the component index.""" + text = str(_symbols_of(var)[0]).rsplit("(", 1)[0] + if getattr(var, "num_components", 1) > 1: + text = text.rsplit("_{", 1)[0] + return text + + +def misfit_duals(misfit, variables, scratch=None): + r"""``dJ/df`` as a dual field on each field the misfit reads. + + ``J = \int misfit`` over the mesh; a field enters through its value and, + for a misfit on a stress or a strain rate, through its gradient. Both + parts are differentiated symbolically and assembled as one load + (:func:`dual_on`), so a misfit written in terms of :math:`\nabla u` needs + no integration by parts by the caller. Returns ``{variable: dual}`` for + the variables that appear; give each dual back to the scratch pool when + done. + """ + scratch = _shared_scratch if scratch is None else scratch + peeled = _peel(misfit) + text = str(peeled) + atoms = set(peeled.atoms(sympy.Function)) + out = {} + for var in variables: + token = _token_of(var) + if token not in text: + continue + symbols = _symbols_of(var) + value = [sympy.diff(peeled, s) for s in symbols] + pattern = re.compile(re.escape(token) + r"_\{ ?(\d*),(\d+)\}\(") + g1 = None + for atom in atoms: + m = pattern.match(str(atom)) + if not m: + continue + if g1 is None: + g1 = sympy.zeros(len(symbols), var.mesh.cdim) + i = int(m.group(1)) if m.group(1) else 0 + g1[i, int(m.group(2))] = sympy.diff(peeled, atom) + if all(v == 0 for v in value) and (g1 is None or g1.is_zero_matrix): + continue + out[var] = dual_on(var, _as_expression(value), g1, scratch) + return out + + _n = [0] @@ -267,10 +314,8 @@ def gradient(self, misfit, parameters: Iterable = (), fields: Iterable = ()): model.load_state(self.final_state) J = float(uw.maths.Integral(self._mesh(), misfit).evaluate()) acc: Dict[str, object] = {} - peeled = _peel(misfit) - for var, symbols in self._fields_in(misfit): - dJ = [sympy.diff(peeled, s) for s in symbols] - self._accumulate(acc, var, dual_on(var, _as_expression(dJ), None, scratch)) + for var, dual in misfit_duals(misfit, self._tokens().values(), scratch).items(): + self._accumulate(acc, var, dual) grad = {p: 0.0 for p in parameters} for p in parameters: @@ -334,24 +379,11 @@ def _tokens(self): A variable prints as its symbol, which need not be its name and can carry nested braces (a history slot is ``{\\psi^{*}_{...}}``), so the - token is taken from the symbol's own text: everything before the - coordinate arguments and, for a vector, before the component index. + token is taken from the symbol's own text (:func:`_token_of`). Matching on the name found the user's fields and silently missed every history, which cut the chain at the first step.""" - out = {} - for var in self.model._variables.values(): - if not hasattr(var, "sym"): - continue - text = str(_symbols_of(var)[0]).rsplit("(", 1)[0] # drop (N.x, N.y) - if getattr(var, "num_components", 1) > 1: - text = text.rsplit("_{", 1)[0] # drop _{ i } - out[text] = var - return out - - def _fields_in(self, expression): - text = str(_peel(expression)) - return [(v, _symbols_of(v)) for token, v in self._tokens().items() - if token in text] + return {_token_of(var): var for var in self.model._variables.values() + if hasattr(var, "sym")} def _reads(self, solver, unknown): """What a solver's residual reads, other than its unknown. diff --git a/tests/test_0021_adjoint_misfit_through_gradient.py b/tests/test_0021_adjoint_misfit_through_gradient.py new file mode 100644 index 000000000..2e67e852e --- /dev/null +++ b/tests/test_0021_adjoint_misfit_through_gradient.py @@ -0,0 +1,89 @@ +"""A misfit on a strain rate: the dual carries a gradient part. + +The fault-segments example in miniature — a weak plane in a transversely +isotropic viscosity, two segments of unknown weak-plane viscosity, the misfit +on the surface uplift rate and on the shear stress near a point. The stress +reads the velocity through its gradient, so ``misfit_duals`` must assemble +the gradient part of the load; without it the adjoint gradient is wrong by +the whole stress term. +""" +import math + +import numpy as np +import pytest +import sympy + +import underworld3 as uw +from underworld3.adjoint import misfit_duals + + +@pytest.mark.level_2 +@pytest.mark.tier_a +def test_segment_strength_gradient_matches_finite_differences(): + mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0.0, 0.0), maxCoords=(2.0, 1.0), + cellSize=1 / 10, qdegree=3) + x, y = mesh.X + v = uw.discretisation.MeshVariable("v", mesh, 2, degree=2) + p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1, continuous=True) + v_obs = uw.discretisation.MeshVariable("v_obs", mesh, 2, degree=2) + + theta = math.radians(45.0) + n_hat = sympy.Matrix([[-math.sin(theta), math.cos(theta)]]) + d = (x - 0.6) * n_hat[0] + y * n_hat[1] + s = (x - 0.6) * math.cos(theta) + y * math.sin(theta) + band = sympy.exp(-(d / 0.15) ** 2) + strengths = [uw.expression(r"\eta_1", 0.1, "segment 1"), + uw.expression(r"\eta_2", 0.1, "segment 2")] + half = 0.5 / math.sin(theta) + seg = [(1 - sympy.tanh((s - half) / 0.15)) / 2, (1 + sympy.tanh((s - half) / 0.15)) / 2] + eta_1 = 1 - band * sum((1 - strengths[k]) * seg[k] for k in range(2)) + + stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + stokes.constitutive_model = uw.constitutive_models.TransverseIsotropicFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = 1 + stokes.constitutive_model.Parameters.shear_viscosity_1 = eta_1 + stokes.constitutive_model.Parameters.director = n_hat + stokes.tolerance = 1e-9 + stokes.add_essential_bc((0.0, 0.0), "Bottom") + stokes.add_essential_bc((0.5, None), "Left") + stokes.add_essential_bc((-0.5, None), "Right") + + def shear(field): + e = mesh.vector.strain_tensor(field.sym) + return 2 * e[0, 1] + + w_top = sympy.exp(-((1 - y) / 0.15) ** 2) + w_pt = sympy.exp(-((x - 1.3) ** 2 + (y - 0.4) ** 2) / 0.3 ** 2) + misfit = (w_top * (v.sym[1] - v_obs.sym[1]) ** 2 + + w_pt * (shear(v) - shear(v_obs)) ** 2) / 2 + + def set_strengths(a, b): + strengths[0].sym, strengths[1].sym = float(a), float(b) + + def J(): + stokes.solve(zero_init_guess=True) + return float(uw.maths.Integral(mesh, misfit).evaluate()) + + set_strengths(0.02, 0.2) + stokes.solve(zero_init_guess=True) + v_obs.array[...] = np.asarray(v.array) + + set_strengths(0.1, 0.1) + J() + dual = misfit_duals(misfit, [v])[v] + dual.array[...] = -np.asarray(dual.array) + mu = uw.discretisation.MeshVariable("mu", mesh, 2, degree=2) + lam = uw.discretisation.MeshVariable("lam", mesh, 1, degree=1) + _, reason = stokes.adjoint_solve((dual, None), target=(mu, lam)) + assert reason > 0 + adjoint = [stokes.sensitivity(mu, e) for e in strengths] + + h = 1e-4 + for k in range(2): + plus = [0.1, 0.1]; plus[k] += h + minus = [0.1, 0.1]; minus[k] -= h + set_strengths(*plus); Jp = J() + set_strengths(*minus); Jm = J() + fd = (Jp - Jm) / (2 * h) + assert adjoint[k] != 0 + assert abs(fd / adjoint[k] - 1) < 2e-3, (k, fd, adjoint[k]) From 459856e3f24b4562ff222e9d78f10dbf835114ed Mon Sep 17 00:00:00 2001 From: lmoresi Date: Wed, 16 Sep 2026 09:12:26 +1000 Subject: [PATCH 03/23] The fault-segments example is listric: a flat decollement, a circular ramp, four strengths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fault leaves a flat at depth horizontally and steepens on a circular ramp to sixty degrees at the surface. Signed distance, position along the fault and the director (vertical on the flat, radial on the ramp) are exact on each piece. Four segments — flat, lower ramp, upper ramp, near surface — with true strengths (0.005, 0.05, 0.02, 0.2); the adjoint gradient matches central differences to 1.00000 on each, and L-BFGS recovers (0.0050, 0.0500, 0.0200, 0.2001) from a uniform 0.1 in twenty evaluations. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- .../adjoint/fault_segments/fault_segments.py | 72 +++++++++++-------- .../fault_segments/plot_fault_segments.py | 5 +- 2 files changed, 46 insertions(+), 31 deletions(-) diff --git a/docs/examples/adjoint/fault_segments/fault_segments.py b/docs/examples/adjoint/fault_segments/fault_segments.py index 3d7db883a..5f92584f5 100644 --- a/docs/examples/adjoint/fault_segments/fault_segments.py +++ b/docs/examples/adjoint/fault_segments/fault_segments.py @@ -1,11 +1,13 @@ -"""Strength of a fault, segment by segment, from surface uplift and stress. +"""Strength of a listric fault, segment by segment, from surface uplift and stress. -A dipping fault under horizontal shortening, represented as a weak plane in a -transversely isotropic viscosity (no cut in the mesh). Its weak-plane -viscosity is different on each of a few segments down dip, and those are the -unknowns. The observations are the uplift rate along the top surface and the -shear stress in the bulk near a handful of points, taken from a run at the -true strengths. +A listric fault under horizontal shortening: a ramp that steepens from a flat +decollement at depth to a dip of sixty degrees at the surface, represented as +a weak plane in a transversely isotropic viscosity (no cut in the mesh). Its +weak-plane viscosity is different on the flat, on the lower and upper parts +of the ramp, and near the surface, and those four strengths are the unknowns. +The observations are the uplift rate along the top surface and the shear +stress in the bulk near a handful of points, taken from a run at the true +strengths. The gradient of the misfit with respect to each segment's strength comes from the solver's own discrete adjoint: a transpose against the Jacobian it @@ -27,12 +29,13 @@ params = uw.Params( cell_size=uw.Param(1 / 24, "mesh cell size (box is 2 x 1)"), - dip=uw.Param(45.0, "fault dip, degrees"), - n_segments=uw.Param(3, "segments of independent strength down dip"), + surface_dip=uw.Param(60.0, "dip of the ramp where it reaches the surface, degrees"), + flat_depth=uw.Param(0.3, "height of the decollement above the base"), + surface_x=uw.Param(1.9, "where the fault reaches the surface"), band=uw.Param(0.08, "half-width of the weak band, in box units"), - true_strengths=uw.Param("0.01,0.1,0.03", "weak-plane viscosity per segment, true run"), + true_strengths=uw.Param("0.005,0.05,0.02,0.2", + "weak-plane viscosity: flat, lower ramp, upper ramp, near surface"), initial_strength=uw.Param(0.1, "starting guess, every segment"), - n_stress_points=uw.Param(4, "shear-stress observation points in the bulk"), check_only=uw.Param(0, "1: gradient check against finite differences, no inversion"), ) @@ -45,27 +48,40 @@ p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1, continuous=True) v_obs = uw.discretisation.MeshVariable("v_obs", mesh, mesh.dim, degree=2) -# The fault: a plane through (0.6, 0) at the given dip, a weak band around it, -# and segments of equal length along it. -theta = math.radians(params.dip) -t_hat = sympy.Matrix([[math.cos(theta), math.sin(theta)]]) # along dip -n_hat = sympy.Matrix([[-math.sin(theta), math.cos(theta)]]) # the director -d = (x - 0.6) * n_hat[0] + y * n_hat[1] # distance from the plane -s = (x - 0.6) * t_hat[0] + y * t_hat[1] # position along it -length = 1.0 / math.sin(theta) +# The fault: a flat decollement at y = flat_depth running from the left wall +# to x = xc, then a circular ramp of radius R about (xc, yc) that leaves the +# flat horizontally and reaches the surface at the given dip. The signed +# distance to it and the position along it are exact on each piece, and the +# director is the normal at the nearest point: vertical on the flat, radial on +# the ramp. +phi_top = math.radians(params.surface_dip) - math.pi / 2 # angle of the surface point about the centre +R = (1.0 - params.flat_depth) / (1.0 + math.sin(phi_top)) +yc = params.flat_depth + R +xc = params.surface_x - R * math.cos(phi_top) +r = sympy.sqrt((x - xc) ** 2 + (y - yc) ** 2) +phi = sympy.atan2(y - yc, x - xc) +on_flat = x < xc +d = sympy.Piecewise((y - params.flat_depth, on_flat), (R - r, True)) +s = sympy.Piecewise((x, on_flat), (xc + R * (phi + sympy.pi / 2), True)) +n_hat = sympy.Matrix([[sympy.Piecewise((0, on_flat), ((x - xc) / r, True)), + sympy.Piecewise((1, on_flat), ((y - yc) / r, True))]]) +ramp = R * (phi_top + math.pi / 2) +length = xc + ramp band = sympy.exp(-(d / params.band) ** 2) -n_seg = int(params.n_segments) -strengths = [uw.expression(rf"\eta_{k + 1}", params.initial_strength, - f"weak-plane viscosity of segment {k + 1}") +# Segments along the fault: the flat, then the ramp in three equal parts. +edges = [0.0, xc, xc + ramp / 3, xc + 2 * ramp / 3, length] +names = ["flat", "lower ramp", "upper ramp", "near surface"] +n_seg = len(names) +strengths = [uw.expression(rf"\eta_{{{k + 1}}}", params.initial_strength, + f"weak-plane viscosity, {names[k]}") for k in range(n_seg)] def segment(k): """A smooth indicator for segment k along the fault, in [0, 1].""" edge = params.band - lo, hi = k * length / n_seg, (k + 1) * length / n_seg - on = 1 if k == 0 else (1 + sympy.tanh((s - lo) / edge)) / 2 - off = 1 if k == n_seg - 1 else (1 - sympy.tanh((s - hi) / edge)) / 2 + on = 1 if k == 0 else (1 + sympy.tanh((s - edges[k]) / edge)) / 2 + off = 1 if k == n_seg - 1 else (1 - sympy.tanh((s - edges[k + 1]) / edge)) / 2 return on * off eta_0 = 1 @@ -86,9 +102,7 @@ def segment(k): # --- the observations ------------------------------------------------------ w_top = sympy.exp(-((1 - y) / params.band) ** 2) -rng = np.random.default_rng(7) -points = [(0.3 + 1.4 * i / max(int(params.n_stress_points) - 1, 1), 0.3 + 0.3 * (i % 2)) - for i in range(int(params.n_stress_points))] +points = [(0.3, 0.65), (0.75, 0.12), (1.25, 0.3), (0.9, 0.6), (1.55, 0.85)] w_points = sum(sympy.exp(-((x - px) ** 2 + (y - py) ** 2) / (2 * params.band) ** 2) for px, py in points) @@ -139,7 +153,7 @@ def J_and_gradient(): stokes.solve(zero_init_guess=True) fd.append(float(uw.maths.Integral(mesh, misfit).evaluate())) fd = (fd[0] - fd[1]) / (2 * h) - uw.pprint(f"segment {k + 1}: adjoint {g0[k]: .6e} finite difference {fd: .6e} " + uw.pprint(f"{names[k]:>13}: adjoint {g0[k]: .6e} finite difference {fd: .6e} " f"ratio {fd / g0[k]:.5f}") set_strengths([params.initial_strength] * n_seg) diff --git a/docs/examples/adjoint/fault_segments/plot_fault_segments.py b/docs/examples/adjoint/fault_segments/plot_fault_segments.py index a379293fc..f8583629d 100644 --- a/docs/examples/adjoint/fault_segments/plot_fault_segments.py +++ b/docs/examples/adjoint/fault_segments/plot_fault_segments.py @@ -6,12 +6,13 @@ d = np.load("fault_segments_data.npz") history, true = d["history"], d["true"] +names = ["flat", "lower ramp", "upper ramp", "near surface"] n_seg = len(true) fig, axes = plt.subplots(1, 3, figsize=(11, 3.2), gridspec_kw={"width_ratios": [2.2, 1.6, 1.4]}) ax = axes[0] -ax.contourf(d["gx"], d["gy"], np.log10(d["eta_1"]), levels=np.linspace(-2.2, 0, 12), cmap="viridis") +ax.contourf(d["gx"], d["gy"], np.log10(d["eta_1"]), levels=np.linspace(-2.4, 0, 13), cmap="viridis") ax.plot(d["points"][:, 0], d["points"][:, 1], "wx", ms=7, mew=1.5) ax.set_aspect("equal") ax.set_xlim(0, 2); ax.set_ylim(0, 1) @@ -28,7 +29,7 @@ ax = axes[2] its = np.arange(len(history)) for k in range(n_seg): - ax.semilogy(its, history[:, 1 + k], "o-", ms=3, lw=1, color=f"C{k}", label=f"segment {k + 1}") + ax.semilogy(its, history[:, 1 + k], "o-", ms=3, lw=1, color=f"C{k}", label=names[k]) ax.axhline(true[k], color=f"C{k}", lw=0.8, ls="--") ax.set_xlabel("misfit evaluation"); ax.set_ylabel("weak-plane viscosity") ax.set_title("strengths (dashed: true)", fontsize=10) From 5058b575cf2716894bf6e62f9cbd0d0aca6d72e0 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Wed, 16 Sep 2026 11:24:25 +1000 Subject: [PATCH 04/23] Friction on the listric fault: a Coulomb yield on the plane, one coefficient per segment fault_friction.py: the plane yields at tau_y = C + mu p under a gravity load, with the plane's viscosity the harmonic combination of the bulk viscosity and tau_y over the resolved shear strain rate, so the residual is nonlinear in the velocity and the pressure. The adjoint is the transpose of the consistent tangent, and matches central differences to 1.00000 on all four segments; L-BFGS recovers (0.0500, 0.1500, 0.2500, 0.4000) from a uniform 0.2 in twenty-one evaluations, the flat last. The plot script takes the data file as an argument. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- .../adjoint/fault_segments/fault_friction.py | 214 ++++++++++++++++++ .../adjoint/fault_segments/fault_segments.py | 1 + .../fault_segments/plot_fault_segments.py | 22 +- 3 files changed, 230 insertions(+), 7 deletions(-) create mode 100644 docs/examples/adjoint/fault_segments/fault_friction.py diff --git a/docs/examples/adjoint/fault_segments/fault_friction.py b/docs/examples/adjoint/fault_segments/fault_friction.py new file mode 100644 index 000000000..21dc80d84 --- /dev/null +++ b/docs/examples/adjoint/fault_segments/fault_friction.py @@ -0,0 +1,214 @@ +"""Friction on a listric fault, segment by segment, from surface uplift and stress. + +A listric fault under horizontal shortening and gravity: a ramp that steepens +from a flat decollement at depth to a dip of sixty degrees at the surface, +represented as a weak plane in a transversely isotropic viscosity (no cut in +the mesh). The plane yields at the Coulomb stress tau_y = C + mu p, where p +is the pressure, and the friction coefficient mu is different on the flat, +on the lower and upper parts of the ramp, and near the surface. Those four +coefficients are the unknowns. The observations are the uplift rate along +the top surface and the shear stress in the bulk near a handful of points, +taken from a run at the true coefficients. + +The residual is nonlinear in the velocity and the pressure, so the forward +solve is Newton, and the adjoint is a transpose of the consistent tangent +the solver assembled. The gradient with respect to each coefficient is the +symbolic derivative of that residual. Nothing here is differenced. + +Run it: + + python fault_friction.py # twin experiment: gradient check, then the inversion + python fault_friction.py -uw_check_only 1 +""" +import math + +import numpy as np +import sympy + +import underworld3 as uw +from underworld3.adjoint import misfit_duals, inner + +params = uw.Params( + cell_size=uw.Param(1 / 24, "mesh cell size (box is 2 x 1)"), + surface_dip=uw.Param(60.0, "dip of the ramp where it reaches the surface, degrees"), + flat_depth=uw.Param(0.3, "height of the decollement above the base"), + surface_x=uw.Param(1.9, "where the fault reaches the surface"), + band=uw.Param(0.08, "half-width of the weak band, in box units"), + true_strengths=uw.Param("0.05,0.15,0.25,0.4", + "friction coefficient: flat, lower ramp, upper ramp, near surface"), + initial_strength=uw.Param(0.2, "starting guess, every segment"), + cohesion=uw.Param(0.05, "cohesion C in tau_y = C + mu p"), + rho_g=uw.Param(10.0, "body force, so the pressure grows with depth"), + check_only=uw.Param(0, "1: gradient check against finite differences, no inversion"), +) + +# --- the model --------------------------------------------------------------- +mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0.0, 0.0), maxCoords=(2.0, 1.0), + cellSize=params.cell_size, qdegree=3) +x, y = mesh.X + +v = uw.discretisation.MeshVariable("v", mesh, mesh.dim, degree=2) +p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1, continuous=True) +v_obs = uw.discretisation.MeshVariable("v_obs", mesh, mesh.dim, degree=2) + +# The fault: a flat decollement at y = flat_depth running from the left wall +# to x = xc, then a circular ramp of radius R about (xc, yc) that leaves the +# flat horizontally and reaches the surface at the given dip. The signed +# distance to it and the position along it are exact on each piece, and the +# director is the normal at the nearest point: vertical on the flat, radial on +# the ramp. +phi_top = math.radians(params.surface_dip) - math.pi / 2 # angle of the surface point about the centre +R = (1.0 - params.flat_depth) / (1.0 + math.sin(phi_top)) +yc = params.flat_depth + R +xc = params.surface_x - R * math.cos(phi_top) +r = sympy.sqrt((x - xc) ** 2 + (y - yc) ** 2) +phi = sympy.atan2(y - yc, x - xc) +on_flat = x < xc +d = sympy.Piecewise((y - params.flat_depth, on_flat), (R - r, True)) +s = sympy.Piecewise((x, on_flat), (xc + R * (phi + sympy.pi / 2), True)) +n_hat = sympy.Matrix([[sympy.Piecewise((0, on_flat), ((x - xc) / r, True)), + sympy.Piecewise((1, on_flat), ((y - yc) / r, True))]]) +ramp = R * (phi_top + math.pi / 2) +length = xc + ramp +band = sympy.exp(-(d / params.band) ** 2) + +# Segments along the fault: the flat, then the ramp in three equal parts. +edges = [0.0, xc, xc + ramp / 3, xc + 2 * ramp / 3, length] +names = ["flat", "lower ramp", "upper ramp", "near surface"] +n_seg = len(names) +strengths = [uw.expression(rf"\mu_{{{k + 1}}}", params.initial_strength, + f"friction coefficient, {names[k]}") + for k in range(n_seg)] + +def segment(k): + """A smooth indicator for segment k along the fault, in [0, 1].""" + edge = params.band + on = 1 if k == 0 else (1 + sympy.tanh((s - edges[k]) / edge)) / 2 + off = 1 if k == n_seg - 1 else (1 - sympy.tanh((s - edges[k + 1]) / edge)) / 2 + return on * off + +eta_0 = 1 + +# Coulomb yield on the plane. The shear strain rate resolved on the plane is +# t.E.n; the plane's viscosity is the harmonic combination of the bulk +# viscosity and the yield stress over that rate, which is smooth everywhere +# and tends to tau_y / (2 e_s) where the plane slips. +E = mesh.vector.strain_tensor(v.sym) +t_hat = sympy.Matrix([[-n_hat[1], n_hat[0]]]) +e_s = sympy.sqrt((t_hat * E * n_hat.T)[0] ** 2 + uw.maths.functions.vanishing) +friction = sum(strengths[k] * segment(k) for k in range(n_seg)) +tau_y = params.cohesion + friction * p.sym[0] +eta_plane = eta_0 * tau_y / (tau_y + 2 * eta_0 * e_s) +eta_1 = eta_0 - band * (eta_0 - eta_plane) + +stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) +stokes.constitutive_model = uw.constitutive_models.TransverseIsotropicFlowModel +stokes.constitutive_model.Parameters.shear_viscosity_0 = eta_0 +stokes.constitutive_model.Parameters.shear_viscosity_1 = eta_1 +stokes.constitutive_model.Parameters.director = n_hat +stokes.tolerance = 1e-8 +stokes.bodyforce = sympy.Matrix([0, -params.rho_g]) + +# Shortening from both sides, a no-slip base, and a free top: the surface +# velocity is the uplift rate. +stokes.add_essential_bc((0.0, 0.0), "Bottom") +stokes.add_essential_bc((0.5, None), "Left") +stokes.add_essential_bc((-0.5, None), "Right") + +# --- the observations ------------------------------------------------------ +w_top = sympy.exp(-((1 - y) / params.band) ** 2) +points = [(0.3, 0.65), (0.75, 0.12), (1.25, 0.3), (0.9, 0.6), (1.55, 0.85)] +w_points = sum(sympy.exp(-((x - px) ** 2 + (y - py) ** 2) / (2 * params.band) ** 2) + for px, py in points) + +def shear_stress(field): + e = mesh.vector.strain_tensor(field.sym) + return 2 * eta_0 * e[0, 1] + +misfit = (w_top * (v.sym[1] - v_obs.sym[1]) ** 2 + + w_points * (shear_stress(v) - shear_stress(v_obs)) ** 2) / 2 + +def set_strengths(values): + for expr, value in zip(strengths, values): + expr.sym = float(value) + +def J_and_gradient(): + """The misfit and dJ/d(log strength) for each segment, by the adjoint.""" + stokes.solve(zero_init_guess=True) + J = float(uw.maths.Integral(mesh, misfit).evaluate()) + dJ_dv = misfit_duals(misfit, [v])[v] + mu = uw.discretisation.MeshVariable(f"mu_{uw.adjoint._counter()}", mesh, mesh.dim, degree=2) + lam = uw.discretisation.MeshVariable(f"lam_{uw.adjoint._counter()}", mesh, 1, degree=1) + dJ_dv.array[...] = -np.asarray(dJ_dv.array) + _, reason = stokes.adjoint_solve((dJ_dv, None), target=(mu, lam)) + assert reason > 0, reason + grad = np.array([stokes.sensitivity(mu, expr) * float(expr.sym) for expr in strengths]) + return J, grad + +# --- the truth, and the twin ------------------------------------------------- +true_values = [float(t) for t in str(params.true_strengths).split(",")][:n_seg] +set_strengths(true_values) +stokes.solve(zero_init_guess=True) +v_obs.array[...] = np.asarray(v.array) +uw.pprint(f"true strengths {true_values}") + +set_strengths([params.initial_strength] * n_seg) +J0, g0 = J_and_gradient() +uw.pprint(f"initial J = {J0:.6e} dJ/dlog eta = {g0}") + +# Gradient check: central differences in each log-strength. +h = 1e-3 +for k in range(n_seg): + base = math.log(params.initial_strength) + fd = [] + for sign in (+1, -1): + vals = [params.initial_strength] * n_seg + vals[k] = math.exp(base + sign * h) + set_strengths(vals) + stokes.solve(zero_init_guess=True) + fd.append(float(uw.maths.Integral(mesh, misfit).evaluate())) + fd = (fd[0] - fd[1]) / (2 * h) + uw.pprint(f"{names[k]:>13}: adjoint {g0[k]: .6e} finite difference {fd: .6e} " + f"ratio {fd / g0[k]:.5f}") +set_strengths([params.initial_strength] * n_seg) + +if int(params.check_only): + raise SystemExit + +# --- the inversion --------------------------------------------------------------- +from scipy.optimize import minimize + +history = [] + +def objective(log_eta): + set_strengths(np.exp(log_eta)) + J, grad = J_and_gradient() + history.append((J, np.exp(log_eta).copy())) + uw.pprint(f" J = {J:.6e} strengths = {np.exp(log_eta)}") + return J, grad + +result = minimize(objective, np.log([params.initial_strength] * n_seg), jac=True, + method="L-BFGS-B", options={"maxiter": 40, "gtol": 1e-10}) +uw.pprint(f"recovered {np.exp(result.x)} true {true_values} " + f"after {len(history)} evaluations") + +# --- what the figure needs ----------------------------------------------------- +# Uplift-rate profiles along the top at the truth, the start and the answer, the +# weak-plane viscosity on a grid, and the path the strengths took. +xs = np.linspace(0.0, 2.0, 161) +top = np.column_stack([xs, np.full_like(xs, 1.0 - 1e-6)]) +profiles = {} +for label, values in (("true", true_values), ("initial", [params.initial_strength] * n_seg), + ("recovered", list(np.exp(result.x)))): + set_strengths(values) + stokes.solve(zero_init_guess=True) + profiles[label] = np.asarray(uw.function.evaluate(v.sym[1], top)).ravel() +set_strengths(true_values) +stokes.solve(zero_init_guess=True) # the field on the grid is the truth's +gx, gy = np.meshgrid(np.linspace(0, 2, 201), np.linspace(0, 1, 101)) +grid = np.column_stack([gx.ravel(), gy.ravel()]) +eta_1_grid = np.asarray(uw.function.evaluate(eta_1, grid)).reshape(gx.shape) +np.savez("fault_friction_data.npz", xs=xs, gx=gx, gy=gy, eta_1=eta_1_grid, + points=np.array(points), true=np.array(true_values), + history=np.array([[J, *vals] for J, vals in history]), + **{f"uplift_{k}": val for k, val in profiles.items()}) diff --git a/docs/examples/adjoint/fault_segments/fault_segments.py b/docs/examples/adjoint/fault_segments/fault_segments.py index 5f92584f5..38e1e21fb 100644 --- a/docs/examples/adjoint/fault_segments/fault_segments.py +++ b/docs/examples/adjoint/fault_segments/fault_segments.py @@ -189,6 +189,7 @@ def objective(log_eta): stokes.solve(zero_init_guess=True) profiles[label] = np.asarray(uw.function.evaluate(v.sym[1], top)).ravel() set_strengths(true_values) +stokes.solve(zero_init_guess=True) # the field on the grid is the truth's gx, gy = np.meshgrid(np.linspace(0, 2, 201), np.linspace(0, 1, 101)) grid = np.column_stack([gx.ravel(), gy.ravel()]) eta_1_grid = np.asarray(uw.function.evaluate(eta_1, grid)).reshape(gx.shape) diff --git a/docs/examples/adjoint/fault_segments/plot_fault_segments.py b/docs/examples/adjoint/fault_segments/plot_fault_segments.py index f8583629d..b8cc2ba5d 100644 --- a/docs/examples/adjoint/fault_segments/plot_fault_segments.py +++ b/docs/examples/adjoint/fault_segments/plot_fault_segments.py @@ -1,10 +1,17 @@ -"""The figure for the fault-segments example, from fault_segments_data.npz.""" +"""The figure for the fault-segments examples. + + python plot_fault_segments.py fault_segments_data.npz # weak-plane viscosity + python plot_fault_segments.py fault_friction_data.npz # friction coefficient +""" +import sys import numpy as np import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt -d = np.load("fault_segments_data.npz") +source = sys.argv[1] if len(sys.argv) > 1 else "fault_segments_data.npz" +friction = "friction" in source +d = np.load(source) history, true = d["history"], d["true"] names = ["flat", "lower ramp", "upper ramp", "near surface"] n_seg = len(true) @@ -31,11 +38,12 @@ for k in range(n_seg): ax.semilogy(its, history[:, 1 + k], "o-", ms=3, lw=1, color=f"C{k}", label=names[k]) ax.axhline(true[k], color=f"C{k}", lw=0.8, ls="--") -ax.set_xlabel("misfit evaluation"); ax.set_ylabel("weak-plane viscosity") -ax.set_title("strengths (dashed: true)", fontsize=10) +ax.set_xlabel("misfit evaluation"); ax.set_ylabel("friction coefficient" if friction else "weak-plane viscosity") +ax.set_title(("friction" if friction else "strengths") + " (dashed: true)", fontsize=10) ax.legend(fontsize=8, frameon=False) fig.tight_layout() -fig.savefig("fault_segments.png", dpi=180) -fig.savefig("fault_segments.pdf") -print("wrote fault_segments.png / .pdf") +stem = source.replace("_data.npz", "") +fig.savefig(stem + ".png", dpi=180) +fig.savefig(stem + ".pdf") +print(f"wrote {stem}.png / .pdf") From b0be142a5251955e8029216a0b29dd5c235a9035 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Wed, 16 Sep 2026 12:10:19 +1000 Subject: [PATCH 05/23] fault_friction: keep the cohesion where the pressure is tensile; a flat colour for the bulk in the figure Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- docs/examples/adjoint/fault_segments/fault_friction.py | 4 +++- .../examples/adjoint/fault_segments/plot_fault_segments.py | 7 +++++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/docs/examples/adjoint/fault_segments/fault_friction.py b/docs/examples/adjoint/fault_segments/fault_friction.py index 21dc80d84..cda53b032 100644 --- a/docs/examples/adjoint/fault_segments/fault_friction.py +++ b/docs/examples/adjoint/fault_segments/fault_friction.py @@ -97,7 +97,9 @@ def segment(k): t_hat = sympy.Matrix([[-n_hat[1], n_hat[0]]]) e_s = sympy.sqrt((t_hat * E * n_hat.T)[0] ** 2 + uw.maths.functions.vanishing) friction = sum(strengths[k] * segment(k) for k in range(n_seg)) -tau_y = params.cohesion + friction * p.sym[0] +# Compression is positive; where the dynamic pressure is tensile the plane +# keeps its cohesion and no more. +tau_y = params.cohesion + friction * sympy.Max(p.sym[0], 0) eta_plane = eta_0 * tau_y / (tau_y + 2 * eta_0 * e_s) eta_1 = eta_0 - band * (eta_0 - eta_plane) diff --git a/docs/examples/adjoint/fault_segments/plot_fault_segments.py b/docs/examples/adjoint/fault_segments/plot_fault_segments.py index b8cc2ba5d..e1e620bd8 100644 --- a/docs/examples/adjoint/fault_segments/plot_fault_segments.py +++ b/docs/examples/adjoint/fault_segments/plot_fault_segments.py @@ -19,12 +19,15 @@ fig, axes = plt.subplots(1, 3, figsize=(11, 3.2), gridspec_kw={"width_ratios": [2.2, 1.6, 1.4]}) ax = axes[0] -ax.contourf(d["gx"], d["gy"], np.log10(d["eta_1"]), levels=np.linspace(-2.4, 0, 13), cmap="viridis") +# The bulk sits at log10(eta_1) = 0 to round-off, so the top level is set a +# little below it and "extend" gives the bulk one flat colour. +ax.contourf(d["gx"], d["gy"], np.log10(d["eta_1"]), levels=np.linspace(-2.4, -0.1, 12), + cmap="viridis", extend="both") ax.plot(d["points"][:, 0], d["points"][:, 1], "wx", ms=7, mew=1.5) ax.set_aspect("equal") ax.set_xlim(0, 2); ax.set_ylim(0, 1) ax.set_xlabel("$x$"); ax.set_ylabel("$y$") -ax.set_title(r"$\log_{10}\eta_1$ at the true strengths; $\times$ stress points", fontsize=10) +ax.set_title(r"$\log_{10}\eta_1$ at the truth; $\times$ stress points", fontsize=10) ax = axes[1] for label, style in (("true", "k-"), ("initial", "C3--"), ("recovered", "C0:")): From 1f86d01ee332ac2ab86dbf75706ed98a95ce8610 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Wed, 16 Sep 2026 14:37:08 +1000 Subject: [PATCH 06/23] =?UTF-8?q?fault=5Ffriction:=20an=20observations=20s?= =?UTF-8?q?witch=20=E2=80=94=20uplift+stress,=20principal-stress=20orienta?= =?UTF-8?q?tion,=20or=20orientation=20along=20the=20surface=20only?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Orientation is a unit vector of the deviatoric strain rate, so there is no angle to wrap. Noise-free, all three recover the four coefficients to four figures; what changes is the sensitivity to the flat, 1e-3 of the near-surface segment's with orientation at the points and the surface, 1e-4 with the surface alone. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- .../adjoint/fault_segments/fault_friction.py | 25 ++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/docs/examples/adjoint/fault_segments/fault_friction.py b/docs/examples/adjoint/fault_segments/fault_friction.py index cda53b032..b55358d84 100644 --- a/docs/examples/adjoint/fault_segments/fault_friction.py +++ b/docs/examples/adjoint/fault_segments/fault_friction.py @@ -40,6 +40,10 @@ cohesion=uw.Param(0.05, "cohesion C in tau_y = C + mu p"), rho_g=uw.Param(10.0, "body force, so the pressure grows with depth"), check_only=uw.Param(0, "1: gradient check against finite differences, no inversion"), + observations=uw.Param("uplift+stress", + "uplift+stress | orientation (principal-stress orientation at the " + "points and along the surface) | orientation_surface (along the " + "surface only)"), ) # --- the model --------------------------------------------------------------- @@ -127,8 +131,23 @@ def shear_stress(field): e = mesh.vector.strain_tensor(field.sym) return 2 * eta_0 * e[0, 1] -misfit = (w_top * (v.sym[1] - v_obs.sym[1]) ** 2 - + w_points * (shear_stress(v) - shear_stress(v_obs)) ** 2) / 2 +def orientation(field): + """The principal-stress orientation as the unit vector (cos 2theta, sin 2theta) + of the deviatoric strain rate, which is the stress orientation in the + isotropic bulk. A unit vector rather than an angle, so there is no wrap.""" + e = mesh.vector.strain_tensor(field.sym) + a, b = e[0, 0] - e[1, 1], 2 * e[0, 1] + norm = sympy.sqrt(a ** 2 + b ** 2 + uw.maths.functions.vanishing) + return sympy.Matrix([[a / norm, b / norm]]) + +what = str(params.observations) +if what == "uplift+stress": + misfit = (w_top * (v.sym[1] - v_obs.sym[1]) ** 2 + + w_points * (shear_stress(v) - shear_stress(v_obs)) ** 2) / 2 +else: + dq = orientation(v) - orientation(v_obs) + weight = w_top if what == "orientation_surface" else w_top + w_points + misfit = weight * (dq[0] ** 2 + dq[1] ** 2) / 2 def set_strengths(values): for expr, value in zip(strengths, values): @@ -210,7 +229,7 @@ def objective(log_eta): gx, gy = np.meshgrid(np.linspace(0, 2, 201), np.linspace(0, 1, 101)) grid = np.column_stack([gx.ravel(), gy.ravel()]) eta_1_grid = np.asarray(uw.function.evaluate(eta_1, grid)).reshape(gx.shape) -np.savez("fault_friction_data.npz", xs=xs, gx=gx, gy=gy, eta_1=eta_1_grid, +np.savez(f"fault_friction_{what}_data.npz", xs=xs, gx=gx, gy=gy, eta_1=eta_1_grid, points=np.array(points), true=np.array(true_values), history=np.array([[J, *vals] for J, vals in history]), **{f"uplift_{k}": val for k, val in profiles.items()}) From bcac5a88f762d2ab8349068d2cd063d0442d8bfa Mon Sep 17 00:00:00 2001 From: lmoresi Date: Wed, 16 Sep 2026 17:27:53 +1000 Subject: [PATCH 07/23] fault_friction: convergence of the three observation sets side by side Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- .../fault_segments/plot_convergence.py | 36 +++++++++++++++++++ 1 file changed, 36 insertions(+) create mode 100644 docs/examples/adjoint/fault_segments/plot_convergence.py diff --git a/docs/examples/adjoint/fault_segments/plot_convergence.py b/docs/examples/adjoint/fault_segments/plot_convergence.py new file mode 100644 index 000000000..79999c080 --- /dev/null +++ b/docs/examples/adjoint/fault_segments/plot_convergence.py @@ -0,0 +1,36 @@ +"""Convergence of the friction inversion under the three observation sets.""" +import numpy as np +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +cases = [("uplift + stress", "fault_friction_data.npz"), + ("orientation, points + surface", "fault_friction_orientation_data.npz"), + ("orientation, surface only", "fault_friction_orientation_surface_data.npz")] +names = ["flat", "lower ramp", "upper ramp", "near surface"] + +fig, axes = plt.subplots(2, 3, figsize=(11, 5.6), sharex="col", + gridspec_kw={"height_ratios": [2.2, 1]}) +for col, (title, path) in enumerate(cases): + d = np.load(path) + history, true = d["history"], d["true"] + its = np.arange(len(history)) + ax = axes[0, col] + for k in range(len(true)): + ax.semilogy(its, history[:, 1 + k], "o-", ms=3, lw=1, color=f"C{k}", label=names[k]) + ax.axhline(true[k], color=f"C{k}", lw=0.8, ls="--") + ax.set_title(title, fontsize=10) + ax.set_ylim(0.03, 0.6) + if col == 0: + ax.set_ylabel("friction coefficient") + ax.legend(fontsize=8, frameon=False, loc="lower left") + ax = axes[1, col] + ax.semilogy(its, history[:, 0] / history[0, 0], "k.-", ms=4, lw=1) + ax.set_xlabel("misfit evaluation") + if col == 0: + ax.set_ylabel("$J / J_0$") + ax.set_ylim(1e-11, 2) +fig.tight_layout() +fig.savefig("fault_friction_convergence.png", dpi=180) +fig.savefig("fault_friction_convergence.pdf") +print("wrote fault_friction_convergence.png / .pdf") From f7d9ccff284b1c225545adf665541c1b54149698 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 17 Sep 2026 10:37:45 +1000 Subject: [PATCH 08/23] fault_friction: PyVista renders of the true state on the mesh's own triangulation Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- .../fault_segments/render_fault_friction.py | 109 ++++++++++++++++++ 1 file changed, 109 insertions(+) create mode 100644 docs/examples/adjoint/fault_segments/render_fault_friction.py diff --git a/docs/examples/adjoint/fault_segments/render_fault_friction.py b/docs/examples/adjoint/fault_segments/render_fault_friction.py new file mode 100644 index 000000000..3b2fdbb56 --- /dev/null +++ b/docs/examples/adjoint/fault_segments/render_fault_friction.py @@ -0,0 +1,109 @@ +"""PyVista renders of the friction example at the true coefficients. + +Writes log10 eta_1 (the plane's viscosity), the velocity, and the pressure +to ~/+Simulations/adjoint_fault_example/, on the mesh's own triangulation. +""" +import math +import os + +import numpy as np +import sympy +import pyvista as pv + +import underworld3 as uw +import underworld3.visualisation as vis + +pv.OFF_SCREEN = True +OUT = os.path.expanduser("~/+Simulations/adjoint_fault_example") + +cell_size, surface_dip, flat_depth, surface_x, band_w = 1 / 24, 60.0, 0.3, 1.9, 0.08 +true_mu, cohesion, rho_g = [0.05, 0.15, 0.25, 0.4], 0.05, 10.0 + +mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0.0, 0.0), maxCoords=(2.0, 1.0), + cellSize=cell_size, qdegree=3) +x, y = mesh.X +v = uw.discretisation.MeshVariable("v", mesh, 2, degree=2) +p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1, continuous=True) + +phi_top = math.radians(surface_dip) - math.pi / 2 +R = (1.0 - flat_depth) / (1.0 + math.sin(phi_top)) +yc = flat_depth + R +xc = surface_x - R * math.cos(phi_top) +r = sympy.sqrt((x - xc) ** 2 + (y - yc) ** 2) +phi = sympy.atan2(y - yc, x - xc) +on_flat = x < xc +d = sympy.Piecewise((y - flat_depth, on_flat), (R - r, True)) +s = sympy.Piecewise((x, on_flat), (xc + R * (phi + sympy.pi / 2), True)) +n_hat = sympy.Matrix([[sympy.Piecewise((0, on_flat), ((x - xc) / r, True)), + sympy.Piecewise((1, on_flat), ((y - yc) / r, True))]]) +ramp = R * (phi_top + math.pi / 2) +edges_s = [0.0, xc, xc + ramp / 3, xc + 2 * ramp / 3, xc + ramp] +band = sympy.exp(-(d / band_w) ** 2) + +def segment(k): + on = 1 if k == 0 else (1 + sympy.tanh((s - edges_s[k]) / band_w)) / 2 + off = 1 if k == 3 else (1 - sympy.tanh((s - edges_s[k + 1]) / band_w)) / 2 + return on * off + +E = mesh.vector.strain_tensor(v.sym) +t_hat = sympy.Matrix([[-n_hat[1], n_hat[0]]]) +e_s = sympy.sqrt((t_hat * E * n_hat.T)[0] ** 2 + uw.maths.functions.vanishing) +friction = sum(true_mu[k] * segment(k) for k in range(4)) +tau_y = cohesion + friction * sympy.Max(p.sym[0], 0) +eta_plane = tau_y / (tau_y + 2 * e_s) +eta_1 = 1 - band * (1 - eta_plane) + +stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) +stokes.constitutive_model = uw.constitutive_models.TransverseIsotropicFlowModel +stokes.constitutive_model.Parameters.shear_viscosity_0 = 1 +stokes.constitutive_model.Parameters.shear_viscosity_1 = eta_1 +stokes.constitutive_model.Parameters.director = n_hat +stokes.tolerance = 1e-8 +stokes.bodyforce = sympy.Matrix([0, -rho_g]) +stokes.add_essential_bc((0.0, 0.0), "Bottom") +stokes.add_essential_bc((0.5, None), "Left") +stokes.add_essential_bc((-0.5, None), "Right") +stokes.solve(zero_init_guess=True) + +# Nodal values on the mesh's own triangulation (skill rule 5): no grid resampling. +pv_mesh = vis.mesh_to_pv_mesh(mesh) +pts = np.asarray(pv_mesh.points[:, :2]) +edges = pv_mesh.extract_all_edges() +pv_mesh.point_data["log10 eta_1"] = np.log10(np.clip(np.asarray(uw.function.evaluate(eta_1, pts)).ravel(), 1e-3, 1.0)) +pv_mesh.point_data["slip rate"] = np.asarray(uw.function.evaluate(band * e_s, pts)).ravel() +pv_mesh.point_data["p"] = np.asarray(uw.function.evaluate(p.sym[0], pts)).ravel() +pv_v = vis.meshVariable_to_pv_mesh_object(v) +vdata = np.asarray(v.data) +pv_v.point_data["|v|"] = np.linalg.norm(vdata, axis=1) +pv_v.point_data["v"] = np.column_stack([vdata, np.zeros(len(vdata))]) +pv_v.point_data["v_y"] = vdata[:, 1] + +def frame(name, obj, scalars, cmap, clim, title, arrows=False): + pl = pv.Plotter(off_screen=True, window_size=(1600, 820)) + pl.set_background("white") + pl.add_mesh(obj, scalars=scalars, cmap=cmap, clim=clim, show_edges=False, lighting=False, + scalar_bar_args=dict(title=title, color="black", vertical=False, position_x=0.3, + position_y=0.04, width=0.4, height=0.06, title_font_size=26, + label_font_size=22)) + pl.add_mesh(edges, color="black", line_width=0.25, lighting=False, opacity=0.2) + if arrows: + # a coarse regular set of points, so the arrows read as a field and not as the mesh + gx, gy = np.meshgrid(np.linspace(0.05, 1.95, 30), np.linspace(0.05, 0.95, 15)) + seeds = np.column_stack([gx.ravel(), gy.ravel()]) + vals = np.asarray(uw.function.evaluate(v.sym, seeds)).reshape(-1, 2) + cloud = pv.PolyData(np.column_stack([seeds, np.zeros(len(seeds))])) + cloud["v"] = np.column_stack([vals, np.zeros(len(vals))]) + cloud["|v|"] = np.linalg.norm(vals, axis=1) + pl.add_mesh(cloud.glyph(orient="v", scale="|v|", factor=0.09), color="black", lighting=False) + pl.view_xy() + pl.camera.parallel_projection = True + pl.camera.focal_point = (1.0, 0.42, 0.0) + pl.camera.parallel_scale = 0.66 + pl.screenshot(os.path.join(OUT, name)) + pl.close() + +frame("fault_eta1.png", pv_mesh, "log10 eta_1", "viridis", (-2.5, 0.0), "log10 plane viscosity") +frame("fault_slip.png", pv_mesh, "slip rate", "magma_r", (0.0, 1.0), "shear strain rate on the plane") +frame("fault_vy.png", pv_v, "v_y", "viridis", (0.0, 0.7), "vertical velocity (uplift rate)", arrows=True) +frame("fault_p.png", pv_mesh, "p", "RdBu_r", (-8.0, 8.0), "pressure") +print("wrote", OUT) From 077af00b966300ae1bf8b040ae13ba440f45361a Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 17 Sep 2026 14:27:56 +1000 Subject: [PATCH 09/23] The fault example records itself; adjoint solves are operators in every transcript view MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An adjoint solve was already recorded as an event; the log, the chart, the figure and the flowchart now show it as "adjoint Stokes(v)" in its own column. The fault example wraps each misfit evaluation in a zero-length step, so its transcript reads "Stokes(v) > adjoint Stokes(v)" per evaluation with the twin's truth run and the finite-difference solves labelled. The mesh is refined once from a coarser base so the velocity block has a multigrid hierarchy. Without one it fell back to gamg and hit its iteration cap on every solve, and the inexact Newton step converged linearly at a fixed rate — forty-odd iterations cold, twenty warm, the same with either tangent. With the hierarchy nothing is capped and Newton is quadratic: six iterations cold, two warm. Evaluations warm-start. The key prints a vector boundary condition as a tuple, with "free" for an unconstrained component, and says of a quantity whose value is itself an expression that it is in the record rather than printing it raw. Both figures mark the surface observation band as well as the points. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- .../adjoint/fault_segments/fault_friction.py | 55 +++++++++------ .../adjoint/fault_segments/fault_segments.py | 17 +++-- .../fault_segments/plot_fault_segments.py | 5 +- .../fault_segments/render_fault_friction.py | 22 ++++-- src/underworld3/model.py | 4 +- .../utilities/transcript_report.py | 67 +++++++++++++++---- 6 files changed, 126 insertions(+), 44 deletions(-) diff --git a/docs/examples/adjoint/fault_segments/fault_friction.py b/docs/examples/adjoint/fault_segments/fault_friction.py index b55358d84..cf3d28685 100644 --- a/docs/examples/adjoint/fault_segments/fault_friction.py +++ b/docs/examples/adjoint/fault_segments/fault_friction.py @@ -29,7 +29,7 @@ from underworld3.adjoint import misfit_duals, inner params = uw.Params( - cell_size=uw.Param(1 / 24, "mesh cell size (box is 2 x 1)"), + cell_size=uw.Param(1 / 12, "base mesh cell size (box is 2 x 1); refined once, so half this"), surface_dip=uw.Param(60.0, "dip of the ramp where it reaches the surface, degrees"), flat_depth=uw.Param(0.3, "height of the decollement above the base"), surface_x=uw.Param(1.9, "where the fault reaches the surface"), @@ -47,8 +47,11 @@ ) # --- the model --------------------------------------------------------------- +# Refined once from the base size: the refinement gives the velocity block a +# multigrid hierarchy. Without one it falls back to gamg, which hits its +# iteration cap on this problem, and an inexact Newton step converges linearly. mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0.0, 0.0), maxCoords=(2.0, 1.0), - cellSize=params.cell_size, qdegree=3) + cellSize=params.cell_size, qdegree=3, refinement=1) x, y = mesh.X v = uw.discretisation.MeshVariable("v", mesh, mesh.dim, degree=2) @@ -149,32 +152,46 @@ def orientation(field): weight = w_top if what == "orientation_surface" else w_top + w_points misfit = weight * (dq[0] ** 2 + dq[1] ** 2) / 2 +model = uw.get_default_model() + def set_strengths(values): for expr, value in zip(strengths, values): expr.sym = float(value) -def J_and_gradient(): - """The misfit and dJ/d(log strength) for each segment, by the adjoint.""" - stokes.solve(zero_init_guess=True) - J = float(uw.maths.Integral(mesh, misfit).evaluate()) - dJ_dv = misfit_duals(misfit, [v])[v] - mu = uw.discretisation.MeshVariable(f"mu_{uw.adjoint._counter()}", mesh, mesh.dim, degree=2) - lam = uw.discretisation.MeshVariable(f"lam_{uw.adjoint._counter()}", mesh, 1, degree=1) - dJ_dv.array[...] = -np.asarray(dJ_dv.array) - _, reason = stokes.adjoint_solve((dJ_dv, None), target=(mu, lam)) - assert reason > 0, reason - grad = np.array([stokes.sensitivity(mu, expr) * float(expr.sym) for expr in strengths]) +evaluations = [0] + +def J_and_gradient(label=None): + """The misfit and dJ/d(log strength) for each segment, by the adjoint. + + Each evaluation is one step of zero length in the model's record, so the + run's transcript lists the forward solve and the adjoint solve it made. + """ + evaluations[0] += 1 + with model.step(0.0, label=label or f"eval {evaluations[0]}"): + stokes.solve(zero_init_guess=False) + J = float(uw.maths.Integral(mesh, misfit).evaluate()) + dJ_dv = misfit_duals(misfit, [v])[v] + mu = uw.discretisation.MeshVariable(f"mu_{uw.adjoint._counter()}", mesh, mesh.dim, degree=2) + lam = uw.discretisation.MeshVariable(f"lam_{uw.adjoint._counter()}", mesh, 1, degree=1) + dJ_dv.array[...] = -np.asarray(dJ_dv.array) + _, reason = stokes.adjoint_solve((dJ_dv, None), target=(mu, lam)) + assert reason > 0, reason + grad = np.array([stokes.sensitivity(mu, expr) * float(expr.sym) for expr in strengths]) return J, grad +def forward(label): + with model.step(0.0, label=label): + stokes.solve(zero_init_guess=False) + # --- the truth, and the twin ------------------------------------------------- true_values = [float(t) for t in str(params.true_strengths).split(",")][:n_seg] set_strengths(true_values) -stokes.solve(zero_init_guess=True) +forward("truth") v_obs.array[...] = np.asarray(v.array) uw.pprint(f"true strengths {true_values}") set_strengths([params.initial_strength] * n_seg) -J0, g0 = J_and_gradient() +J0, g0 = J_and_gradient("start") uw.pprint(f"initial J = {J0:.6e} dJ/dlog eta = {g0}") # Gradient check: central differences in each log-strength. @@ -186,7 +203,7 @@ def J_and_gradient(): vals = [params.initial_strength] * n_seg vals[k] = math.exp(base + sign * h) set_strengths(vals) - stokes.solve(zero_init_guess=True) + forward(f"fd {names[k]} {'+' if sign > 0 else '-'}h") fd.append(float(uw.maths.Integral(mesh, misfit).evaluate())) fd = (fd[0] - fd[1]) / (2 * h) uw.pprint(f"{names[k]:>13}: adjoint {g0[k]: .6e} finite difference {fd: .6e} " @@ -222,14 +239,14 @@ def objective(log_eta): for label, values in (("true", true_values), ("initial", [params.initial_strength] * n_seg), ("recovered", list(np.exp(result.x)))): set_strengths(values) - stokes.solve(zero_init_guess=True) + forward(f"profile {label}") profiles[label] = np.asarray(uw.function.evaluate(v.sym[1], top)).ravel() set_strengths(true_values) -stokes.solve(zero_init_guess=True) # the field on the grid is the truth's +forward("truth again") # the field on the grid is the truth's gx, gy = np.meshgrid(np.linspace(0, 2, 201), np.linspace(0, 1, 101)) grid = np.column_stack([gx.ravel(), gy.ravel()]) eta_1_grid = np.asarray(uw.function.evaluate(eta_1, grid)).reshape(gx.shape) np.savez(f"fault_friction_{what}_data.npz", xs=xs, gx=gx, gy=gy, eta_1=eta_1_grid, - points=np.array(points), true=np.array(true_values), + points=np.array(points), true=np.array(true_values), band=params.band, history=np.array([[J, *vals] for J, vals in history]), **{f"uplift_{k}": val for k, val in profiles.items()}) diff --git a/docs/examples/adjoint/fault_segments/fault_segments.py b/docs/examples/adjoint/fault_segments/fault_segments.py index 38e1e21fb..0fd574c6d 100644 --- a/docs/examples/adjoint/fault_segments/fault_segments.py +++ b/docs/examples/adjoint/fault_segments/fault_segments.py @@ -28,7 +28,7 @@ from underworld3.adjoint import misfit_duals, inner params = uw.Params( - cell_size=uw.Param(1 / 24, "mesh cell size (box is 2 x 1)"), + cell_size=uw.Param(1 / 12, "base mesh cell size (box is 2 x 1); refined once, so half this"), surface_dip=uw.Param(60.0, "dip of the ramp where it reaches the surface, degrees"), flat_depth=uw.Param(0.3, "height of the decollement above the base"), surface_x=uw.Param(1.9, "where the fault reaches the surface"), @@ -40,8 +40,11 @@ ) # --- the model --------------------------------------------------------------- +# Refined once from the base size: the refinement gives the velocity block a +# multigrid hierarchy. Without one it falls back to gamg, which hits its +# iteration cap on this problem, and an inexact Newton step converges linearly. mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0.0, 0.0), maxCoords=(2.0, 1.0), - cellSize=params.cell_size, qdegree=3) + cellSize=params.cell_size, qdegree=3, refinement=1) x, y = mesh.X v = uw.discretisation.MeshVariable("v", mesh, mesh.dim, degree=2) @@ -119,7 +122,7 @@ def set_strengths(values): def J_and_gradient(): """The misfit and dJ/d(log strength) for each segment, by the adjoint.""" - stokes.solve(zero_init_guess=True) + stokes.solve(zero_init_guess=False) J = float(uw.maths.Integral(mesh, misfit).evaluate()) dJ_dv = misfit_duals(misfit, [v])[v] mu = uw.discretisation.MeshVariable(f"mu_{uw.adjoint._counter()}", mesh, mesh.dim, degree=2) @@ -133,7 +136,7 @@ def J_and_gradient(): # --- the truth, and the twin ------------------------------------------------- true_values = [float(t) for t in str(params.true_strengths).split(",")][:n_seg] set_strengths(true_values) -stokes.solve(zero_init_guess=True) +stokes.solve(zero_init_guess=False) v_obs.array[...] = np.asarray(v.array) uw.pprint(f"true strengths {true_values}") @@ -150,7 +153,7 @@ def J_and_gradient(): vals = [params.initial_strength] * n_seg vals[k] = math.exp(base + sign * h) set_strengths(vals) - stokes.solve(zero_init_guess=True) + stokes.solve(zero_init_guess=False) fd.append(float(uw.maths.Integral(mesh, misfit).evaluate())) fd = (fd[0] - fd[1]) / (2 * h) uw.pprint(f"{names[k]:>13}: adjoint {g0[k]: .6e} finite difference {fd: .6e} " @@ -186,10 +189,10 @@ def objective(log_eta): for label, values in (("true", true_values), ("initial", [params.initial_strength] * n_seg), ("recovered", list(np.exp(result.x)))): set_strengths(values) - stokes.solve(zero_init_guess=True) + stokes.solve(zero_init_guess=False) profiles[label] = np.asarray(uw.function.evaluate(v.sym[1], top)).ravel() set_strengths(true_values) -stokes.solve(zero_init_guess=True) # the field on the grid is the truth's +stokes.solve(zero_init_guess=False) # the field on the grid is the truth's gx, gy = np.meshgrid(np.linspace(0, 2, 201), np.linspace(0, 1, 101)) grid = np.column_stack([gx.ravel(), gy.ravel()]) eta_1_grid = np.asarray(uw.function.evaluate(eta_1, grid)).reshape(gx.shape) diff --git a/docs/examples/adjoint/fault_segments/plot_fault_segments.py b/docs/examples/adjoint/fault_segments/plot_fault_segments.py index e1e620bd8..2067dbacd 100644 --- a/docs/examples/adjoint/fault_segments/plot_fault_segments.py +++ b/docs/examples/adjoint/fault_segments/plot_fault_segments.py @@ -23,11 +23,14 @@ # little below it and "extend" gives the bulk one flat colour. ax.contourf(d["gx"], d["gy"], np.log10(d["eta_1"]), levels=np.linspace(-2.4, -0.1, 12), cmap="viridis", extend="both") +band = float(d["band"]) if "band" in d else 0.08 +# the surface band under the top, where the uplift rate (or the orientation) is read +ax.axhspan(1 - 2 * band, 1.0, color="white", alpha=0.35, lw=0) ax.plot(d["points"][:, 0], d["points"][:, 1], "wx", ms=7, mew=1.5) ax.set_aspect("equal") ax.set_xlim(0, 2); ax.set_ylim(0, 1) ax.set_xlabel("$x$"); ax.set_ylabel("$y$") -ax.set_title(r"$\log_{10}\eta_1$ at the truth; $\times$ stress points", fontsize=10) +ax.set_title(r"$\log_{10}\eta_1$ at the truth; $\times$ points, white band: surface observations", fontsize=10) ax = axes[1] for label, style in (("true", "k-"), ("initial", "C3--"), ("recovered", "C0:")): diff --git a/docs/examples/adjoint/fault_segments/render_fault_friction.py b/docs/examples/adjoint/fault_segments/render_fault_friction.py index 3b2fdbb56..e23e3aecf 100644 --- a/docs/examples/adjoint/fault_segments/render_fault_friction.py +++ b/docs/examples/adjoint/fault_segments/render_fault_friction.py @@ -16,11 +16,11 @@ pv.OFF_SCREEN = True OUT = os.path.expanduser("~/+Simulations/adjoint_fault_example") -cell_size, surface_dip, flat_depth, surface_x, band_w = 1 / 24, 60.0, 0.3, 1.9, 0.08 +cell_size, surface_dip, flat_depth, surface_x, band_w = 1 / 12, 60.0, 0.3, 1.9, 0.08 true_mu, cohesion, rho_g = [0.05, 0.15, 0.25, 0.4], 0.05, 10.0 mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0.0, 0.0), maxCoords=(2.0, 1.0), - cellSize=cell_size, qdegree=3) + cellSize=cell_size, qdegree=3, refinement=1) x, y = mesh.X v = uw.discretisation.MeshVariable("v", mesh, 2, degree=2) p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1, continuous=True) @@ -78,7 +78,19 @@ def segment(k): pv_v.point_data["v"] = np.column_stack([vdata, np.zeros(len(vdata))]) pv_v.point_data["v_y"] = vdata[:, 1] -def frame(name, obj, scalars, cmap, clim, title, arrows=False): +points = [(0.3, 0.65), (0.75, 0.12), (1.25, 0.3), (0.9, 0.6), (1.55, 0.85)] + +def observations(pl): + """Where the observations are read: the five points, and the band under the + surface (two band widths deep, where its Gaussian weight is above 2%).""" + strip = pv.Rectangle([[0.0, 1 - 2 * band_w, 0.001], [2.0, 1 - 2 * band_w, 0.001], [2.0, 1.0, 0.001]]) + pl.add_mesh(strip, color="white", opacity=0.45, lighting=False) + pl.add_mesh(pv.Line([0, 1 - 2 * band_w, 0.002], [2, 1 - 2 * band_w, 0.002]), color="white", line_width=2, lighting=False) + for px, py in points: + pl.add_mesh(pv.Disc(center=(px, py, 0.003), inner=0.0, outer=0.022, normal=(0, 0, 1)), color="white", lighting=False) + pl.add_mesh(pv.Disc(center=(px, py, 0.004), inner=0.0, outer=0.013, normal=(0, 0, 1)), color="black", lighting=False) + +def frame(name, obj, scalars, cmap, clim, title, arrows=False, observed=False): pl = pv.Plotter(off_screen=True, window_size=(1600, 820)) pl.set_background("white") pl.add_mesh(obj, scalars=scalars, cmap=cmap, clim=clim, show_edges=False, lighting=False, @@ -95,6 +107,8 @@ def frame(name, obj, scalars, cmap, clim, title, arrows=False): cloud["v"] = np.column_stack([vals, np.zeros(len(vals))]) cloud["|v|"] = np.linalg.norm(vals, axis=1) pl.add_mesh(cloud.glyph(orient="v", scale="|v|", factor=0.09), color="black", lighting=False) + if observed: + observations(pl) pl.view_xy() pl.camera.parallel_projection = True pl.camera.focal_point = (1.0, 0.42, 0.0) @@ -102,7 +116,7 @@ def frame(name, obj, scalars, cmap, clim, title, arrows=False): pl.screenshot(os.path.join(OUT, name)) pl.close() -frame("fault_eta1.png", pv_mesh, "log10 eta_1", "viridis", (-2.5, 0.0), "log10 plane viscosity") +frame("fault_eta1.png", pv_mesh, "log10 eta_1", "viridis", (-2.5, 0.0), "log10 plane viscosity", observed=True) frame("fault_slip.png", pv_mesh, "slip rate", "magma_r", (0.0, 1.0), "shear strain rate on the plane") frame("fault_vy.png", pv_v, "v_y", "viridis", (0.0, 0.7), "vertical velocity (uplift rate)", arrows=True) frame("fault_p.png", pv_mesh, "p", "RdBu_r", (-8.0, 8.0), "pressure") diff --git a/src/underworld3/model.py b/src/underworld3/model.py index 1098be38f..96c0d2f9b 100644 --- a/src/underworld3/model.py +++ b/src/underworld3/model.py @@ -243,6 +243,8 @@ def _operator_text(event): return name if kind == "history_shift": return f"shift {name}" + if kind == "adjoint_solve": + return f"adjoint {name}" return f"{kind}:{name}" @@ -1350,7 +1352,7 @@ def _render_transcript_text(self, payload): # older transcript may carry is not an operator and is left out. operators = " > ".join( _operator_text(e) for e in events - if e.get("kind") in ("solve", "history_shift") + if e.get("kind") in ("solve", "history_shift", "adjoint_solve") ) or "(nothing)" notes = [] diff --git a/src/underworld3/utilities/transcript_report.py b/src/underworld3/utilities/transcript_report.py index daeb6e730..63d5d565a 100644 --- a/src/underworld3/utilities/transcript_report.py +++ b/src/underworld3/utilities/transcript_report.py @@ -163,6 +163,17 @@ def _magnitude_and_unit(value_text, units): m = re.match(r"^\s*Matrix\(\[\[([-+0-9.eE]+)\]\]\)\s*$", text) if m: return _compact_number(m.group(1)), _plain_unit(units) + # A vector of numbers, as a boundary condition often is: (0, 0), or + # (0.5, free) where a component is left unconstrained. + m = re.match(r"^\s*Matrix\(\[(\[.*\])\]\)\s*$", text) + if m and "(" not in m.group(1): + entries = re.findall(r"\[([^\[\]]*)\]", m.group(1)) + parts = [] + for e in entries: + e = e.strip() + parts.append("free" if e in ("oo", "zoo", "nan") else _compact_number(e)) + if parts and all(re.match(r"^[-+0-9.eE]+$|^free$", q) for q in parts): + return "(" + ", ".join(parts) + ")", _plain_unit(units) if re.match(r"^\s*[-+0-9.eE]+\s*$", text): return _compact_number(text), _plain_unit(units) return text, _plain_unit(units) @@ -185,7 +196,13 @@ def _latex_value(value_latex, value_text, units): if re.match(r"^[-+0-9.eE]+$", number or ""): latex = _number_latex(number) return latex + (rf"\ \mathrm{{{_plain_unit_latex(unit)}}}" if unit else "") - return str(value_latex) if value_latex not in (None, "") else str(value_text) + if number and number.startswith("("): + return number + text = str(value_latex) if value_latex not in (None, "") else str(value_text) + # A quantity whose value is itself an expression of the fields — a + # yield-limited viscosity, say — is not a number to print; the full + # expression is in the record. + return text if len(text) <= 60 else r"\text{(an expression; in the record)}" def _number_latex(number): @@ -595,18 +612,18 @@ def _signature(step): return tuple( (event["kind"], _short_operator(event["name"])) for event in step.get("events", []) - if event.get("kind") in ("solve", "history_shift") + if event.get("kind") in _OPERATOR_KINDS ) def _describe(signature): return " ".join( - name if kind == "solve" else f"shift {name}" for kind, name in signature + _kind_text(kind, name) for kind, name in signature ) or "(nothing)" def _sequence_text(signature): - parts = [name if kind == "solve" else f"shift {name}" + parts = [_kind_text(kind, name) for kind, name in signature] return " > ".join(parts) or "(nothing)" @@ -1532,6 +1549,33 @@ def transcript_flowchart(source, run=-1, out=None): # The transcript as a chart: parts across the page, steps down it # --------------------------------------------------------------------------- +_OPERATOR_KINDS = ("solve", "history_shift", "adjoint_solve") + + +def _kind_text(kind, name): + if kind == "history_shift": + return f"shift {name}" + if kind == "adjoint_solve": + return f"adjoint {name}" + return name + + +def _part_key(event): + """Which column an event belongs to. An adjoint solve is the same solver + as its forward solve and a different operator, so it gets its own.""" + key = event.get("part") or event.get("name") + return f"{key}/adjoint" if event.get("kind") == "adjoint_solve" else key + + +def _part_label(event): + name = _short_operator(event.get("name", event.get("part", "?"))) + if event.get("kind") == "history_shift": + return name + if event.get("kind") == "adjoint_solve": + return f"adjoint {name}" + return name + + def _parts_of(steps): """The roster, in a stable order, from what actually played. @@ -1544,14 +1588,14 @@ def _parts_of(steps): order, labels = [], {} for step in steps: for event in step.get("events", []): - if event.get("kind") not in ("solve", "history_shift"): + if event.get("kind") not in _OPERATOR_KINDS: continue - key = event.get("part") or event.get("name") + key = _part_key(event) if key not in labels: order.append(key) - labels[key] = _short_operator(event.get("name", key)) + labels[key] = _part_label(event) elif event.get("kind") == "history_shift": - labels[key] = _short_operator(event.get("name", key)) + labels[key] = _part_label(event) # In the order they first ran. A step is drawn as a bar whose events # descend in the order they ran, so with the columns in that same order # the usual step reads as a staircase down and to the right, and any step @@ -1573,7 +1617,7 @@ def _outcome(event): it is neither a clean convergence nor a failure, and reading it as either loses the thing worth seeing. """ - if event.get("kind") != "solve" or "converged" not in event: + if event.get("kind") not in ("solve", "adjoint_solve") or "converged" not in event: return None if not event.get("converged"): return "diverged" @@ -1592,9 +1636,8 @@ def _step_cells(step, parts): """ played = [] for event in step.get("events", []): - if event.get("kind") in ("solve", "history_shift"): - played.append((event.get("part") or event.get("name"), - _outcome(event))) + if event.get("kind") in _OPERATOR_KINDS: + played.append((_part_key(event), _outcome(event))) cells = [] for key, _ in parts: hits = tuple((i + 1, outcome) for i, (k, outcome) in enumerate(played) From b1d068230bfb88d56aa6695f15b90183051525f1 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 17 Sep 2026 14:39:18 +1000 Subject: [PATCH 10/23] adjoint_solve transposes the matrices and solves forwards with the solver's own KSP MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit KSP.solveTranspose applies the transpose of the preconditioner, and PETSc refuses that for a multigrid whose smoother is a one-sided SOR sweep — which is the velocity block of every Stokes solve that has a mesh hierarchy. So the fault example's adjoint worked only on the mesh whose velocity block had fallen back to gamg, and failed the moment the mesh was refined and FMG could be built. J and P are now transposed explicitly, the null space carried across, and the same KSP solves the transposed system forwards, so every preconditioner the forward solve can use the adjoint can use. The forward operators are put back after. On the refined mesh the gradient check reads 1.00000 on all four segments with cold starts, and each recorded step takes about 2.5 s against 40–60 s on the capped mesh. Warm starts were tried and dropped: the warm-started Newton stops on its step criterion a little early and the finite-difference side lost the last digits (0.99976). Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- .../adjoint/fault_segments/fault_friction.py | 6 ++- .../adjoint/fault_segments/fault_segments.py | 12 +++--- .../cython/petsc_generic_snes_solvers.pyx | 42 +++++++++++++++---- 3 files changed, 45 insertions(+), 15 deletions(-) diff --git a/docs/examples/adjoint/fault_segments/fault_friction.py b/docs/examples/adjoint/fault_segments/fault_friction.py index cf3d28685..e2a02f1fc 100644 --- a/docs/examples/adjoint/fault_segments/fault_friction.py +++ b/docs/examples/adjoint/fault_segments/fault_friction.py @@ -50,6 +50,8 @@ # Refined once from the base size: the refinement gives the velocity block a # multigrid hierarchy. Without one it falls back to gamg, which hits its # iteration cap on this problem, and an inexact Newton step converges linearly. +# Every solve starts cold: six Newton iterations, and a misfit that does not +# depend on the previous evaluation, which the finite-difference check needs. mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0.0, 0.0), maxCoords=(2.0, 1.0), cellSize=params.cell_size, qdegree=3, refinement=1) x, y = mesh.X @@ -168,7 +170,7 @@ def J_and_gradient(label=None): """ evaluations[0] += 1 with model.step(0.0, label=label or f"eval {evaluations[0]}"): - stokes.solve(zero_init_guess=False) + stokes.solve(zero_init_guess=True) J = float(uw.maths.Integral(mesh, misfit).evaluate()) dJ_dv = misfit_duals(misfit, [v])[v] mu = uw.discretisation.MeshVariable(f"mu_{uw.adjoint._counter()}", mesh, mesh.dim, degree=2) @@ -181,7 +183,7 @@ def J_and_gradient(label=None): def forward(label): with model.step(0.0, label=label): - stokes.solve(zero_init_guess=False) + stokes.solve(zero_init_guess=True) # --- the truth, and the twin ------------------------------------------------- true_values = [float(t) for t in str(params.true_strengths).split(",")][:n_seg] diff --git a/docs/examples/adjoint/fault_segments/fault_segments.py b/docs/examples/adjoint/fault_segments/fault_segments.py index 0fd574c6d..635e8f620 100644 --- a/docs/examples/adjoint/fault_segments/fault_segments.py +++ b/docs/examples/adjoint/fault_segments/fault_segments.py @@ -43,6 +43,8 @@ # Refined once from the base size: the refinement gives the velocity block a # multigrid hierarchy. Without one it falls back to gamg, which hits its # iteration cap on this problem, and an inexact Newton step converges linearly. +# Every solve starts cold: six Newton iterations, and a misfit that does not +# depend on the previous evaluation, which the finite-difference check needs. mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0.0, 0.0), maxCoords=(2.0, 1.0), cellSize=params.cell_size, qdegree=3, refinement=1) x, y = mesh.X @@ -122,7 +124,7 @@ def set_strengths(values): def J_and_gradient(): """The misfit and dJ/d(log strength) for each segment, by the adjoint.""" - stokes.solve(zero_init_guess=False) + stokes.solve(zero_init_guess=True) J = float(uw.maths.Integral(mesh, misfit).evaluate()) dJ_dv = misfit_duals(misfit, [v])[v] mu = uw.discretisation.MeshVariable(f"mu_{uw.adjoint._counter()}", mesh, mesh.dim, degree=2) @@ -136,7 +138,7 @@ def J_and_gradient(): # --- the truth, and the twin ------------------------------------------------- true_values = [float(t) for t in str(params.true_strengths).split(",")][:n_seg] set_strengths(true_values) -stokes.solve(zero_init_guess=False) +stokes.solve(zero_init_guess=True) v_obs.array[...] = np.asarray(v.array) uw.pprint(f"true strengths {true_values}") @@ -153,7 +155,7 @@ def J_and_gradient(): vals = [params.initial_strength] * n_seg vals[k] = math.exp(base + sign * h) set_strengths(vals) - stokes.solve(zero_init_guess=False) + stokes.solve(zero_init_guess=True) fd.append(float(uw.maths.Integral(mesh, misfit).evaluate())) fd = (fd[0] - fd[1]) / (2 * h) uw.pprint(f"{names[k]:>13}: adjoint {g0[k]: .6e} finite difference {fd: .6e} " @@ -189,10 +191,10 @@ def objective(log_eta): for label, values in (("true", true_values), ("initial", [params.initial_strength] * n_seg), ("recovered", list(np.exp(result.x)))): set_strengths(values) - stokes.solve(zero_init_guess=False) + stokes.solve(zero_init_guess=True) profiles[label] = np.asarray(uw.function.evaluate(v.sym[1], top)).ravel() set_strengths(true_values) -stokes.solve(zero_init_guess=False) # the field on the grid is the truth's +stokes.solve(zero_init_guess=True) # the field on the grid is the truth's gx, gy = np.meshgrid(np.linspace(0, 2, 201), np.linspace(0, 1, 101)) grid = np.column_stack([gx.ravel(), gy.ravel()]) eta_1_grid = np.asarray(uw.function.evaluate(eta_1, grid)).reshape(gx.shape) diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index 38dc9275c..a0f1b1108 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -53,6 +53,38 @@ expression = lambda *x, **X: public_expression(*x, _unique_name_generation=True, from underworld3.function.expressions import unwrap_expression as _unwrap_expression +def _solve_transposed(ksp, J, P, b, x): + """Solve :math:`J^T x = b` with the solver's OWN KSP and preconditioner. + + ``KSP.solveTranspose`` applies the transpose of the preconditioner, which + PETSc refuses for a multigrid whose smoother is a one-sided SOR sweep — + the velocity block of every Stokes solve with a mesh hierarchy. So the + matrices are transposed explicitly and the same KSP solves the + transposed system forwards: every preconditioner the forward solve can + use, the adjoint can use, set up on :math:`J^T` exactly as it was on + :math:`J`. The forward operators are put back afterwards. A null space + attached to ``J`` (the pressure constant on an enclosed domain) is the + transpose null space of ``J^T``, and is carried across. + """ + Jt = J.transpose() + Pt = Jt if P.handle == J.handle else P.transpose() + for source, put in ((J.getNullSpace(), Jt.setTransposeNullSpace), + (J.getTransposeNullSpace(), Jt.setNullSpace), + (J.getNearNullSpace(), Jt.setNearNullSpace)): + if source is not None and source.handle != 0: + put(source) + ksp.setOperators(Jt, Pt) + try: + ksp.solve(b, x) + reason = int(ksp.getConvergedReason()) + finally: + ksp.setOperators(J, P) + if Pt is not Jt: + Pt.destroy() + Jt.destroy() + return reason + + def _jacobian_unwrap(expr): """Expand UWexpressions down to (but NOT including) constant atoms, for use as the input to a Jacobian derivative (``derive_by_array`` / ``diff``). @@ -1551,10 +1583,7 @@ class SolverBaseClass(uw_object): x = gvec.duplicate() x.set(0.0) - ksp = self.snes.getKSP() - ksp.setOperators(J, P) - ksp.solveTranspose(b, x) - reason = int(ksp.getConvergedReason()) + reason = _solve_transposed(self.snes.getKSP(), J, P, b, x) self._restore_tangent(tangent) if target is not None: @@ -10170,10 +10199,7 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): x = gvec.duplicate() x.set(0.0) - ksp = self.snes.getKSP() - ksp.setOperators(J, P) - ksp.solveTranspose(b, x) - reason = int(ksp.getConvergedReason()) + reason = _solve_transposed(self.snes.getKSP(), J, P, b, x) self._restore_tangent(tangent) if target is not None: From 57fcdf35f8776b91ec80f1cbf24d0cdfc397ebc4 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 17 Sep 2026 15:06:15 +1000 Subject: [PATCH 11/23] sensitivity: leave numeric constants as named expressions when peeling, so the kernel is not recompiled per value Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- src/underworld3/adjoint.py | 3 ++- src/underworld3/cython/petsc_generic_snes_solvers.pyx | 8 +++++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/src/underworld3/adjoint.py b/src/underworld3/adjoint.py index 8e90e9e46..6b7cb0287 100644 --- a/src/underworld3/adjoint.py +++ b/src/underworld3/adjoint.py @@ -502,7 +502,8 @@ def _peel_except(expression, wrt, depth=8): the derivative of a number is zero.""" for _ in range(depth): named = [e for e in uw.function.fn_extract_expressions(expression) - if e is not wrt and e != wrt] + if e is not wrt and e != wrt + and not getattr(getattr(e, "sym", None), "is_Number", False)] if not named: break expression = expression.subs({e: e.sym for e in named}) diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index a0f1b1108..581a5766d 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -1837,9 +1837,15 @@ class SolverBaseClass(uw_object): definition, one level at a time, and stops at ``wrt`` so the chain rule has something to hold on to. """ + # A named expression whose value is a bare number is left as it is: + # the JIT passes it as a run-time constant, so the kernel compiled + # for the integrand survives a change of its value. Substituting the + # number recompiled every sensitivity at every evaluation of an + # inversion (four compiles per evaluation on the fault example). for _ in range(depth): named = [e for e in uw.function.fn_extract_expressions(expression) - if e is not wrt and e != wrt] + if e is not wrt and e != wrt + and not getattr(getattr(e, "sym", None), "is_Number", False)] if not named: break expression = expression.subs({e: e.sym for e in named}) From 45688997b48dab28389a49986f68b4020c227981 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 17 Sep 2026 15:32:41 +1000 Subject: [PATCH 12/23] plot_convergence: read the per-observation data files Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- docs/examples/adjoint/fault_segments/plot_convergence.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/examples/adjoint/fault_segments/plot_convergence.py b/docs/examples/adjoint/fault_segments/plot_convergence.py index 79999c080..655fabf23 100644 --- a/docs/examples/adjoint/fault_segments/plot_convergence.py +++ b/docs/examples/adjoint/fault_segments/plot_convergence.py @@ -4,7 +4,7 @@ matplotlib.use("Agg") import matplotlib.pyplot as plt -cases = [("uplift + stress", "fault_friction_data.npz"), +cases = [("uplift + stress", "fault_friction_uplift+stress_data.npz"), ("orientation, points + surface", "fault_friction_orientation_data.npz"), ("orientation, surface only", "fault_friction_orientation_surface_data.npz")] names = ["flat", "lower ramp", "upper ramp", "near surface"] From 52deffe59762cabf2ee07b7d054f7851167ed5f4 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 17 Sep 2026 19:40:03 +1000 Subject: [PATCH 13/23] The adjoint operator is assembled from the transposed pointwise kernels MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PETSc assembles a Jacobian block from four pointwise kernels, g0..g3, which the JIT already produces by differentiating the residual. The transposed bilinear form has the same four with trial and test exchanged: g0 and g3 transposed on their paired indices, g1 and g2 swapped. _transpose_kernels does that relabelling on the SymPy matrices, in PETSc's flat layout; each solver swaps its kernels when _adjoint_kernels is set (the saddle point swaps uu in place and exchanges up and pu, registering all four kernels of the transposed (p,u) block); and adjoint_solve installs the set with a rewire, lets the SNES assemble K^T at the converged state, and solves it forwards with the solver's own KSP. Nothing is differentiated again and no matrix is transposed. test_0022 pins the assembled operator against the explicit transpose: 7e-17 on a non-symmetric SUPG step, 5e-17 on a nonlinear Stokes with pressure in the viscosity. The gradient tests (0019-0021) pass through the new path, in half the time. The explicit-transpose route stays as the fallback for a solver with boundary Jacobian kernels, whose swap is not written yet — and it is fixed: Mat.transpose() with no target transposes in place, which had been transposing the forward Jacobian under the SNES. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- .../cython/petsc_generic_snes_solvers.pyx | 176 +++++++++++++++++- ...test_0022_adjoint_operator_from_kernels.py | 89 +++++++++ 2 files changed, 255 insertions(+), 10 deletions(-) create mode 100644 tests/test_0022_adjoint_operator_from_kernels.py diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index 581a5766d..5a9256ad7 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -53,6 +53,57 @@ expression = lambda *x, **X: public_expression(*x, _unique_name_generation=True, from underworld3.function.expressions import unwrap_expression as _unwrap_expression +def _transpose_kernels(G0, G1, G2, G3, nc_test, nc_trial, dim): + """The pointwise Jacobian kernels of the TRANSPOSED bilinear form. + + PETSc reads a Jacobian block (test field with ``nc_test`` components, + trial field with ``nc_trial``) from four flat buffers laid out + ``g0[fc, gc]``, ``g1[fc, gc, d]``, ``g2[fc, gc, d]``, ``g3[fc, gc, d, e]`` + — ``fc`` the test component, ``gc`` the trial component, ``d`` the + test-side derivative direction in g2/g3 and the trial-side one in g1, + ``e`` the trial-side direction in g3. The form is + + phi_fc g0 u_gc + phi_fc g1 d_d u_gc + d_d phi_fc g2 u_gc + d_d phi_fc g3 d_e u_gc. + + Swapping trial and test gives the block of K^T with the roles of the + two fields exchanged: h0 is g0 transposed, h1 and h2 are g2 and g1 with + the components exchanged, and h3 is g3 with both index pairs exchanged. + Nothing is differentiated again; the derivatives are the forward ones, + evaluated at the same state. Matrices are returned whose row-major + flattening is PETSc's layout for the transposed block. A ``None`` + kernel is an absent (zero) one. + """ + import sympy + + def flat(G, n): + if G is None: + return [sympy.Integer(0)] * n + values = list(G) + if len(values) != n: + raise ValueError(f"kernel has {len(values)} entries, expected {n}") + return values + + nt, nu = nc_test, nc_trial + g0 = flat(G0, nt * nu) + g1 = flat(G1, nt * nu * dim) + g2 = flat(G2, nt * nu * dim) + g3 = flat(G3, nt * nu * dim * dim) + H0 = sympy.zeros(nu, nt) + H1 = sympy.zeros(nu * nt, dim) + H2 = sympy.zeros(nu * nt, dim) + H3 = sympy.zeros(nu * nt, dim * dim) + for fc in range(nt): + for gc in range(nu): + H0[gc, fc] = g0[fc * nu + gc] + for d in range(dim): + H1[gc * nt + fc, d] = g2[(fc * nu + gc) * dim + d] + H2[gc * nt + fc, d] = g1[(fc * nu + gc) * dim + d] + for e in range(dim): + H3[gc * nt + fc, e * dim + d] = g3[((fc * nu + gc) * dim + d) * dim + e] + return (sympy.ImmutableMatrix(H0), sympy.ImmutableMatrix(H1), + sympy.ImmutableMatrix(H2), sympy.ImmutableMatrix(H3)) + + def _solve_transposed(ksp, J, P, b, x): """Solve :math:`J^T x = b` with the solver's OWN KSP and preconditioner. @@ -66,8 +117,14 @@ def _solve_transposed(ksp, J, P, b, x): attached to ``J`` (the pressure constant on an enclosed domain) is the transpose null space of ``J^T``, and is carried across. """ - Jt = J.transpose() - Pt = Jt if P.handle == J.handle else P.transpose() + # Mat.transpose() with no target transposes IN PLACE; give it a new Mat. + Jt = PETSc.Mat() + J.transpose(Jt) + if P.handle == J.handle: + Pt = Jt + else: + Pt = PETSc.Mat() + P.transpose(Pt) for source, put in ((J.getNullSpace(), Jt.setTransposeNullSpace), (J.getTransposeNullSpace(), Jt.setNullSpace), (J.getNearNullSpace(), Jt.setNearNullSpace)): @@ -1549,7 +1606,9 @@ class SolverBaseClass(uw_object): # The kernel first: if the forward ran Picard, the rebuild below may # replace the DS the SNES assembles with, so every handle taken from # the DM must be taken AFTER it. - tangent = self._consistent_tangent_for_adjoint() + by_kernels = self._adjoint_by_kernels() + tangent = (self._install_adjoint_kernels() if by_kernels + else self._consistent_tangent_for_adjoint()) dm = self.dm # The Jacobian at the state the forward solve ended in. @@ -1583,8 +1642,16 @@ class SolverBaseClass(uw_object): x = gvec.duplicate() x.set(0.0) - reason = _solve_transposed(self.snes.getKSP(), J, P, b, x) - self._restore_tangent(tangent) + if by_kernels: + # J IS K^T: assembled from the transposed kernels. Solve forwards. + ksp = self.snes.getKSP() + ksp.setOperators(J, P) + ksp.solve(b, x) + reason = int(ksp.getConvergedReason()) + self._uninstall_adjoint_kernels(tangent) + else: + reason = _solve_transposed(self.snes.getKSP(), J, P, b, x) + self._restore_tangent(tangent) if target is not None: # Homogeneous constraints: the local vector is zeroed before the @@ -1721,6 +1788,41 @@ class SolverBaseClass(uw_object): out = out + d0[i] * mu_sym[i] return out + uw.maths.tensor.rank2_inner_product(d1, grad_mu) + def _adjoint_by_kernels(self): + """Whether the adjoint operator is ASSEMBLED from the transposed + kernels rather than obtained by transposing the assembled matrix. + The kernel route covers the volume terms; a boundary Jacobian (a + natural or Nitsche condition with a tangent) is not swapped yet, so + such a solver takes the matrix route.""" + return not (getattr(self, "natural_bcs", None) or []) + + def _install_adjoint_kernels(self): + """Rewire the solver to the kernels of the transposed form. + + The Jacobian the SNES then assembles IS K^T at the state the forward + solve ended in: the consistent tangent's derivatives with trial and + test exchanged (:func:`_transpose_kernels`), compiled and cached like + any other kernel set. Returns what to hand back to + :meth:`_uninstall_adjoint_kernels`. + """ + previous = self._consistent_jacobian + if self.consistent_jacobian is False and not self._residual_is_linear_in_unknown(): + self._consistent_jacobian = True # transpose dR/du, not the Picard kernel + self._adjoint_kernels = True + self._needs_function_rewire = True + self._build(False, False, None) + self.snes.setUp() # a direct computeJacobian needs it + return previous + + def _uninstall_adjoint_kernels(self, previous): + """Mark the forward kernels as wanted again. The next forward solve's + own build rewires; nothing is torn down here, so a second adjoint on + the same forward state reuses the installed set.""" + self._consistent_jacobian = previous + self._adjoint_kernels = False + self._adjoint_kernel_installed = True + self._needs_function_rewire = True + def _consistent_tangent_for_adjoint(self): """Make sure the Jacobian kernel the adjoint assembles is dR/du. @@ -4520,6 +4622,12 @@ class SNES_Scalar(SolverBaseClass): self._G1 = sympy.ImmutableMatrix(G1) self._G2 = sympy.ImmutableMatrix(G2) self._G3 = sympy.ImmutableMatrix(G3) + if getattr(self, "_adjoint_kernels", False): + # The adjoint's operator: the same derivatives, trial and test + # exchanged. Compiled as its own kernel set and cached. + _nc = int(self._G0.shape[0]) + self._G0, self._G1, self._G2, self._G3 = _transpose_kernels( + self._G0, self._G1, self._G2, self._G3, _nc, _nc, self.mesh.cdim) ################## @@ -5521,6 +5629,12 @@ class SNES_Vector(SolverBaseClass): self._G1 = sympy.ImmutableMatrix(G1) self._G2 = sympy.ImmutableMatrix(G2) self._G3 = sympy.ImmutableMatrix(G3) + if getattr(self, "_adjoint_kernels", False): + # The adjoint's operator: the same derivatives, trial and test + # exchanged. Compiled as its own kernel set and cached. + _nc = int(self._G0.shape[0]) + self._G0, self._G1, self._G2, self._G3 = _transpose_kernels( + self._G0, self._G1, self._G2, self._G3, _nc, _nc, self.mesh.cdim) ################## @@ -6319,6 +6433,12 @@ class SNES_MultiComponent(SolverBaseClass): self._G1 = sympy.ImmutableMatrix(G1) self._G2 = sympy.ImmutableMatrix(G2) self._G3 = sympy.ImmutableMatrix(G3) + if getattr(self, "_adjoint_kernels", False): + # The adjoint's operator: the same derivatives, trial and test + # exchanged. Compiled as its own kernel set and cached. + _nc = int(self._G0.shape[0]) + self._G0, self._G1, self._G2, self._G3 = _transpose_kernels( + self._G0, self._G1, self._G2, self._G3, _nc, _nc, self.mesh.cdim) fns_jacobian = (self._G0, self._G1, self._G2, self._G3) @@ -8767,6 +8887,25 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): fns_jacobian.append(self._pp_G0) + if getattr(self, "_adjoint_kernels", False): + # K^T block by block: uu transposed in place; the (u,p) block of + # K^T is the transpose of K's (p,u) block and vice versa; pp is a + # scalar. Same derivatives, trial and test exchanged. + forward = [self._uu_G0, self._uu_G1, self._uu_G2, self._uu_G3, + self._up_G0, self._up_G1, self._up_G2, self._up_G3, + self._pu_G0, self._pu_G1] + up_from_pu = _transpose_kernels(self._pu_G0, self._pu_G1, None, None, 1, dim, dim) + pu_from_up = _transpose_kernels(self._up_G0, self._up_G1, self._up_G2, self._up_G3, + dim, 1, dim) + self._uu_G0, self._uu_G1, self._uu_G2, self._uu_G3 = _transpose_kernels( + self._uu_G0, self._uu_G1, self._uu_G2, self._uu_G3, dim, dim, dim) + self._up_G0, self._up_G1, self._up_G2, self._up_G3 = up_from_pu + self._pu_G0, self._pu_G1, self._pu_G2, self._pu_G3 = pu_from_up + fns_jacobian = [f for f in fns_jacobian if not any(f is g for g in forward)] + fns_jacobian += [self._uu_G0, self._uu_G1, self._uu_G2, self._uu_G3, + self._up_G0, self._up_G1, self._up_G2, self._up_G3, + self._pu_G0, self._pu_G1, self._pu_G2, self._pu_G3] + ## 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 @@ -9442,10 +9581,17 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): PetscDSSetJacobian( ds.ds, 0, 0, ext.fns_jacobian[i_jac[self._uu_G0]], ext.fns_jacobian[i_jac[self._uu_G1]], ext.fns_jacobian[i_jac[self._uu_G2]], ext.fns_jacobian[i_jac[self._uu_G3]]) PetscDSSetJacobian( ds.ds, 0, 1, ext.fns_jacobian[i_jac[self._up_G0]], ext.fns_jacobian[i_jac[self._up_G1]], ext.fns_jacobian[i_jac[self._up_G2]], ext.fns_jacobian[i_jac[self._up_G3]]) - PetscDSSetJacobian( ds.ds, 1, 0, ext.fns_jacobian[i_jac[self._pu_G0]], ext.fns_jacobian[i_jac[self._pu_G1]], NULL, NULL) + if getattr(self, "_adjoint_kernels", False): + # the transposed (p,u) block carries the flux derivatives of K's (u,p) block + PetscDSSetJacobian( ds.ds, 1, 0, ext.fns_jacobian[i_jac[self._pu_G0]], ext.fns_jacobian[i_jac[self._pu_G1]], ext.fns_jacobian[i_jac[self._pu_G2]], ext.fns_jacobian[i_jac[self._pu_G3]]) + else: + PetscDSSetJacobian( 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, 0, 0, ext.fns_jacobian[i_jac[self._uu_G0]], ext.fns_jacobian[i_jac[self._uu_G1]], ext.fns_jacobian[i_jac[self._uu_G2]], ext.fns_jacobian[i_jac[self._uu_G3]]) PetscDSSetJacobianPreconditioner(ds.ds, 0, 1, ext.fns_jacobian[i_jac[self._up_G0]], ext.fns_jacobian[i_jac[self._up_G1]], ext.fns_jacobian[i_jac[self._up_G2]], ext.fns_jacobian[i_jac[self._up_G3]]) - PetscDSSetJacobianPreconditioner(ds.ds, 1, 0, ext.fns_jacobian[i_jac[self._pu_G0]], ext.fns_jacobian[i_jac[self._pu_G1]], NULL, NULL) + if getattr(self, "_adjoint_kernels", False): + PetscDSSetJacobianPreconditioner(ds.ds, 1, 0, ext.fns_jacobian[i_jac[self._pu_G0]], ext.fns_jacobian[i_jac[self._pu_G1]], ext.fns_jacobian[i_jac[self._pu_G2]], ext.fns_jacobian[i_jac[self._pu_G3]]) + else: + 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 @@ -10167,7 +10313,9 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): import numpy as np - tangent = self._consistent_tangent_for_adjoint() # before any DM vector + by_kernels = self._adjoint_by_kernels() # before any DM vector + tangent = (self._install_adjoint_kernels() if by_kernels + else self._consistent_tangent_for_adjoint()) gvec = self.dm.getGlobalVec() gvec.setArray(0.0) self._gather_fields_to_global(gvec) @@ -10205,8 +10353,16 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): x = gvec.duplicate() x.set(0.0) - reason = _solve_transposed(self.snes.getKSP(), J, P, b, x) - self._restore_tangent(tangent) + if by_kernels: + # J IS K^T: assembled from the transposed kernels. Solve forwards. + ksp = self.snes.getKSP() + ksp.setOperators(J, P) + ksp.solve(b, x) + reason = int(ksp.getConvergedReason()) + self._uninstall_adjoint_kernels(tangent) + else: + reason = _solve_transposed(self.snes.getKSP(), J, P, b, x) + self._restore_tangent(tangent) if target is not None: u_adj, p_adj = target diff --git a/tests/test_0022_adjoint_operator_from_kernels.py b/tests/test_0022_adjoint_operator_from_kernels.py new file mode 100644 index 000000000..bd4c530d5 --- /dev/null +++ b/tests/test_0022_adjoint_operator_from_kernels.py @@ -0,0 +1,89 @@ +"""The adjoint operator assembled from the transposed kernels IS K^T. + +``adjoint_solve`` no longer transposes the assembled Jacobian: it rewires the +solver to the pointwise kernels of the transposed bilinear form (g0 and g3 +transposed on their paired indices, g1 and g2 exchanged) and lets the SNES +assemble. This checks that matrix against the explicit transpose to machine +precision, on a non-symmetric operator and on a nonlinear saddle point with +pressure in the viscosity, so every block and every index order is exercised. +""" +import numpy as np +import pytest +import sympy + +import underworld3 as uw +from petsc4py import PETSc + + +def _adjoint_versus_transpose(solver): + def assemble(): + dm = solver.dm + g = dm.getGlobalVec() + if hasattr(solver, "_gather_fields_to_global"): + solver._gather_fields_to_global(g) + else: + dm.localToGlobal(solver.u.vec, g) + solver.mesh.update_lvec() + dm.setAuxiliaryVec(solver.mesh.lvec, None) + J, P = solver.snes.getJacobian()[:2] + solver.snes.computeJacobian(g, J, P) + out = J.copy() + dm.restoreGlobalVec(g) + return out + + K = assemble() + token = solver._install_adjoint_kernels() + K_adj = assemble() + solver._uninstall_adjoint_kernels(token) + Kt = PETSc.Mat() + K.transpose(Kt) # a new matrix; the no-argument form transposes in place + D = K_adj.copy() + D.axpy(-1.0, Kt) + S = K.copy() + S.axpy(-1.0, Kt) + return D.norm() / K.norm(), S.norm() / K.norm() + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_supg_step_adjoint_kernels_give_the_transpose(): + mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0, 0), maxCoords=(1, 1), + cellSize=1 / 8, qdegree=3) + T = uw.discretisation.MeshVariable("T", mesh, 1, degree=2) + adv = uw.systems.AdvDiffusion(mesh, u_Field=T, V_fn=sympy.Matrix([[1.0, 0.3]]), theta=1.0) + adv.constitutive_model = uw.constitutive_models.DiffusionModel + adv.constitutive_model.Parameters.diffusivity = 0.05 + adv.add_essential_bc(1.0, "Bottom") + adv.add_essential_bc(0.0, "Top") + adv.solve(timestep=0.1, zero_init_guess=True) + error, asymmetry = _adjoint_versus_transpose(adv) + assert asymmetry > 1e-3, "the operator must be non-symmetric for this to test anything" + assert error < 1e-12 + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_nonlinear_stokes_adjoint_kernels_give_the_transpose(): + mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0, 0), maxCoords=(1, 1), + cellSize=1 / 6, qdegree=3) + x, y = mesh.X + v = uw.discretisation.MeshVariable("v", mesh, 2, degree=2) + p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1, continuous=True) + stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + # viscosity depending on the strain rate AND the pressure: every block of + # the Jacobian is non-trivial, and (u,p) and (p,u) are not transposes of + # each other + edot = stokes.constitutive_model.Parameters.strainrate_inv_II if hasattr( + stokes.constitutive_model.Parameters, "strainrate_inv_II") else None + E = mesh.vector.strain_tensor(v.sym) + e2 = (E[0, 0] ** 2 + E[1, 1] ** 2 + 2 * E[0, 1] ** 2) / 2 + uw.maths.functions.vanishing + stokes.constitutive_model.Parameters.shear_viscosity_0 = (1 + sympy.Max(p.sym[0], 0)) / (1 + sympy.sqrt(e2)) + stokes.bodyforce = sympy.Matrix([0, -sympy.sin(3 * x)]) + stokes.add_essential_bc((0.0, 0.0), "Bottom") + stokes.add_essential_bc((0.5, None), "Left") + stokes.add_essential_bc((-0.5, None), "Right") + stokes.solve(zero_init_guess=True) + error, asymmetry = _adjoint_versus_transpose(stokes) + assert asymmetry > 1e-3 + assert error < 1e-12 From 557417b87f1cf89726a7265b13472a4e91544d2f Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 17 Sep 2026 19:56:13 +1000 Subject: [PATCH 14/23] solver.gradient(misfit, parameters, fields): the steady adjoint as one call, and the adjoint problem in writing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit uw.adjoint.gradient(solver, misfit, parameters=, fields=) does the three steps — the misfit's dual on the unknown, the transposed solve, the sensitivities — and returns {"J", "parameters", "fields"} like the transcript driver does for a run; solver.gradient is the method form. A field control that the residual reads through a history slot (the initial condition of a transport step) is routed from the slot to the field, the unknown included. The driver shares the read-detection with it (field_duals). Introspection: adjoint_kernels() gives the transposed pointwise kernels per block; adjoint_templates() writes the adjoint problem in the residual template language, f0_adj = g0^T mu + g2^T:grad mu, f1_adj = g1^T mu + g3^T grad mu, per block for a saddle point; adjoint_view() typesets it in a notebook. test_0023 checks the one-call gradient against central differences in the diffusivity and in the initial field on a non-symmetric SUPG step, and the templates against the kernels. The fault example's evaluation is now stokes.gradient(misfit, parameters=strengths); its check still reads 1.00000 on every segment. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- .../adjoint/fault_segments/fault_friction.py | 22 ++- src/underworld3/adjoint.py | 149 +++++++++++++++--- .../cython/petsc_generic_snes_solvers.pyx | 105 ++++++++++++ ...lver_gradient_and_adjoint_introspection.py | 89 +++++++++++ 4 files changed, 328 insertions(+), 37 deletions(-) create mode 100644 tests/test_0023_solver_gradient_and_adjoint_introspection.py diff --git a/docs/examples/adjoint/fault_segments/fault_friction.py b/docs/examples/adjoint/fault_segments/fault_friction.py index e2a02f1fc..78b9e364a 100644 --- a/docs/examples/adjoint/fault_segments/fault_friction.py +++ b/docs/examples/adjoint/fault_segments/fault_friction.py @@ -11,9 +11,11 @@ taken from a run at the true coefficients. The residual is nonlinear in the velocity and the pressure, so the forward -solve is Newton, and the adjoint is a transpose of the consistent tangent -the solver assembled. The gradient with respect to each coefficient is the -symbolic derivative of that residual. Nothing here is differenced. +solve is Newton. The gradient with respect to each coefficient comes from +one call, stokes.gradient(misfit, parameters=...): the solver assembles the +adjoint operator from its own Jacobian kernels with trial and test +exchanged, solves it, and differentiates its residual symbolically with +respect to each named coefficient. Nothing here is differenced. Run it: @@ -26,7 +28,6 @@ import sympy import underworld3 as uw -from underworld3.adjoint import misfit_duals, inner params = uw.Params( cell_size=uw.Param(1 / 12, "base mesh cell size (box is 2 x 1); refined once, so half this"), @@ -171,15 +172,10 @@ def J_and_gradient(label=None): evaluations[0] += 1 with model.step(0.0, label=label or f"eval {evaluations[0]}"): stokes.solve(zero_init_guess=True) - J = float(uw.maths.Integral(mesh, misfit).evaluate()) - dJ_dv = misfit_duals(misfit, [v])[v] - mu = uw.discretisation.MeshVariable(f"mu_{uw.adjoint._counter()}", mesh, mesh.dim, degree=2) - lam = uw.discretisation.MeshVariable(f"lam_{uw.adjoint._counter()}", mesh, 1, degree=1) - dJ_dv.array[...] = -np.asarray(dJ_dv.array) - _, reason = stokes.adjoint_solve((dJ_dv, None), target=(mu, lam)) - assert reason > 0, reason - grad = np.array([stokes.sensitivity(mu, expr) * float(expr.sym) for expr in strengths]) - return J, grad + out = stokes.gradient(misfit, parameters=strengths) + # d/d(log mu) = mu d/d(mu) + grad = np.array([out["parameters"][expr] * float(expr.sym) for expr in strengths]) + return out["J"], grad def forward(label): with model.step(0.0, label=label): diff --git a/src/underworld3/adjoint.py b/src/underworld3/adjoint.py index 6b7cb0287..83f71f41a 100644 --- a/src/underworld3/adjoint.py +++ b/src/underworld3/adjoint.py @@ -249,6 +249,130 @@ def misfit_duals(misfit, variables, scratch=None): return out +def _reads_of(solver, unknown, tokens): + """What a solver's residual reads, other than its unknown. + + ``(variable, value symbols, {(component, direction): derivative atom})`` + per variable. A component prints as ``{v}_{ 0 }``; a derivative carries + a comma — ``{v}_{ 0,1}`` for a vector, ``{T}_{,1}`` for a scalar — and + is read through the gradient part of the load. + """ + f0 = _peel(solver.F0.sym) + f1 = _peel(solver.F1.sym) + text = str(f0) + str(f1) + atoms = set(f0.atoms(sympy.Function)) | set(f1.atoms(sympy.Function)) + found = [] + for token, var in tokens.items(): + if var is unknown or token not in text: + continue + derivatives = {} + pattern = re.compile(re.escape(token) + r"_\{ ?(\d*),(\d+)\}\(") + for atom in atoms: + m = pattern.match(str(atom)) + if m: + i = int(m.group(1)) if m.group(1) else 0 + derivatives[(i, int(m.group(2)))] = atom + found.append((var, _symbols_of(var), derivatives)) + return found + + +def _tokens_of(variables): + return {_token_of(var): var for var in variables if hasattr(var, "sym")} + + +def field_duals(solver, mu, variables, scratch=None): + r"""``(\partial R/\partial f)^T \mu`` as a dual on each field ``f`` the + solver's residual reads, among ``variables``. + + A field read through its value gives the value part of the load; one + read through its gradient (a Crank–Nicolson step reads the old flux) + gives the gradient part. Both come from :meth:`adjoint_integrand`, the + symbolic derivative of the residual, and are assembled as one load. + """ + scratch = _shared_scratch if scratch is None else scratch + out = {} + for var, symbols, derivatives in _reads_of(solver, solver.u, _tokens_of(variables)): + value = [solver.adjoint_integrand(mu, s) for s in symbols] + g1 = None + if derivatives: + cdim = var.mesh.cdim + g1 = sympy.zeros(len(symbols), cdim) + for (i, k), atom in derivatives.items(): + g1[i, k] = solver.adjoint_integrand(mu, atom) + out[var] = dual_on(var, _as_expression(value), g1, scratch) + return out + + +def gradient(solver, misfit, parameters=(), fields=(), scratch=None): + r"""``dJ/dm`` and the duals on fields, by the adjoint of ONE solve. + + For :math:`J = \int` ``misfit`` over the mesh, evaluated in the state the + solver ended in: the dual of :math:`J` on the unknown is assembled + (:func:`misfit_duals`), the adjoint system :math:`K^T\mu = -\partial J/ + \partial u` is solved (:meth:`adjoint_solve`), and each parameter gets + :math:`\partial J/\partial m + \mu^T\partial R/\partial m` + (:meth:`sensitivity`). Each requested field gets :math:`\partial J/ + \partial f + (\partial R/\partial f)^T\mu` as a dual on its own space + — the derivative through a field the residual reads, an initial + condition or a coefficient field. + + Returns ``{"J": float, "parameters": {expr: float}, "fields": {var: dual}}``; + the duals are NumPy copies, safe to keep. + """ + scratch = _shared_scratch if scratch is None else scratch + parameters, fields = list(parameters), list(fields) + u = solver.u + mesh = u.mesh + J = float(uw.maths.Integral(mesh, misfit).evaluate()) + duals = misfit_duals(misfit, [u] + [f for f in fields if f is not u], scratch) + grad = {} + for p in parameters: + explicit = sympy.diff(_peel_except(misfit, p), p) + grad[p] = 0.0 if explicit == 0 else float(uw.maths.Integral(mesh, explicit).evaluate()) + # The explicit part on a field the misfit reads directly — never on the + # unknown, whose misfit dual is the adjoint's right-hand side. + out_fields = {var: (np.array(duals[var].array, copy=True) if (var in duals and var is not u) + else np.zeros_like(np.asarray(var.array))) for var in fields} + # A history slot the residual reads (psi_star[0]) holds the tracked field + # at the solve's input, so its dual is the derivative with respect to + # that field: read the slot, route the dual to the field. + route = {} + for history in (getattr(solver, "DuDt", None), getattr(solver, "DFDt", None)): + if history is None or not getattr(history, "psi_star", None): + continue + text = str(history.psi_fn) + for field in fields: + # the unknown itself is a control through its INPUT level, which + # is what the slot holds + if _token_of(field) in text: + route[history.psi_star[0]] = field + read_vars = [f for f in fields if f is not u] + list(route) + if u in duals: + rhs = duals.pop(u) + rhs.array[...] = -np.asarray(rhs.array) + mu = scratch.take(u) + if getattr(solver, "p", None) is not None and hasattr(solver, "_subdict"): + lam = scratch.take(solver.p) + _, reason = solver.adjoint_solve((rhs, None), target=(mu, lam)) + scratch.give(lam) + else: + _, reason = solver.adjoint_solve(rhs, target=mu) + scratch.give(rhs) + if reason <= 0: + raise RuntimeError(f"gradient: the adjoint of {type(solver).__name__}({u.name}) " + f"did not converge ({reason})") + for p in parameters: + grad[p] += solver.sensitivity(mu, p) + for var, dual in field_duals(solver, mu, read_vars, scratch).items(): + target = route.get(var, var) + out_fields[target] = out_fields[target] + np.asarray(dual.array) + scratch.give(dual) + scratch.give(mu) + for dual in duals.values(): + scratch.give(dual) + return {"J": J, "parameters": grad, "fields": out_fields} + + _n = [0] @@ -386,30 +510,7 @@ def _tokens(self): if hasattr(var, "sym")} def _reads(self, solver, unknown): - """What a solver's residual reads, other than its unknown. - - ``(variable, value symbols, {(component, direction): derivative atom})`` - per variable. A component prints as ``{v}_{ 0 }``; a derivative - carries a comma — ``{v}_{ 0,1}`` for a vector, ``{T}_{,1}`` for a - scalar — and is read through the gradient part of the load. - """ - f0 = _peel(solver.F0.sym) - f1 = _peel(solver.F1.sym) - text = str(f0) + str(f1) - atoms = set(f0.atoms(sympy.Function)) | set(f1.atoms(sympy.Function)) - found = [] - for token, var in self._tokens().items(): - if var is unknown or token not in text: - continue - derivatives = {} - pattern = re.compile(re.escape(token) + r"_\{ ?(\d*),(\d+)\}\(") - for atom in atoms: - m = pattern.match(str(atom)) - if m: - i = int(m.group(1)) if m.group(1) else 0 - derivatives[(i, int(m.group(2)))] = atom - found.append((var, _symbols_of(var), derivatives)) - return found + return _reads_of(solver, unknown, self._tokens()) def _linearise_at(self, step, solves, j): """Restore the step's snapshot, replay solves 0..j, and put each diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index 5a9256ad7..b788ffaaf 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -1788,6 +1788,111 @@ class SolverBaseClass(uw_object): out = out + d0[i] * mu_sym[i] return out + uw.maths.tensor.rank2_inner_product(d1, grad_mu) + def gradient(self, misfit, parameters=(), fields=()): + r"""``dJ/dm`` for each parameter, and the dual on each field, by the + adjoint of this solve. + + ``misfit`` is :math:`J` as an integrand over the mesh in the state + the solve ended in. The three steps — the dual of :math:`J` on the + unknown, the transposed solve, the sensitivities — are + :func:`underworld3.adjoint.gradient`; this is the method form of it. + Returns ``{"J", "parameters": {expr: dJ/dm}, "fields": {var: dual}}``. + """ + from underworld3.adjoint import gradient as _gradient + return _gradient(self, misfit, parameters=parameters, fields=fields) + + def adjoint_kernels(self): + """The pointwise kernels of the adjoint operator, as SymPy matrices. + + ``{"F0_u", "F0_grad_u", "F1_u", "F1_grad_u"}`` per block, in + PETSc's flat layout: the forward Jacobian kernels with trial and + test exchanged (:func:`_transpose_kernels`), which is what + :meth:`adjoint_solve` assembles. Read-only; the solver must have + been built. + """ + if getattr(self, "_G0", None) is None and getattr(self, "_uu_G0", None) is None: + self._build(False, False, None) + dim = self.mesh.cdim + if getattr(self, "_uu_G0", None) is not None: + uu = _transpose_kernels(self._uu_G0, self._uu_G1, self._uu_G2, self._uu_G3, dim, dim, dim) + up = _transpose_kernels(self._pu_G0, self._pu_G1, None, None, 1, dim, dim) + pu = _transpose_kernels(self._up_G0, self._up_G1, self._up_G2, self._up_G3, dim, 1, dim) + names = ("F0_u", "F0_grad_u", "F1_u", "F1_grad_u") + return {"uu": dict(zip(names, uu)), "up": dict(zip(names, up)), "pu": dict(zip(names, pu))} + nc = int(self._G0.shape[0]) + H = _transpose_kernels(self._G0, self._G1, self._G2, self._G3, nc, nc, dim) + return dict(zip(("F0_u", "F0_grad_u", "F1_u", "F1_grad_u"), H)) + + def adjoint_templates(self): + r"""The adjoint problem in the residual template form. + + :math:`K^T\mu = b` is linear in :math:`\mu`, so it is a solver in + the same language as the forward one: + + .. math:: + + f_0^{\rm adj} = g_0^T\mu + g_2^T\!:\!\nabla\mu, \qquad + \mathbf f_1^{\rm adj} = g_1^T\mu + g_3^T\nabla\mu, + + the forward Jacobian kernels with trial and test exchanged, applied + to the adjoint variable. Returns ``(F0_adj, F1_adj)`` as SymPy + expressions in :math:`\mu` (a scalar or a row vector of the + unknown's size) for a single-field solver, and the same per block + for a saddle point, which :meth:`adjoint_view` typesets. + """ + import sympy + kernels = self.adjoint_kernels() + dim = self.mesh.cdim + + def apply(block, nc, mu, grad_mu): + H0, H1, H2, H3 = (block["F0_u"], block["F0_grad_u"], block["F1_u"], block["F1_grad_u"]) + nt = int(H0.shape[0]) # test components of the adjoint block + f0 = sympy.zeros(nt, 1) + f1 = sympy.zeros(nt, dim) + for a in range(nt): + for c in range(nc): + f0[a] += H0[a, c] * mu[c] + for d in range(dim): + f0[a] += H1[a * nc + c, d] * grad_mu[c, d] + f1[a, d] += H2[a * nc + c, d] * mu[c] + for e in range(dim): + f1[a, d] += H3[a * nc + c, d * dim + e] * grad_mu[c, e] + return f0, f1 + + if "uu" in kernels: + mu = sympy.Matrix([sympy.Symbol(f"\\mu_{{{i}}}") for i in range(dim)]) + gmu = sympy.Matrix(dim, dim, lambda i, j: sympy.Symbol(f"\\mu_{{{i},{j}}}")) + lam = sympy.Matrix([sympy.Symbol(r"\lambda")]) + glam = sympy.Matrix(1, dim, lambda i, j: sympy.Symbol(f"\\lambda_{{,{j}}}")) + f0_uu, f1_uu = apply(kernels["uu"], dim, mu, gmu) + f0_up, f1_up = apply(kernels["up"], 1, lam, glam) + f0_pu, f1_pu = apply(kernels["pu"], dim, mu, gmu) + return {"u": (f0_uu + f0_up, f1_uu + f1_up), "p": (f0_pu, f1_pu)} + nc = int(kernels["F0_u"].shape[0]) + if nc == 1: + mu = sympy.Matrix([sympy.Symbol(r"\mu")]) + gmu = sympy.Matrix(1, dim, lambda i, j: sympy.Symbol(f"\\mu_{{,{j}}}")) + else: + mu = sympy.Matrix([sympy.Symbol(f"\\mu_{{{i}}}") for i in range(nc)]) + gmu = sympy.Matrix(nc, dim, lambda i, j: sympy.Symbol(f"\\mu_{{{i},{j}}}")) + return apply(kernels, nc, mu, gmu) + + def adjoint_view(self): + """Typeset the adjoint problem this solver assembles, in a notebook.""" + from IPython.display import Latex, Markdown, display + import sympy + templates = self.adjoint_templates() + display(Markdown("### Adjoint problem, as assembled")) + display(Markdown(r"$\int \phi\, f_0^{\rm adj}(\mu,\nabla\mu) + \nabla\phi\cdot\mathbf f_1^{\rm adj}(\mu,\nabla\mu) = b(\phi)$ with")) + if isinstance(templates, dict): + for block, (f0, f1) in templates.items(): + display(Latex(f"$f_0^{{\\rm adj}}[{block}] = {sympy.latex(f0)}$")) + display(Latex(f"$\\mathbf f_1^{{\\rm adj}}[{block}] = {sympy.latex(f1)}$")) + else: + f0, f1 = templates + display(Latex(f"$f_0^{{\\rm adj}} = {sympy.latex(f0)}$")) + display(Latex(f"$\\mathbf f_1^{{\\rm adj}} = {sympy.latex(f1)}$")) + def _adjoint_by_kernels(self): """Whether the adjoint operator is ASSEMBLED from the transposed kernels rather than obtained by transposing the assembled matrix. diff --git a/tests/test_0023_solver_gradient_and_adjoint_introspection.py b/tests/test_0023_solver_gradient_and_adjoint_introspection.py new file mode 100644 index 000000000..8e62f0483 --- /dev/null +++ b/tests/test_0023_solver_gradient_and_adjoint_introspection.py @@ -0,0 +1,89 @@ +"""One call for the steady adjoint, and the adjoint problem it assembles, in writing. + +``solver.gradient(misfit, parameters=..., fields=...)`` is the three-call +route — dual of the misfit, transposed solve, sensitivities — as one method, +so a steady inversion reads like the time-dependent one. Checked against a +central difference on a non-symmetric SUPG step, in the parameter and in the +initial field. ``adjoint_templates()`` writes the adjoint residual in the +same template language as the forward one; the check here is that it is +the forward Jacobian kernels with trial and test exchanged, by evaluating +both at a point. +""" +import numpy as np +import pytest +import sympy + +import underworld3 as uw + + +def _supg(cell=1 / 8): + mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0, 0), maxCoords=(1, 1), + cellSize=cell, qdegree=3) + x, y = mesh.X + T = uw.discretisation.MeshVariable("T", mesh, 1, degree=2) + T_obs = uw.discretisation.MeshVariable("T_obs", mesh, 1, degree=2) + kappa = uw.expression(r"\kappa", 0.05, "diffusivity") + adv = uw.systems.AdvDiffusion(mesh, u_Field=T, V_fn=sympy.Matrix([[1.0, 0.3]]), theta=1.0) + adv.constitutive_model = uw.constitutive_models.DiffusionModel + adv.constitutive_model.Parameters.diffusivity = kappa + adv.add_essential_bc(1.0, "Bottom") + adv.add_essential_bc(0.0, "Top") + return mesh, T, T_obs, kappa, adv + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_gradient_method_matches_finite_differences_in_parameter_and_initial_field(): + mesh, T, T_obs, kappa, adv = _supg() + x, y = mesh.X + dt = 0.05 + + def initial(c=0.5): + X = np.asarray(T.coords) + return np.exp(-((X[:, 0] - c) ** 2 + (X[:, 1] - 0.5) ** 2) / 0.05) + + def run(k, T0): + kappa.sym = float(k) + T.array[:, 0, 0] = T0 + adv.DuDt.initialise_history() if hasattr(adv.DuDt, "initialise_history") else None + adv.solve(timestep=dt, zero_init_guess=True) + return float(uw.maths.Integral(mesh, misfit).evaluate()) + + misfit = (T.sym[0] - T_obs.sym[0]) ** 2 / 2 + run(0.05, initial(0.55)) + T_obs.array[...] = np.asarray(T.array) + T0 = initial(0.5) + J0 = run(0.05, T0) + adv.DuDt.psi_star[0].array[:, 0, 0] = T0 # the step's input, not its output + out = adv.gradient(misfit, parameters=[kappa], fields=[T]) + assert abs(out["J"] - J0) < 1e-12 + + h = 1e-4 + fd = (run(0.05 + h, T0) - run(0.05 - h, T0)) / (2 * h) + assert abs(fd / out["parameters"][kappa] - 1) < 1e-3, (fd, out["parameters"][kappa]) + + # the initial field enters through the history slot the solve reads + direction = initial(0.45) - initial(0.5) + dual = out["fields"][T] + adjoint = uw.adjoint.inner(T, dual, direction) + fd = (run(0.05, T0 + h * direction) - run(0.05, T0 - h * direction)) / (2 * h) + assert abs(fd / adjoint - 1) < 1e-3, (fd, adjoint) + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_adjoint_templates_are_the_forward_kernels_with_trial_and_test_exchanged(): + mesh, T, T_obs, kappa, adv = _supg(cell=1 / 4) + adv.solve(timestep=0.05, zero_init_guess=True) + f0, f1 = adv.adjoint_templates() + kernels = adv.adjoint_kernels() + # value part: the coefficient of mu in f0 is H0, of mu_{,d} is H1 + mu = sympy.Symbol(r"\mu") + assert sympy.simplify(sympy.diff(f0[0], mu) - kernels["F0_u"][0, 0]) == 0 + for d in range(mesh.cdim): + md = sympy.Symbol(f"\\mu_{{,{d}}}") + assert sympy.simplify(sympy.diff(f0[0], md) - kernels["F0_grad_u"][0, d]) == 0 + assert sympy.simplify(sympy.diff(f1[0, d], mu) - kernels["F1_u"][0, d]) == 0 + # and the forward kernels' transpose: the SUPG term makes g3 non-symmetric only + # through the velocity; g1 <-> g2 is where a wrong swap would show + assert kernels["F0_grad_u"] == sympy.ImmutableMatrix(adv._G2).reshape(1, mesh.cdim) From 30b09385c0f3613a8f23bf6cb32b2d44a6df25b4 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 17 Sep 2026 21:33:12 +1000 Subject: [PATCH 15/23] A misfit on a boundary, and a parameter that enters through one MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit dual_on(..., boundary=) assembles the facet load int_Gamma v phi_j as a natural condition of the load assembler with zero volume templates; misfit_duals and gradient take the same boundary= and integrate J over it with BdIntegral. So surface observations — the uplift rate along a free top — are a boundary integral in the misfit, not a band beneath the surface. A gradient part on a boundary is not assembled yet and raises. sensitivity adds the facet part of mu^T dR/dm for every natural condition whose expression carries the parameter: a prescribed traction or flux is differentiated on its boundary and paired with mu there. test_0024 checks both against central differences on Stokes with a free top: the viscosity through a misfit on the top's uplift rate, and a traction amplitude on the top through a volume misfit. Twenty adjoint tests pass. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- src/underworld3/adjoint.py | 74 +++++++++++----- .../cython/petsc_generic_snes_solvers.pyx | 23 ++++- ...oundary_misfit_and_boundary_sensitivity.py | 85 +++++++++++++++++++ 3 files changed, 159 insertions(+), 23 deletions(-) create mode 100644 tests/test_0024_adjoint_boundary_misfit_and_boundary_sensitivity.py diff --git a/src/underworld3/adjoint.py b/src/underworld3/adjoint.py index 83f71f41a..ac97f2d20 100644 --- a/src/underworld3/adjoint.py +++ b/src/underworld3/adjoint.py @@ -127,7 +127,7 @@ def assembler(self, like): _shared_scratch = _Scratch() -def dual_on(variable, value, grad=None, scratch=None): +def dual_on(variable, value, grad=None, scratch=None, boundary=None): r"""The dual of a load on ``variable``'s space, held as a field. :math:`b_j = \int v\,\phi_j + \mathbf g\cdot\nabla\phi_j` for every basis @@ -139,18 +139,40 @@ def dual_on(variable, value, grad=None, scratch=None): solve and no integration by parts. The returned field comes from ``scratch`` (a :class:`_Scratch` pool; the module's shared one by default) — give it back with ``scratch.give(field)`` when done. + + With ``boundary`` (a mesh boundary label), the load is the facet + integral :math:`b_j = \int_\Gamma v\,\phi_j` instead — a misfit on + surface observations — assembled as a natural condition of the same + generic solver with zero volume templates. A gradient part on a boundary + is not assembled yet and raises. """ scratch = _shared_scratch if scratch is None else scratch mesh = variable.mesh n = getattr(variable, "num_components", 1) dim, cdim = mesh.dim, mesh.cdim asm = scratch.assembler(variable) - if grad is None: - grad = sympy.zeros(1, cdim) if n == 1 else sympy.zeros(dim, cdim) - asm._g0 = value - asm._g1 = sympy.Matrix(grad) - asm._needs_function_rewire = True # the templates re-evaluate - asm._build(False, False, None) + if boundary is not None: + if grad is not None and not sympy.Matrix(grad).is_zero_matrix: + raise NotImplementedError( + "dual_on: a load read through the gradient on a boundary is " + "not assembled yet; only the value part is") + # zero volume templates; the natural condition IS the load + asm._g0 = sympy.zeros(1, 1) if n == 1 else sympy.zeros(1, dim) + asm._g1 = sympy.zeros(1, cdim) if n == 1 else sympy.zeros(dim, cdim) + asm.natural_bcs.clear() + asm.add_natural_bc(value, boundary) + asm.is_setup = False # the facet kernels are registered on build + asm._build(False, False, None) + else: + if grad is None: + grad = sympy.zeros(1, cdim) if n == 1 else sympy.zeros(dim, cdim) + if asm.natural_bcs: + asm.natural_bcs.clear() + asm.is_setup = False + asm._g0 = value + asm._g1 = sympy.Matrix(grad) + asm._needs_function_rewire = True # the templates re-evaluate + asm._build(False, False, None) out_var = scratch.take(variable) gvec = asm.dm.getGlobalVec() gvec.set(0.0) @@ -211,16 +233,17 @@ def _token_of(var): return text -def misfit_duals(misfit, variables, scratch=None): +def misfit_duals(misfit, variables, scratch=None, boundary=None): r"""``dJ/df`` as a dual field on each field the misfit reads. - ``J = \int misfit`` over the mesh; a field enters through its value and, - for a misfit on a stress or a strain rate, through its gradient. Both - parts are differentiated symbolically and assembled as one load - (:func:`dual_on`), so a misfit written in terms of :math:`\nabla u` needs - no integration by parts by the caller. Returns ``{variable: dual}`` for - the variables that appear; give each dual back to the scratch pool when - done. + ``J = \int misfit`` over the mesh — or over the boundary ``boundary`` + when one is named, for a misfit on surface observations. A field enters + through its value and, for a misfit on a stress or a strain rate, + through its gradient. Both parts are differentiated symbolically and + assembled as one load (:func:`dual_on`), so a misfit written in terms of + :math:`\nabla u` needs no integration by parts by the caller. Returns + ``{variable: dual}`` for the variables that appear; give each dual back + to the scratch pool when done. """ scratch = _shared_scratch if scratch is None else scratch peeled = _peel(misfit) @@ -245,10 +268,17 @@ def misfit_duals(misfit, variables, scratch=None): g1[i, int(m.group(2))] = sympy.diff(peeled, atom) if all(v == 0 for v in value) and (g1 is None or g1.is_zero_matrix): continue - out[var] = dual_on(var, _as_expression(value), g1, scratch) + out[var] = dual_on(var, _as_expression(value), g1, scratch, boundary=boundary) return out +def integral(mesh, expression, boundary=None): + """``float(∫ expression)`` over the mesh, or over ``boundary`` if named.""" + if boundary is None: + return float(uw.maths.Integral(mesh, expression).evaluate()) + return float(uw.maths.BdIntegral(mesh, expression, boundary).evaluate()) + + def _reads_of(solver, unknown, tokens): """What a solver's residual reads, other than its unknown. @@ -303,7 +333,7 @@ def field_duals(solver, mu, variables, scratch=None): return out -def gradient(solver, misfit, parameters=(), fields=(), scratch=None): +def gradient(solver, misfit, parameters=(), fields=(), scratch=None, boundary=None): r"""``dJ/dm`` and the duals on fields, by the adjoint of ONE solve. For :math:`J = \int` ``misfit`` over the mesh, evaluated in the state the @@ -316,6 +346,9 @@ def gradient(solver, misfit, parameters=(), fields=(), scratch=None): — the derivative through a field the residual reads, an initial condition or a coefficient field. + ``boundary`` names a mesh boundary label over which the misfit is + integrated instead of the volume: surface observations. + Returns ``{"J": float, "parameters": {expr: float}, "fields": {var: dual}}``; the duals are NumPy copies, safe to keep. """ @@ -323,12 +356,13 @@ def gradient(solver, misfit, parameters=(), fields=(), scratch=None): parameters, fields = list(parameters), list(fields) u = solver.u mesh = u.mesh - J = float(uw.maths.Integral(mesh, misfit).evaluate()) - duals = misfit_duals(misfit, [u] + [f for f in fields if f is not u], scratch) + J = integral(mesh, misfit, boundary) + duals = misfit_duals(misfit, [u] + [f for f in fields if f is not u], scratch, + boundary=boundary) grad = {} for p in parameters: explicit = sympy.diff(_peel_except(misfit, p), p) - grad[p] = 0.0 if explicit == 0 else float(uw.maths.Integral(mesh, explicit).evaluate()) + grad[p] = 0.0 if explicit == 0 else integral(mesh, explicit, boundary) # The explicit part on a field the misfit reads directly — never on the # unknown, whose misfit dual is the adjoint's right-hand side. out_fields = {var: (np.array(duals[var].array, copy=True) if (var in duals and var is not u) diff --git a/src/underworld3/cython/petsc_generic_snes_solvers.pyx b/src/underworld3/cython/petsc_generic_snes_solvers.pyx index b788ffaaf..5c4299b03 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -1788,7 +1788,7 @@ class SolverBaseClass(uw_object): out = out + d0[i] * mu_sym[i] return out + uw.maths.tensor.rank2_inner_product(d1, grad_mu) - def gradient(self, misfit, parameters=(), fields=()): + def gradient(self, misfit, parameters=(), fields=(), boundary=None): r"""``dJ/dm`` for each parameter, and the dual on each field, by the adjoint of this solve. @@ -1799,7 +1799,8 @@ class SolverBaseClass(uw_object): Returns ``{"J", "parameters": {expr: dJ/dm}, "fields": {var: dual}}``. """ from underworld3.adjoint import gradient as _gradient - return _gradient(self, misfit, parameters=parameters, fields=fields) + return _gradient(self, misfit, parameters=parameters, fields=fields, + boundary=boundary) def adjoint_kernels(self): """The pointwise kernels of the adjoint operator, as SymPy matrices. @@ -2066,7 +2067,23 @@ class SolverBaseClass(uw_object): implicit part of the gradient; add :math:`\partial J/\partial m` if the misfit depends on the parameter directly. """ - return float(uw.maths.Integral(self.mesh, self.adjoint_integrand(mu, wrt)).evaluate()) + total = float(uw.maths.Integral(self.mesh, self.adjoint_integrand(mu, wrt)).evaluate()) + # A parameter that enters through a natural condition — a prescribed + # traction or flux — has a facet part: (d bd_F0 / dm) . mu on that boundary. + import sympy + mu_sym = mu.sym + for bc in (getattr(self, "natural_bcs", None) or []): + fn = getattr(bc, "fn_f", None) + if fn is None: + continue + d = sympy.diff(self._peel_except(sympy.Matrix(fn), wrt), wrt) + if d.is_zero_matrix: + continue + values = list(d) + comps = getattr(mu, "num_components", 1) + integrand = sum(values[i] * mu_sym[i] for i in range(min(len(values), comps))) + total += float(uw.maths.BdIntegral(self.mesh, integrand, bc.boundary).evaluate()) + return total def _constraint_mechanisms(self): """Every way a constraint can have been put on this solver. diff --git a/tests/test_0024_adjoint_boundary_misfit_and_boundary_sensitivity.py b/tests/test_0024_adjoint_boundary_misfit_and_boundary_sensitivity.py new file mode 100644 index 000000000..62f761942 --- /dev/null +++ b/tests/test_0024_adjoint_boundary_misfit_and_boundary_sensitivity.py @@ -0,0 +1,85 @@ +"""Surface observations, and a parameter that lives on a boundary. + +Two things an adjoint paper derives by hand and this machinery must get +without derivation. A misfit integrated over a boundary — the uplift rate +along a free surface — whose dual is a facet load. And a parameter that +enters the residual through a natural condition — the traction on a +boundary — whose sensitivity has a facet part. Both against central +differences, on a Stokes problem small enough to run in seconds. +""" +import numpy as np +import pytest +import sympy + +import underworld3 as uw + + +def _box(): + mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0, 0), maxCoords=(1, 1), + cellSize=1 / 6, qdegree=3) + v = uw.discretisation.MeshVariable("v", mesh, 2, degree=2) + p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1, continuous=True) + v_obs = uw.discretisation.MeshVariable("v_obs", mesh, 2, degree=2) + return mesh, v, p, v_obs + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_misfit_on_the_free_surface_matches_finite_differences(): + mesh, v, p, v_obs = _box() + x, y = mesh.X + eta = uw.expression(r"\eta", 1.0, "viscosity") + stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = eta * (1 + 2 * (x - 0.5) ** 2) + stokes.bodyforce = sympy.Matrix([0, -sympy.sin(3 * x)]) + stokes.add_essential_bc((0.0, 0.0), "Bottom") + stokes.add_essential_bc((0.3, None), "Left") + stokes.add_essential_bc((-0.3, None), "Right") # free top + stokes.tolerance = 1e-10 # the gradient is small; the difference must be clean + + misfit = (v.sym[1] - v_obs.sym[1]) ** 2 / 2 # uplift rate along the top + + def J(value): + eta.sym = float(value) + stokes.solve(zero_init_guess=True) + return float(uw.maths.BdIntegral(mesh, misfit, "Top").evaluate()) + + J(2.5); v_obs.array[...] = np.asarray(v.array) + J0 = J(1.0) + out = stokes.gradient(misfit, parameters=[eta], boundary="Top") + assert abs(out["J"] - J0) < 1e-12 + h = 1e-4 + fd = (J(1.0 + h) - J(1.0 - h)) / (2 * h) + assert abs(fd / out["parameters"][eta] - 1) < 1e-3, (fd, out["parameters"][eta]) + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_parameter_in_a_traction_condition_matches_finite_differences(): + mesh, v, p, v_obs = _box() + x, y = mesh.X + tau = uw.expression(r"\tau", 0.5, "traction on the top") + stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = 1 + 2 * (x - 0.5) ** 2 + stokes.add_essential_bc((0.0, 0.0), "Bottom") + stokes.add_essential_bc((0.0, None), "Left") + stokes.add_essential_bc((0.0, None), "Right") + stokes.add_natural_bc(sympy.Matrix([[tau * sympy.sin(sympy.pi * x), 0.0]]), "Top") + stokes.tolerance = 1e-10 + + misfit = (v.sym[0] ** 2 + v.sym[1] ** 2) / 2 # a volume misfit + + def J(value): + tau.sym = float(value) + stokes.solve(zero_init_guess=True) + return float(uw.maths.Integral(mesh, misfit).evaluate()) + + J0 = J(0.5) + out = stokes.gradient(misfit, parameters=[tau]) + assert abs(out["J"] - J0) < 1e-12 + h = 1e-4 + fd = (J(0.5 + h) - J(0.5 - h)) / (2 * h) + assert out["parameters"][tau] != 0.0 + assert abs(fd / out["parameters"][tau] - 1) < 1e-3, (fd, out["parameters"][tau]) From 3ccc2b63efb51c9cfeec14664675a109b5bd93d4 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 17 Sep 2026 21:58:43 +1000 Subject: [PATCH 16/23] A misfit with terms on several domains; the fault example's surface terms are boundary integrals gradient() takes {None: volume integrand, "Top": surface integrand}; J, the duals and the explicit parameter parts are sums over the terms. dual_on(boundary=) assembles a gradient part on a facet for a vector space through the natural condition's flux slot, the one a Nitsche condition uses for its symmetry term, so a stress orientation read on the surface is a true boundary integral too. The fault example's uplift and surface-orientation terms are boundary integrals over "Top" rather than a Gaussian band beneath it; the check reads 1.00000 on every segment for all three observation sets. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- .../adjoint/fault_segments/fault_friction.py | 21 ++++++--- src/underworld3/adjoint.py | 43 ++++++++++++++----- 2 files changed, 46 insertions(+), 18 deletions(-) diff --git a/docs/examples/adjoint/fault_segments/fault_friction.py b/docs/examples/adjoint/fault_segments/fault_friction.py index 78b9e364a..1e7503ca8 100644 --- a/docs/examples/adjoint/fault_segments/fault_friction.py +++ b/docs/examples/adjoint/fault_segments/fault_friction.py @@ -128,7 +128,6 @@ def segment(k): stokes.add_essential_bc((-0.5, None), "Right") # --- the observations ------------------------------------------------------ -w_top = sympy.exp(-((1 - y) / params.band) ** 2) points = [(0.3, 0.65), (0.75, 0.12), (1.25, 0.3), (0.9, 0.6), (1.55, 0.85)] w_points = sum(sympy.exp(-((x - px) ** 2 + (y - py) ** 2) / (2 * params.band) ** 2) for px, py in points) @@ -146,14 +145,22 @@ def orientation(field): norm = sympy.sqrt(a ** 2 + b ** 2 + uw.maths.functions.vanishing) return sympy.Matrix([[a / norm, b / norm]]) +# The misfit has a term on the top surface — a true boundary integral of +# the uplift rate, or of the stress orientation — and a term in the volume +# around the stress points. gradient() takes them as {domain: integrand}. what = str(params.observations) +dq = orientation(v) - orientation(v_obs) if what == "uplift+stress": - misfit = (w_top * (v.sym[1] - v_obs.sym[1]) ** 2 - + w_points * (shear_stress(v) - shear_stress(v_obs)) ** 2) / 2 + misfit = {"Top": (v.sym[1] - v_obs.sym[1]) ** 2 / 2, + None: w_points * (shear_stress(v) - shear_stress(v_obs)) ** 2 / 2} +elif what == "orientation": + misfit = {"Top": (dq[0] ** 2 + dq[1] ** 2) / 2, + None: w_points * (dq[0] ** 2 + dq[1] ** 2) / 2} else: - dq = orientation(v) - orientation(v_obs) - weight = w_top if what == "orientation_surface" else w_top + w_points - misfit = weight * (dq[0] ** 2 + dq[1] ** 2) / 2 + misfit = {"Top": (dq[0] ** 2 + dq[1] ** 2) / 2} + +def misfit_value(): + return sum(uw.adjoint.integral(mesh, term, where) for where, term in misfit.items()) model = uw.get_default_model() @@ -202,7 +209,7 @@ def forward(label): vals[k] = math.exp(base + sign * h) set_strengths(vals) forward(f"fd {names[k]} {'+' if sign > 0 else '-'}h") - fd.append(float(uw.maths.Integral(mesh, misfit).evaluate())) + fd.append(misfit_value()) fd = (fd[0] - fd[1]) / (2 * h) uw.pprint(f"{names[k]:>13}: adjoint {g0[k]: .6e} finite difference {fd: .6e} " f"ratio {fd / g0[k]:.5f}") diff --git a/src/underworld3/adjoint.py b/src/underworld3/adjoint.py index ac97f2d20..2dcd1af2f 100644 --- a/src/underworld3/adjoint.py +++ b/src/underworld3/adjoint.py @@ -152,15 +152,22 @@ def dual_on(variable, value, grad=None, scratch=None, boundary=None): dim, cdim = mesh.dim, mesh.cdim asm = scratch.assembler(variable) if boundary is not None: - if grad is not None and not sympy.Matrix(grad).is_zero_matrix: + has_grad = grad is not None and not sympy.Matrix(grad).is_zero_matrix + if has_grad and n == 1: raise NotImplementedError( - "dual_on: a load read through the gradient on a boundary is " - "not assembled yet; only the value part is") + "dual_on: a scalar load read through the gradient on a boundary " + "is not assembled yet (the scalar assembler has no facet flux " + "term); a vector one is") # zero volume templates; the natural condition IS the load asm._g0 = sympy.zeros(1, 1) if n == 1 else sympy.zeros(1, dim) asm._g1 = sympy.zeros(1, cdim) if n == 1 else sympy.zeros(dim, cdim) asm.natural_bcs.clear() asm.add_natural_bc(value, boundary) + if has_grad: + # the facet flux part int_Gamma g . grad(phi_j): the same slot a + # Nitsche condition uses for its symmetry term + bc = asm.natural_bcs[-1] + asm.natural_bcs[-1] = bc._replace(fn_F=sympy.Matrix(grad).as_immutable()) asm.is_setup = False # the facet kernels are registered on build asm._build(False, False, None) else: @@ -347,7 +354,9 @@ def gradient(solver, misfit, parameters=(), fields=(), scratch=None, boundary=No condition or a coefficient field. ``boundary`` names a mesh boundary label over which the misfit is - integrated instead of the volume: surface observations. + integrated instead of the volume: surface observations. A misfit with + terms on several domains is a dict ``{None: volume integrand, "Top": + surface integrand}``; then ``boundary`` is ignored. Returns ``{"J": float, "parameters": {expr: float}, "fields": {var: dual}}``; the duals are NumPy copies, safe to keep. @@ -356,13 +365,25 @@ def gradient(solver, misfit, parameters=(), fields=(), scratch=None, boundary=No parameters, fields = list(parameters), list(fields) u = solver.u mesh = u.mesh - J = integral(mesh, misfit, boundary) - duals = misfit_duals(misfit, [u] + [f for f in fields if f is not u], scratch, - boundary=boundary) - grad = {} - for p in parameters: - explicit = sympy.diff(_peel_except(misfit, p), p) - grad[p] = 0.0 if explicit == 0 else integral(mesh, explicit, boundary) + # One misfit, or several terms on different domains: {None: volume + # integrand, "Top": surface integrand, ...}. J and every dual are sums. + terms = dict(misfit) if isinstance(misfit, dict) else {boundary: misfit} + J = 0.0 + duals = {} + grad = {p: 0.0 for p in parameters} + read = [u] + [f for f in fields if f is not u] + for where, term in terms.items(): + J += integral(mesh, term, where) + for var, dual in misfit_duals(term, read, scratch, boundary=where).items(): + if var in duals: + duals[var].array[...] = np.asarray(duals[var].array) + np.asarray(dual.array) + scratch.give(dual) + else: + duals[var] = dual + for p in parameters: + explicit = sympy.diff(_peel_except(term, p), p) + if explicit != 0: + grad[p] += integral(mesh, explicit, where) # The explicit part on a field the misfit reads directly — never on the # unknown, whose misfit dual is the adjoint's right-hand side. out_fields = {var: (np.array(duals[var].array, copy=True) if (var in duals and var is not u) From c3adaf8ab0f3ccfaa1e587aad23a554aa28eeb4a Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 17 Sep 2026 23:34:16 +1000 Subject: [PATCH 17/23] fault_friction: the optimiser sees the misfit relative to its starting value L-BFGS-B stops on the absolute decrease of its objective, and a surface integral of a velocity misfit is a small number: unscaled, the surface-only orientation inversion stopped after seven evaluations with the deep segments untouched. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- docs/examples/adjoint/fault_segments/fault_friction.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/docs/examples/adjoint/fault_segments/fault_friction.py b/docs/examples/adjoint/fault_segments/fault_friction.py index 1e7503ca8..ca680c384 100644 --- a/docs/examples/adjoint/fault_segments/fault_friction.py +++ b/docs/examples/adjoint/fault_segments/fault_friction.py @@ -223,12 +223,17 @@ def forward(label): history = [] +# The optimiser sees the misfit relative to its starting value: L-BFGS-B stops +# on the absolute decrease of its objective, and a surface integral of a +# velocity misfit is a small number. +J_scale = J0 + def objective(log_eta): set_strengths(np.exp(log_eta)) J, grad = J_and_gradient() history.append((J, np.exp(log_eta).copy())) uw.pprint(f" J = {J:.6e} strengths = {np.exp(log_eta)}") - return J, grad + return J / J_scale, grad / J_scale result = minimize(objective, np.log([params.initial_strength] * n_seg), jac=True, method="L-BFGS-B", options={"maxiter": 40, "gtol": 1e-10}) From 503989c0cf89da14d0c4f6595d688a32c3f97569 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Fri, 18 Sep 2026 07:28:25 +1000 Subject: [PATCH 18/23] fault_friction: orientation is an interior observable; the surface gives the strain rate On a traction-free surface the shear strain rate vanishes, so a stress orientation read there is a sign, and the surface-only orientation case had been fitting the discrete strain rate's departure from that. The observation sets are now uplift+stress, the principal-stress orientation at the five interior points only, and the surface strain rate dv_x/dx along the top only. All three recover the four coefficients to five figures (22, 19 and 40 evaluations); the checks read 1.00000. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- .../adjoint/fault_segments/fault_friction.py | 19 ++++++++++++------- .../fault_segments/plot_convergence.py | 4 ++-- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/docs/examples/adjoint/fault_segments/fault_friction.py b/docs/examples/adjoint/fault_segments/fault_friction.py index ca680c384..0720d02ee 100644 --- a/docs/examples/adjoint/fault_segments/fault_friction.py +++ b/docs/examples/adjoint/fault_segments/fault_friction.py @@ -42,9 +42,9 @@ rho_g=uw.Param(10.0, "body force, so the pressure grows with depth"), check_only=uw.Param(0, "1: gradient check against finite differences, no inversion"), observations=uw.Param("uplift+stress", - "uplift+stress | orientation (principal-stress orientation at the " - "points and along the surface) | orientation_surface (along the " - "surface only)"), + "uplift+stress | orientation_points (principal-stress orientation " + "at the five interior points only) | surface_strain (the surface " + "strain rate d v_x / d x along the top only)"), ) # --- the model --------------------------------------------------------------- @@ -148,16 +148,21 @@ def orientation(field): # The misfit has a term on the top surface — a true boundary integral of # the uplift rate, or of the stress orientation — and a term in the volume # around the stress points. gradient() takes them as {domain: integrand}. +# On a traction-free surface the shear strain rate vanishes, so a stress +# orientation read there is only a sign; orientation is an interior +# observable (boreholes, focal mechanisms), and the surface gives velocities +# and their tangential derivative — the geodetic strain rate. what = str(params.observations) dq = orientation(v) - orientation(v_obs) if what == "uplift+stress": misfit = {"Top": (v.sym[1] - v_obs.sym[1]) ** 2 / 2, None: w_points * (shear_stress(v) - shear_stress(v_obs)) ** 2 / 2} -elif what == "orientation": - misfit = {"Top": (dq[0] ** 2 + dq[1] ** 2) / 2, - None: w_points * (dq[0] ** 2 + dq[1] ** 2) / 2} +elif what == "orientation_points": + misfit = {None: w_points * (dq[0] ** 2 + dq[1] ** 2) / 2} +elif what == "surface_strain": + misfit = {"Top": (v.sym[0].diff(x) - v_obs.sym[0].diff(x)) ** 2 / 2} else: - misfit = {"Top": (dq[0] ** 2 + dq[1] ** 2) / 2} + raise ValueError(f"observations: {what!r}") def misfit_value(): return sum(uw.adjoint.integral(mesh, term, where) for where, term in misfit.items()) diff --git a/docs/examples/adjoint/fault_segments/plot_convergence.py b/docs/examples/adjoint/fault_segments/plot_convergence.py index 655fabf23..7b458006c 100644 --- a/docs/examples/adjoint/fault_segments/plot_convergence.py +++ b/docs/examples/adjoint/fault_segments/plot_convergence.py @@ -5,8 +5,8 @@ import matplotlib.pyplot as plt cases = [("uplift + stress", "fault_friction_uplift+stress_data.npz"), - ("orientation, points + surface", "fault_friction_orientation_data.npz"), - ("orientation, surface only", "fault_friction_orientation_surface_data.npz")] + ("stress orientation at five interior points", "fault_friction_orientation_points_data.npz"), + ("surface strain rate only", "fault_friction_surface_strain_data.npz")] names = ["flat", "lower ramp", "upper ramp", "near surface"] fig, axes = plt.subplots(2, 3, figsize=(11, 5.6), sharex="col", From ab4b7b0723a97aa1fe1b2a99dfda35d8855b1e28 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Fri, 18 Sep 2026 12:19:13 +1000 Subject: [PATCH 19/23] uw.adjoint.minimise: PETSc TAO as the inversion driver objective(x) -> (J, dJ/dx) driven by TAO's limited-memory quasi-Newton method (lmvm, or blmvm with bounds), its own line search choosing the step. A few scalar controls are replicated on every rank; the objective's collective solves keep the ranks in step. Returns the solution, the iteration count, the converged reason and the (J, x) history. The fault example takes -uw_optimiser tao|scipy and defaults to tao. test_0025 checks the driver on a quadratic, bounded and not. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- .../adjoint/fault_segments/fault_friction.py | 17 +++++-- src/underworld3/adjoint.py | 51 +++++++++++++++++++ tests/test_0025_tao_driver.py | 29 +++++++++++ 3 files changed, 93 insertions(+), 4 deletions(-) create mode 100644 tests/test_0025_tao_driver.py diff --git a/docs/examples/adjoint/fault_segments/fault_friction.py b/docs/examples/adjoint/fault_segments/fault_friction.py index 0720d02ee..20b420043 100644 --- a/docs/examples/adjoint/fault_segments/fault_friction.py +++ b/docs/examples/adjoint/fault_segments/fault_friction.py @@ -41,6 +41,7 @@ cohesion=uw.Param(0.05, "cohesion C in tau_y = C + mu p"), rho_g=uw.Param(10.0, "body force, so the pressure grows with depth"), check_only=uw.Param(0, "1: gradient check against finite differences, no inversion"), + optimiser=uw.Param("tao", "tao (PETSc, limited-memory quasi-Newton) | scipy (L-BFGS-B)"), observations=uw.Param("uplift+stress", "uplift+stress | orientation_points (principal-stress orientation " "at the five interior points only) | surface_strain (the surface " @@ -240,9 +241,17 @@ def objective(log_eta): uw.pprint(f" J = {J:.6e} strengths = {np.exp(log_eta)}") return J / J_scale, grad / J_scale -result = minimize(objective, np.log([params.initial_strength] * n_seg), jac=True, - method="L-BFGS-B", options={"maxiter": 40, "gtol": 1e-10}) -uw.pprint(f"recovered {np.exp(result.x)} true {true_values} " +x0 = np.log([params.initial_strength] * n_seg) +if str(params.optimiser) == "tao": + # PETSc's own driver: the same objective and gradient, TAO's quasi-Newton + # update and line search. + x_best, info = uw.adjoint.minimise(objective, x0, max_evaluations=60, + gradient_tolerance=1e-10) +else: + result = minimize(objective, x0, jac=True, method="L-BFGS-B", + options={"maxiter": 40, "gtol": 1e-10}) + x_best = result.x +uw.pprint(f"recovered {np.exp(x_best)} true {true_values} " f"after {len(history)} evaluations") # --- what the figure needs ----------------------------------------------------- @@ -252,7 +261,7 @@ def objective(log_eta): top = np.column_stack([xs, np.full_like(xs, 1.0 - 1e-6)]) profiles = {} for label, values in (("true", true_values), ("initial", [params.initial_strength] * n_seg), - ("recovered", list(np.exp(result.x)))): + ("recovered", list(np.exp(x_best)))): set_strengths(values) forward(f"profile {label}") profiles[label] = np.asarray(uw.function.evaluate(v.sym[1], top)).ravel() diff --git a/src/underworld3/adjoint.py b/src/underworld3/adjoint.py index 2dcd1af2f..d15da2ebe 100644 --- a/src/underworld3/adjoint.py +++ b/src/underworld3/adjoint.py @@ -428,6 +428,57 @@ def gradient(solver, misfit, parameters=(), fields=(), scratch=None, boundary=No return {"J": J, "parameters": grad, "fields": out_fields} +def minimise(objective, x0, bounds=None, method="lmvm", options=None, max_evaluations=100, + gradient_tolerance=1e-8, callback=None): + r"""Minimise ``objective(x) -> (J, dJ/dx)`` with PETSc TAO. + + The driver for an inversion: the misfit and its gradient come from the + adjoint, the step along the gradient is the line search's, and the + quasi-Newton update is TAO's limited-memory one (``"lmvm"``; ``"blmvm"`` + honours ``bounds``, a pair of arrays). ``x0`` is a NumPy array; a few + scalar controls are replicated on every rank, each rank's TAO doing the + same arithmetic on the same numbers, and the objective's own collective + solves keep the ranks in step. Returns ``(x, info)`` with the iterations, + the converged reason and the history of ``(J, x)`` per evaluation. + """ + from petsc4py import PETSc + x0 = np.asarray(x0, dtype=float).ravel() + x = PETSc.Vec().createSeq(x0.size, comm=PETSc.COMM_SELF) + x.setArray(x0) + history = [] + + def fg(tao, xv, g): + values = np.array(xv.getArray(readonly=True), copy=True) + J, grad = objective(values) + g.setArray(np.asarray(grad, dtype=float).ravel()) + history.append((float(J), values)) + if callback is not None: + callback(J, values) + return float(J) + + tao = PETSc.TAO().create(comm=PETSc.COMM_SELF) + tao.setType(method) + tao.setObjectiveGradient(fg, None) + if bounds is not None: + lo, hi = bounds + lower = x.duplicate(); lower.setArray(np.asarray(lo, dtype=float).ravel()) + upper = x.duplicate(); upper.setArray(np.asarray(hi, dtype=float).ravel()) + tao.setVariableBounds(lower, upper) + tao.setMaximumFunctionEvaluations(int(max_evaluations)) + tao.setTolerances(gatol=gradient_tolerance) + for key, value in (options or {}).items(): + PETSc.Options().setValue(key, value) + tao.setFromOptions() + tao.setSolution(x) + tao.solve() + info = {"iterations": int(tao.getIterationNumber()), + "reason": int(tao.getConvergedReason()), + "history": history} + out = np.array(x.getArray(readonly=True), copy=True) + tao.destroy() + return out, info + + _n = [0] diff --git a/tests/test_0025_tao_driver.py b/tests/test_0025_tao_driver.py new file mode 100644 index 000000000..beea9c280 --- /dev/null +++ b/tests/test_0025_tao_driver.py @@ -0,0 +1,29 @@ +"""uw.adjoint.minimise drives an objective with PETSc TAO: a quadratic, bounded and not.""" +import numpy as np +import pytest + +import underworld3 as uw + + +def _quadratic(x): + c = np.array([1.0, -0.5, 2.0]) + return float(np.sum((x - c) ** 2 * [1, 2, 3])), 2 * (x - c) * [1, 2, 3] + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_lmvm_finds_the_minimum(): + x, info = uw.adjoint.minimise(_quadratic, np.zeros(3), gradient_tolerance=1e-12) + assert np.allclose(x, [1.0, -0.5, 2.0], atol=1e-6) + assert info["reason"] > 0 + assert len(info["history"]) == len(info["history"]) and info["history"][0][0] > info["history"][-1][0] + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_blmvm_honours_bounds(): + x, info = uw.adjoint.minimise(_quadratic, np.zeros(3), method="blmvm", + bounds=(np.array([0.0, 0.0, 0.0]), np.array([0.5, 1.0, 1.0])), + gradient_tolerance=1e-12) + assert np.allclose(x, [0.5, 0.0, 1.0], atol=1e-6) + assert info["reason"] > 0 From 53417bbc205bec413f65e3b757595304a91a276f Mon Sep 17 00:00:00 2001 From: lmoresi Date: Fri, 18 Sep 2026 12:39:43 +1000 Subject: [PATCH 20/23] fault_friction: scipy stays the default driver, since the note's figures come from it; tao is the switch Under TAO the uplift-and-stress inversion recovers (0.05, 0.15, 0.25, 0.4) to machine precision in 26 evaluations against L-BFGS-B's 22 to five figures; the first trial step overshoots in log space and the line search recovers. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- docs/examples/adjoint/fault_segments/fault_friction.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/examples/adjoint/fault_segments/fault_friction.py b/docs/examples/adjoint/fault_segments/fault_friction.py index 20b420043..82e461fea 100644 --- a/docs/examples/adjoint/fault_segments/fault_friction.py +++ b/docs/examples/adjoint/fault_segments/fault_friction.py @@ -41,7 +41,7 @@ cohesion=uw.Param(0.05, "cohesion C in tau_y = C + mu p"), rho_g=uw.Param(10.0, "body force, so the pressure grows with depth"), check_only=uw.Param(0, "1: gradient check against finite differences, no inversion"), - optimiser=uw.Param("tao", "tao (PETSc, limited-memory quasi-Newton) | scipy (L-BFGS-B)"), + optimiser=uw.Param("scipy", "scipy (L-BFGS-B; the figures in the note are from it) | tao (PETSc, limited-memory quasi-Newton)"), observations=uw.Param("uplift+stress", "uplift+stress | orientation_points (principal-stress orientation " "at the five interior points only) | surface_strain (the surface " From 183ea41288b81a4f594c546d30f600ef5caad3af Mon Sep 17 00:00:00 2001 From: lmoresi Date: Fri, 18 Sep 2026 14:51:33 +1000 Subject: [PATCH 21/23] fault_friction: noise on the observations, a Tikhonov term, and bounds for TAO's blmvm -uw_noise adds Gaussian noise to the observed velocity field at a fraction of each component's rms, drawn once; every observation set inherits it. -uw_regularisation adds alpha (log mu - log mu_start)^2 to the scaled objective; -uw_bounds lo,hi switches the TAO driver to blmvm on the log-strengths. TAO is the default driver. plot_noise.py draws the recovered friction against the noise level. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- .../adjoint/fault_segments/fault_friction.py | 34 +++++++++++++++--- .../adjoint/fault_segments/plot_noise.py | 36 +++++++++++++++++++ 2 files changed, 65 insertions(+), 5 deletions(-) create mode 100644 docs/examples/adjoint/fault_segments/plot_noise.py diff --git a/docs/examples/adjoint/fault_segments/fault_friction.py b/docs/examples/adjoint/fault_segments/fault_friction.py index 82e461fea..2fe9ef561 100644 --- a/docs/examples/adjoint/fault_segments/fault_friction.py +++ b/docs/examples/adjoint/fault_segments/fault_friction.py @@ -41,7 +41,11 @@ cohesion=uw.Param(0.05, "cohesion C in tau_y = C + mu p"), rho_g=uw.Param(10.0, "body force, so the pressure grows with depth"), check_only=uw.Param(0, "1: gradient check against finite differences, no inversion"), - optimiser=uw.Param("scipy", "scipy (L-BFGS-B; the figures in the note are from it) | tao (PETSc, limited-memory quasi-Newton)"), + optimiser=uw.Param("tao", "tao (PETSc, limited-memory quasi-Newton; blmvm when bounds are set) | scipy (L-BFGS-B)"), + noise=uw.Param(0.0, "Gaussian noise on the observed velocity, as a fraction of its rms, per component"), + seed=uw.Param(7, "seed for the noise"), + regularisation=uw.Param(0.0, "Tikhonov weight on (log mu - log mu_start)^2, relative to the starting misfit"), + bounds=uw.Param("", "friction bounds lo,hi for TAO's blmvm; empty for none"), observations=uw.Param("uplift+stress", "uplift+stress | orientation_points (principal-stress orientation " "at the five interior points only) | surface_strain (the surface " @@ -199,6 +203,15 @@ def forward(label): set_strengths(true_values) forward("truth") v_obs.array[...] = np.asarray(v.array) +if float(params.noise) > 0: + # Noise on the observed velocity field, drawn once, at a fraction of each + # component's rms. Every observation set reads v_obs, so the uplift, the + # stress and the strain rate all inherit it. + rng = np.random.default_rng(int(params.seed)) + obs = np.asarray(v_obs.array) + rms = np.sqrt(np.mean(obs ** 2, axis=0, keepdims=True)) + v_obs.array[...] = obs + float(params.noise) * rms * rng.standard_normal(obs.shape) + uw.pprint(f"noise: {float(params.noise):.3f} of the rms per component, seed {int(params.seed)}") uw.pprint(f"true strengths {true_values}") set_strengths([params.initial_strength] * n_seg) @@ -234,19 +247,29 @@ def forward(label): # velocity misfit is a small number. J_scale = J0 +alpha = float(params.regularisation) +log_prior = np.log([params.initial_strength] * n_seg) + def objective(log_eta): + """The misfit relative to its start, plus a Tikhonov term on the log-strengths.""" set_strengths(np.exp(log_eta)) J, grad = J_and_gradient() history.append((J, np.exp(log_eta).copy())) uw.pprint(f" J = {J:.6e} strengths = {np.exp(log_eta)}") - return J / J_scale, grad / J_scale + penalty = alpha * np.sum((log_eta - log_prior) ** 2) + return J / J_scale + penalty, grad / J_scale + 2 * alpha * (log_eta - log_prior) x0 = np.log([params.initial_strength] * n_seg) if str(params.optimiser) == "tao": # PETSc's own driver: the same objective and gradient, TAO's quasi-Newton - # update and line search. + # update and line search; the bounded variant when bounds are given. + bounds = None + if str(params.bounds).strip(): + lo, hi = (float(b) for b in str(params.bounds).split(",")) + bounds = (np.log([lo] * n_seg), np.log([hi] * n_seg)) x_best, info = uw.adjoint.minimise(objective, x0, max_evaluations=60, - gradient_tolerance=1e-10) + gradient_tolerance=1e-10, bounds=bounds, + method="blmvm" if bounds else "lmvm") else: result = minimize(objective, x0, jac=True, method="L-BFGS-B", options={"maxiter": 40, "gtol": 1e-10}) @@ -270,7 +293,8 @@ def objective(log_eta): gx, gy = np.meshgrid(np.linspace(0, 2, 201), np.linspace(0, 1, 101)) grid = np.column_stack([gx.ravel(), gy.ravel()]) eta_1_grid = np.asarray(uw.function.evaluate(eta_1, grid)).reshape(gx.shape) -np.savez(f"fault_friction_{what}_data.npz", xs=xs, gx=gx, gy=gy, eta_1=eta_1_grid, +tag = f"{what}" + (f"_noise{float(params.noise):g}" if float(params.noise) > 0 else "") + (f"_reg{alpha:g}" if alpha > 0 else "") +np.savez(f"fault_friction_{tag}_data.npz", xs=xs, gx=gx, gy=gy, eta_1=eta_1_grid, points=np.array(points), true=np.array(true_values), band=params.band, history=np.array([[J, *vals] for J, vals in history]), **{f"uplift_{k}": val for k, val in profiles.items()}) diff --git a/docs/examples/adjoint/fault_segments/plot_noise.py b/docs/examples/adjoint/fault_segments/plot_noise.py new file mode 100644 index 000000000..8457a8aa2 --- /dev/null +++ b/docs/examples/adjoint/fault_segments/plot_noise.py @@ -0,0 +1,36 @@ +"""Recovered friction against the noise on the observations, from the data files.""" +import glob +import re + +import numpy as np +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +names = ["flat", "lower ramp", "upper ramp", "near surface"] +runs = [] +for path in sorted(glob.glob("fault_friction_uplift+stress*_data.npz")): + m = re.search(r"_noise([0-9.]+)", path) + level = float(m.group(1)) if m else 0.0 + d = np.load(path) + runs.append((level, d["history"][-1, 1:], d["true"])) +runs.sort(key=lambda r: r[0]) +levels = np.array([r[0] for r in runs]) +true = runs[0][2] + +fig, ax = plt.subplots(figsize=(6.2, 3.6)) +x = np.arange(len(levels)) +w = 0.18 +for k in range(len(names)): + ax.bar(x + (k - 1.5) * w, [r[1][k] for r in runs], width=w, color=f"C{k}", label=names[k]) + ax.hlines(true[k], -0.5, len(levels) - 0.5, colors=f"C{k}", linestyles="--", lw=0.8) +ax.set_yscale("log") +ax.set_xticks(x) +ax.set_xticklabels([f"{l:g}" for l in levels]) +ax.set_xlabel("noise on the observed velocity, fraction of its rms") +ax.set_ylabel("recovered friction (dashed: true)") +ax.legend(fontsize=8, frameon=False, ncol=2) +fig.tight_layout() +fig.savefig("fault_friction_noise.png", dpi=180) +fig.savefig("fault_friction_noise.pdf") +print("wrote fault_friction_noise.png / .pdf") From 3be3f18a21cec16480dc682f44e59260d14674e1 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Fri, 18 Sep 2026 18:44:48 +1000 Subject: [PATCH 22/23] fault_friction: the objective is a negative log posterior, with the prior given by its width The data term is chi-squared/2: the misfit scaled by its expected value at the truth under the noise (read directly in the twin) times the number of independent data, the surface nodes and the nodes under the point weights. The prior term is (log mu - log mu_start)^2 / 2 sigma_m^2 with -uw_prior_sigma in log units. The weight between them is then the noise and the prior, not a number to tune; which coefficients the data move is decided by their sensitivities against it. Replaces the Tikhonov weight on the J/J0-scaled misfit, whose range under ten percent noise was parts in ten thousand, so that any weight above 1e-4 pinned every coefficient at the prior. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- .../adjoint/fault_segments/fault_friction.py | 38 +++++++++++++++---- .../adjoint/fault_segments/plot_noise.py | 24 +++++++----- 2 files changed, 45 insertions(+), 17 deletions(-) diff --git a/docs/examples/adjoint/fault_segments/fault_friction.py b/docs/examples/adjoint/fault_segments/fault_friction.py index 2fe9ef561..e40002a83 100644 --- a/docs/examples/adjoint/fault_segments/fault_friction.py +++ b/docs/examples/adjoint/fault_segments/fault_friction.py @@ -44,7 +44,7 @@ optimiser=uw.Param("tao", "tao (PETSc, limited-memory quasi-Newton; blmvm when bounds are set) | scipy (L-BFGS-B)"), noise=uw.Param(0.0, "Gaussian noise on the observed velocity, as a fraction of its rms, per component"), seed=uw.Param(7, "seed for the noise"), - regularisation=uw.Param(0.0, "Tikhonov weight on (log mu - log mu_start)^2, relative to the starting misfit"), + prior_sigma=uw.Param(0.0, "width of a Gaussian prior on log mu about the start, in log units; 0 for none"), bounds=uw.Param("", "friction bounds lo,hi for TAO's blmvm; empty for none"), observations=uw.Param("uplift+stress", "uplift+stress | orientation_points (principal-stress orientation " @@ -213,7 +213,9 @@ def forward(label): v_obs.array[...] = obs + float(params.noise) * rms * rng.standard_normal(obs.shape) uw.pprint(f"noise: {float(params.noise):.3f} of the rms per component, seed {int(params.seed)}") uw.pprint(f"true strengths {true_values}") +J_truth = None +J_truth = misfit_value() # v still holds the truth's velocity set_strengths([params.initial_strength] * n_seg) J0, g0 = J_and_gradient("start") uw.pprint(f"initial J = {J0:.6e} dJ/dlog eta = {g0}") @@ -247,17 +249,39 @@ def forward(label): # velocity misfit is a small number. J_scale = J0 -alpha = float(params.regularisation) +# The objective as a negative log posterior. The data term is chi-squared/2: +# the misfit scaled by its expected value at the truth under the noise, which +# a twin experiment can read directly, times the number of independent data — +# the surface nodes and the nodes under the point weights. The prior term is +# (log mu - log mu_start)^2 / (2 sigma_m^2) with sigma_m in log units. The +# weight between them is then a statement about the noise and the prior, not +# a number to tune, and which coefficients the data move is decided by their +# sensitivities against that. +sigma_m = float(params.prior_sigma) log_prior = np.log([params.initial_strength] * n_seg) +X = np.asarray(v.coords) +on_surface = X[:, 1] > 1.0 - 1e-6 +near_points = np.zeros(len(X), dtype=bool) +for px, py in points: + near_points |= (X[:, 0] - px) ** 2 + (X[:, 1] - py) ** 2 < (2 * params.band) ** 2 +N_eff = {"uplift+stress": on_surface.sum() + near_points.sum(), + "orientation_points": near_points.sum(), + "surface_strain": on_surface.sum()}[what] +J_floor = J_truth if float(params.noise) > 0 else J0 +chi2_scale = N_eff / J_floor # chi^2 = J * chi2_scale +uw.pprint(f"N_eff = {N_eff}, misfit floor at the truth = {J_floor:.4e}") def objective(log_eta): - """The misfit relative to its start, plus a Tikhonov term on the log-strengths.""" set_strengths(np.exp(log_eta)) J, grad = J_and_gradient() history.append((J, np.exp(log_eta).copy())) - uw.pprint(f" J = {J:.6e} strengths = {np.exp(log_eta)}") - penalty = alpha * np.sum((log_eta - log_prior) ** 2) - return J / J_scale + penalty, grad / J_scale + 2 * alpha * (log_eta - log_prior) + uw.pprint(f" J = {J:.6e} chi2/N = {J * chi2_scale / N_eff:.4f} strengths = {np.exp(log_eta)}") + value = J * chi2_scale / 2 + g = grad * chi2_scale / 2 + if sigma_m > 0: + value += np.sum((log_eta - log_prior) ** 2) / (2 * sigma_m ** 2) + g = g + (log_eta - log_prior) / sigma_m ** 2 + return value, g x0 = np.log([params.initial_strength] * n_seg) if str(params.optimiser) == "tao": @@ -293,7 +317,7 @@ def objective(log_eta): gx, gy = np.meshgrid(np.linspace(0, 2, 201), np.linspace(0, 1, 101)) grid = np.column_stack([gx.ravel(), gy.ravel()]) eta_1_grid = np.asarray(uw.function.evaluate(eta_1, grid)).reshape(gx.shape) -tag = f"{what}" + (f"_noise{float(params.noise):g}" if float(params.noise) > 0 else "") + (f"_reg{alpha:g}" if alpha > 0 else "") +tag = f"{what}" + (f"_noise{float(params.noise):g}" if float(params.noise) > 0 else "") + (f"_prior{sigma_m:g}" if sigma_m > 0 else "") np.savez(f"fault_friction_{tag}_data.npz", xs=xs, gx=gx, gy=gy, eta_1=eta_1_grid, points=np.array(points), true=np.array(true_values), band=params.band, history=np.array([[J, *vals] for J, vals in history]), diff --git a/docs/examples/adjoint/fault_segments/plot_noise.py b/docs/examples/adjoint/fault_segments/plot_noise.py index 8457a8aa2..1d88b9858 100644 --- a/docs/examples/adjoint/fault_segments/plot_noise.py +++ b/docs/examples/adjoint/fault_segments/plot_noise.py @@ -12,23 +12,27 @@ for path in sorted(glob.glob("fault_friction_uplift+stress*_data.npz")): m = re.search(r"_noise([0-9.]+)", path) level = float(m.group(1)) if m else 0.0 + r = re.search(r"_prior([0-9.]+)", path) + alpha = float(r.group(1)) if r else 0.0 d = np.load(path) - runs.append((level, d["history"][-1, 1:], d["true"])) -runs.sort(key=lambda r: r[0]) -levels = np.array([r[0] for r in runs]) -true = runs[0][2] + runs.append((level, alpha, d["history"][-1, 1:], d["true"])) +runs.sort(key=lambda r: (r[0], r[1])) +labels = [f"{r[0]:g}" + (f"\nprior σ={r[1]:g}" if r[1] > 0 else "") for r in runs] +true = runs[0][3] -fig, ax = plt.subplots(figsize=(6.2, 3.6)) -x = np.arange(len(levels)) +fig, ax = plt.subplots(figsize=(7.2, 3.8)) +x = np.arange(len(runs)) w = 0.18 for k in range(len(names)): - ax.bar(x + (k - 1.5) * w, [r[1][k] for r in runs], width=w, color=f"C{k}", label=names[k]) - ax.hlines(true[k], -0.5, len(levels) - 0.5, colors=f"C{k}", linestyles="--", lw=0.8) + ax.bar(x + (k - 1.5) * w, [r[2][k] for r in runs], width=w, color=f"C{k}", label=names[k]) + ax.hlines(true[k], -0.5, len(runs) - 0.5, colors=f"C{k}", linestyles="--", lw=0.8) +ax.axhline(1.0, color="0.3", lw=0.8, ls=":") +ax.axhline(0.005, color="0.3", lw=0.8, ls=":") ax.set_yscale("log") ax.set_xticks(x) -ax.set_xticklabels([f"{l:g}" for l in levels]) +ax.set_xticklabels(labels, fontsize=8) ax.set_xlabel("noise on the observed velocity, fraction of its rms") -ax.set_ylabel("recovered friction (dashed: true)") +ax.set_ylabel("recovered friction (dashed: true; dotted: bounds)") ax.legend(fontsize=8, frameon=False, ncol=2) fig.tight_layout() fig.savefig("fault_friction_noise.png", dpi=180) From b019c7511940c113b89cf9e2cac878616afb940f Mon Sep 17 00:00:00 2001 From: lmoresi Date: Sat, 19 Sep 2026 18:15:54 +1000 Subject: [PATCH 23/23] fault_friction: the example as a notebook, in kilometres, with the run header in the declared units The listric-fault friction example is now a jupytext percent-format script with a notebook made from it, named constants ahead of a uw.Params block, and the orchestration model declared first with three reference quantities: the depth of 10 km, the bulk viscosity of 1e21 Pa s, and the convergence rate, 8.4 mm/yr, fixed by the one number the Coulomb problem depends on, the lithostatic pressure at the base over the viscous stress of the shortening. The nondimensional problem that reaches the solver is the one that ran before, to every printed digit of the misfit and the gradient check, so the figures and the numbers in the note stand; the axes now read in km and mm/yr. The default run was broken: the negative-log-posterior objective went in for the noise study and was never run at zero noise, where its scale is a thousand times larger and TAO's first trial step overflows the exponential. Without noise the objective is J/J0 again, as it was for the published runs, and a prior without a noise level is refused. The run header in the transcript, the figure and the table now reports the reference quantities as they were declared, in their own units, beside the fundamental scales the record keeps; composite pint unit strings abbreviate to "Pa s" and "mm/yr", and a quantity declared as an expression of others is reduced to base units. The units system's status lines print only with verbose=True, and without emoji. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Na7qBenCp67rDTZhFGTh5V --- .../examples/adjoint/fault_segments/README.md | 24 + .../fault_segments/fault_friction.ipynb | 676 ++++++++++++++++++ .../adjoint/fault_segments/fault_friction.py | 537 +++++++++----- .../fault_segments/fault_friction_key.md | 38 + .../fault_friction_transcript.svg | 390 ++++++++++ .../fault_segments/plot_fault_segments.py | 13 +- .../fault_segments/render_fault_friction.py | 7 +- src/underworld3/model.py | 63 +- .../utilities/transcript_report.py | 36 +- 9 files changed, 1556 insertions(+), 228 deletions(-) create mode 100644 docs/examples/adjoint/fault_segments/README.md create mode 100644 docs/examples/adjoint/fault_segments/fault_friction.ipynb create mode 100644 docs/examples/adjoint/fault_segments/fault_friction_key.md create mode 100644 docs/examples/adjoint/fault_segments/fault_friction_transcript.svg diff --git a/docs/examples/adjoint/fault_segments/README.md b/docs/examples/adjoint/fault_segments/README.md new file mode 100644 index 000000000..48d5610f6 --- /dev/null +++ b/docs/examples/adjoint/fault_segments/README.md @@ -0,0 +1,24 @@ +# Friction on a listric fault from surface observations + +The example for the technical note *A Discrete Adjoint from the Run Record* +(UWTN 2026-020). Four friction coefficients on a listric fault are recovered +from the surface uplift rate and the shear stress at five interior points, +with the gradient from the solver's own adjoint. + +| file | what it does | +|---|---| +| `fault_friction.ipynb` | the example as a notebook | +| `fault_friction.py` | the same, as a script in jupytext percent format; `python fault_friction.py -uw_check_only 1` runs the gradient check alone | +| `plot_fault_segments.py` | the figure: the weak plane, the uplift profiles, and the path of the coefficients, from the `_data.npz` a run writes | +| `plot_convergence.py` | convergence under the three observation sets | +| `plot_noise.py` | recovered coefficients against the noise, the bounds and the prior | +| `render_fault_friction.py` | PyVista renders of the truth: the plane's viscosity, the slip rate, the uplift and the pressure | + +The gradient check takes about three minutes and the inversion about ten on +a laptop. The transcript of a run is written to `transcripts//`, +and `uw.transcript_figure(...)` draws it. + +The problem is stated in kilometres, pascal seconds and millimetres per year, +and the solver works in units of the depth, the viscosity and the +convergence rate. The observation sets, the noise and the prior are switches +in the parameter block at the top of the notebook. diff --git a/docs/examples/adjoint/fault_segments/fault_friction.ipynb b/docs/examples/adjoint/fault_segments/fault_friction.ipynb new file mode 100644 index 000000000..e43827190 --- /dev/null +++ b/docs/examples/adjoint/fault_segments/fault_friction.ipynb @@ -0,0 +1,676 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "5a56e0e3", + "metadata": {}, + "source": [ + "# Friction on a listric fault from surface observations\n", + "\n", + "A listric fault under horizontal shortening and gravity: a flat décollement\n", + "at depth that steepens through a circular ramp to a dip of sixty degrees at\n", + "the surface. The fault is a weak plane in a transversely isotropic viscosity,\n", + "with no cut in the mesh. The plane yields at the Coulomb stress\n", + "$\\tau_y = C + \\mu p$, and the friction coefficient $\\mu$ takes a different\n", + "value on the flat, on the lower and upper parts of the ramp, and near the\n", + "surface. Those four coefficients are the unknowns.\n", + "\n", + "The observations are the uplift rate along the top surface and the shear\n", + "stress near a handful of interior points, read from a run at the true\n", + "coefficients. The gradient of the misfit with respect to each coefficient\n", + "comes from one call, `stokes.gradient(misfit, parameters=...)`: the solver\n", + "assembles the adjoint operator from its own Jacobian kernels with trial and\n", + "test functions exchanged, solves it, and differentiates its residual\n", + "symbolically with respect to each named coefficient. Nothing is differenced.\n", + "\n", + "The notebook checks that gradient against central differences, then hands\n", + "the misfit and its gradient to PETSc's TAO to recover the four\n", + "coefficients. Each misfit evaluation is one step of zero length in the\n", + "model's record, so the run's transcript lists the forward solve and the\n", + "adjoint solve it made." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5a51a9d0", + "metadata": {}, + "outputs": [], + "source": [ + "import math\n", + "\n", + "import numpy as np\n", + "import sympy\n", + "\n", + "import underworld3 as uw" + ] + }, + { + "cell_type": "markdown", + "id": "033a1223", + "metadata": {}, + "source": [ + "### Configurable parameters\n", + "\n", + "Default values are defined as named constants below. From the command line,\n", + "override them with PETSc-style flags:\n", + "\n", + "```bash\n", + "python fault_friction.py -uw_check_only 1\n", + "python fault_friction.py -uw_observations surface_strain\n", + "python fault_friction.py -uw_noise 0.03 -uw_prior_sigma 1 -uw_bounds 0.005,1\n", + "```\n", + "\n", + "The problem depends on one dimensionless number: the lithostatic pressure at\n", + "the base of the model over the viscous stress of the shortening,\n", + "$\\rho g L^2 / (\\eta V)$. We fix the depth, the viscosity and the density,\n", + "and that ratio sets the convergence rate." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0d0d83bb", + "metadata": {}, + "outputs": [], + "source": [ + "# --- Scales ---\n", + "DOMAIN_DEPTH = uw.quantity(10, \"km\") # the box is 20 km wide and 10 km deep\n", + "REF_VISCOSITY = uw.quantity(1e21, \"Pa*s\") # the bulk viscosity\n", + "DENSITY = uw.quantity(2700, \"kg/m**3\")\n", + "GRAVITY = uw.quantity(9.81, \"m/s**2\")\n", + "PRESSURE_TO_STRESS = 10 # rho g L^2 / (eta V)\n", + "CONVERGENCE_RATE = (DENSITY * GRAVITY * DOMAIN_DEPTH**2 / (PRESSURE_TO_STRESS * REF_VISCOSITY)).to(\"mm/yr\") # 8.4 mm/yr\n", + "REF_STRESS = (REF_VISCOSITY * CONVERGENCE_RATE / DOMAIN_DEPTH).to(\"MPa\") # 26 MPa\n", + "\n", + "# --- Geometry ---\n", + "CELL_SIZE = DOMAIN_DEPTH / 12 # base mesh; refined once, so half this\n", + "SURFACE_DIP = 60.0 # degrees, where the ramp reaches the surface\n", + "FLAT_DEPTH = uw.quantity(3, \"km\") # height of the décollement above the base\n", + "SURFACE_X = uw.quantity(19, \"km\") # where the fault reaches the surface\n", + "BAND_HALF_WIDTH = uw.quantity(0.8, \"km\") # of the weak plane\n", + "\n", + "# --- Rheology ---\n", + "COHESION = (0.05 * REF_STRESS).to(\"MPa\") # C in tau_y = C + mu p; 1.3 MPa\n", + "TRUE_FRICTION = \"0.05,0.15,0.25,0.4\" # flat, lower ramp, upper ramp, near surface\n", + "INITIAL_FRICTION = 0.2 # the starting guess, every segment\n", + "\n", + "# --- The inversion ---\n", + "CHECK_ONLY = 0 # 1: the gradient check, no inversion\n", + "OPTIMISER = \"tao\" # tao (PETSc quasi-Newton) | scipy (L-BFGS-B)\n", + "OBSERVATIONS = \"uplift+stress\" # uplift+stress | orientation_points | surface_strain\n", + "NOISE = 0.0 # on the observed velocity, as a fraction of its rms\n", + "SEED = 7\n", + "PRIOR_SIGMA = 0.0 # width of a Gaussian prior on log mu; 0 for none\n", + "BOUNDS = \"\" # friction bounds lo,hi for TAO's blmvm; empty for none\n", + "\n", + "params = uw.Params(\n", + " uw_cell_size=uw.Param(CELL_SIZE, description=\"base mesh cell size\"),\n", + " uw_surface_dip=uw.Param(SURFACE_DIP, description=\"dip of the ramp at the surface, degrees\"),\n", + " uw_flat_depth=uw.Param(FLAT_DEPTH, description=\"height of the décollement above the base\"),\n", + " uw_surface_x=uw.Param(SURFACE_X, description=\"where the fault reaches the surface\"),\n", + " uw_band=uw.Param(BAND_HALF_WIDTH, description=\"half-width of the weak plane\"),\n", + " uw_cohesion=uw.Param(COHESION, description=\"cohesion C in tau_y = C + mu p\"),\n", + " uw_true_friction=uw.Param(TRUE_FRICTION, description=\"friction: flat, lower ramp, upper ramp, near surface\"),\n", + " uw_initial_friction=uw.Param(INITIAL_FRICTION, description=\"starting guess, every segment\"),\n", + " uw_check_only=uw.Param(CHECK_ONLY, description=\"1: gradient check only\"),\n", + " uw_optimiser=uw.Param(OPTIMISER, description=\"tao | scipy\"),\n", + " uw_observations=uw.Param(OBSERVATIONS, description=\"uplift+stress | orientation_points | surface_strain\"),\n", + " uw_noise=uw.Param(NOISE, description=\"noise on the observed velocity, fraction of its rms\"),\n", + " uw_seed=uw.Param(SEED, description=\"seed for the noise\"),\n", + " uw_prior_sigma=uw.Param(PRIOR_SIGMA, description=\"width of the prior on log mu; 0 for none\"),\n", + " uw_bounds=uw.Param(BOUNDS, description=\"friction bounds lo,hi; empty for none\"),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "df86a530", + "metadata": {}, + "source": [ + "## The model and its scales\n", + "\n", + "The model is declared first, with the three reference quantities that fix\n", + "the scaling: the depth, the viscosity and the convergence rate. The solver\n", + "then works in units of those, so the box is $2 \\times 1$, the bulk viscosity\n", + "is one, the walls close at one, and the body force is the pressure-to-stress\n", + "ratio. `coords` on a variable come back in physical units, and `coords_nd`\n", + "are the model's own." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a9c02f07", + "metadata": {}, + "outputs": [], + "source": [ + "orchestration_model = uw.get_default_model()\n", + "orchestration_model.set_reference_quantities(\n", + " domain_depth=DOMAIN_DEPTH,\n", + " viscosity=REF_VISCOSITY,\n", + " convergence_rate=CONVERGENCE_RATE,\n", + ")\n", + "\n", + "\n", + "def _nd(quantity):\n", + " \"\"\"The plain number a dimensionless quantity stands for.\"\"\"\n", + " try:\n", + " return float(quantity.to(\"dimensionless\").magnitude)\n", + " except AttributeError:\n", + " return float(quantity)\n", + "\n", + "\n", + "def _lengths(quantity):\n", + " \"\"\"A length in units of the model depth.\"\"\"\n", + " return _nd(quantity / DOMAIN_DEPTH)\n", + "\n", + "\n", + "mm_per_yr = float(CONVERGENCE_RATE.to(\"mm/yr\").magnitude) # one model velocity unit\n", + "km = float(DOMAIN_DEPTH.to(\"km\").magnitude) # one model length unit" + ] + }, + { + "cell_type": "markdown", + "id": "db10286a", + "metadata": {}, + "source": [ + "## The mesh\n", + "\n", + "Refined once from the base cell size, which gives the velocity block a\n", + "multigrid hierarchy. Without one it falls back to an algebraic coarsening\n", + "that hits its iteration cap on this problem, and an inexact Newton step then\n", + "converges only linearly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8a456b31", + "metadata": {}, + "outputs": [], + "source": [ + "mesh = uw.meshing.UnstructuredSimplexBox(\n", + " minCoords=(0.0, 0.0), maxCoords=(2.0, 1.0),\n", + " cellSize=_lengths(params.uw_cell_size), qdegree=3, refinement=1,\n", + ")\n", + "x, y = mesh.X\n", + "\n", + "v = uw.discretisation.MeshVariable(\"v\", mesh, mesh.dim, degree=2)\n", + "p = uw.discretisation.MeshVariable(\"p\", mesh, 1, degree=1, continuous=True)\n", + "v_obs = uw.discretisation.MeshVariable(\"v_obs\", mesh, mesh.dim, degree=2)" + ] + }, + { + "cell_type": "markdown", + "id": "32c94ffc", + "metadata": {}, + "source": [ + "## The fault\n", + "\n", + "A flat décollement runs from the left wall to $x = x_c$, then a circular\n", + "ramp of radius $R$ about $(x_c, y_c)$ leaves the flat horizontally and\n", + "reaches the surface at the given dip. The signed distance to the fault and\n", + "the position along it are exact on each piece. The director is the normal\n", + "at the nearest point: vertical on the flat, radial on the ramp." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "35be5a5f", + "metadata": {}, + "outputs": [], + "source": [ + "flat_depth = _lengths(params.uw_flat_depth)\n", + "band_width = _lengths(params.uw_band)\n", + "phi_top = math.radians(params.uw_surface_dip) - math.pi / 2 # angle of the surface point about the centre\n", + "R = (1.0 - flat_depth) / (1.0 + math.sin(phi_top))\n", + "yc = flat_depth + R\n", + "xc = _lengths(params.uw_surface_x) - R * math.cos(phi_top)\n", + "r = sympy.sqrt((x - xc) ** 2 + (y - yc) ** 2)\n", + "phi = sympy.atan2(y - yc, x - xc)\n", + "on_flat = x < xc\n", + "d = sympy.Piecewise((y - flat_depth, on_flat), (R - r, True))\n", + "s = sympy.Piecewise((x, on_flat), (xc + R * (phi + sympy.pi / 2), True))\n", + "n_hat = sympy.Matrix([[sympy.Piecewise((0, on_flat), ((x - xc) / r, True)),\n", + " sympy.Piecewise((1, on_flat), ((y - yc) / r, True))]])\n", + "ramp = R * (phi_top + math.pi / 2)\n", + "length = xc + ramp\n", + "band = sympy.exp(-(d / band_width) ** 2)" + ] + }, + { + "cell_type": "markdown", + "id": "67982371", + "metadata": {}, + "source": [ + "## The segments and their friction coefficients\n", + "\n", + "Four segments along the fault: the flat, then the ramp in three equal\n", + "parts. Each coefficient is a named expression, so the solver can\n", + "differentiate its residual with respect to it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "113b01c3", + "metadata": {}, + "outputs": [], + "source": [ + "edges = [0.0, xc, xc + ramp / 3, xc + 2 * ramp / 3, length]\n", + "names = [\"flat\", \"lower ramp\", \"upper ramp\", \"near surface\"]\n", + "n_seg = len(names)\n", + "friction_coefficients = [\n", + " uw.expression(rf\"\\mu_{{{k + 1}}}\", params.uw_initial_friction, f\"friction coefficient, {names[k]}\")\n", + " for k in range(n_seg)\n", + "]\n", + "\n", + "\n", + "def segment(k):\n", + " \"\"\"A smooth indicator for segment k along the fault, in [0, 1].\"\"\"\n", + " on = 1 if k == 0 else (1 + sympy.tanh((s - edges[k]) / band_width)) / 2\n", + " off = 1 if k == n_seg - 1 else (1 - sympy.tanh((s - edges[k + 1]) / band_width)) / 2\n", + " return on * off\n", + "\n", + "\n", + "def set_friction(values):\n", + " for coefficient, value in zip(friction_coefficients, values):\n", + " coefficient.sym = float(value)" + ] + }, + { + "cell_type": "markdown", + "id": "49faa869", + "metadata": {}, + "source": [ + "## Coulomb yield on the plane\n", + "\n", + "The shear strain rate resolved on the plane is $\\hat t \\cdot E \\cdot \\hat n$.\n", + "The plane's viscosity is the harmonic combination of the bulk viscosity and\n", + "the yield stress over that rate, which is smooth everywhere and tends to\n", + "$\\tau_y / 2\\dot\\varepsilon_s$ where the plane slips. Compression is positive.\n", + "Where the dynamic pressure is tensile the plane keeps its cohesion and no\n", + "more." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "db143519", + "metadata": {}, + "outputs": [], + "source": [ + "eta_0 = 1\n", + "cohesion = _nd(params.uw_cohesion / REF_STRESS)\n", + "rho_g = _nd(DENSITY * GRAVITY * DOMAIN_DEPTH / REF_STRESS)\n", + "\n", + "E = mesh.vector.strain_tensor(v.sym)\n", + "t_hat = sympy.Matrix([[-n_hat[1], n_hat[0]]])\n", + "e_s = sympy.sqrt((t_hat * E * n_hat.T)[0] ** 2 + uw.maths.functions.vanishing)\n", + "friction = sum(friction_coefficients[k] * segment(k) for k in range(n_seg))\n", + "tau_y = cohesion + friction * sympy.Max(p.sym[0], 0)\n", + "eta_plane = eta_0 * tau_y / (tau_y + 2 * eta_0 * e_s)\n", + "eta_1 = eta_0 - band * (eta_0 - eta_plane)" + ] + }, + { + "cell_type": "markdown", + "id": "11d02db3", + "metadata": {}, + "source": [ + "## The Stokes solver\n", + "\n", + "Shortening from both sides, a no-slip base, and a free top: the surface\n", + "velocity is the uplift rate. The residual is nonlinear in the velocity and\n", + "the pressure, so the forward solve is Newton." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c3812d14", + "metadata": {}, + "outputs": [], + "source": [ + "stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p)\n", + "stokes.constitutive_model = uw.constitutive_models.TransverseIsotropicFlowModel\n", + "stokes.constitutive_model.Parameters.shear_viscosity_0 = eta_0\n", + "stokes.constitutive_model.Parameters.shear_viscosity_1 = eta_1\n", + "stokes.constitutive_model.Parameters.director = n_hat\n", + "stokes.tolerance = 1e-8\n", + "stokes.bodyforce = sympy.Matrix([0, -rho_g])\n", + "\n", + "stokes.add_essential_bc((0.0, 0.0), \"Bottom\")\n", + "stokes.add_essential_bc((0.5, None), \"Left\")\n", + "stokes.add_essential_bc((-0.5, None), \"Right\")" + ] + }, + { + "cell_type": "markdown", + "id": "4c5f2e6e", + "metadata": {}, + "source": [ + "## The observations\n", + "\n", + "The misfit has a term on the top surface, a boundary integral of the uplift\n", + "rate, and a term in the volume around five interior points where the shear\n", + "stress is read. `gradient()` takes the terms as `{domain: integrand}`.\n", + "\n", + "On a traction-free surface the shear strain rate vanishes, so a stress\n", + "orientation read there is only a sign. Orientation is an interior\n", + "observable (boreholes, focal mechanisms). The surface gives velocities and\n", + "their tangential derivative, the geodetic strain rate. Both alternatives\n", + "are available through `uw_observations`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c934afd6", + "metadata": {}, + "outputs": [], + "source": [ + "points = [(0.3, 0.65), (0.75, 0.12), (1.25, 0.3), (0.9, 0.6), (1.55, 0.85)] # in units of the depth\n", + "w_points = sum(sympy.exp(-((x - px) ** 2 + (y - py) ** 2) / (2 * band_width) ** 2)\n", + " for px, py in points)\n", + "\n", + "\n", + "def shear_stress(field):\n", + " e = mesh.vector.strain_tensor(field.sym)\n", + " return 2 * eta_0 * e[0, 1]\n", + "\n", + "\n", + "def orientation(field):\n", + " \"\"\"The principal-stress orientation as the unit vector (cos 2theta, sin 2theta)\n", + " of the deviatoric strain rate, which is the stress orientation in the\n", + " isotropic bulk. A unit vector rather than an angle, so there is no wrap.\"\"\"\n", + " e = mesh.vector.strain_tensor(field.sym)\n", + " a, b = e[0, 0] - e[1, 1], 2 * e[0, 1]\n", + " norm = sympy.sqrt(a ** 2 + b ** 2 + uw.maths.functions.vanishing)\n", + " return sympy.Matrix([[a / norm, b / norm]])\n", + "\n", + "\n", + "what = str(params.uw_observations)\n", + "dq = orientation(v) - orientation(v_obs)\n", + "if what == \"uplift+stress\":\n", + " misfit = {\"Top\": (v.sym[1] - v_obs.sym[1]) ** 2 / 2,\n", + " None: w_points * (shear_stress(v) - shear_stress(v_obs)) ** 2 / 2}\n", + "elif what == \"orientation_points\":\n", + " misfit = {None: w_points * (dq[0] ** 2 + dq[1] ** 2) / 2}\n", + "elif what == \"surface_strain\":\n", + " misfit = {\"Top\": (v.sym[0].diff(x) - v_obs.sym[0].diff(x)) ** 2 / 2}\n", + "else:\n", + " raise ValueError(f\"observations: {what!r}\")\n", + "\n", + "\n", + "def misfit_value():\n", + " return sum(uw.adjoint.integral(mesh, term, where) for where, term in misfit.items())" + ] + }, + { + "cell_type": "markdown", + "id": "6b77d313", + "metadata": {}, + "source": [ + "## The forward solve and the gradient\n", + "\n", + "Every solve starts cold: six Newton iterations, and a misfit that does not\n", + "depend on the previous evaluation, which the finite-difference check needs.\n", + "Each evaluation is one step of zero length in the model's record." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b47600f1", + "metadata": {}, + "outputs": [], + "source": [ + "evaluations = [0]\n", + "\n", + "\n", + "def forward(label):\n", + " with orchestration_model.step(0.0, label=label):\n", + " stokes.solve(zero_init_guess=True)\n", + "\n", + "\n", + "def misfit_and_gradient(label=None):\n", + " \"\"\"The misfit and dJ/d(log mu) for each segment, by the adjoint.\"\"\"\n", + " evaluations[0] += 1\n", + " with orchestration_model.step(0.0, label=label or f\"eval {evaluations[0]}\"):\n", + " stokes.solve(zero_init_guess=True)\n", + " out = stokes.gradient(misfit, parameters=friction_coefficients)\n", + " # d/d(log mu) = mu d/d(mu)\n", + " grad = np.array([out[\"parameters\"][c] * float(c.sym) for c in friction_coefficients])\n", + " return out[\"J\"], grad" + ] + }, + { + "cell_type": "markdown", + "id": "0ca6e48b", + "metadata": {}, + "source": [ + "## The twin\n", + "\n", + "The observations are the velocity field at the true coefficients, with\n", + "optional noise drawn once at a fraction of each component's rms. Every\n", + "observation set reads `v_obs`, so the uplift, the stress and the strain\n", + "rate all inherit the same noise." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "82a72cd4", + "metadata": {}, + "outputs": [], + "source": [ + "true_values = [float(t) for t in str(params.uw_true_friction).split(\",\")][:n_seg]\n", + "set_friction(true_values)\n", + "forward(\"truth\")\n", + "v_obs.array[...] = np.asarray(v.array)\n", + "if float(params.uw_noise) > 0:\n", + " rng = np.random.default_rng(int(params.uw_seed))\n", + " obs = np.asarray(v_obs.array)\n", + " rms = np.sqrt(np.mean(obs ** 2, axis=0, keepdims=True))\n", + " v_obs.array[...] = obs + float(params.uw_noise) * rms * rng.standard_normal(obs.shape)\n", + " print(f\"noise: {float(params.uw_noise):.3f} of the rms per component, seed {int(params.uw_seed)}\")\n", + "J_truth = misfit_value() # v still holds the truth's velocity\n", + "\n", + "set_friction([params.uw_initial_friction] * n_seg)\n", + "J0, g0 = misfit_and_gradient(\"start\")\n", + "print(f\"true friction {true_values}\")\n", + "print(f\"initial J = {J0:.6e} dJ/dlog mu = {g0}\")" + ] + }, + { + "cell_type": "markdown", + "id": "2f1cc0f0", + "metadata": {}, + "source": [ + "## The gradient check\n", + "\n", + "Central differences in each log-coefficient, against the adjoint gradient." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d5a3c746", + "metadata": {}, + "outputs": [], + "source": [ + "h = 1e-3\n", + "for k in range(n_seg):\n", + " base = math.log(params.uw_initial_friction)\n", + " fd = []\n", + " for sign in (+1, -1):\n", + " values = [params.uw_initial_friction] * n_seg\n", + " values[k] = math.exp(base + sign * h)\n", + " set_friction(values)\n", + " forward(f\"fd {names[k]} {'+' if sign > 0 else '-'}h\")\n", + " fd.append(misfit_value())\n", + " fd = (fd[0] - fd[1]) / (2 * h)\n", + " print(f\"{names[k]:>13}: adjoint {g0[k]: .6e} finite difference {fd: .6e} ratio {fd / g0[k]:.5f}\")\n", + "set_friction([params.uw_initial_friction] * n_seg)\n", + "\n", + "if int(params.uw_check_only):\n", + " raise SystemExit" + ] + }, + { + "cell_type": "markdown", + "id": "91bb23cf", + "metadata": {}, + "source": [ + "## The objective\n", + "\n", + "Without noise the misfit has no natural scale and the optimiser sees it\n", + "relative to its starting value, $J/J_0$. With noise the objective is a\n", + "negative log posterior. The data term is $\\chi^2/2$: the misfit scaled by\n", + "its expected value at the truth under the noise, which a twin experiment\n", + "can read directly, times the number of independent data, the surface nodes\n", + "and the nodes under the point weights. The prior term is\n", + "$(\\log\\mu - \\log\\mu_0)^2 / 2\\sigma_m^2$ with $\\sigma_m$ in log units. The\n", + "weight between the two is then a statement about the noise and the prior,\n", + "and which coefficients the data move is decided by their sensitivities\n", + "against that." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e95ede9e", + "metadata": {}, + "outputs": [], + "source": [ + "history = []\n", + "noisy = float(params.uw_noise) > 0\n", + "sigma_m = float(params.uw_prior_sigma)\n", + "log_prior = np.log([params.uw_initial_friction] * n_seg)\n", + "if sigma_m > 0 and not noisy:\n", + " raise ValueError(\"a prior is weighed against the noise: give uw_noise as well\")\n", + "\n", + "if noisy:\n", + " X = np.asarray(v.coords_nd)\n", + " on_surface = X[:, 1] > 1.0 - 1e-6\n", + " near_points = np.zeros(len(X), dtype=bool)\n", + " for px, py in points:\n", + " near_points |= (X[:, 0] - px) ** 2 + (X[:, 1] - py) ** 2 < (2 * band_width) ** 2\n", + " N_eff = {\"uplift+stress\": on_surface.sum() + near_points.sum(),\n", + " \"orientation_points\": near_points.sum(),\n", + " \"surface_strain\": on_surface.sum()}[what]\n", + " chi2_scale = N_eff / J_truth # chi^2 = J * chi2_scale\n", + " print(f\"N_eff = {N_eff}, misfit floor at the truth = {J_truth:.4e}\")\n", + "\n", + "\n", + "def objective(log_mu):\n", + " set_friction(np.exp(log_mu))\n", + " J, grad = misfit_and_gradient()\n", + " history.append((J, np.exp(log_mu).copy()))\n", + " if not noisy:\n", + " print(f\" J/J0 = {J / J0:.6e} friction = {np.exp(log_mu)}\")\n", + " return J / J0, grad / J0\n", + " print(f\" J = {J:.6e} chi2/N = {J * chi2_scale / N_eff:.4f} friction = {np.exp(log_mu)}\")\n", + " value = J * chi2_scale / 2\n", + " g = grad * chi2_scale / 2\n", + " if sigma_m > 0:\n", + " value += np.sum((log_mu - log_prior) ** 2) / (2 * sigma_m ** 2)\n", + " g = g + (log_mu - log_prior) / sigma_m ** 2\n", + " return value, g" + ] + }, + { + "cell_type": "markdown", + "id": "fed874e2", + "metadata": {}, + "source": [ + "## The inversion\n", + "\n", + "PETSc's TAO drives it by default: the same objective and gradient, a\n", + "limited-memory quasi-Newton update and TAO's line search, and the bounded\n", + "variant when bounds are given. SciPy's L-BFGS-B is the alternative." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b3f436f0", + "metadata": {}, + "outputs": [], + "source": [ + "x0 = np.log([params.uw_initial_friction] * n_seg)\n", + "if str(params.uw_optimiser) == \"tao\":\n", + " bounds = None\n", + " if str(params.uw_bounds).strip():\n", + " lo, hi = (float(b) for b in str(params.uw_bounds).split(\",\"))\n", + " bounds = (np.log([lo] * n_seg), np.log([hi] * n_seg))\n", + " x_best, info = uw.adjoint.minimise(objective, x0, max_evaluations=60,\n", + " gradient_tolerance=1e-10, bounds=bounds,\n", + " method=\"blmvm\" if bounds else \"lmvm\")\n", + "else:\n", + " from scipy.optimize import minimize\n", + " result = minimize(objective, x0, jac=True, method=\"L-BFGS-B\",\n", + " options={\"maxiter\": 40, \"gtol\": 1e-10})\n", + " x_best = result.x\n", + "print(f\"recovered {np.exp(x_best)} true {true_values} after {len(history)} evaluations\")" + ] + }, + { + "cell_type": "markdown", + "id": "795a2cdf", + "metadata": {}, + "source": [ + "## What the figure needs\n", + "\n", + "Uplift-rate profiles along the top at the truth, the start and the answer,\n", + "the weak-plane viscosity on a grid, and the path the coefficients took.\n", + "Lengths are saved in kilometres and velocities in millimetres per year." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0d4d8b90", + "metadata": {}, + "outputs": [], + "source": [ + "xs = np.linspace(0.0, 2.0, 161)\n", + "top = np.column_stack([xs, np.full_like(xs, 1.0 - 1e-6)])\n", + "profiles = {}\n", + "for label, values in ((\"true\", true_values), (\"initial\", [params.uw_initial_friction] * n_seg),\n", + " (\"recovered\", list(np.exp(x_best)))):\n", + " set_friction(values)\n", + " forward(f\"profile {label}\")\n", + " profiles[label] = np.asarray(uw.function.evaluate(v.sym[1], top)).ravel() * mm_per_yr\n", + "set_friction(true_values)\n", + "forward(\"truth again\") # the field on the grid is the truth's\n", + "gx, gy = np.meshgrid(np.linspace(0, 2, 201), np.linspace(0, 1, 101))\n", + "grid = np.column_stack([gx.ravel(), gy.ravel()])\n", + "eta_1_grid = np.asarray(uw.function.evaluate(eta_1, grid)).reshape(gx.shape)\n", + "\n", + "tag = (f\"{what}\" + (f\"_noise{float(params.uw_noise):g}\" if float(params.uw_noise) > 0 else \"\")\n", + " + (f\"_prior{sigma_m:g}\" if sigma_m > 0 else \"\"))\n", + "np.savez(f\"fault_friction_{tag}_data.npz\", xs=xs * km, gx=gx * km, gy=gy * km, eta_1=eta_1_grid,\n", + " points=np.array(points) * km, true=np.array(true_values), band=band_width * km,\n", + " length_unit=\"km\", velocity_unit=\"mm/yr\",\n", + " history=np.array([[J, *vals] for J, vals in history]),\n", + " **{f\"uplift_{k}\": val for k, val in profiles.items()})" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/examples/adjoint/fault_segments/fault_friction.py b/docs/examples/adjoint/fault_segments/fault_friction.py index e40002a83..c177fe8b9 100644 --- a/docs/examples/adjoint/fault_segments/fault_friction.py +++ b/docs/examples/adjoint/fault_segments/fault_friction.py @@ -1,27 +1,29 @@ -"""Friction on a listric fault, segment by segment, from surface uplift and stress. - -A listric fault under horizontal shortening and gravity: a ramp that steepens -from a flat decollement at depth to a dip of sixty degrees at the surface, -represented as a weak plane in a transversely isotropic viscosity (no cut in -the mesh). The plane yields at the Coulomb stress tau_y = C + mu p, where p -is the pressure, and the friction coefficient mu is different on the flat, -on the lower and upper parts of the ramp, and near the surface. Those four -coefficients are the unknowns. The observations are the uplift rate along -the top surface and the shear stress in the bulk near a handful of points, -taken from a run at the true coefficients. - -The residual is nonlinear in the velocity and the pressure, so the forward -solve is Newton. The gradient with respect to each coefficient comes from -one call, stokes.gradient(misfit, parameters=...): the solver assembles the -adjoint operator from its own Jacobian kernels with trial and test -exchanged, solves it, and differentiates its residual symbolically with -respect to each named coefficient. Nothing here is differenced. - -Run it: - - python fault_friction.py # twin experiment: gradient check, then the inversion - python fault_friction.py -uw_check_only 1 -""" +# %% [markdown] +# # Friction on a listric fault from surface observations +# +# A listric fault under horizontal shortening and gravity: a flat décollement +# at depth that steepens through a circular ramp to a dip of sixty degrees at +# the surface. The fault is a weak plane in a transversely isotropic viscosity, +# with no cut in the mesh. The plane yields at the Coulomb stress +# $\tau_y = C + \mu p$, and the friction coefficient $\mu$ takes a different +# value on the flat, on the lower and upper parts of the ramp, and near the +# surface. Those four coefficients are the unknowns. +# +# The observations are the uplift rate along the top surface and the shear +# stress near a handful of interior points, read from a run at the true +# coefficients. The gradient of the misfit with respect to each coefficient +# comes from one call, `stokes.gradient(misfit, parameters=...)`: the solver +# assembles the adjoint operator from its own Jacobian kernels with trial and +# test functions exchanged, solves it, and differentiates its residual +# symbolically with respect to each named coefficient. Nothing is differenced. +# +# The notebook checks that gradient against central differences, then hands +# the misfit and its gradient to PETSc's TAO to recover the four +# coefficients. Each misfit evaluation is one step of zero length in the +# model's record, so the run's transcript lists the forward solve and the +# adjoint solve it made. + +# %% import math import numpy as np @@ -29,118 +31,248 @@ import underworld3 as uw +# %% [markdown] +# ### Configurable parameters +# +# Default values are defined as named constants below. From the command line, +# override them with PETSc-style flags: +# +# ```bash +# python fault_friction.py -uw_check_only 1 +# python fault_friction.py -uw_observations surface_strain +# python fault_friction.py -uw_noise 0.03 -uw_prior_sigma 1 -uw_bounds 0.005,1 +# ``` +# +# The problem depends on one dimensionless number: the lithostatic pressure at +# the base of the model over the viscous stress of the shortening, +# $\rho g L^2 / (\eta V)$. We fix the depth, the viscosity and the density, +# and that ratio sets the convergence rate. + +# %% +# --- Scales --- +DOMAIN_DEPTH = uw.quantity(10, "km") # the box is 20 km wide and 10 km deep +REF_VISCOSITY = uw.quantity(1e21, "Pa*s") # the bulk viscosity +DENSITY = uw.quantity(2700, "kg/m**3") +GRAVITY = uw.quantity(9.81, "m/s**2") +PRESSURE_TO_STRESS = 10 # rho g L^2 / (eta V) +CONVERGENCE_RATE = (DENSITY * GRAVITY * DOMAIN_DEPTH**2 / (PRESSURE_TO_STRESS * REF_VISCOSITY)).to("mm/yr") # 8.4 mm/yr +REF_STRESS = (REF_VISCOSITY * CONVERGENCE_RATE / DOMAIN_DEPTH).to("MPa") # 26 MPa + +# --- Geometry --- +CELL_SIZE = DOMAIN_DEPTH / 12 # base mesh; refined once, so half this +SURFACE_DIP = 60.0 # degrees, where the ramp reaches the surface +FLAT_DEPTH = uw.quantity(3, "km") # height of the décollement above the base +SURFACE_X = uw.quantity(19, "km") # where the fault reaches the surface +BAND_HALF_WIDTH = uw.quantity(0.8, "km") # of the weak plane + +# --- Rheology --- +COHESION = (0.05 * REF_STRESS).to("MPa") # C in tau_y = C + mu p; 1.3 MPa +TRUE_FRICTION = "0.05,0.15,0.25,0.4" # flat, lower ramp, upper ramp, near surface +INITIAL_FRICTION = 0.2 # the starting guess, every segment + +# --- The inversion --- +CHECK_ONLY = 0 # 1: the gradient check, no inversion +OPTIMISER = "tao" # tao (PETSc quasi-Newton) | scipy (L-BFGS-B) +OBSERVATIONS = "uplift+stress" # uplift+stress | orientation_points | surface_strain +NOISE = 0.0 # on the observed velocity, as a fraction of its rms +SEED = 7 +PRIOR_SIGMA = 0.0 # width of a Gaussian prior on log mu; 0 for none +BOUNDS = "" # friction bounds lo,hi for TAO's blmvm; empty for none + params = uw.Params( - cell_size=uw.Param(1 / 12, "base mesh cell size (box is 2 x 1); refined once, so half this"), - surface_dip=uw.Param(60.0, "dip of the ramp where it reaches the surface, degrees"), - flat_depth=uw.Param(0.3, "height of the decollement above the base"), - surface_x=uw.Param(1.9, "where the fault reaches the surface"), - band=uw.Param(0.08, "half-width of the weak band, in box units"), - true_strengths=uw.Param("0.05,0.15,0.25,0.4", - "friction coefficient: flat, lower ramp, upper ramp, near surface"), - initial_strength=uw.Param(0.2, "starting guess, every segment"), - cohesion=uw.Param(0.05, "cohesion C in tau_y = C + mu p"), - rho_g=uw.Param(10.0, "body force, so the pressure grows with depth"), - check_only=uw.Param(0, "1: gradient check against finite differences, no inversion"), - optimiser=uw.Param("tao", "tao (PETSc, limited-memory quasi-Newton; blmvm when bounds are set) | scipy (L-BFGS-B)"), - noise=uw.Param(0.0, "Gaussian noise on the observed velocity, as a fraction of its rms, per component"), - seed=uw.Param(7, "seed for the noise"), - prior_sigma=uw.Param(0.0, "width of a Gaussian prior on log mu about the start, in log units; 0 for none"), - bounds=uw.Param("", "friction bounds lo,hi for TAO's blmvm; empty for none"), - observations=uw.Param("uplift+stress", - "uplift+stress | orientation_points (principal-stress orientation " - "at the five interior points only) | surface_strain (the surface " - "strain rate d v_x / d x along the top only)"), + uw_cell_size=uw.Param(CELL_SIZE, description="base mesh cell size"), + uw_surface_dip=uw.Param(SURFACE_DIP, description="dip of the ramp at the surface, degrees"), + uw_flat_depth=uw.Param(FLAT_DEPTH, description="height of the décollement above the base"), + uw_surface_x=uw.Param(SURFACE_X, description="where the fault reaches the surface"), + uw_band=uw.Param(BAND_HALF_WIDTH, description="half-width of the weak plane"), + uw_cohesion=uw.Param(COHESION, description="cohesion C in tau_y = C + mu p"), + uw_true_friction=uw.Param(TRUE_FRICTION, description="friction: flat, lower ramp, upper ramp, near surface"), + uw_initial_friction=uw.Param(INITIAL_FRICTION, description="starting guess, every segment"), + uw_check_only=uw.Param(CHECK_ONLY, description="1: gradient check only"), + uw_optimiser=uw.Param(OPTIMISER, description="tao | scipy"), + uw_observations=uw.Param(OBSERVATIONS, description="uplift+stress | orientation_points | surface_strain"), + uw_noise=uw.Param(NOISE, description="noise on the observed velocity, fraction of its rms"), + uw_seed=uw.Param(SEED, description="seed for the noise"), + uw_prior_sigma=uw.Param(PRIOR_SIGMA, description="width of the prior on log mu; 0 for none"), + uw_bounds=uw.Param(BOUNDS, description="friction bounds lo,hi; empty for none"), ) -# --- the model --------------------------------------------------------------- -# Refined once from the base size: the refinement gives the velocity block a -# multigrid hierarchy. Without one it falls back to gamg, which hits its -# iteration cap on this problem, and an inexact Newton step converges linearly. -# Every solve starts cold: six Newton iterations, and a misfit that does not -# depend on the previous evaluation, which the finite-difference check needs. -mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0.0, 0.0), maxCoords=(2.0, 1.0), - cellSize=params.cell_size, qdegree=3, refinement=1) +# %% [markdown] +# ## The model and its scales +# +# The model is declared first, with the three reference quantities that fix +# the scaling: the depth, the viscosity and the convergence rate. The solver +# then works in units of those, so the box is $2 \times 1$, the bulk viscosity +# is one, the walls close at one, and the body force is the pressure-to-stress +# ratio. `coords` on a variable come back in physical units, and `coords_nd` +# are the model's own. + +# %% +orchestration_model = uw.get_default_model() +orchestration_model.set_reference_quantities( + domain_depth=DOMAIN_DEPTH, + viscosity=REF_VISCOSITY, + convergence_rate=CONVERGENCE_RATE, +) + + +def _nd(quantity): + """The plain number a dimensionless quantity stands for.""" + try: + return float(quantity.to("dimensionless").magnitude) + except AttributeError: + return float(quantity) + + +def _lengths(quantity): + """A length in units of the model depth.""" + return _nd(quantity / DOMAIN_DEPTH) + + +mm_per_yr = float(CONVERGENCE_RATE.to("mm/yr").magnitude) # one model velocity unit +km = float(DOMAIN_DEPTH.to("km").magnitude) # one model length unit + +# %% [markdown] +# ## The mesh +# +# Refined once from the base cell size, which gives the velocity block a +# multigrid hierarchy. Without one it falls back to an algebraic coarsening +# that hits its iteration cap on this problem, and an inexact Newton step then +# converges only linearly. + +# %% +mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(2.0, 1.0), + cellSize=_lengths(params.uw_cell_size), qdegree=3, refinement=1, +) x, y = mesh.X v = uw.discretisation.MeshVariable("v", mesh, mesh.dim, degree=2) p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1, continuous=True) v_obs = uw.discretisation.MeshVariable("v_obs", mesh, mesh.dim, degree=2) -# The fault: a flat decollement at y = flat_depth running from the left wall -# to x = xc, then a circular ramp of radius R about (xc, yc) that leaves the -# flat horizontally and reaches the surface at the given dip. The signed -# distance to it and the position along it are exact on each piece, and the -# director is the normal at the nearest point: vertical on the flat, radial on -# the ramp. -phi_top = math.radians(params.surface_dip) - math.pi / 2 # angle of the surface point about the centre -R = (1.0 - params.flat_depth) / (1.0 + math.sin(phi_top)) -yc = params.flat_depth + R -xc = params.surface_x - R * math.cos(phi_top) +# %% [markdown] +# ## The fault +# +# A flat décollement runs from the left wall to $x = x_c$, then a circular +# ramp of radius $R$ about $(x_c, y_c)$ leaves the flat horizontally and +# reaches the surface at the given dip. The signed distance to the fault and +# the position along it are exact on each piece. The director is the normal +# at the nearest point: vertical on the flat, radial on the ramp. + +# %% +flat_depth = _lengths(params.uw_flat_depth) +band_width = _lengths(params.uw_band) +phi_top = math.radians(params.uw_surface_dip) - math.pi / 2 # angle of the surface point about the centre +R = (1.0 - flat_depth) / (1.0 + math.sin(phi_top)) +yc = flat_depth + R +xc = _lengths(params.uw_surface_x) - R * math.cos(phi_top) r = sympy.sqrt((x - xc) ** 2 + (y - yc) ** 2) phi = sympy.atan2(y - yc, x - xc) on_flat = x < xc -d = sympy.Piecewise((y - params.flat_depth, on_flat), (R - r, True)) +d = sympy.Piecewise((y - flat_depth, on_flat), (R - r, True)) s = sympy.Piecewise((x, on_flat), (xc + R * (phi + sympy.pi / 2), True)) n_hat = sympy.Matrix([[sympy.Piecewise((0, on_flat), ((x - xc) / r, True)), sympy.Piecewise((1, on_flat), ((y - yc) / r, True))]]) ramp = R * (phi_top + math.pi / 2) length = xc + ramp -band = sympy.exp(-(d / params.band) ** 2) +band = sympy.exp(-(d / band_width) ** 2) + +# %% [markdown] +# ## The segments and their friction coefficients +# +# Four segments along the fault: the flat, then the ramp in three equal +# parts. Each coefficient is a named expression, so the solver can +# differentiate its residual with respect to it. -# Segments along the fault: the flat, then the ramp in three equal parts. +# %% edges = [0.0, xc, xc + ramp / 3, xc + 2 * ramp / 3, length] names = ["flat", "lower ramp", "upper ramp", "near surface"] n_seg = len(names) -strengths = [uw.expression(rf"\mu_{{{k + 1}}}", params.initial_strength, - f"friction coefficient, {names[k]}") - for k in range(n_seg)] +friction_coefficients = [ + uw.expression(rf"\mu_{{{k + 1}}}", params.uw_initial_friction, f"friction coefficient, {names[k]}") + for k in range(n_seg) +] + def segment(k): """A smooth indicator for segment k along the fault, in [0, 1].""" - edge = params.band - on = 1 if k == 0 else (1 + sympy.tanh((s - edges[k]) / edge)) / 2 - off = 1 if k == n_seg - 1 else (1 - sympy.tanh((s - edges[k + 1]) / edge)) / 2 + on = 1 if k == 0 else (1 + sympy.tanh((s - edges[k]) / band_width)) / 2 + off = 1 if k == n_seg - 1 else (1 - sympy.tanh((s - edges[k + 1]) / band_width)) / 2 return on * off + +def set_friction(values): + for coefficient, value in zip(friction_coefficients, values): + coefficient.sym = float(value) + +# %% [markdown] +# ## Coulomb yield on the plane +# +# The shear strain rate resolved on the plane is $\hat t \cdot E \cdot \hat n$. +# The plane's viscosity is the harmonic combination of the bulk viscosity and +# the yield stress over that rate, which is smooth everywhere and tends to +# $\tau_y / 2\dot\varepsilon_s$ where the plane slips. Compression is positive. +# Where the dynamic pressure is tensile the plane keeps its cohesion and no +# more. + +# %% eta_0 = 1 +cohesion = _nd(params.uw_cohesion / REF_STRESS) +rho_g = _nd(DENSITY * GRAVITY * DOMAIN_DEPTH / REF_STRESS) -# Coulomb yield on the plane. The shear strain rate resolved on the plane is -# t.E.n; the plane's viscosity is the harmonic combination of the bulk -# viscosity and the yield stress over that rate, which is smooth everywhere -# and tends to tau_y / (2 e_s) where the plane slips. E = mesh.vector.strain_tensor(v.sym) t_hat = sympy.Matrix([[-n_hat[1], n_hat[0]]]) e_s = sympy.sqrt((t_hat * E * n_hat.T)[0] ** 2 + uw.maths.functions.vanishing) -friction = sum(strengths[k] * segment(k) for k in range(n_seg)) -# Compression is positive; where the dynamic pressure is tensile the plane -# keeps its cohesion and no more. -tau_y = params.cohesion + friction * sympy.Max(p.sym[0], 0) +friction = sum(friction_coefficients[k] * segment(k) for k in range(n_seg)) +tau_y = cohesion + friction * sympy.Max(p.sym[0], 0) eta_plane = eta_0 * tau_y / (tau_y + 2 * eta_0 * e_s) eta_1 = eta_0 - band * (eta_0 - eta_plane) +# %% [markdown] +# ## The Stokes solver +# +# Shortening from both sides, a no-slip base, and a free top: the surface +# velocity is the uplift rate. The residual is nonlinear in the velocity and +# the pressure, so the forward solve is Newton. + +# %% stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) stokes.constitutive_model = uw.constitutive_models.TransverseIsotropicFlowModel stokes.constitutive_model.Parameters.shear_viscosity_0 = eta_0 stokes.constitutive_model.Parameters.shear_viscosity_1 = eta_1 stokes.constitutive_model.Parameters.director = n_hat stokes.tolerance = 1e-8 -stokes.bodyforce = sympy.Matrix([0, -params.rho_g]) +stokes.bodyforce = sympy.Matrix([0, -rho_g]) -# Shortening from both sides, a no-slip base, and a free top: the surface -# velocity is the uplift rate. stokes.add_essential_bc((0.0, 0.0), "Bottom") stokes.add_essential_bc((0.5, None), "Left") stokes.add_essential_bc((-0.5, None), "Right") -# --- the observations ------------------------------------------------------ -points = [(0.3, 0.65), (0.75, 0.12), (1.25, 0.3), (0.9, 0.6), (1.55, 0.85)] -w_points = sum(sympy.exp(-((x - px) ** 2 + (y - py) ** 2) / (2 * params.band) ** 2) +# %% [markdown] +# ## The observations +# +# The misfit has a term on the top surface, a boundary integral of the uplift +# rate, and a term in the volume around five interior points where the shear +# stress is read. `gradient()` takes the terms as `{domain: integrand}`. +# +# On a traction-free surface the shear strain rate vanishes, so a stress +# orientation read there is only a sign. Orientation is an interior +# observable (boreholes, focal mechanisms). The surface gives velocities and +# their tangential derivative, the geodetic strain rate. Both alternatives +# are available through `uw_observations`. + +# %% +points = [(0.3, 0.65), (0.75, 0.12), (1.25, 0.3), (0.9, 0.6), (1.55, 0.85)] # in units of the depth +w_points = sum(sympy.exp(-((x - px) ** 2 + (y - py) ** 2) / (2 * band_width) ** 2) for px, py in points) + def shear_stress(field): e = mesh.vector.strain_tensor(field.sym) return 2 * eta_0 * e[0, 1] + def orientation(field): """The principal-stress orientation as the unit vector (cos 2theta, sin 2theta) of the deviatoric strain rate, which is the stress orientation in the @@ -150,14 +282,8 @@ def orientation(field): norm = sympy.sqrt(a ** 2 + b ** 2 + uw.maths.functions.vanishing) return sympy.Matrix([[a / norm, b / norm]]) -# The misfit has a term on the top surface — a true boundary integral of -# the uplift rate, or of the stress orientation — and a term in the volume -# around the stress points. gradient() takes them as {domain: integrand}. -# On a traction-free surface the shear strain rate vanishes, so a stress -# orientation read there is only a sign; orientation is an interior -# observable (boreholes, focal mechanisms), and the surface gives velocities -# and their tangential derivative — the geodetic strain rate. -what = str(params.observations) + +what = str(params.uw_observations) dq = orientation(v) - orientation(v_obs) if what == "uplift+stress": misfit = {"Top": (v.sym[1] - v_obs.sym[1]) ** 2 / 2, @@ -169,156 +295,185 @@ def orientation(field): else: raise ValueError(f"observations: {what!r}") + def misfit_value(): return sum(uw.adjoint.integral(mesh, term, where) for where, term in misfit.items()) -model = uw.get_default_model() - -def set_strengths(values): - for expr, value in zip(strengths, values): - expr.sym = float(value) +# %% [markdown] +# ## The forward solve and the gradient +# +# Every solve starts cold: six Newton iterations, and a misfit that does not +# depend on the previous evaluation, which the finite-difference check needs. +# Each evaluation is one step of zero length in the model's record. +# %% evaluations = [0] -def J_and_gradient(label=None): - """The misfit and dJ/d(log strength) for each segment, by the adjoint. - Each evaluation is one step of zero length in the model's record, so the - run's transcript lists the forward solve and the adjoint solve it made. - """ +def forward(label): + with orchestration_model.step(0.0, label=label): + stokes.solve(zero_init_guess=True) + + +def misfit_and_gradient(label=None): + """The misfit and dJ/d(log mu) for each segment, by the adjoint.""" evaluations[0] += 1 - with model.step(0.0, label=label or f"eval {evaluations[0]}"): + with orchestration_model.step(0.0, label=label or f"eval {evaluations[0]}"): stokes.solve(zero_init_guess=True) - out = stokes.gradient(misfit, parameters=strengths) + out = stokes.gradient(misfit, parameters=friction_coefficients) # d/d(log mu) = mu d/d(mu) - grad = np.array([out["parameters"][expr] * float(expr.sym) for expr in strengths]) + grad = np.array([out["parameters"][c] * float(c.sym) for c in friction_coefficients]) return out["J"], grad -def forward(label): - with model.step(0.0, label=label): - stokes.solve(zero_init_guess=True) - -# --- the truth, and the twin ------------------------------------------------- -true_values = [float(t) for t in str(params.true_strengths).split(",")][:n_seg] -set_strengths(true_values) +# %% [markdown] +# ## The twin +# +# The observations are the velocity field at the true coefficients, with +# optional noise drawn once at a fraction of each component's rms. Every +# observation set reads `v_obs`, so the uplift, the stress and the strain +# rate all inherit the same noise. + +# %% +true_values = [float(t) for t in str(params.uw_true_friction).split(",")][:n_seg] +set_friction(true_values) forward("truth") v_obs.array[...] = np.asarray(v.array) -if float(params.noise) > 0: - # Noise on the observed velocity field, drawn once, at a fraction of each - # component's rms. Every observation set reads v_obs, so the uplift, the - # stress and the strain rate all inherit it. - rng = np.random.default_rng(int(params.seed)) +if float(params.uw_noise) > 0: + rng = np.random.default_rng(int(params.uw_seed)) obs = np.asarray(v_obs.array) rms = np.sqrt(np.mean(obs ** 2, axis=0, keepdims=True)) - v_obs.array[...] = obs + float(params.noise) * rms * rng.standard_normal(obs.shape) - uw.pprint(f"noise: {float(params.noise):.3f} of the rms per component, seed {int(params.seed)}") -uw.pprint(f"true strengths {true_values}") -J_truth = None - + v_obs.array[...] = obs + float(params.uw_noise) * rms * rng.standard_normal(obs.shape) + print(f"noise: {float(params.uw_noise):.3f} of the rms per component, seed {int(params.uw_seed)}") J_truth = misfit_value() # v still holds the truth's velocity -set_strengths([params.initial_strength] * n_seg) -J0, g0 = J_and_gradient("start") -uw.pprint(f"initial J = {J0:.6e} dJ/dlog eta = {g0}") -# Gradient check: central differences in each log-strength. +set_friction([params.uw_initial_friction] * n_seg) +J0, g0 = misfit_and_gradient("start") +print(f"true friction {true_values}") +print(f"initial J = {J0:.6e} dJ/dlog mu = {g0}") + +# %% [markdown] +# ## The gradient check +# +# Central differences in each log-coefficient, against the adjoint gradient. + +# %% h = 1e-3 for k in range(n_seg): - base = math.log(params.initial_strength) + base = math.log(params.uw_initial_friction) fd = [] for sign in (+1, -1): - vals = [params.initial_strength] * n_seg - vals[k] = math.exp(base + sign * h) - set_strengths(vals) + values = [params.uw_initial_friction] * n_seg + values[k] = math.exp(base + sign * h) + set_friction(values) forward(f"fd {names[k]} {'+' if sign > 0 else '-'}h") fd.append(misfit_value()) fd = (fd[0] - fd[1]) / (2 * h) - uw.pprint(f"{names[k]:>13}: adjoint {g0[k]: .6e} finite difference {fd: .6e} " - f"ratio {fd / g0[k]:.5f}") -set_strengths([params.initial_strength] * n_seg) + print(f"{names[k]:>13}: adjoint {g0[k]: .6e} finite difference {fd: .6e} ratio {fd / g0[k]:.5f}") +set_friction([params.uw_initial_friction] * n_seg) -if int(params.check_only): +if int(params.uw_check_only): raise SystemExit -# --- the inversion --------------------------------------------------------------- -from scipy.optimize import minimize - +# %% [markdown] +# ## The objective +# +# Without noise the misfit has no natural scale and the optimiser sees it +# relative to its starting value, $J/J_0$. With noise the objective is a +# negative log posterior. The data term is $\chi^2/2$: the misfit scaled by +# its expected value at the truth under the noise, which a twin experiment +# can read directly, times the number of independent data, the surface nodes +# and the nodes under the point weights. The prior term is +# $(\log\mu - \log\mu_0)^2 / 2\sigma_m^2$ with $\sigma_m$ in log units. The +# weight between the two is then a statement about the noise and the prior, +# and which coefficients the data move is decided by their sensitivities +# against that. + +# %% history = [] - -# The optimiser sees the misfit relative to its starting value: L-BFGS-B stops -# on the absolute decrease of its objective, and a surface integral of a -# velocity misfit is a small number. -J_scale = J0 - -# The objective as a negative log posterior. The data term is chi-squared/2: -# the misfit scaled by its expected value at the truth under the noise, which -# a twin experiment can read directly, times the number of independent data — -# the surface nodes and the nodes under the point weights. The prior term is -# (log mu - log mu_start)^2 / (2 sigma_m^2) with sigma_m in log units. The -# weight between them is then a statement about the noise and the prior, not -# a number to tune, and which coefficients the data move is decided by their -# sensitivities against that. -sigma_m = float(params.prior_sigma) -log_prior = np.log([params.initial_strength] * n_seg) -X = np.asarray(v.coords) -on_surface = X[:, 1] > 1.0 - 1e-6 -near_points = np.zeros(len(X), dtype=bool) -for px, py in points: - near_points |= (X[:, 0] - px) ** 2 + (X[:, 1] - py) ** 2 < (2 * params.band) ** 2 -N_eff = {"uplift+stress": on_surface.sum() + near_points.sum(), - "orientation_points": near_points.sum(), - "surface_strain": on_surface.sum()}[what] -J_floor = J_truth if float(params.noise) > 0 else J0 -chi2_scale = N_eff / J_floor # chi^2 = J * chi2_scale -uw.pprint(f"N_eff = {N_eff}, misfit floor at the truth = {J_floor:.4e}") - -def objective(log_eta): - set_strengths(np.exp(log_eta)) - J, grad = J_and_gradient() - history.append((J, np.exp(log_eta).copy())) - uw.pprint(f" J = {J:.6e} chi2/N = {J * chi2_scale / N_eff:.4f} strengths = {np.exp(log_eta)}") +noisy = float(params.uw_noise) > 0 +sigma_m = float(params.uw_prior_sigma) +log_prior = np.log([params.uw_initial_friction] * n_seg) +if sigma_m > 0 and not noisy: + raise ValueError("a prior is weighed against the noise: give uw_noise as well") + +if noisy: + X = np.asarray(v.coords_nd) + on_surface = X[:, 1] > 1.0 - 1e-6 + near_points = np.zeros(len(X), dtype=bool) + for px, py in points: + near_points |= (X[:, 0] - px) ** 2 + (X[:, 1] - py) ** 2 < (2 * band_width) ** 2 + N_eff = {"uplift+stress": on_surface.sum() + near_points.sum(), + "orientation_points": near_points.sum(), + "surface_strain": on_surface.sum()}[what] + chi2_scale = N_eff / J_truth # chi^2 = J * chi2_scale + print(f"N_eff = {N_eff}, misfit floor at the truth = {J_truth:.4e}") + + +def objective(log_mu): + set_friction(np.exp(log_mu)) + J, grad = misfit_and_gradient() + history.append((J, np.exp(log_mu).copy())) + if not noisy: + print(f" J/J0 = {J / J0:.6e} friction = {np.exp(log_mu)}") + return J / J0, grad / J0 + print(f" J = {J:.6e} chi2/N = {J * chi2_scale / N_eff:.4f} friction = {np.exp(log_mu)}") value = J * chi2_scale / 2 g = grad * chi2_scale / 2 if sigma_m > 0: - value += np.sum((log_eta - log_prior) ** 2) / (2 * sigma_m ** 2) - g = g + (log_eta - log_prior) / sigma_m ** 2 + value += np.sum((log_mu - log_prior) ** 2) / (2 * sigma_m ** 2) + g = g + (log_mu - log_prior) / sigma_m ** 2 return value, g -x0 = np.log([params.initial_strength] * n_seg) -if str(params.optimiser) == "tao": - # PETSc's own driver: the same objective and gradient, TAO's quasi-Newton - # update and line search; the bounded variant when bounds are given. +# %% [markdown] +# ## The inversion +# +# PETSc's TAO drives it by default: the same objective and gradient, a +# limited-memory quasi-Newton update and TAO's line search, and the bounded +# variant when bounds are given. SciPy's L-BFGS-B is the alternative. + +# %% +x0 = np.log([params.uw_initial_friction] * n_seg) +if str(params.uw_optimiser) == "tao": bounds = None - if str(params.bounds).strip(): - lo, hi = (float(b) for b in str(params.bounds).split(",")) + if str(params.uw_bounds).strip(): + lo, hi = (float(b) for b in str(params.uw_bounds).split(",")) bounds = (np.log([lo] * n_seg), np.log([hi] * n_seg)) x_best, info = uw.adjoint.minimise(objective, x0, max_evaluations=60, gradient_tolerance=1e-10, bounds=bounds, method="blmvm" if bounds else "lmvm") else: + from scipy.optimize import minimize result = minimize(objective, x0, jac=True, method="L-BFGS-B", options={"maxiter": 40, "gtol": 1e-10}) x_best = result.x -uw.pprint(f"recovered {np.exp(x_best)} true {true_values} " - f"after {len(history)} evaluations") +print(f"recovered {np.exp(x_best)} true {true_values} after {len(history)} evaluations") -# --- what the figure needs ----------------------------------------------------- -# Uplift-rate profiles along the top at the truth, the start and the answer, the -# weak-plane viscosity on a grid, and the path the strengths took. +# %% [markdown] +# ## What the figure needs +# +# Uplift-rate profiles along the top at the truth, the start and the answer, +# the weak-plane viscosity on a grid, and the path the coefficients took. +# Lengths are saved in kilometres and velocities in millimetres per year. + +# %% xs = np.linspace(0.0, 2.0, 161) top = np.column_stack([xs, np.full_like(xs, 1.0 - 1e-6)]) profiles = {} -for label, values in (("true", true_values), ("initial", [params.initial_strength] * n_seg), +for label, values in (("true", true_values), ("initial", [params.uw_initial_friction] * n_seg), ("recovered", list(np.exp(x_best)))): - set_strengths(values) + set_friction(values) forward(f"profile {label}") - profiles[label] = np.asarray(uw.function.evaluate(v.sym[1], top)).ravel() -set_strengths(true_values) + profiles[label] = np.asarray(uw.function.evaluate(v.sym[1], top)).ravel() * mm_per_yr +set_friction(true_values) forward("truth again") # the field on the grid is the truth's gx, gy = np.meshgrid(np.linspace(0, 2, 201), np.linspace(0, 1, 101)) grid = np.column_stack([gx.ravel(), gy.ravel()]) eta_1_grid = np.asarray(uw.function.evaluate(eta_1, grid)).reshape(gx.shape) -tag = f"{what}" + (f"_noise{float(params.noise):g}" if float(params.noise) > 0 else "") + (f"_prior{sigma_m:g}" if sigma_m > 0 else "") -np.savez(f"fault_friction_{tag}_data.npz", xs=xs, gx=gx, gy=gy, eta_1=eta_1_grid, - points=np.array(points), true=np.array(true_values), band=params.band, + +tag = (f"{what}" + (f"_noise{float(params.uw_noise):g}" if float(params.uw_noise) > 0 else "") + + (f"_prior{sigma_m:g}" if sigma_m > 0 else "")) +np.savez(f"fault_friction_{tag}_data.npz", xs=xs * km, gx=gx * km, gy=gy * km, eta_1=eta_1_grid, + points=np.array(points) * km, true=np.array(true_values), band=band_width * km, + length_unit="km", velocity_unit="mm/yr", history=np.array([[J, *vals] for J, vals in history]), **{f"uplift_{k}": val for k, val in profiles.items()}) diff --git a/docs/examples/adjoint/fault_segments/fault_friction_key.md b/docs/examples/adjoint/fault_segments/fault_friction_key.md new file mode 100644 index 000000000..dabbd0255 --- /dev/null +++ b/docs/examples/adjoint/fault_segments/fault_friction_key.md @@ -0,0 +1,38 @@ +## Key — what each part solved in fault_friction +*run started 2026-09-19T17:44:38+10:00* + +### Stokes(v) +`SNES_Stokes`, unknown `v`, 2-D; recorded at step 0 + +Residual $\int F_0\,\phi + F_1 \cdot \nabla\phi = 0$ with + +$$\mathbf{f}_0\left( \mathbf{u} \right) = \left[\begin{matrix}0\\10.0\end{matrix}\right]$$ +*Velocity equation body force term (pointwise).* + +$$\mathbf{F}_1\left( \mathbf{u} \right) = \left[\begin{matrix}\eta_0 \uplambda \left({v}_{ 0,0}(\mathbf{x}) + {v}_{ 1,1}(\mathbf{x})\right) - 2 \left(2 \eta_0 - 2 \eta_1\right) \left(- 2 \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) + \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right)\right) \left(\frac{{v}_{ 0,1}(\mathbf{x})}{2} + \frac{{v}_{ 1,0}(\mathbf{x})}{2}\right) + 2 \left(2 \eta_0 - 2 \eta_1\right) \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) {v}_{ 1,1}(\mathbf{x}) + \left(2 \eta_0 - \left(2 \eta_0 - 2 \eta_1\right) \left(- 2 \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) + 2 \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right)\right)\right) {v}_{ 0,0}(\mathbf{x}) - {p}(\mathbf{x}) & - \left(2 \eta_0 - 2 \eta_1\right) \left(- 2 \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) + \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right)\right) {v}_{ 0,0}(\mathbf{x}) - \left(2 \eta_0 - 2 \eta_1\right) \left(- 2 \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) + \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right)\right) {v}_{ 1,1}(\mathbf{x}) + 2 \left(\eta_0 - \left(2 \eta_0 - 2 \eta_1\right) \left(- 2 \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) + \frac{\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}}{2} + \frac{\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}}{2}\right)\right) \left(\frac{{v}_{ 0,1}(\mathbf{x})}{2} + \frac{{v}_{ 1,0}(\mathbf{x})}{2}\right)\\- \left(2 \eta_0 - 2 \eta_1\right) \left(- 2 \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) + \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right)\right) {v}_{ 0,0}(\mathbf{x}) - \left(2 \eta_0 - 2 \eta_1\right) \left(- 2 \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) + \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right)\right) {v}_{ 1,1}(\mathbf{x}) + 2 \left(\eta_0 - \left(2 \eta_0 - 2 \eta_1\right) \left(- 2 \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) + \frac{\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}}{2} + \frac{\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}}{2}\right)\right) \left(\frac{{v}_{ 0,1}(\mathbf{x})}{2} + \frac{{v}_{ 1,0}(\mathbf{x})}{2}\right) & \eta_0 \uplambda \left({v}_{ 0,0}(\mathbf{x}) + {v}_{ 1,1}(\mathbf{x})\right) - 2 \left(2 \eta_0 - 2 \eta_1\right) \left(- 2 \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) + \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right)\right) \left(\frac{{v}_{ 0,1}(\mathbf{x})}{2} + \frac{{v}_{ 1,0}(\mathbf{x})}{2}\right) + 2 \left(2 \eta_0 - 2 \eta_1\right) \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) {v}_{ 0,0}(\mathbf{x}) + \left(2 \eta_0 - \left(2 \eta_0 - 2 \eta_1\right) \left(- 2 \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) + 2 \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right)\right)\right) {v}_{ 1,1}(\mathbf{x}) - {p}(\mathbf{x})\end{matrix}\right]$$ +*Velocity equation flux/stress term (pointwise).* + +where +- $\eta_0$ $= 1$ — Shear viscosity +- $\eta_1$ $= \text{(an expression; in the record)}$ — Second viscosity + - $\mu_{1}$ $= 0.05$ — friction coefficient, flat + - $\mu_{2}$ $= 0.15$ — friction coefficient, lower ramp + - $\mu_{3}$ $= 0.25$ — friction coefficient, upper ramp + - $\mu_{4}$ $= 0.4$ — friction coefficient, near surface + - $\varepsilon$ $= 10^{-18}$ — vanishingly small value +- $\uplambda$ $= 0$ — Numerical Penalty + +$$\mathbf{h}_0\left( \mathbf{p} \right) = \left[\begin{matrix}{v}_{ 0,0}(\mathbf{x}) + {v}_{ 1,1}(\mathbf{x})\end{matrix}\right]$$ +*Pressure equation constraint term (continuity).* + +Boundary conditions: +- essential on Bottom: $\left[\begin{matrix}0.0\\0.0\end{matrix}\right]$ +- essential on Left: $\left[\begin{matrix}0.5\\\infty\end{matrix}\right]$ +- essential on Right: $\left[\begin{matrix}-0.5\\\infty\end{matrix}\right]$ + +Given: +- `bodyforce` $= (0, -10)$ — body force per unit volume; F0 is its negative +- `penalty` $= 0$ — augmented-Lagrangian grad-div penalty (0 = off) +- `TransverseIsotropicFlowModel.director` $= \text{(an expression; in the record)}$ — constitutive parameter +- `TransverseIsotropicFlowModel.shear_viscosity_0` $= 1$ — constitutive parameter +- `TransverseIsotropicFlowModel.shear_viscosity_1` $= \text{(an expression; in the record)}$ — constitutive parameter diff --git a/docs/examples/adjoint/fault_segments/fault_friction_transcript.svg b/docs/examples/adjoint/fault_segments/fault_friction_transcript.svg new file mode 100644 index 000000000..873b97401 --- /dev/null +++ b/docs/examples/adjoint/fault_segments/fault_friction_transcript.svg @@ -0,0 +1,390 @@ + + +Transcript — fault_friction +started 2026-09-19T17:44:38+10:00 · complete, 40 steps · 2 parts +scales: domain_depth 10 km viscosity 1e+21 Pa s convergence_rate 8.359 mm/yr +step +t +dt +Stokes(v) +adjoint Stokes(v) + +0 +0 +0 + +✅ + +1 +0 +0 + + + +✅ +✅ + +2 +0 +0 + +✅ + +3 +0 +0 + +✅ + +4 +0 +0 + +✅ + +5 +0 +0 + +✅ + +6 +0 +0 + +✅ + +7 +0 +0 + +✅ + +8 +0 +0 + +✅ + +9 +0 +0 + +✅ + +10 +0 +0 + + + +✅ +✅ + +11 +0 +0 + + + +✅ +✅ + +12 +0 +0 + + + +✅ +✅ + +13 +0 +0 + + + +✅ +✅ + +14 +0 +0 + + + +✅ +✅ + +15 +0 +0 + + + +✅ +✅ + +16 +0 +0 + + + +✅ +✅ + +17 +0 +0 + + + +✅ +✅ + +18 +0 +0 + + + +✅ +✅ + +19 +0 +0 + + + +✅ +✅ + +20 +0 +0 + + + +✅ +✅ + +21 +0 +0 + + + +✅ +✅ + +22 +0 +0 + + + +✅ +✅ + +23 +0 +0 + + + +✅ +✅ + +24 +0 +0 + + + +✅ +✅ + +25 +0 +0 + + + +✅ +✅ + +26 +0 +0 + + + +✅ +✅ + +27 +0 +0 + + + +✅ +✅ + +28 +0 +0 + + + +✅ +✅ + +29 +0 +0 + + + +✅ +✅ + +30 +0 +0 + + + +✅ +✅ + +31 +0 +0 + + + +✅ +✅ + +32 +0 +0 + + + +✅ +✅ + +33 +0 +0 + + + +✅ +✅ + +34 +0 +0 + + + +✅ +✅ + +35 +0 +0 + + + +✅ +✅ + +36 +0 +0 + +✅ + +37 +0 +0 + +✅ + +38 +0 +0 + +✅ + +39 +0 +0 + +✅ + + + + + +ran; a step reads down, the path joining what ran in the order it ran +✅ +solved, and converged +⚠️ +converged, but a fieldsplit block hit its iteration cap — that block did not solve +❌ +did not converge; the reason is in the record + +ran in a step that was then abandoned + +did nothing in this step + + +the steps between did exactly this, unchanged; first and last values are shown where they differ +Key — the parts, and what is in their residuals +Stokes(v) +recorded at step 0 + +— with +Velocity equation body force term (pointwise). + + + + + + + + + +Velocity equation flux/stress term (pointwise). +(the form is in the record; uw.transcript_key renders it) + +— Shear viscosity + +— Second viscosity + +— friction coefficient, flat + +— friction coefficient, lower ramp + +— friction coefficient, upper ramp + +— friction coefficient, near surface + +— vanishingly small value + +— Numerical Penalty +Pressure equation constraint term (continuity). + + +essential on Bottom: (0, 0) +essential on Left: (0.5, free) +essential on Right: (-0.5, free) + diff --git a/docs/examples/adjoint/fault_segments/plot_fault_segments.py b/docs/examples/adjoint/fault_segments/plot_fault_segments.py index 2067dbacd..8da98b596 100644 --- a/docs/examples/adjoint/fault_segments/plot_fault_segments.py +++ b/docs/examples/adjoint/fault_segments/plot_fault_segments.py @@ -13,6 +13,9 @@ friction = "friction" in source d = np.load(source) history, true = d["history"], d["true"] +length_unit = str(d["length_unit"]) if "length_unit" in d else "" +velocity_unit = str(d["velocity_unit"]) if "velocity_unit" in d else "" +x_max, y_max = float(d["gx"].max()), float(d["gy"].max()) names = ["flat", "lower ramp", "upper ramp", "near surface"] n_seg = len(true) @@ -25,17 +28,19 @@ cmap="viridis", extend="both") band = float(d["band"]) if "band" in d else 0.08 # the surface band under the top, where the uplift rate (or the orientation) is read -ax.axhspan(1 - 2 * band, 1.0, color="white", alpha=0.35, lw=0) +ax.axhspan(y_max - 2 * band, y_max, color="white", alpha=0.35, lw=0) ax.plot(d["points"][:, 0], d["points"][:, 1], "wx", ms=7, mew=1.5) ax.set_aspect("equal") -ax.set_xlim(0, 2); ax.set_ylim(0, 1) -ax.set_xlabel("$x$"); ax.set_ylabel("$y$") +ax.set_xlim(0, x_max); ax.set_ylim(0, y_max) +ax.set_xlabel(f"$x$ ({length_unit})" if length_unit else "$x$") +ax.set_ylabel(f"$y$ ({length_unit})" if length_unit else "$y$") ax.set_title(r"$\log_{10}\eta_1$ at the truth; $\times$ points, white band: surface observations", fontsize=10) ax = axes[1] for label, style in (("true", "k-"), ("initial", "C3--"), ("recovered", "C0:")): ax.plot(d["xs"], d[f"uplift_{label}"], style, lw=1.6, label=label) -ax.set_xlabel("$x$ along the surface"); ax.set_ylabel("uplift rate $v_y$") +ax.set_xlabel(f"$x$ along the surface ({length_unit})" if length_unit else "$x$ along the surface") +ax.set_ylabel(f"uplift rate $v_y$ ({velocity_unit})" if velocity_unit else "uplift rate $v_y$") ax.set_title("surface uplift rate", fontsize=10) ax.legend(fontsize=8, frameon=False) diff --git a/docs/examples/adjoint/fault_segments/render_fault_friction.py b/docs/examples/adjoint/fault_segments/render_fault_friction.py index e23e3aecf..8aa3fc875 100644 --- a/docs/examples/adjoint/fault_segments/render_fault_friction.py +++ b/docs/examples/adjoint/fault_segments/render_fault_friction.py @@ -1,7 +1,10 @@ """PyVista renders of the friction example at the true coefficients. -Writes log10 eta_1 (the plane's viscosity), the velocity, and the pressure -to ~/+Simulations/adjoint_fault_example/, on the mesh's own triangulation. +Writes log10 eta_1 (the plane's viscosity), the slip rate, the velocity and +the pressure to ~/+Simulations/adjoint_fault_example/, on the mesh's own +triangulation. The model is the notebook's in its own units: lengths in +units of the depth (10 km), stresses in units of the viscous stress of the +shortening (26 MPa), velocities in units of the convergence rate (8.4 mm/yr). """ import math import os diff --git a/src/underworld3/model.py b/src/underworld3/model.py index 96c0d2f9b..1dd66a908 100644 --- a/src/underworld3/model.py +++ b/src/underworld3/model.py @@ -31,6 +31,7 @@ # Import the Pint-native implementation import os +import re import sys sys.path.append(os.path.dirname(__file__)) @@ -261,15 +262,23 @@ def _quantity_parts(value): return None, None +_UNIT_SYMBOLS = { + "second": "s", "minute": "min", "hour": "hr", "day": "d", + "year": "yr", "kiloyear": "kyr", "megayear": "Myr", "gigayear": "Gyr", + "meter": "m", "kilometer": "km", "centimeter": "cm", "millimeter": "mm", + "kelvin": "K", "kilogram": "kg", "pascal": "Pa", "newton": "N", + "joule": "J", "watt": "W", +} + + def _abbreviate_unit(unit): - """A short unit name for a column header. Falls back to the full name.""" + """A short unit name for a column header or a scale: each unit name in a + pint unit string by its symbol, "pascal * second" as "Pa s" and + "millimeter / year" as "mm/yr". Names without a symbol stay as they are.""" if unit is None: return "" - return { - "second": "s", "minute": "min", "hour": "hr", "day": "d", - "year": "yr", "kiloyear": "kyr", "megayear": "Myr", "gigayear": "Gyr", - "meter": "m", "kilometer": "km", "kelvin": "K", "kilogram": "kg", - }.get(str(unit), str(unit)) + text = re.sub(r"[A-Za-z_]+", lambda m: _UNIT_SYMBOLS.get(m.group(0), m.group(0)), str(unit)) + return text.replace(" * ", " ").replace(" / ", "/") def _in_units_of(value, unit): @@ -1256,12 +1265,31 @@ def _run_header(self): """The record that opens a run in the log, so the file is self-describing.""" from datetime import datetime, timezone - scales = {} + # The record keeps the fundamental scales, and beside them the + # reference quantities as they were declared, in the units they were + # quoted in. The readable header reports the declaration when there is + # one: "domain_depth 10 km", not the fundamental length in metres. + scales, reference = {}, {} try: for name, scale in (self.get_fundamental_scales() or {}).items(): scales[str(name)] = _jsonable_quantity(scale) except Exception: scales = {} + try: + from .scaling import units as ureg + + for name, quantity in (self.get_reference_quantities() or {}).items(): + if not (isinstance(quantity, dict) and "magnitude" in quantity): + continue + magnitude, unit = float(quantity["magnitude"]), str(quantity.get("units")) + # A quantity declared as an expression of others carries the + # raw composite of their units; in SI base units it reads. + if "**" in unit or unit.count("/") > 1: + base = ureg.Quantity(magnitude, unit).to_base_units() + magnitude, unit = float(base.magnitude), str(base.units) + reference[str(name)] = {"magnitude": magnitude, "units": unit} + except Exception: + reference = {} script = None try: entry = sys.argv[0] if sys.argv else "" @@ -1275,6 +1303,7 @@ def _run_header(self): "script": script, "started": datetime.now().astimezone().isoformat(timespec="seconds"), "scales": scales, + "reference": reference, } # ------------------------------------------------------------------ @@ -1287,7 +1316,7 @@ def _render_transcript_text(self, payload): kind = payload.get("kind") if kind == "run": - scales = payload.get("scales") or {} + scales = payload.get("reference") or payload.get("scales") or {} summary = " | ".join( f"{name} {value['magnitude']:.4g} {_abbreviate_unit(value['units'])}" for name, value in scales.items() @@ -2220,12 +2249,12 @@ def set_reference_quantities(self, verbose=False, nondimensional_scaling=True, * # Enable/disable non-dimensionalization based on parameter import underworld3 as uw - if nondimensional_scaling: - uw.use_nondimensional_scaling(True) - uw.pprint("✓ Units system active with automatic non-dimensionalization", proc=0) - else: - uw.use_nondimensional_scaling(False) - uw.pprint("⚠ Expert mode: Units active WITHOUT non-dimensionalization", proc=0) + uw.use_nondimensional_scaling(bool(nondimensional_scaling)) + if verbose: + if nondimensional_scaling: + uw.pprint("Units system active with automatic non-dimensionalization", proc=0) + else: + uw.pprint("Units active without non-dimensionalization", proc=0) uw.pprint(" (Warning: This mode may have numerical conditioning issues)", proc=0) def get_reference_quantities(self): @@ -2535,9 +2564,9 @@ def _solve_available_dimensions(self, matrix, names, fundamental_dims, ureg): import underworld3 as uw # Informational message about missing dimensions (not an error!) - if missing_dims: - uw.pprint(f"ℹ️ Dimensional coverage: {covered_dims}", proc=0) - uw.pprint(f" (Not covered: {missing_dims} - will fail only if needed)", proc=0) + if missing_dims and getattr(self, "_verbose_units", False): + uw.pprint(f"Dimensional coverage: {covered_dims}", proc=0) + uw.pprint(f" (not covered: {missing_dims}; needed only if a quantity uses them)", proc=0) # Extract sub-matrix for covered dimensions only sub_matrix = matrix[:, covered_indices] diff --git a/src/underworld3/utilities/transcript_report.py b/src/underworld3/utilities/transcript_report.py index 63d5d565a..5647f953e 100644 --- a/src/underworld3/utilities/transcript_report.py +++ b/src/underworld3/utilities/transcript_report.py @@ -523,13 +523,19 @@ def _unit(value): def _short_unit(unit): - if unit is None: - return "" - return { - "second": "s", "minute": "min", "hour": "hr", "day": "d", "year": "yr", - "kiloyear": "kyr", "megayear": "Myr", "gigayear": "Gyr", - "meter": "m", "kilometer": "km", "kelvin": "K", "kilogram": "kg", - }.get(str(unit), str(unit)) + """The unit symbol the log uses, so the figure and the log agree.""" + from underworld3.model import _abbreviate_unit + return _abbreviate_unit(unit) + + +def _scales_line(header, sep=" "): + """The run's scales as one line: the reference quantities as they were + declared, in their own units, or the fundamental scales for a run that + declared none. Empty for a nondimensional run.""" + scales = header.get("reference") or header.get("scales") or {} + items = [f"{name} {value['magnitude']:.4g} {_short_unit(value['units'])}" + for name, value in scales.items() if isinstance(value, dict)] + return "scales: " + sep.join(items) if items else "" def _converted(value, unit): @@ -870,12 +876,8 @@ def draw_header(y, first): bits.append(f"{len(notes)} backtrack(s)") canvas.text(_MARGIN, y + 8, " · ".join(bits), size=8.5, fill=_MUTED) y += 13 - scales = header.get("scales") or {} - if scales: - canvas.text(_MARGIN, y + 8, "scales: " + " ".join( - f"{name} {value['magnitude']:.4g} {_short_unit(value['units'])}" - for name, value in scales.items() if isinstance(value, dict) - ), size=8, fill=_MUTED) + if _scales_line(header): + canvas.text(_MARGIN, y + 8, _scales_line(header), size=8, fill=_MUTED) y += 12 y += 10 # column captions @@ -1790,7 +1792,11 @@ def _transcript_layout(header, steps, notes, entry, title=None, width=PAGE_W, bits.append("no terminator: still running, or interrupted") bits.append(f"{len(parts)} parts") canvas.text(_MARGIN, y + 8, " · ".join(bits), size=8.5, fill=_MUTED) - y += 22 + y += 13 + if _scales_line(header): + canvas.text(_MARGIN, y + 8, _scales_line(header), size=8, fill=_MUTED) + y += 12 + y += 9 # --- columns --- gutter = _MARGIN + 46.0 @@ -2099,6 +2105,8 @@ def transcript_table(source, run=-1, width=11, collapse=True): out.append(f"transcript · {_run_title(header, fallback='')}".rstrip(" ·")) if header.get("started"): out.append(f"started {header['started']}") + if _scales_line(header, sep=" | "): + out.append(_scales_line(header, sep=" | ")) if ended: out.append(f"complete — {ended.get('steps', len(steps))} step(s)") elif entry.get("live"):