diff --git a/CLAUDE.md b/CLAUDE.md index 82745c01..f2648f50 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 58ad6026..66a7ec41 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 c5fcd662..5dfeb9a4 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 c9e45368..08f31294 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 00000000..80384778 --- /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 00000000..f6cad0e4 --- /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 00000000..fd93a9e8 --- /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 00000000..25ce8bc8 --- /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 00000000..321b3047 --- /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 00000000..da5dfa2a --- /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 bf3faf59..03ee873b 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 00000000..8f1b11ba --- /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 2647afb1..38dc9275 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 8d0845db..1098be38 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 fe9ae544..fd20c04e 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 5e3c13ed..888d08ac 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 612b519d..daeb6e73 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 00000000..0bbfda38 --- /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 00000000..31d5e267 --- /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 00000000..b102ca11 --- /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 b143ef89..87fdb2bd 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 8bd74fd6..6606ce8e 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())