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/fault_segments/README.md b/docs/examples/adjoint/fault_segments/README.md new file mode 100644 index 00000000..48d5610f --- /dev/null +++ b/docs/examples/adjoint/fault_segments/README.md @@ -0,0 +1,24 @@ +# Friction on a listric fault from surface observations + +The example for the technical note *A Discrete Adjoint from the Run Record* +(UWTN 2026-020). Four friction coefficients on a listric fault are recovered +from the surface uplift rate and the shear stress at five interior points, +with the gradient from the solver's own adjoint. + +| file | what it does | +|---|---| +| `fault_friction.ipynb` | the example as a notebook | +| `fault_friction.py` | the same, as a script in jupytext percent format; `python fault_friction.py -uw_check_only 1` runs the gradient check alone | +| `plot_fault_segments.py` | the figure: the weak plane, the uplift profiles, and the path of the coefficients, from the `_data.npz` a run writes | +| `plot_convergence.py` | convergence under the three observation sets | +| `plot_noise.py` | recovered coefficients against the noise, the bounds and the prior | +| `render_fault_friction.py` | PyVista renders of the truth: the plane's viscosity, the slip rate, the uplift and the pressure | + +The gradient check takes about three minutes and the inversion about ten on +a laptop. The transcript of a run is written to `transcripts//`, +and `uw.transcript_figure(...)` draws it. + +The problem is stated in kilometres, pascal seconds and millimetres per year, +and the solver works in units of the depth, the viscosity and the +convergence rate. The observation sets, the noise and the prior are switches +in the parameter block at the top of the notebook. diff --git a/docs/examples/adjoint/fault_segments/fault_friction.ipynb b/docs/examples/adjoint/fault_segments/fault_friction.ipynb new file mode 100644 index 00000000..e4382719 --- /dev/null +++ b/docs/examples/adjoint/fault_segments/fault_friction.ipynb @@ -0,0 +1,676 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "5a56e0e3", + "metadata": {}, + "source": [ + "# Friction on a listric fault from surface observations\n", + "\n", + "A listric fault under horizontal shortening and gravity: a flat décollement\n", + "at depth that steepens through a circular ramp to a dip of sixty degrees at\n", + "the surface. The fault is a weak plane in a transversely isotropic viscosity,\n", + "with no cut in the mesh. The plane yields at the Coulomb stress\n", + "$\\tau_y = C + \\mu p$, and the friction coefficient $\\mu$ takes a different\n", + "value on the flat, on the lower and upper parts of the ramp, and near the\n", + "surface. Those four coefficients are the unknowns.\n", + "\n", + "The observations are the uplift rate along the top surface and the shear\n", + "stress near a handful of interior points, read from a run at the true\n", + "coefficients. The gradient of the misfit with respect to each coefficient\n", + "comes from one call, `stokes.gradient(misfit, parameters=...)`: the solver\n", + "assembles the adjoint operator from its own Jacobian kernels with trial and\n", + "test functions exchanged, solves it, and differentiates its residual\n", + "symbolically with respect to each named coefficient. Nothing is differenced.\n", + "\n", + "The notebook checks that gradient against central differences, then hands\n", + "the misfit and its gradient to PETSc's TAO to recover the four\n", + "coefficients. Each misfit evaluation is one step of zero length in the\n", + "model's record, so the run's transcript lists the forward solve and the\n", + "adjoint solve it made." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5a51a9d0", + "metadata": {}, + "outputs": [], + "source": [ + "import math\n", + "\n", + "import numpy as np\n", + "import sympy\n", + "\n", + "import underworld3 as uw" + ] + }, + { + "cell_type": "markdown", + "id": "033a1223", + "metadata": {}, + "source": [ + "### Configurable parameters\n", + "\n", + "Default values are defined as named constants below. From the command line,\n", + "override them with PETSc-style flags:\n", + "\n", + "```bash\n", + "python fault_friction.py -uw_check_only 1\n", + "python fault_friction.py -uw_observations surface_strain\n", + "python fault_friction.py -uw_noise 0.03 -uw_prior_sigma 1 -uw_bounds 0.005,1\n", + "```\n", + "\n", + "The problem depends on one dimensionless number: the lithostatic pressure at\n", + "the base of the model over the viscous stress of the shortening,\n", + "$\\rho g L^2 / (\\eta V)$. We fix the depth, the viscosity and the density,\n", + "and that ratio sets the convergence rate." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0d0d83bb", + "metadata": {}, + "outputs": [], + "source": [ + "# --- Scales ---\n", + "DOMAIN_DEPTH = uw.quantity(10, \"km\") # the box is 20 km wide and 10 km deep\n", + "REF_VISCOSITY = uw.quantity(1e21, \"Pa*s\") # the bulk viscosity\n", + "DENSITY = uw.quantity(2700, \"kg/m**3\")\n", + "GRAVITY = uw.quantity(9.81, \"m/s**2\")\n", + "PRESSURE_TO_STRESS = 10 # rho g L^2 / (eta V)\n", + "CONVERGENCE_RATE = (DENSITY * GRAVITY * DOMAIN_DEPTH**2 / (PRESSURE_TO_STRESS * REF_VISCOSITY)).to(\"mm/yr\") # 8.4 mm/yr\n", + "REF_STRESS = (REF_VISCOSITY * CONVERGENCE_RATE / DOMAIN_DEPTH).to(\"MPa\") # 26 MPa\n", + "\n", + "# --- Geometry ---\n", + "CELL_SIZE = DOMAIN_DEPTH / 12 # base mesh; refined once, so half this\n", + "SURFACE_DIP = 60.0 # degrees, where the ramp reaches the surface\n", + "FLAT_DEPTH = uw.quantity(3, \"km\") # height of the décollement above the base\n", + "SURFACE_X = uw.quantity(19, \"km\") # where the fault reaches the surface\n", + "BAND_HALF_WIDTH = uw.quantity(0.8, \"km\") # of the weak plane\n", + "\n", + "# --- Rheology ---\n", + "COHESION = (0.05 * REF_STRESS).to(\"MPa\") # C in tau_y = C + mu p; 1.3 MPa\n", + "TRUE_FRICTION = \"0.05,0.15,0.25,0.4\" # flat, lower ramp, upper ramp, near surface\n", + "INITIAL_FRICTION = 0.2 # the starting guess, every segment\n", + "\n", + "# --- The inversion ---\n", + "CHECK_ONLY = 0 # 1: the gradient check, no inversion\n", + "OPTIMISER = \"tao\" # tao (PETSc quasi-Newton) | scipy (L-BFGS-B)\n", + "OBSERVATIONS = \"uplift+stress\" # uplift+stress | orientation_points | surface_strain\n", + "NOISE = 0.0 # on the observed velocity, as a fraction of its rms\n", + "SEED = 7\n", + "PRIOR_SIGMA = 0.0 # width of a Gaussian prior on log mu; 0 for none\n", + "BOUNDS = \"\" # friction bounds lo,hi for TAO's blmvm; empty for none\n", + "\n", + "params = uw.Params(\n", + " uw_cell_size=uw.Param(CELL_SIZE, description=\"base mesh cell size\"),\n", + " uw_surface_dip=uw.Param(SURFACE_DIP, description=\"dip of the ramp at the surface, degrees\"),\n", + " uw_flat_depth=uw.Param(FLAT_DEPTH, description=\"height of the décollement above the base\"),\n", + " uw_surface_x=uw.Param(SURFACE_X, description=\"where the fault reaches the surface\"),\n", + " uw_band=uw.Param(BAND_HALF_WIDTH, description=\"half-width of the weak plane\"),\n", + " uw_cohesion=uw.Param(COHESION, description=\"cohesion C in tau_y = C + mu p\"),\n", + " uw_true_friction=uw.Param(TRUE_FRICTION, description=\"friction: flat, lower ramp, upper ramp, near surface\"),\n", + " uw_initial_friction=uw.Param(INITIAL_FRICTION, description=\"starting guess, every segment\"),\n", + " uw_check_only=uw.Param(CHECK_ONLY, description=\"1: gradient check only\"),\n", + " uw_optimiser=uw.Param(OPTIMISER, description=\"tao | scipy\"),\n", + " uw_observations=uw.Param(OBSERVATIONS, description=\"uplift+stress | orientation_points | surface_strain\"),\n", + " uw_noise=uw.Param(NOISE, description=\"noise on the observed velocity, fraction of its rms\"),\n", + " uw_seed=uw.Param(SEED, description=\"seed for the noise\"),\n", + " uw_prior_sigma=uw.Param(PRIOR_SIGMA, description=\"width of the prior on log mu; 0 for none\"),\n", + " uw_bounds=uw.Param(BOUNDS, description=\"friction bounds lo,hi; empty for none\"),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "df86a530", + "metadata": {}, + "source": [ + "## The model and its scales\n", + "\n", + "The model is declared first, with the three reference quantities that fix\n", + "the scaling: the depth, the viscosity and the convergence rate. The solver\n", + "then works in units of those, so the box is $2 \\times 1$, the bulk viscosity\n", + "is one, the walls close at one, and the body force is the pressure-to-stress\n", + "ratio. `coords` on a variable come back in physical units, and `coords_nd`\n", + "are the model's own." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a9c02f07", + "metadata": {}, + "outputs": [], + "source": [ + "orchestration_model = uw.get_default_model()\n", + "orchestration_model.set_reference_quantities(\n", + " domain_depth=DOMAIN_DEPTH,\n", + " viscosity=REF_VISCOSITY,\n", + " convergence_rate=CONVERGENCE_RATE,\n", + ")\n", + "\n", + "\n", + "def _nd(quantity):\n", + " \"\"\"The plain number a dimensionless quantity stands for.\"\"\"\n", + " try:\n", + " return float(quantity.to(\"dimensionless\").magnitude)\n", + " except AttributeError:\n", + " return float(quantity)\n", + "\n", + "\n", + "def _lengths(quantity):\n", + " \"\"\"A length in units of the model depth.\"\"\"\n", + " return _nd(quantity / DOMAIN_DEPTH)\n", + "\n", + "\n", + "mm_per_yr = float(CONVERGENCE_RATE.to(\"mm/yr\").magnitude) # one model velocity unit\n", + "km = float(DOMAIN_DEPTH.to(\"km\").magnitude) # one model length unit" + ] + }, + { + "cell_type": "markdown", + "id": "db10286a", + "metadata": {}, + "source": [ + "## The mesh\n", + "\n", + "Refined once from the base cell size, which gives the velocity block a\n", + "multigrid hierarchy. Without one it falls back to an algebraic coarsening\n", + "that hits its iteration cap on this problem, and an inexact Newton step then\n", + "converges only linearly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8a456b31", + "metadata": {}, + "outputs": [], + "source": [ + "mesh = uw.meshing.UnstructuredSimplexBox(\n", + " minCoords=(0.0, 0.0), maxCoords=(2.0, 1.0),\n", + " cellSize=_lengths(params.uw_cell_size), qdegree=3, refinement=1,\n", + ")\n", + "x, y = mesh.X\n", + "\n", + "v = uw.discretisation.MeshVariable(\"v\", mesh, mesh.dim, degree=2)\n", + "p = uw.discretisation.MeshVariable(\"p\", mesh, 1, degree=1, continuous=True)\n", + "v_obs = uw.discretisation.MeshVariable(\"v_obs\", mesh, mesh.dim, degree=2)" + ] + }, + { + "cell_type": "markdown", + "id": "32c94ffc", + "metadata": {}, + "source": [ + "## The fault\n", + "\n", + "A flat décollement runs from the left wall to $x = x_c$, then a circular\n", + "ramp of radius $R$ about $(x_c, y_c)$ leaves the flat horizontally and\n", + "reaches the surface at the given dip. The signed distance to the fault and\n", + "the position along it are exact on each piece. The director is the normal\n", + "at the nearest point: vertical on the flat, radial on the ramp." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "35be5a5f", + "metadata": {}, + "outputs": [], + "source": [ + "flat_depth = _lengths(params.uw_flat_depth)\n", + "band_width = _lengths(params.uw_band)\n", + "phi_top = math.radians(params.uw_surface_dip) - math.pi / 2 # angle of the surface point about the centre\n", + "R = (1.0 - flat_depth) / (1.0 + math.sin(phi_top))\n", + "yc = flat_depth + R\n", + "xc = _lengths(params.uw_surface_x) - R * math.cos(phi_top)\n", + "r = sympy.sqrt((x - xc) ** 2 + (y - yc) ** 2)\n", + "phi = sympy.atan2(y - yc, x - xc)\n", + "on_flat = x < xc\n", + "d = sympy.Piecewise((y - flat_depth, on_flat), (R - r, True))\n", + "s = sympy.Piecewise((x, on_flat), (xc + R * (phi + sympy.pi / 2), True))\n", + "n_hat = sympy.Matrix([[sympy.Piecewise((0, on_flat), ((x - xc) / r, True)),\n", + " sympy.Piecewise((1, on_flat), ((y - yc) / r, True))]])\n", + "ramp = R * (phi_top + math.pi / 2)\n", + "length = xc + ramp\n", + "band = sympy.exp(-(d / band_width) ** 2)" + ] + }, + { + "cell_type": "markdown", + "id": "67982371", + "metadata": {}, + "source": [ + "## The segments and their friction coefficients\n", + "\n", + "Four segments along the fault: the flat, then the ramp in three equal\n", + "parts. Each coefficient is a named expression, so the solver can\n", + "differentiate its residual with respect to it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "113b01c3", + "metadata": {}, + "outputs": [], + "source": [ + "edges = [0.0, xc, xc + ramp / 3, xc + 2 * ramp / 3, length]\n", + "names = [\"flat\", \"lower ramp\", \"upper ramp\", \"near surface\"]\n", + "n_seg = len(names)\n", + "friction_coefficients = [\n", + " uw.expression(rf\"\\mu_{{{k + 1}}}\", params.uw_initial_friction, f\"friction coefficient, {names[k]}\")\n", + " for k in range(n_seg)\n", + "]\n", + "\n", + "\n", + "def segment(k):\n", + " \"\"\"A smooth indicator for segment k along the fault, in [0, 1].\"\"\"\n", + " on = 1 if k == 0 else (1 + sympy.tanh((s - edges[k]) / band_width)) / 2\n", + " off = 1 if k == n_seg - 1 else (1 - sympy.tanh((s - edges[k + 1]) / band_width)) / 2\n", + " return on * off\n", + "\n", + "\n", + "def set_friction(values):\n", + " for coefficient, value in zip(friction_coefficients, values):\n", + " coefficient.sym = float(value)" + ] + }, + { + "cell_type": "markdown", + "id": "49faa869", + "metadata": {}, + "source": [ + "## Coulomb yield on the plane\n", + "\n", + "The shear strain rate resolved on the plane is $\\hat t \\cdot E \\cdot \\hat n$.\n", + "The plane's viscosity is the harmonic combination of the bulk viscosity and\n", + "the yield stress over that rate, which is smooth everywhere and tends to\n", + "$\\tau_y / 2\\dot\\varepsilon_s$ where the plane slips. Compression is positive.\n", + "Where the dynamic pressure is tensile the plane keeps its cohesion and no\n", + "more." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "db143519", + "metadata": {}, + "outputs": [], + "source": [ + "eta_0 = 1\n", + "cohesion = _nd(params.uw_cohesion / REF_STRESS)\n", + "rho_g = _nd(DENSITY * GRAVITY * DOMAIN_DEPTH / REF_STRESS)\n", + "\n", + "E = mesh.vector.strain_tensor(v.sym)\n", + "t_hat = sympy.Matrix([[-n_hat[1], n_hat[0]]])\n", + "e_s = sympy.sqrt((t_hat * E * n_hat.T)[0] ** 2 + uw.maths.functions.vanishing)\n", + "friction = sum(friction_coefficients[k] * segment(k) for k in range(n_seg))\n", + "tau_y = cohesion + friction * sympy.Max(p.sym[0], 0)\n", + "eta_plane = eta_0 * tau_y / (tau_y + 2 * eta_0 * e_s)\n", + "eta_1 = eta_0 - band * (eta_0 - eta_plane)" + ] + }, + { + "cell_type": "markdown", + "id": "11d02db3", + "metadata": {}, + "source": [ + "## The Stokes solver\n", + "\n", + "Shortening from both sides, a no-slip base, and a free top: the surface\n", + "velocity is the uplift rate. The residual is nonlinear in the velocity and\n", + "the pressure, so the forward solve is Newton." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c3812d14", + "metadata": {}, + "outputs": [], + "source": [ + "stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p)\n", + "stokes.constitutive_model = uw.constitutive_models.TransverseIsotropicFlowModel\n", + "stokes.constitutive_model.Parameters.shear_viscosity_0 = eta_0\n", + "stokes.constitutive_model.Parameters.shear_viscosity_1 = eta_1\n", + "stokes.constitutive_model.Parameters.director = n_hat\n", + "stokes.tolerance = 1e-8\n", + "stokes.bodyforce = sympy.Matrix([0, -rho_g])\n", + "\n", + "stokes.add_essential_bc((0.0, 0.0), \"Bottom\")\n", + "stokes.add_essential_bc((0.5, None), \"Left\")\n", + "stokes.add_essential_bc((-0.5, None), \"Right\")" + ] + }, + { + "cell_type": "markdown", + "id": "4c5f2e6e", + "metadata": {}, + "source": [ + "## The observations\n", + "\n", + "The misfit has a term on the top surface, a boundary integral of the uplift\n", + "rate, and a term in the volume around five interior points where the shear\n", + "stress is read. `gradient()` takes the terms as `{domain: integrand}`.\n", + "\n", + "On a traction-free surface the shear strain rate vanishes, so a stress\n", + "orientation read there is only a sign. Orientation is an interior\n", + "observable (boreholes, focal mechanisms). The surface gives velocities and\n", + "their tangential derivative, the geodetic strain rate. Both alternatives\n", + "are available through `uw_observations`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c934afd6", + "metadata": {}, + "outputs": [], + "source": [ + "points = [(0.3, 0.65), (0.75, 0.12), (1.25, 0.3), (0.9, 0.6), (1.55, 0.85)] # in units of the depth\n", + "w_points = sum(sympy.exp(-((x - px) ** 2 + (y - py) ** 2) / (2 * band_width) ** 2)\n", + " for px, py in points)\n", + "\n", + "\n", + "def shear_stress(field):\n", + " e = mesh.vector.strain_tensor(field.sym)\n", + " return 2 * eta_0 * e[0, 1]\n", + "\n", + "\n", + "def orientation(field):\n", + " \"\"\"The principal-stress orientation as the unit vector (cos 2theta, sin 2theta)\n", + " of the deviatoric strain rate, which is the stress orientation in the\n", + " isotropic bulk. A unit vector rather than an angle, so there is no wrap.\"\"\"\n", + " e = mesh.vector.strain_tensor(field.sym)\n", + " a, b = e[0, 0] - e[1, 1], 2 * e[0, 1]\n", + " norm = sympy.sqrt(a ** 2 + b ** 2 + uw.maths.functions.vanishing)\n", + " return sympy.Matrix([[a / norm, b / norm]])\n", + "\n", + "\n", + "what = str(params.uw_observations)\n", + "dq = orientation(v) - orientation(v_obs)\n", + "if what == \"uplift+stress\":\n", + " misfit = {\"Top\": (v.sym[1] - v_obs.sym[1]) ** 2 / 2,\n", + " None: w_points * (shear_stress(v) - shear_stress(v_obs)) ** 2 / 2}\n", + "elif what == \"orientation_points\":\n", + " misfit = {None: w_points * (dq[0] ** 2 + dq[1] ** 2) / 2}\n", + "elif what == \"surface_strain\":\n", + " misfit = {\"Top\": (v.sym[0].diff(x) - v_obs.sym[0].diff(x)) ** 2 / 2}\n", + "else:\n", + " raise ValueError(f\"observations: {what!r}\")\n", + "\n", + "\n", + "def misfit_value():\n", + " return sum(uw.adjoint.integral(mesh, term, where) for where, term in misfit.items())" + ] + }, + { + "cell_type": "markdown", + "id": "6b77d313", + "metadata": {}, + "source": [ + "## The forward solve and the gradient\n", + "\n", + "Every solve starts cold: six Newton iterations, and a misfit that does not\n", + "depend on the previous evaluation, which the finite-difference check needs.\n", + "Each evaluation is one step of zero length in the model's record." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b47600f1", + "metadata": {}, + "outputs": [], + "source": [ + "evaluations = [0]\n", + "\n", + "\n", + "def forward(label):\n", + " with orchestration_model.step(0.0, label=label):\n", + " stokes.solve(zero_init_guess=True)\n", + "\n", + "\n", + "def misfit_and_gradient(label=None):\n", + " \"\"\"The misfit and dJ/d(log mu) for each segment, by the adjoint.\"\"\"\n", + " evaluations[0] += 1\n", + " with orchestration_model.step(0.0, label=label or f\"eval {evaluations[0]}\"):\n", + " stokes.solve(zero_init_guess=True)\n", + " out = stokes.gradient(misfit, parameters=friction_coefficients)\n", + " # d/d(log mu) = mu d/d(mu)\n", + " grad = np.array([out[\"parameters\"][c] * float(c.sym) for c in friction_coefficients])\n", + " return out[\"J\"], grad" + ] + }, + { + "cell_type": "markdown", + "id": "0ca6e48b", + "metadata": {}, + "source": [ + "## The twin\n", + "\n", + "The observations are the velocity field at the true coefficients, with\n", + "optional noise drawn once at a fraction of each component's rms. Every\n", + "observation set reads `v_obs`, so the uplift, the stress and the strain\n", + "rate all inherit the same noise." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "82a72cd4", + "metadata": {}, + "outputs": [], + "source": [ + "true_values = [float(t) for t in str(params.uw_true_friction).split(\",\")][:n_seg]\n", + "set_friction(true_values)\n", + "forward(\"truth\")\n", + "v_obs.array[...] = np.asarray(v.array)\n", + "if float(params.uw_noise) > 0:\n", + " rng = np.random.default_rng(int(params.uw_seed))\n", + " obs = np.asarray(v_obs.array)\n", + " rms = np.sqrt(np.mean(obs ** 2, axis=0, keepdims=True))\n", + " v_obs.array[...] = obs + float(params.uw_noise) * rms * rng.standard_normal(obs.shape)\n", + " print(f\"noise: {float(params.uw_noise):.3f} of the rms per component, seed {int(params.uw_seed)}\")\n", + "J_truth = misfit_value() # v still holds the truth's velocity\n", + "\n", + "set_friction([params.uw_initial_friction] * n_seg)\n", + "J0, g0 = misfit_and_gradient(\"start\")\n", + "print(f\"true friction {true_values}\")\n", + "print(f\"initial J = {J0:.6e} dJ/dlog mu = {g0}\")" + ] + }, + { + "cell_type": "markdown", + "id": "2f1cc0f0", + "metadata": {}, + "source": [ + "## The gradient check\n", + "\n", + "Central differences in each log-coefficient, against the adjoint gradient." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d5a3c746", + "metadata": {}, + "outputs": [], + "source": [ + "h = 1e-3\n", + "for k in range(n_seg):\n", + " base = math.log(params.uw_initial_friction)\n", + " fd = []\n", + " for sign in (+1, -1):\n", + " values = [params.uw_initial_friction] * n_seg\n", + " values[k] = math.exp(base + sign * h)\n", + " set_friction(values)\n", + " forward(f\"fd {names[k]} {'+' if sign > 0 else '-'}h\")\n", + " fd.append(misfit_value())\n", + " fd = (fd[0] - fd[1]) / (2 * h)\n", + " print(f\"{names[k]:>13}: adjoint {g0[k]: .6e} finite difference {fd: .6e} ratio {fd / g0[k]:.5f}\")\n", + "set_friction([params.uw_initial_friction] * n_seg)\n", + "\n", + "if int(params.uw_check_only):\n", + " raise SystemExit" + ] + }, + { + "cell_type": "markdown", + "id": "91bb23cf", + "metadata": {}, + "source": [ + "## The objective\n", + "\n", + "Without noise the misfit has no natural scale and the optimiser sees it\n", + "relative to its starting value, $J/J_0$. With noise the objective is a\n", + "negative log posterior. The data term is $\\chi^2/2$: the misfit scaled by\n", + "its expected value at the truth under the noise, which a twin experiment\n", + "can read directly, times the number of independent data, the surface nodes\n", + "and the nodes under the point weights. The prior term is\n", + "$(\\log\\mu - \\log\\mu_0)^2 / 2\\sigma_m^2$ with $\\sigma_m$ in log units. The\n", + "weight between the two is then a statement about the noise and the prior,\n", + "and which coefficients the data move is decided by their sensitivities\n", + "against that." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e95ede9e", + "metadata": {}, + "outputs": [], + "source": [ + "history = []\n", + "noisy = float(params.uw_noise) > 0\n", + "sigma_m = float(params.uw_prior_sigma)\n", + "log_prior = np.log([params.uw_initial_friction] * n_seg)\n", + "if sigma_m > 0 and not noisy:\n", + " raise ValueError(\"a prior is weighed against the noise: give uw_noise as well\")\n", + "\n", + "if noisy:\n", + " X = np.asarray(v.coords_nd)\n", + " on_surface = X[:, 1] > 1.0 - 1e-6\n", + " near_points = np.zeros(len(X), dtype=bool)\n", + " for px, py in points:\n", + " near_points |= (X[:, 0] - px) ** 2 + (X[:, 1] - py) ** 2 < (2 * band_width) ** 2\n", + " N_eff = {\"uplift+stress\": on_surface.sum() + near_points.sum(),\n", + " \"orientation_points\": near_points.sum(),\n", + " \"surface_strain\": on_surface.sum()}[what]\n", + " chi2_scale = N_eff / J_truth # chi^2 = J * chi2_scale\n", + " print(f\"N_eff = {N_eff}, misfit floor at the truth = {J_truth:.4e}\")\n", + "\n", + "\n", + "def objective(log_mu):\n", + " set_friction(np.exp(log_mu))\n", + " J, grad = misfit_and_gradient()\n", + " history.append((J, np.exp(log_mu).copy()))\n", + " if not noisy:\n", + " print(f\" J/J0 = {J / J0:.6e} friction = {np.exp(log_mu)}\")\n", + " return J / J0, grad / J0\n", + " print(f\" J = {J:.6e} chi2/N = {J * chi2_scale / N_eff:.4f} friction = {np.exp(log_mu)}\")\n", + " value = J * chi2_scale / 2\n", + " g = grad * chi2_scale / 2\n", + " if sigma_m > 0:\n", + " value += np.sum((log_mu - log_prior) ** 2) / (2 * sigma_m ** 2)\n", + " g = g + (log_mu - log_prior) / sigma_m ** 2\n", + " return value, g" + ] + }, + { + "cell_type": "markdown", + "id": "fed874e2", + "metadata": {}, + "source": [ + "## The inversion\n", + "\n", + "PETSc's TAO drives it by default: the same objective and gradient, a\n", + "limited-memory quasi-Newton update and TAO's line search, and the bounded\n", + "variant when bounds are given. SciPy's L-BFGS-B is the alternative." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b3f436f0", + "metadata": {}, + "outputs": [], + "source": [ + "x0 = np.log([params.uw_initial_friction] * n_seg)\n", + "if str(params.uw_optimiser) == \"tao\":\n", + " bounds = None\n", + " if str(params.uw_bounds).strip():\n", + " lo, hi = (float(b) for b in str(params.uw_bounds).split(\",\"))\n", + " bounds = (np.log([lo] * n_seg), np.log([hi] * n_seg))\n", + " x_best, info = uw.adjoint.minimise(objective, x0, max_evaluations=60,\n", + " gradient_tolerance=1e-10, bounds=bounds,\n", + " method=\"blmvm\" if bounds else \"lmvm\")\n", + "else:\n", + " from scipy.optimize import minimize\n", + " result = minimize(objective, x0, jac=True, method=\"L-BFGS-B\",\n", + " options={\"maxiter\": 40, \"gtol\": 1e-10})\n", + " x_best = result.x\n", + "print(f\"recovered {np.exp(x_best)} true {true_values} after {len(history)} evaluations\")" + ] + }, + { + "cell_type": "markdown", + "id": "795a2cdf", + "metadata": {}, + "source": [ + "## What the figure needs\n", + "\n", + "Uplift-rate profiles along the top at the truth, the start and the answer,\n", + "the weak-plane viscosity on a grid, and the path the coefficients took.\n", + "Lengths are saved in kilometres and velocities in millimetres per year." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0d4d8b90", + "metadata": {}, + "outputs": [], + "source": [ + "xs = np.linspace(0.0, 2.0, 161)\n", + "top = np.column_stack([xs, np.full_like(xs, 1.0 - 1e-6)])\n", + "profiles = {}\n", + "for label, values in ((\"true\", true_values), (\"initial\", [params.uw_initial_friction] * n_seg),\n", + " (\"recovered\", list(np.exp(x_best)))):\n", + " set_friction(values)\n", + " forward(f\"profile {label}\")\n", + " profiles[label] = np.asarray(uw.function.evaluate(v.sym[1], top)).ravel() * mm_per_yr\n", + "set_friction(true_values)\n", + "forward(\"truth again\") # the field on the grid is the truth's\n", + "gx, gy = np.meshgrid(np.linspace(0, 2, 201), np.linspace(0, 1, 101))\n", + "grid = np.column_stack([gx.ravel(), gy.ravel()])\n", + "eta_1_grid = np.asarray(uw.function.evaluate(eta_1, grid)).reshape(gx.shape)\n", + "\n", + "tag = (f\"{what}\" + (f\"_noise{float(params.uw_noise):g}\" if float(params.uw_noise) > 0 else \"\")\n", + " + (f\"_prior{sigma_m:g}\" if sigma_m > 0 else \"\"))\n", + "np.savez(f\"fault_friction_{tag}_data.npz\", xs=xs * km, gx=gx * km, gy=gy * km, eta_1=eta_1_grid,\n", + " points=np.array(points) * km, true=np.array(true_values), band=band_width * km,\n", + " length_unit=\"km\", velocity_unit=\"mm/yr\",\n", + " history=np.array([[J, *vals] for J, vals in history]),\n", + " **{f\"uplift_{k}\": val for k, val in profiles.items()})" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/examples/adjoint/fault_segments/fault_friction.py b/docs/examples/adjoint/fault_segments/fault_friction.py new file mode 100644 index 00000000..c177fe8b --- /dev/null +++ b/docs/examples/adjoint/fault_segments/fault_friction.py @@ -0,0 +1,479 @@ +# %% [markdown] +# # Friction on a listric fault from surface observations +# +# A listric fault under horizontal shortening and gravity: a flat décollement +# at depth that steepens through a circular ramp to a dip of sixty degrees at +# the surface. The fault is a weak plane in a transversely isotropic viscosity, +# with no cut in the mesh. The plane yields at the Coulomb stress +# $\tau_y = C + \mu p$, and the friction coefficient $\mu$ takes a different +# value on the flat, on the lower and upper parts of the ramp, and near the +# surface. Those four coefficients are the unknowns. +# +# The observations are the uplift rate along the top surface and the shear +# stress near a handful of interior points, read from a run at the true +# coefficients. The gradient of the misfit with respect to each coefficient +# comes from one call, `stokes.gradient(misfit, parameters=...)`: the solver +# assembles the adjoint operator from its own Jacobian kernels with trial and +# test functions exchanged, solves it, and differentiates its residual +# symbolically with respect to each named coefficient. Nothing is differenced. +# +# The notebook checks that gradient against central differences, then hands +# the misfit and its gradient to PETSc's TAO to recover the four +# coefficients. Each misfit evaluation is one step of zero length in the +# model's record, so the run's transcript lists the forward solve and the +# adjoint solve it made. + +# %% +import math + +import numpy as np +import sympy + +import underworld3 as uw + +# %% [markdown] +# ### Configurable parameters +# +# Default values are defined as named constants below. From the command line, +# override them with PETSc-style flags: +# +# ```bash +# python fault_friction.py -uw_check_only 1 +# python fault_friction.py -uw_observations surface_strain +# python fault_friction.py -uw_noise 0.03 -uw_prior_sigma 1 -uw_bounds 0.005,1 +# ``` +# +# The problem depends on one dimensionless number: the lithostatic pressure at +# the base of the model over the viscous stress of the shortening, +# $\rho g L^2 / (\eta V)$. We fix the depth, the viscosity and the density, +# and that ratio sets the convergence rate. + +# %% +# --- Scales --- +DOMAIN_DEPTH = uw.quantity(10, "km") # the box is 20 km wide and 10 km deep +REF_VISCOSITY = uw.quantity(1e21, "Pa*s") # the bulk viscosity +DENSITY = uw.quantity(2700, "kg/m**3") +GRAVITY = uw.quantity(9.81, "m/s**2") +PRESSURE_TO_STRESS = 10 # rho g L^2 / (eta V) +CONVERGENCE_RATE = (DENSITY * GRAVITY * DOMAIN_DEPTH**2 / (PRESSURE_TO_STRESS * REF_VISCOSITY)).to("mm/yr") # 8.4 mm/yr +REF_STRESS = (REF_VISCOSITY * CONVERGENCE_RATE / DOMAIN_DEPTH).to("MPa") # 26 MPa + +# --- Geometry --- +CELL_SIZE = DOMAIN_DEPTH / 12 # base mesh; refined once, so half this +SURFACE_DIP = 60.0 # degrees, where the ramp reaches the surface +FLAT_DEPTH = uw.quantity(3, "km") # height of the décollement above the base +SURFACE_X = uw.quantity(19, "km") # where the fault reaches the surface +BAND_HALF_WIDTH = uw.quantity(0.8, "km") # of the weak plane + +# --- Rheology --- +COHESION = (0.05 * REF_STRESS).to("MPa") # C in tau_y = C + mu p; 1.3 MPa +TRUE_FRICTION = "0.05,0.15,0.25,0.4" # flat, lower ramp, upper ramp, near surface +INITIAL_FRICTION = 0.2 # the starting guess, every segment + +# --- The inversion --- +CHECK_ONLY = 0 # 1: the gradient check, no inversion +OPTIMISER = "tao" # tao (PETSc quasi-Newton) | scipy (L-BFGS-B) +OBSERVATIONS = "uplift+stress" # uplift+stress | orientation_points | surface_strain +NOISE = 0.0 # on the observed velocity, as a fraction of its rms +SEED = 7 +PRIOR_SIGMA = 0.0 # width of a Gaussian prior on log mu; 0 for none +BOUNDS = "" # friction bounds lo,hi for TAO's blmvm; empty for none + +params = uw.Params( + uw_cell_size=uw.Param(CELL_SIZE, description="base mesh cell size"), + uw_surface_dip=uw.Param(SURFACE_DIP, description="dip of the ramp at the surface, degrees"), + uw_flat_depth=uw.Param(FLAT_DEPTH, description="height of the décollement above the base"), + uw_surface_x=uw.Param(SURFACE_X, description="where the fault reaches the surface"), + uw_band=uw.Param(BAND_HALF_WIDTH, description="half-width of the weak plane"), + uw_cohesion=uw.Param(COHESION, description="cohesion C in tau_y = C + mu p"), + uw_true_friction=uw.Param(TRUE_FRICTION, description="friction: flat, lower ramp, upper ramp, near surface"), + uw_initial_friction=uw.Param(INITIAL_FRICTION, description="starting guess, every segment"), + uw_check_only=uw.Param(CHECK_ONLY, description="1: gradient check only"), + uw_optimiser=uw.Param(OPTIMISER, description="tao | scipy"), + uw_observations=uw.Param(OBSERVATIONS, description="uplift+stress | orientation_points | surface_strain"), + uw_noise=uw.Param(NOISE, description="noise on the observed velocity, fraction of its rms"), + uw_seed=uw.Param(SEED, description="seed for the noise"), + uw_prior_sigma=uw.Param(PRIOR_SIGMA, description="width of the prior on log mu; 0 for none"), + uw_bounds=uw.Param(BOUNDS, description="friction bounds lo,hi; empty for none"), +) + +# %% [markdown] +# ## The model and its scales +# +# The model is declared first, with the three reference quantities that fix +# the scaling: the depth, the viscosity and the convergence rate. The solver +# then works in units of those, so the box is $2 \times 1$, the bulk viscosity +# is one, the walls close at one, and the body force is the pressure-to-stress +# ratio. `coords` on a variable come back in physical units, and `coords_nd` +# are the model's own. + +# %% +orchestration_model = uw.get_default_model() +orchestration_model.set_reference_quantities( + domain_depth=DOMAIN_DEPTH, + viscosity=REF_VISCOSITY, + convergence_rate=CONVERGENCE_RATE, +) + + +def _nd(quantity): + """The plain number a dimensionless quantity stands for.""" + try: + return float(quantity.to("dimensionless").magnitude) + except AttributeError: + return float(quantity) + + +def _lengths(quantity): + """A length in units of the model depth.""" + return _nd(quantity / DOMAIN_DEPTH) + + +mm_per_yr = float(CONVERGENCE_RATE.to("mm/yr").magnitude) # one model velocity unit +km = float(DOMAIN_DEPTH.to("km").magnitude) # one model length unit + +# %% [markdown] +# ## The mesh +# +# Refined once from the base cell size, which gives the velocity block a +# multigrid hierarchy. Without one it falls back to an algebraic coarsening +# that hits its iteration cap on this problem, and an inexact Newton step then +# converges only linearly. + +# %% +mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(2.0, 1.0), + cellSize=_lengths(params.uw_cell_size), qdegree=3, refinement=1, +) +x, y = mesh.X + +v = uw.discretisation.MeshVariable("v", mesh, mesh.dim, degree=2) +p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1, continuous=True) +v_obs = uw.discretisation.MeshVariable("v_obs", mesh, mesh.dim, degree=2) + +# %% [markdown] +# ## The fault +# +# A flat décollement runs from the left wall to $x = x_c$, then a circular +# ramp of radius $R$ about $(x_c, y_c)$ leaves the flat horizontally and +# reaches the surface at the given dip. The signed distance to the fault and +# the position along it are exact on each piece. The director is the normal +# at the nearest point: vertical on the flat, radial on the ramp. + +# %% +flat_depth = _lengths(params.uw_flat_depth) +band_width = _lengths(params.uw_band) +phi_top = math.radians(params.uw_surface_dip) - math.pi / 2 # angle of the surface point about the centre +R = (1.0 - flat_depth) / (1.0 + math.sin(phi_top)) +yc = flat_depth + R +xc = _lengths(params.uw_surface_x) - R * math.cos(phi_top) +r = sympy.sqrt((x - xc) ** 2 + (y - yc) ** 2) +phi = sympy.atan2(y - yc, x - xc) +on_flat = x < xc +d = sympy.Piecewise((y - flat_depth, on_flat), (R - r, True)) +s = sympy.Piecewise((x, on_flat), (xc + R * (phi + sympy.pi / 2), True)) +n_hat = sympy.Matrix([[sympy.Piecewise((0, on_flat), ((x - xc) / r, True)), + sympy.Piecewise((1, on_flat), ((y - yc) / r, True))]]) +ramp = R * (phi_top + math.pi / 2) +length = xc + ramp +band = sympy.exp(-(d / band_width) ** 2) + +# %% [markdown] +# ## The segments and their friction coefficients +# +# Four segments along the fault: the flat, then the ramp in three equal +# parts. Each coefficient is a named expression, so the solver can +# differentiate its residual with respect to it. + +# %% +edges = [0.0, xc, xc + ramp / 3, xc + 2 * ramp / 3, length] +names = ["flat", "lower ramp", "upper ramp", "near surface"] +n_seg = len(names) +friction_coefficients = [ + uw.expression(rf"\mu_{{{k + 1}}}", params.uw_initial_friction, f"friction coefficient, {names[k]}") + for k in range(n_seg) +] + + +def segment(k): + """A smooth indicator for segment k along the fault, in [0, 1].""" + on = 1 if k == 0 else (1 + sympy.tanh((s - edges[k]) / band_width)) / 2 + off = 1 if k == n_seg - 1 else (1 - sympy.tanh((s - edges[k + 1]) / band_width)) / 2 + return on * off + + +def set_friction(values): + for coefficient, value in zip(friction_coefficients, values): + coefficient.sym = float(value) + +# %% [markdown] +# ## Coulomb yield on the plane +# +# The shear strain rate resolved on the plane is $\hat t \cdot E \cdot \hat n$. +# The plane's viscosity is the harmonic combination of the bulk viscosity and +# the yield stress over that rate, which is smooth everywhere and tends to +# $\tau_y / 2\dot\varepsilon_s$ where the plane slips. Compression is positive. +# Where the dynamic pressure is tensile the plane keeps its cohesion and no +# more. + +# %% +eta_0 = 1 +cohesion = _nd(params.uw_cohesion / REF_STRESS) +rho_g = _nd(DENSITY * GRAVITY * DOMAIN_DEPTH / REF_STRESS) + +E = mesh.vector.strain_tensor(v.sym) +t_hat = sympy.Matrix([[-n_hat[1], n_hat[0]]]) +e_s = sympy.sqrt((t_hat * E * n_hat.T)[0] ** 2 + uw.maths.functions.vanishing) +friction = sum(friction_coefficients[k] * segment(k) for k in range(n_seg)) +tau_y = cohesion + friction * sympy.Max(p.sym[0], 0) +eta_plane = eta_0 * tau_y / (tau_y + 2 * eta_0 * e_s) +eta_1 = eta_0 - band * (eta_0 - eta_plane) + +# %% [markdown] +# ## The Stokes solver +# +# Shortening from both sides, a no-slip base, and a free top: the surface +# velocity is the uplift rate. The residual is nonlinear in the velocity and +# the pressure, so the forward solve is Newton. + +# %% +stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) +stokes.constitutive_model = uw.constitutive_models.TransverseIsotropicFlowModel +stokes.constitutive_model.Parameters.shear_viscosity_0 = eta_0 +stokes.constitutive_model.Parameters.shear_viscosity_1 = eta_1 +stokes.constitutive_model.Parameters.director = n_hat +stokes.tolerance = 1e-8 +stokes.bodyforce = sympy.Matrix([0, -rho_g]) + +stokes.add_essential_bc((0.0, 0.0), "Bottom") +stokes.add_essential_bc((0.5, None), "Left") +stokes.add_essential_bc((-0.5, None), "Right") + +# %% [markdown] +# ## The observations +# +# The misfit has a term on the top surface, a boundary integral of the uplift +# rate, and a term in the volume around five interior points where the shear +# stress is read. `gradient()` takes the terms as `{domain: integrand}`. +# +# On a traction-free surface the shear strain rate vanishes, so a stress +# orientation read there is only a sign. Orientation is an interior +# observable (boreholes, focal mechanisms). The surface gives velocities and +# their tangential derivative, the geodetic strain rate. Both alternatives +# are available through `uw_observations`. + +# %% +points = [(0.3, 0.65), (0.75, 0.12), (1.25, 0.3), (0.9, 0.6), (1.55, 0.85)] # in units of the depth +w_points = sum(sympy.exp(-((x - px) ** 2 + (y - py) ** 2) / (2 * band_width) ** 2) + for px, py in points) + + +def shear_stress(field): + e = mesh.vector.strain_tensor(field.sym) + return 2 * eta_0 * e[0, 1] + + +def orientation(field): + """The principal-stress orientation as the unit vector (cos 2theta, sin 2theta) + of the deviatoric strain rate, which is the stress orientation in the + isotropic bulk. A unit vector rather than an angle, so there is no wrap.""" + e = mesh.vector.strain_tensor(field.sym) + a, b = e[0, 0] - e[1, 1], 2 * e[0, 1] + norm = sympy.sqrt(a ** 2 + b ** 2 + uw.maths.functions.vanishing) + return sympy.Matrix([[a / norm, b / norm]]) + + +what = str(params.uw_observations) +dq = orientation(v) - orientation(v_obs) +if what == "uplift+stress": + misfit = {"Top": (v.sym[1] - v_obs.sym[1]) ** 2 / 2, + None: w_points * (shear_stress(v) - shear_stress(v_obs)) ** 2 / 2} +elif what == "orientation_points": + misfit = {None: w_points * (dq[0] ** 2 + dq[1] ** 2) / 2} +elif what == "surface_strain": + misfit = {"Top": (v.sym[0].diff(x) - v_obs.sym[0].diff(x)) ** 2 / 2} +else: + raise ValueError(f"observations: {what!r}") + + +def misfit_value(): + return sum(uw.adjoint.integral(mesh, term, where) for where, term in misfit.items()) + +# %% [markdown] +# ## The forward solve and the gradient +# +# Every solve starts cold: six Newton iterations, and a misfit that does not +# depend on the previous evaluation, which the finite-difference check needs. +# Each evaluation is one step of zero length in the model's record. + +# %% +evaluations = [0] + + +def forward(label): + with orchestration_model.step(0.0, label=label): + stokes.solve(zero_init_guess=True) + + +def misfit_and_gradient(label=None): + """The misfit and dJ/d(log mu) for each segment, by the adjoint.""" + evaluations[0] += 1 + with orchestration_model.step(0.0, label=label or f"eval {evaluations[0]}"): + stokes.solve(zero_init_guess=True) + out = stokes.gradient(misfit, parameters=friction_coefficients) + # d/d(log mu) = mu d/d(mu) + grad = np.array([out["parameters"][c] * float(c.sym) for c in friction_coefficients]) + return out["J"], grad + +# %% [markdown] +# ## The twin +# +# The observations are the velocity field at the true coefficients, with +# optional noise drawn once at a fraction of each component's rms. Every +# observation set reads `v_obs`, so the uplift, the stress and the strain +# rate all inherit the same noise. + +# %% +true_values = [float(t) for t in str(params.uw_true_friction).split(",")][:n_seg] +set_friction(true_values) +forward("truth") +v_obs.array[...] = np.asarray(v.array) +if float(params.uw_noise) > 0: + rng = np.random.default_rng(int(params.uw_seed)) + obs = np.asarray(v_obs.array) + rms = np.sqrt(np.mean(obs ** 2, axis=0, keepdims=True)) + v_obs.array[...] = obs + float(params.uw_noise) * rms * rng.standard_normal(obs.shape) + print(f"noise: {float(params.uw_noise):.3f} of the rms per component, seed {int(params.uw_seed)}") +J_truth = misfit_value() # v still holds the truth's velocity + +set_friction([params.uw_initial_friction] * n_seg) +J0, g0 = misfit_and_gradient("start") +print(f"true friction {true_values}") +print(f"initial J = {J0:.6e} dJ/dlog mu = {g0}") + +# %% [markdown] +# ## The gradient check +# +# Central differences in each log-coefficient, against the adjoint gradient. + +# %% +h = 1e-3 +for k in range(n_seg): + base = math.log(params.uw_initial_friction) + fd = [] + for sign in (+1, -1): + values = [params.uw_initial_friction] * n_seg + values[k] = math.exp(base + sign * h) + set_friction(values) + forward(f"fd {names[k]} {'+' if sign > 0 else '-'}h") + fd.append(misfit_value()) + fd = (fd[0] - fd[1]) / (2 * h) + print(f"{names[k]:>13}: adjoint {g0[k]: .6e} finite difference {fd: .6e} ratio {fd / g0[k]:.5f}") +set_friction([params.uw_initial_friction] * n_seg) + +if int(params.uw_check_only): + raise SystemExit + +# %% [markdown] +# ## The objective +# +# Without noise the misfit has no natural scale and the optimiser sees it +# relative to its starting value, $J/J_0$. With noise the objective is a +# negative log posterior. The data term is $\chi^2/2$: the misfit scaled by +# its expected value at the truth under the noise, which a twin experiment +# can read directly, times the number of independent data, the surface nodes +# and the nodes under the point weights. The prior term is +# $(\log\mu - \log\mu_0)^2 / 2\sigma_m^2$ with $\sigma_m$ in log units. The +# weight between the two is then a statement about the noise and the prior, +# and which coefficients the data move is decided by their sensitivities +# against that. + +# %% +history = [] +noisy = float(params.uw_noise) > 0 +sigma_m = float(params.uw_prior_sigma) +log_prior = np.log([params.uw_initial_friction] * n_seg) +if sigma_m > 0 and not noisy: + raise ValueError("a prior is weighed against the noise: give uw_noise as well") + +if noisy: + X = np.asarray(v.coords_nd) + on_surface = X[:, 1] > 1.0 - 1e-6 + near_points = np.zeros(len(X), dtype=bool) + for px, py in points: + near_points |= (X[:, 0] - px) ** 2 + (X[:, 1] - py) ** 2 < (2 * band_width) ** 2 + N_eff = {"uplift+stress": on_surface.sum() + near_points.sum(), + "orientation_points": near_points.sum(), + "surface_strain": on_surface.sum()}[what] + chi2_scale = N_eff / J_truth # chi^2 = J * chi2_scale + print(f"N_eff = {N_eff}, misfit floor at the truth = {J_truth:.4e}") + + +def objective(log_mu): + set_friction(np.exp(log_mu)) + J, grad = misfit_and_gradient() + history.append((J, np.exp(log_mu).copy())) + if not noisy: + print(f" J/J0 = {J / J0:.6e} friction = {np.exp(log_mu)}") + return J / J0, grad / J0 + print(f" J = {J:.6e} chi2/N = {J * chi2_scale / N_eff:.4f} friction = {np.exp(log_mu)}") + value = J * chi2_scale / 2 + g = grad * chi2_scale / 2 + if sigma_m > 0: + value += np.sum((log_mu - log_prior) ** 2) / (2 * sigma_m ** 2) + g = g + (log_mu - log_prior) / sigma_m ** 2 + return value, g + +# %% [markdown] +# ## The inversion +# +# PETSc's TAO drives it by default: the same objective and gradient, a +# limited-memory quasi-Newton update and TAO's line search, and the bounded +# variant when bounds are given. SciPy's L-BFGS-B is the alternative. + +# %% +x0 = np.log([params.uw_initial_friction] * n_seg) +if str(params.uw_optimiser) == "tao": + bounds = None + if str(params.uw_bounds).strip(): + lo, hi = (float(b) for b in str(params.uw_bounds).split(",")) + bounds = (np.log([lo] * n_seg), np.log([hi] * n_seg)) + x_best, info = uw.adjoint.minimise(objective, x0, max_evaluations=60, + gradient_tolerance=1e-10, bounds=bounds, + method="blmvm" if bounds else "lmvm") +else: + from scipy.optimize import minimize + result = minimize(objective, x0, jac=True, method="L-BFGS-B", + options={"maxiter": 40, "gtol": 1e-10}) + x_best = result.x +print(f"recovered {np.exp(x_best)} true {true_values} after {len(history)} evaluations") + +# %% [markdown] +# ## What the figure needs +# +# Uplift-rate profiles along the top at the truth, the start and the answer, +# the weak-plane viscosity on a grid, and the path the coefficients took. +# Lengths are saved in kilometres and velocities in millimetres per year. + +# %% +xs = np.linspace(0.0, 2.0, 161) +top = np.column_stack([xs, np.full_like(xs, 1.0 - 1e-6)]) +profiles = {} +for label, values in (("true", true_values), ("initial", [params.uw_initial_friction] * n_seg), + ("recovered", list(np.exp(x_best)))): + set_friction(values) + forward(f"profile {label}") + profiles[label] = np.asarray(uw.function.evaluate(v.sym[1], top)).ravel() * mm_per_yr +set_friction(true_values) +forward("truth again") # the field on the grid is the truth's +gx, gy = np.meshgrid(np.linspace(0, 2, 201), np.linspace(0, 1, 101)) +grid = np.column_stack([gx.ravel(), gy.ravel()]) +eta_1_grid = np.asarray(uw.function.evaluate(eta_1, grid)).reshape(gx.shape) + +tag = (f"{what}" + (f"_noise{float(params.uw_noise):g}" if float(params.uw_noise) > 0 else "") + + (f"_prior{sigma_m:g}" if sigma_m > 0 else "")) +np.savez(f"fault_friction_{tag}_data.npz", xs=xs * km, gx=gx * km, gy=gy * km, eta_1=eta_1_grid, + points=np.array(points) * km, true=np.array(true_values), band=band_width * km, + length_unit="km", velocity_unit="mm/yr", + history=np.array([[J, *vals] for J, vals in history]), + **{f"uplift_{k}": val for k, val in profiles.items()}) diff --git a/docs/examples/adjoint/fault_segments/fault_friction_key.md b/docs/examples/adjoint/fault_segments/fault_friction_key.md new file mode 100644 index 00000000..dabbd025 --- /dev/null +++ b/docs/examples/adjoint/fault_segments/fault_friction_key.md @@ -0,0 +1,38 @@ +## Key — what each part solved in fault_friction +*run started 2026-09-19T17:44:38+10:00* + +### Stokes(v) +`SNES_Stokes`, unknown `v`, 2-D; recorded at step 0 + +Residual $\int F_0\,\phi + F_1 \cdot \nabla\phi = 0$ with + +$$\mathbf{f}_0\left( \mathbf{u} \right) = \left[\begin{matrix}0\\10.0\end{matrix}\right]$$ +*Velocity equation body force term (pointwise).* + +$$\mathbf{F}_1\left( \mathbf{u} \right) = \left[\begin{matrix}\eta_0 \uplambda \left({v}_{ 0,0}(\mathbf{x}) + {v}_{ 1,1}(\mathbf{x})\right) - 2 \left(2 \eta_0 - 2 \eta_1\right) \left(- 2 \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) + \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right)\right) \left(\frac{{v}_{ 0,1}(\mathbf{x})}{2} + \frac{{v}_{ 1,0}(\mathbf{x})}{2}\right) + 2 \left(2 \eta_0 - 2 \eta_1\right) \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) {v}_{ 1,1}(\mathbf{x}) + \left(2 \eta_0 - \left(2 \eta_0 - 2 \eta_1\right) \left(- 2 \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) + 2 \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right)\right)\right) {v}_{ 0,0}(\mathbf{x}) - {p}(\mathbf{x}) & - \left(2 \eta_0 - 2 \eta_1\right) \left(- 2 \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) + \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right)\right) {v}_{ 0,0}(\mathbf{x}) - \left(2 \eta_0 - 2 \eta_1\right) \left(- 2 \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) + \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right)\right) {v}_{ 1,1}(\mathbf{x}) + 2 \left(\eta_0 - \left(2 \eta_0 - 2 \eta_1\right) \left(- 2 \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) + \frac{\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}}{2} + \frac{\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}}{2}\right)\right) \left(\frac{{v}_{ 0,1}(\mathbf{x})}{2} + \frac{{v}_{ 1,0}(\mathbf{x})}{2}\right)\\- \left(2 \eta_0 - 2 \eta_1\right) \left(- 2 \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) + \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right)\right) {v}_{ 0,0}(\mathbf{x}) - \left(2 \eta_0 - 2 \eta_1\right) \left(- 2 \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) + \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right)\right) {v}_{ 1,1}(\mathbf{x}) + 2 \left(\eta_0 - \left(2 \eta_0 - 2 \eta_1\right) \left(- 2 \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) + \frac{\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}}{2} + \frac{\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}}{2}\right)\right) \left(\frac{{v}_{ 0,1}(\mathbf{x})}{2} + \frac{{v}_{ 1,0}(\mathbf{x})}{2}\right) & \eta_0 \uplambda \left({v}_{ 0,0}(\mathbf{x}) + {v}_{ 1,1}(\mathbf{x})\right) - 2 \left(2 \eta_0 - 2 \eta_1\right) \left(- 2 \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) + \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{x} - 0.687564434701786}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\mathrm{y} - 1.7}{\sqrt{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}}} & \text{otherwise} \end{cases}\right)\right) \left(\frac{{v}_{ 0,1}(\mathbf{x})}{2} + \frac{{v}_{ 1,0}(\mathbf{x})}{2}\right) + 2 \left(2 \eta_0 - 2 \eta_1\right) \left(\begin{cases} 0 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{x} - 0.687564434701786\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) {v}_{ 0,0}(\mathbf{x}) + \left(2 \eta_0 - \left(2 \eta_0 - 2 \eta_1\right) \left(- 2 \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right) + 2 \left(\begin{cases} 1 & \text{for}\: \mathrm{x} < 0.687564434701786 \\\frac{\left(\mathrm{y} - 1.7\right)^{2}}{\left(\mathrm{x} - 0.687564434701786\right)^{2} + \left(\mathrm{y} - 1.7\right)^{2}} & \text{otherwise} \end{cases}\right)\right)\right) {v}_{ 1,1}(\mathbf{x}) - {p}(\mathbf{x})\end{matrix}\right]$$ +*Velocity equation flux/stress term (pointwise).* + +where +- $\eta_0$ $= 1$ — Shear viscosity +- $\eta_1$ $= \text{(an expression; in the record)}$ — Second viscosity + - $\mu_{1}$ $= 0.05$ — friction coefficient, flat + - $\mu_{2}$ $= 0.15$ — friction coefficient, lower ramp + - $\mu_{3}$ $= 0.25$ — friction coefficient, upper ramp + - $\mu_{4}$ $= 0.4$ — friction coefficient, near surface + - $\varepsilon$ $= 10^{-18}$ — vanishingly small value +- $\uplambda$ $= 0$ — Numerical Penalty + +$$\mathbf{h}_0\left( \mathbf{p} \right) = \left[\begin{matrix}{v}_{ 0,0}(\mathbf{x}) + {v}_{ 1,1}(\mathbf{x})\end{matrix}\right]$$ +*Pressure equation constraint term (continuity).* + +Boundary conditions: +- essential on Bottom: $\left[\begin{matrix}0.0\\0.0\end{matrix}\right]$ +- essential on Left: $\left[\begin{matrix}0.5\\\infty\end{matrix}\right]$ +- essential on Right: $\left[\begin{matrix}-0.5\\\infty\end{matrix}\right]$ + +Given: +- `bodyforce` $= (0, -10)$ — body force per unit volume; F0 is its negative +- `penalty` $= 0$ — augmented-Lagrangian grad-div penalty (0 = off) +- `TransverseIsotropicFlowModel.director` $= \text{(an expression; in the record)}$ — constitutive parameter +- `TransverseIsotropicFlowModel.shear_viscosity_0` $= 1$ — constitutive parameter +- `TransverseIsotropicFlowModel.shear_viscosity_1` $= \text{(an expression; in the record)}$ — constitutive parameter diff --git a/docs/examples/adjoint/fault_segments/fault_friction_transcript.svg b/docs/examples/adjoint/fault_segments/fault_friction_transcript.svg new file mode 100644 index 00000000..873b9740 --- /dev/null +++ b/docs/examples/adjoint/fault_segments/fault_friction_transcript.svg @@ -0,0 +1,390 @@ + + +Transcript — fault_friction +started 2026-09-19T17:44:38+10:00 · complete, 40 steps · 2 parts +scales: domain_depth 10 km viscosity 1e+21 Pa s convergence_rate 8.359 mm/yr +step +t +dt +Stokes(v) +adjoint Stokes(v) + +0 +0 +0 + +✅ + +1 +0 +0 + + + +✅ +✅ + +2 +0 +0 + +✅ + +3 +0 +0 + +✅ + +4 +0 +0 + +✅ + +5 +0 +0 + +✅ + +6 +0 +0 + +✅ + +7 +0 +0 + +✅ + +8 +0 +0 + +✅ + +9 +0 +0 + +✅ + +10 +0 +0 + + + +✅ +✅ + +11 +0 +0 + + + +✅ +✅ + +12 +0 +0 + + + +✅ +✅ + +13 +0 +0 + + + +✅ +✅ + +14 +0 +0 + + + +✅ +✅ + +15 +0 +0 + + + +✅ +✅ + +16 +0 +0 + + + +✅ +✅ + +17 +0 +0 + + + +✅ +✅ + +18 +0 +0 + + + +✅ +✅ + +19 +0 +0 + + + +✅ +✅ + +20 +0 +0 + + + +✅ +✅ + +21 +0 +0 + + + +✅ +✅ + +22 +0 +0 + + + +✅ +✅ + +23 +0 +0 + + + +✅ +✅ + +24 +0 +0 + + + +✅ +✅ + +25 +0 +0 + + + +✅ +✅ + +26 +0 +0 + + + +✅ +✅ + +27 +0 +0 + + + +✅ +✅ + +28 +0 +0 + + + +✅ +✅ + +29 +0 +0 + + + +✅ +✅ + +30 +0 +0 + + + +✅ +✅ + +31 +0 +0 + + + +✅ +✅ + +32 +0 +0 + + + +✅ +✅ + +33 +0 +0 + + + +✅ +✅ + +34 +0 +0 + + + +✅ +✅ + +35 +0 +0 + + + +✅ +✅ + +36 +0 +0 + +✅ + +37 +0 +0 + +✅ + +38 +0 +0 + +✅ + +39 +0 +0 + +✅ + + + + + +ran; a step reads down, the path joining what ran in the order it ran +✅ +solved, and converged +⚠️ +converged, but a fieldsplit block hit its iteration cap — that block did not solve +❌ +did not converge; the reason is in the record + +ran in a step that was then abandoned + +did nothing in this step + + +the steps between did exactly this, unchanged; first and last values are shown where they differ +Key — the parts, and what is in their residuals +Stokes(v) +recorded at step 0 + +— with +Velocity equation body force term (pointwise). + + + + + + + + + +Velocity equation flux/stress term (pointwise). +(the form is in the record; uw.transcript_key renders it) + +— Shear viscosity + +— Second viscosity + +— friction coefficient, flat + +— friction coefficient, lower ramp + +— friction coefficient, upper ramp + +— friction coefficient, near surface + +— vanishingly small value + +— Numerical Penalty +Pressure equation constraint term (continuity). + + +essential on Bottom: (0, 0) +essential on Left: (0.5, free) +essential on Right: (-0.5, free) + diff --git a/docs/examples/adjoint/fault_segments/fault_segments.py b/docs/examples/adjoint/fault_segments/fault_segments.py new file mode 100644 index 00000000..635e8f62 --- /dev/null +++ b/docs/examples/adjoint/fault_segments/fault_segments.py @@ -0,0 +1,204 @@ +"""Strength of a listric fault, segment by segment, from surface uplift and stress. + +A listric fault under horizontal shortening: a ramp that steepens from a flat +decollement at depth to a dip of sixty degrees at the surface, represented as +a weak plane in a transversely isotropic viscosity (no cut in the mesh). Its +weak-plane viscosity is different on the flat, on the lower and upper parts +of the ramp, and near the surface, and those four strengths are the unknowns. +The observations are the uplift rate along the top surface and the shear +stress in the bulk near a handful of points, taken from a run at the true +strengths. + +The gradient of the misfit with respect to each segment's strength comes +from the solver's own discrete adjoint: a transpose against the Jacobian it +assembled, and the symbolic derivative of its residual with respect to the +named strength. Nothing here is differenced. + +Run it: + + python fault_segments.py # twin experiment: gradient check, then the inversion + python fault_segments.py -uw_check_only 1 +""" +import math + +import numpy as np +import sympy + +import underworld3 as uw +from underworld3.adjoint import misfit_duals, inner + +params = uw.Params( + cell_size=uw.Param(1 / 12, "base mesh cell size (box is 2 x 1); refined once, so half this"), + surface_dip=uw.Param(60.0, "dip of the ramp where it reaches the surface, degrees"), + flat_depth=uw.Param(0.3, "height of the decollement above the base"), + surface_x=uw.Param(1.9, "where the fault reaches the surface"), + band=uw.Param(0.08, "half-width of the weak band, in box units"), + true_strengths=uw.Param("0.005,0.05,0.02,0.2", + "weak-plane viscosity: flat, lower ramp, upper ramp, near surface"), + initial_strength=uw.Param(0.1, "starting guess, every segment"), + check_only=uw.Param(0, "1: gradient check against finite differences, no inversion"), +) + +# --- the model --------------------------------------------------------------- +# Refined once from the base size: the refinement gives the velocity block a +# multigrid hierarchy. Without one it falls back to gamg, which hits its +# iteration cap on this problem, and an inexact Newton step converges linearly. +# Every solve starts cold: six Newton iterations, and a misfit that does not +# depend on the previous evaluation, which the finite-difference check needs. +mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0.0, 0.0), maxCoords=(2.0, 1.0), + cellSize=params.cell_size, qdegree=3, refinement=1) +x, y = mesh.X + +v = uw.discretisation.MeshVariable("v", mesh, mesh.dim, degree=2) +p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1, continuous=True) +v_obs = uw.discretisation.MeshVariable("v_obs", mesh, mesh.dim, degree=2) + +# The fault: a flat decollement at y = flat_depth running from the left wall +# to x = xc, then a circular ramp of radius R about (xc, yc) that leaves the +# flat horizontally and reaches the surface at the given dip. The signed +# distance to it and the position along it are exact on each piece, and the +# director is the normal at the nearest point: vertical on the flat, radial on +# the ramp. +phi_top = math.radians(params.surface_dip) - math.pi / 2 # angle of the surface point about the centre +R = (1.0 - params.flat_depth) / (1.0 + math.sin(phi_top)) +yc = params.flat_depth + R +xc = params.surface_x - R * math.cos(phi_top) +r = sympy.sqrt((x - xc) ** 2 + (y - yc) ** 2) +phi = sympy.atan2(y - yc, x - xc) +on_flat = x < xc +d = sympy.Piecewise((y - params.flat_depth, on_flat), (R - r, True)) +s = sympy.Piecewise((x, on_flat), (xc + R * (phi + sympy.pi / 2), True)) +n_hat = sympy.Matrix([[sympy.Piecewise((0, on_flat), ((x - xc) / r, True)), + sympy.Piecewise((1, on_flat), ((y - yc) / r, True))]]) +ramp = R * (phi_top + math.pi / 2) +length = xc + ramp +band = sympy.exp(-(d / params.band) ** 2) + +# Segments along the fault: the flat, then the ramp in three equal parts. +edges = [0.0, xc, xc + ramp / 3, xc + 2 * ramp / 3, length] +names = ["flat", "lower ramp", "upper ramp", "near surface"] +n_seg = len(names) +strengths = [uw.expression(rf"\eta_{{{k + 1}}}", params.initial_strength, + f"weak-plane viscosity, {names[k]}") + for k in range(n_seg)] + +def segment(k): + """A smooth indicator for segment k along the fault, in [0, 1].""" + edge = params.band + on = 1 if k == 0 else (1 + sympy.tanh((s - edges[k]) / edge)) / 2 + off = 1 if k == n_seg - 1 else (1 - sympy.tanh((s - edges[k + 1]) / edge)) / 2 + return on * off + +eta_0 = 1 +eta_1 = eta_0 - band * sum((eta_0 - strengths[k]) * segment(k) for k in range(n_seg)) + +stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) +stokes.constitutive_model = uw.constitutive_models.TransverseIsotropicFlowModel +stokes.constitutive_model.Parameters.shear_viscosity_0 = eta_0 +stokes.constitutive_model.Parameters.shear_viscosity_1 = eta_1 +stokes.constitutive_model.Parameters.director = n_hat +stokes.tolerance = 1e-8 + +# Shortening from both sides, a no-slip base, and a free top: the surface +# velocity is the uplift rate. +stokes.add_essential_bc((0.0, 0.0), "Bottom") +stokes.add_essential_bc((0.5, None), "Left") +stokes.add_essential_bc((-0.5, None), "Right") + +# --- the observations ------------------------------------------------------ +w_top = sympy.exp(-((1 - y) / params.band) ** 2) +points = [(0.3, 0.65), (0.75, 0.12), (1.25, 0.3), (0.9, 0.6), (1.55, 0.85)] +w_points = sum(sympy.exp(-((x - px) ** 2 + (y - py) ** 2) / (2 * params.band) ** 2) + for px, py in points) + +def shear_stress(field): + e = mesh.vector.strain_tensor(field.sym) + return 2 * eta_0 * e[0, 1] + +misfit = (w_top * (v.sym[1] - v_obs.sym[1]) ** 2 + + w_points * (shear_stress(v) - shear_stress(v_obs)) ** 2) / 2 + +def set_strengths(values): + for expr, value in zip(strengths, values): + expr.sym = float(value) + +def J_and_gradient(): + """The misfit and dJ/d(log strength) for each segment, by the adjoint.""" + stokes.solve(zero_init_guess=True) + J = float(uw.maths.Integral(mesh, misfit).evaluate()) + dJ_dv = misfit_duals(misfit, [v])[v] + mu = uw.discretisation.MeshVariable(f"mu_{uw.adjoint._counter()}", mesh, mesh.dim, degree=2) + lam = uw.discretisation.MeshVariable(f"lam_{uw.adjoint._counter()}", mesh, 1, degree=1) + dJ_dv.array[...] = -np.asarray(dJ_dv.array) + _, reason = stokes.adjoint_solve((dJ_dv, None), target=(mu, lam)) + assert reason > 0, reason + grad = np.array([stokes.sensitivity(mu, expr) * float(expr.sym) for expr in strengths]) + return J, grad + +# --- the truth, and the twin ------------------------------------------------- +true_values = [float(t) for t in str(params.true_strengths).split(",")][:n_seg] +set_strengths(true_values) +stokes.solve(zero_init_guess=True) +v_obs.array[...] = np.asarray(v.array) +uw.pprint(f"true strengths {true_values}") + +set_strengths([params.initial_strength] * n_seg) +J0, g0 = J_and_gradient() +uw.pprint(f"initial J = {J0:.6e} dJ/dlog eta = {g0}") + +# Gradient check: central differences in each log-strength. +h = 1e-3 +for k in range(n_seg): + base = math.log(params.initial_strength) + fd = [] + for sign in (+1, -1): + vals = [params.initial_strength] * n_seg + vals[k] = math.exp(base + sign * h) + set_strengths(vals) + stokes.solve(zero_init_guess=True) + fd.append(float(uw.maths.Integral(mesh, misfit).evaluate())) + fd = (fd[0] - fd[1]) / (2 * h) + uw.pprint(f"{names[k]:>13}: adjoint {g0[k]: .6e} finite difference {fd: .6e} " + f"ratio {fd / g0[k]:.5f}") +set_strengths([params.initial_strength] * n_seg) + +if int(params.check_only): + raise SystemExit + +# --- the inversion --------------------------------------------------------------- +from scipy.optimize import minimize + +history = [] + +def objective(log_eta): + set_strengths(np.exp(log_eta)) + J, grad = J_and_gradient() + history.append((J, np.exp(log_eta).copy())) + uw.pprint(f" J = {J:.6e} strengths = {np.exp(log_eta)}") + return J, grad + +result = minimize(objective, np.log([params.initial_strength] * n_seg), jac=True, + method="L-BFGS-B", options={"maxiter": 40, "gtol": 1e-10}) +uw.pprint(f"recovered {np.exp(result.x)} true {true_values} " + f"after {len(history)} evaluations") + +# --- what the figure needs ----------------------------------------------------- +# Uplift-rate profiles along the top at the truth, the start and the answer, the +# weak-plane viscosity on a grid, and the path the strengths took. +xs = np.linspace(0.0, 2.0, 161) +top = np.column_stack([xs, np.full_like(xs, 1.0 - 1e-6)]) +profiles = {} +for label, values in (("true", true_values), ("initial", [params.initial_strength] * n_seg), + ("recovered", list(np.exp(result.x)))): + set_strengths(values) + stokes.solve(zero_init_guess=True) + profiles[label] = np.asarray(uw.function.evaluate(v.sym[1], top)).ravel() +set_strengths(true_values) +stokes.solve(zero_init_guess=True) # the field on the grid is the truth's +gx, gy = np.meshgrid(np.linspace(0, 2, 201), np.linspace(0, 1, 101)) +grid = np.column_stack([gx.ravel(), gy.ravel()]) +eta_1_grid = np.asarray(uw.function.evaluate(eta_1, grid)).reshape(gx.shape) +np.savez("fault_segments_data.npz", xs=xs, gx=gx, gy=gy, eta_1=eta_1_grid, + points=np.array(points), true=np.array(true_values), + history=np.array([[J, *vals] for J, vals in history]), + **{f"uplift_{k}": val for k, val in profiles.items()}) diff --git a/docs/examples/adjoint/fault_segments/plot_convergence.py b/docs/examples/adjoint/fault_segments/plot_convergence.py new file mode 100644 index 00000000..7b458006 --- /dev/null +++ b/docs/examples/adjoint/fault_segments/plot_convergence.py @@ -0,0 +1,36 @@ +"""Convergence of the friction inversion under the three observation sets.""" +import numpy as np +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +cases = [("uplift + stress", "fault_friction_uplift+stress_data.npz"), + ("stress orientation at five interior points", "fault_friction_orientation_points_data.npz"), + ("surface strain rate only", "fault_friction_surface_strain_data.npz")] +names = ["flat", "lower ramp", "upper ramp", "near surface"] + +fig, axes = plt.subplots(2, 3, figsize=(11, 5.6), sharex="col", + gridspec_kw={"height_ratios": [2.2, 1]}) +for col, (title, path) in enumerate(cases): + d = np.load(path) + history, true = d["history"], d["true"] + its = np.arange(len(history)) + ax = axes[0, col] + for k in range(len(true)): + ax.semilogy(its, history[:, 1 + k], "o-", ms=3, lw=1, color=f"C{k}", label=names[k]) + ax.axhline(true[k], color=f"C{k}", lw=0.8, ls="--") + ax.set_title(title, fontsize=10) + ax.set_ylim(0.03, 0.6) + if col == 0: + ax.set_ylabel("friction coefficient") + ax.legend(fontsize=8, frameon=False, loc="lower left") + ax = axes[1, col] + ax.semilogy(its, history[:, 0] / history[0, 0], "k.-", ms=4, lw=1) + ax.set_xlabel("misfit evaluation") + if col == 0: + ax.set_ylabel("$J / J_0$") + ax.set_ylim(1e-11, 2) +fig.tight_layout() +fig.savefig("fault_friction_convergence.png", dpi=180) +fig.savefig("fault_friction_convergence.pdf") +print("wrote fault_friction_convergence.png / .pdf") diff --git a/docs/examples/adjoint/fault_segments/plot_fault_segments.py b/docs/examples/adjoint/fault_segments/plot_fault_segments.py new file mode 100644 index 00000000..8da98b59 --- /dev/null +++ b/docs/examples/adjoint/fault_segments/plot_fault_segments.py @@ -0,0 +1,60 @@ +"""The figure for the fault-segments examples. + + python plot_fault_segments.py fault_segments_data.npz # weak-plane viscosity + python plot_fault_segments.py fault_friction_data.npz # friction coefficient +""" +import sys +import numpy as np +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +source = sys.argv[1] if len(sys.argv) > 1 else "fault_segments_data.npz" +friction = "friction" in source +d = np.load(source) +history, true = d["history"], d["true"] +length_unit = str(d["length_unit"]) if "length_unit" in d else "" +velocity_unit = str(d["velocity_unit"]) if "velocity_unit" in d else "" +x_max, y_max = float(d["gx"].max()), float(d["gy"].max()) +names = ["flat", "lower ramp", "upper ramp", "near surface"] +n_seg = len(true) + +fig, axes = plt.subplots(1, 3, figsize=(11, 3.2), gridspec_kw={"width_ratios": [2.2, 1.6, 1.4]}) + +ax = axes[0] +# The bulk sits at log10(eta_1) = 0 to round-off, so the top level is set a +# little below it and "extend" gives the bulk one flat colour. +ax.contourf(d["gx"], d["gy"], np.log10(d["eta_1"]), levels=np.linspace(-2.4, -0.1, 12), + cmap="viridis", extend="both") +band = float(d["band"]) if "band" in d else 0.08 +# the surface band under the top, where the uplift rate (or the orientation) is read +ax.axhspan(y_max - 2 * band, y_max, color="white", alpha=0.35, lw=0) +ax.plot(d["points"][:, 0], d["points"][:, 1], "wx", ms=7, mew=1.5) +ax.set_aspect("equal") +ax.set_xlim(0, x_max); ax.set_ylim(0, y_max) +ax.set_xlabel(f"$x$ ({length_unit})" if length_unit else "$x$") +ax.set_ylabel(f"$y$ ({length_unit})" if length_unit else "$y$") +ax.set_title(r"$\log_{10}\eta_1$ at the truth; $\times$ points, white band: surface observations", fontsize=10) + +ax = axes[1] +for label, style in (("true", "k-"), ("initial", "C3--"), ("recovered", "C0:")): + ax.plot(d["xs"], d[f"uplift_{label}"], style, lw=1.6, label=label) +ax.set_xlabel(f"$x$ along the surface ({length_unit})" if length_unit else "$x$ along the surface") +ax.set_ylabel(f"uplift rate $v_y$ ({velocity_unit})" if velocity_unit else "uplift rate $v_y$") +ax.set_title("surface uplift rate", fontsize=10) +ax.legend(fontsize=8, frameon=False) + +ax = axes[2] +its = np.arange(len(history)) +for k in range(n_seg): + ax.semilogy(its, history[:, 1 + k], "o-", ms=3, lw=1, color=f"C{k}", label=names[k]) + ax.axhline(true[k], color=f"C{k}", lw=0.8, ls="--") +ax.set_xlabel("misfit evaluation"); ax.set_ylabel("friction coefficient" if friction else "weak-plane viscosity") +ax.set_title(("friction" if friction else "strengths") + " (dashed: true)", fontsize=10) +ax.legend(fontsize=8, frameon=False) + +fig.tight_layout() +stem = source.replace("_data.npz", "") +fig.savefig(stem + ".png", dpi=180) +fig.savefig(stem + ".pdf") +print(f"wrote {stem}.png / .pdf") diff --git a/docs/examples/adjoint/fault_segments/plot_noise.py b/docs/examples/adjoint/fault_segments/plot_noise.py new file mode 100644 index 00000000..1d88b985 --- /dev/null +++ b/docs/examples/adjoint/fault_segments/plot_noise.py @@ -0,0 +1,40 @@ +"""Recovered friction against the noise on the observations, from the data files.""" +import glob +import re + +import numpy as np +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +names = ["flat", "lower ramp", "upper ramp", "near surface"] +runs = [] +for path in sorted(glob.glob("fault_friction_uplift+stress*_data.npz")): + m = re.search(r"_noise([0-9.]+)", path) + level = float(m.group(1)) if m else 0.0 + r = re.search(r"_prior([0-9.]+)", path) + alpha = float(r.group(1)) if r else 0.0 + d = np.load(path) + runs.append((level, alpha, d["history"][-1, 1:], d["true"])) +runs.sort(key=lambda r: (r[0], r[1])) +labels = [f"{r[0]:g}" + (f"\nprior σ={r[1]:g}" if r[1] > 0 else "") for r in runs] +true = runs[0][3] + +fig, ax = plt.subplots(figsize=(7.2, 3.8)) +x = np.arange(len(runs)) +w = 0.18 +for k in range(len(names)): + ax.bar(x + (k - 1.5) * w, [r[2][k] for r in runs], width=w, color=f"C{k}", label=names[k]) + ax.hlines(true[k], -0.5, len(runs) - 0.5, colors=f"C{k}", linestyles="--", lw=0.8) +ax.axhline(1.0, color="0.3", lw=0.8, ls=":") +ax.axhline(0.005, color="0.3", lw=0.8, ls=":") +ax.set_yscale("log") +ax.set_xticks(x) +ax.set_xticklabels(labels, fontsize=8) +ax.set_xlabel("noise on the observed velocity, fraction of its rms") +ax.set_ylabel("recovered friction (dashed: true; dotted: bounds)") +ax.legend(fontsize=8, frameon=False, ncol=2) +fig.tight_layout() +fig.savefig("fault_friction_noise.png", dpi=180) +fig.savefig("fault_friction_noise.pdf") +print("wrote fault_friction_noise.png / .pdf") diff --git a/docs/examples/adjoint/fault_segments/render_fault_friction.py b/docs/examples/adjoint/fault_segments/render_fault_friction.py new file mode 100644 index 00000000..8aa3fc87 --- /dev/null +++ b/docs/examples/adjoint/fault_segments/render_fault_friction.py @@ -0,0 +1,126 @@ +"""PyVista renders of the friction example at the true coefficients. + +Writes log10 eta_1 (the plane's viscosity), the slip rate, the velocity and +the pressure to ~/+Simulations/adjoint_fault_example/, on the mesh's own +triangulation. The model is the notebook's in its own units: lengths in +units of the depth (10 km), stresses in units of the viscous stress of the +shortening (26 MPa), velocities in units of the convergence rate (8.4 mm/yr). +""" +import math +import os + +import numpy as np +import sympy +import pyvista as pv + +import underworld3 as uw +import underworld3.visualisation as vis + +pv.OFF_SCREEN = True +OUT = os.path.expanduser("~/+Simulations/adjoint_fault_example") + +cell_size, surface_dip, flat_depth, surface_x, band_w = 1 / 12, 60.0, 0.3, 1.9, 0.08 +true_mu, cohesion, rho_g = [0.05, 0.15, 0.25, 0.4], 0.05, 10.0 + +mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0.0, 0.0), maxCoords=(2.0, 1.0), + cellSize=cell_size, qdegree=3, refinement=1) +x, y = mesh.X +v = uw.discretisation.MeshVariable("v", mesh, 2, degree=2) +p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1, continuous=True) + +phi_top = math.radians(surface_dip) - math.pi / 2 +R = (1.0 - flat_depth) / (1.0 + math.sin(phi_top)) +yc = flat_depth + R +xc = surface_x - R * math.cos(phi_top) +r = sympy.sqrt((x - xc) ** 2 + (y - yc) ** 2) +phi = sympy.atan2(y - yc, x - xc) +on_flat = x < xc +d = sympy.Piecewise((y - flat_depth, on_flat), (R - r, True)) +s = sympy.Piecewise((x, on_flat), (xc + R * (phi + sympy.pi / 2), True)) +n_hat = sympy.Matrix([[sympy.Piecewise((0, on_flat), ((x - xc) / r, True)), + sympy.Piecewise((1, on_flat), ((y - yc) / r, True))]]) +ramp = R * (phi_top + math.pi / 2) +edges_s = [0.0, xc, xc + ramp / 3, xc + 2 * ramp / 3, xc + ramp] +band = sympy.exp(-(d / band_w) ** 2) + +def segment(k): + on = 1 if k == 0 else (1 + sympy.tanh((s - edges_s[k]) / band_w)) / 2 + off = 1 if k == 3 else (1 - sympy.tanh((s - edges_s[k + 1]) / band_w)) / 2 + return on * off + +E = mesh.vector.strain_tensor(v.sym) +t_hat = sympy.Matrix([[-n_hat[1], n_hat[0]]]) +e_s = sympy.sqrt((t_hat * E * n_hat.T)[0] ** 2 + uw.maths.functions.vanishing) +friction = sum(true_mu[k] * segment(k) for k in range(4)) +tau_y = cohesion + friction * sympy.Max(p.sym[0], 0) +eta_plane = tau_y / (tau_y + 2 * e_s) +eta_1 = 1 - band * (1 - eta_plane) + +stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) +stokes.constitutive_model = uw.constitutive_models.TransverseIsotropicFlowModel +stokes.constitutive_model.Parameters.shear_viscosity_0 = 1 +stokes.constitutive_model.Parameters.shear_viscosity_1 = eta_1 +stokes.constitutive_model.Parameters.director = n_hat +stokes.tolerance = 1e-8 +stokes.bodyforce = sympy.Matrix([0, -rho_g]) +stokes.add_essential_bc((0.0, 0.0), "Bottom") +stokes.add_essential_bc((0.5, None), "Left") +stokes.add_essential_bc((-0.5, None), "Right") +stokes.solve(zero_init_guess=True) + +# Nodal values on the mesh's own triangulation (skill rule 5): no grid resampling. +pv_mesh = vis.mesh_to_pv_mesh(mesh) +pts = np.asarray(pv_mesh.points[:, :2]) +edges = pv_mesh.extract_all_edges() +pv_mesh.point_data["log10 eta_1"] = np.log10(np.clip(np.asarray(uw.function.evaluate(eta_1, pts)).ravel(), 1e-3, 1.0)) +pv_mesh.point_data["slip rate"] = np.asarray(uw.function.evaluate(band * e_s, pts)).ravel() +pv_mesh.point_data["p"] = np.asarray(uw.function.evaluate(p.sym[0], pts)).ravel() +pv_v = vis.meshVariable_to_pv_mesh_object(v) +vdata = np.asarray(v.data) +pv_v.point_data["|v|"] = np.linalg.norm(vdata, axis=1) +pv_v.point_data["v"] = np.column_stack([vdata, np.zeros(len(vdata))]) +pv_v.point_data["v_y"] = vdata[:, 1] + +points = [(0.3, 0.65), (0.75, 0.12), (1.25, 0.3), (0.9, 0.6), (1.55, 0.85)] + +def observations(pl): + """Where the observations are read: the five points, and the band under the + surface (two band widths deep, where its Gaussian weight is above 2%).""" + strip = pv.Rectangle([[0.0, 1 - 2 * band_w, 0.001], [2.0, 1 - 2 * band_w, 0.001], [2.0, 1.0, 0.001]]) + pl.add_mesh(strip, color="white", opacity=0.45, lighting=False) + pl.add_mesh(pv.Line([0, 1 - 2 * band_w, 0.002], [2, 1 - 2 * band_w, 0.002]), color="white", line_width=2, lighting=False) + for px, py in points: + pl.add_mesh(pv.Disc(center=(px, py, 0.003), inner=0.0, outer=0.022, normal=(0, 0, 1)), color="white", lighting=False) + pl.add_mesh(pv.Disc(center=(px, py, 0.004), inner=0.0, outer=0.013, normal=(0, 0, 1)), color="black", lighting=False) + +def frame(name, obj, scalars, cmap, clim, title, arrows=False, observed=False): + pl = pv.Plotter(off_screen=True, window_size=(1600, 820)) + pl.set_background("white") + pl.add_mesh(obj, scalars=scalars, cmap=cmap, clim=clim, show_edges=False, lighting=False, + scalar_bar_args=dict(title=title, color="black", vertical=False, position_x=0.3, + position_y=0.04, width=0.4, height=0.06, title_font_size=26, + label_font_size=22)) + pl.add_mesh(edges, color="black", line_width=0.25, lighting=False, opacity=0.2) + if arrows: + # a coarse regular set of points, so the arrows read as a field and not as the mesh + gx, gy = np.meshgrid(np.linspace(0.05, 1.95, 30), np.linspace(0.05, 0.95, 15)) + seeds = np.column_stack([gx.ravel(), gy.ravel()]) + vals = np.asarray(uw.function.evaluate(v.sym, seeds)).reshape(-1, 2) + cloud = pv.PolyData(np.column_stack([seeds, np.zeros(len(seeds))])) + cloud["v"] = np.column_stack([vals, np.zeros(len(vals))]) + cloud["|v|"] = np.linalg.norm(vals, axis=1) + pl.add_mesh(cloud.glyph(orient="v", scale="|v|", factor=0.09), color="black", lighting=False) + if observed: + observations(pl) + pl.view_xy() + pl.camera.parallel_projection = True + pl.camera.focal_point = (1.0, 0.42, 0.0) + pl.camera.parallel_scale = 0.66 + pl.screenshot(os.path.join(OUT, name)) + pl.close() + +frame("fault_eta1.png", pv_mesh, "log10 eta_1", "viridis", (-2.5, 0.0), "log10 plane viscosity", observed=True) +frame("fault_slip.png", pv_mesh, "slip rate", "magma_r", (0.0, 1.0), "shear strain rate on the plane") +frame("fault_vy.png", pv_v, "v_y", "viridis", (0.0, 0.7), "vertical velocity (uplift rate)", arrows=True) +frame("fault_p.png", pv_mesh, "p", "RdBu_r", (-8.0, 8.0), "pressure") +print("wrote", OUT) 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..d15da2eb --- /dev/null +++ b/src/underworld3/adjoint.py @@ -0,0 +1,726 @@ +"""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, boundary=None): + r"""The dual of a load on ``variable``'s space, held as a field. + + :math:`b_j = \int v\,\phi_j + \mathbf g\cdot\nabla\phi_j` for every basis + 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. + + With ``boundary`` (a mesh boundary label), the load is the facet + integral :math:`b_j = \int_\Gamma v\,\phi_j` instead — a misfit on + surface observations — assembled as a natural condition of the same + generic solver with zero volume templates. A gradient part on a boundary + is not assembled yet and raises. + """ + scratch = _shared_scratch if scratch is None else scratch + mesh = variable.mesh + n = getattr(variable, "num_components", 1) + dim, cdim = mesh.dim, mesh.cdim + asm = scratch.assembler(variable) + if boundary is not None: + has_grad = grad is not None and not sympy.Matrix(grad).is_zero_matrix + if has_grad and n == 1: + raise NotImplementedError( + "dual_on: a scalar load read through the gradient on a boundary " + "is not assembled yet (the scalar assembler has no facet flux " + "term); a vector one is") + # zero volume templates; the natural condition IS the load + asm._g0 = sympy.zeros(1, 1) if n == 1 else sympy.zeros(1, dim) + asm._g1 = sympy.zeros(1, cdim) if n == 1 else sympy.zeros(dim, cdim) + asm.natural_bcs.clear() + asm.add_natural_bc(value, boundary) + if has_grad: + # the facet flux part int_Gamma g . grad(phi_j): the same slot a + # Nitsche condition uses for its symmetry term + bc = asm.natural_bcs[-1] + asm.natural_bcs[-1] = bc._replace(fn_F=sympy.Matrix(grad).as_immutable()) + asm.is_setup = False # the facet kernels are registered on build + asm._build(False, False, None) + else: + if grad is None: + grad = sympy.zeros(1, cdim) if n == 1 else sympy.zeros(dim, cdim) + if asm.natural_bcs: + asm.natural_bcs.clear() + asm.is_setup = False + asm._g0 = value + asm._g1 = sympy.Matrix(grad) + asm._needs_function_rewire = True # the templates re-evaluate + asm._build(False, False, None) + out_var = scratch.take(variable) + gvec = asm.dm.getGlobalVec() + gvec.set(0.0) + 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 + + +def _token_of(var): + """How ``var`` prints inside an expression: everything before the + coordinate arguments and, for a vector, before the component index.""" + text = str(_symbols_of(var)[0]).rsplit("(", 1)[0] + if getattr(var, "num_components", 1) > 1: + text = text.rsplit("_{", 1)[0] + return text + + +def misfit_duals(misfit, variables, scratch=None, boundary=None): + r"""``dJ/df`` as a dual field on each field the misfit reads. + + ``J = \int misfit`` over the mesh — or over the boundary ``boundary`` + when one is named, for a misfit on surface observations. A field enters + through its value and, for a misfit on a stress or a strain rate, + through its gradient. Both parts are differentiated symbolically and + assembled as one load (:func:`dual_on`), so a misfit written in terms of + :math:`\nabla u` needs no integration by parts by the caller. Returns + ``{variable: dual}`` for the variables that appear; give each dual back + to the scratch pool when done. + """ + scratch = _shared_scratch if scratch is None else scratch + peeled = _peel(misfit) + text = str(peeled) + atoms = set(peeled.atoms(sympy.Function)) + out = {} + for var in variables: + token = _token_of(var) + if token not in text: + continue + symbols = _symbols_of(var) + value = [sympy.diff(peeled, s) for s in symbols] + pattern = re.compile(re.escape(token) + r"_\{ ?(\d*),(\d+)\}\(") + g1 = None + for atom in atoms: + m = pattern.match(str(atom)) + if not m: + continue + if g1 is None: + g1 = sympy.zeros(len(symbols), var.mesh.cdim) + i = int(m.group(1)) if m.group(1) else 0 + g1[i, int(m.group(2))] = sympy.diff(peeled, atom) + if all(v == 0 for v in value) and (g1 is None or g1.is_zero_matrix): + continue + out[var] = dual_on(var, _as_expression(value), g1, scratch, boundary=boundary) + return out + + +def integral(mesh, expression, boundary=None): + """``float(∫ expression)`` over the mesh, or over ``boundary`` if named.""" + if boundary is None: + return float(uw.maths.Integral(mesh, expression).evaluate()) + return float(uw.maths.BdIntegral(mesh, expression, boundary).evaluate()) + + +def _reads_of(solver, unknown, tokens): + """What a solver's residual reads, other than its unknown. + + ``(variable, value symbols, {(component, direction): derivative atom})`` + per variable. A component prints as ``{v}_{ 0 }``; a derivative carries + a comma — ``{v}_{ 0,1}`` for a vector, ``{T}_{,1}`` for a scalar — and + is read through the gradient part of the load. + """ + f0 = _peel(solver.F0.sym) + f1 = _peel(solver.F1.sym) + text = str(f0) + str(f1) + atoms = set(f0.atoms(sympy.Function)) | set(f1.atoms(sympy.Function)) + found = [] + for token, var in tokens.items(): + if var is unknown or token not in text: + continue + derivatives = {} + pattern = re.compile(re.escape(token) + r"_\{ ?(\d*),(\d+)\}\(") + for atom in atoms: + m = pattern.match(str(atom)) + if m: + i = int(m.group(1)) if m.group(1) else 0 + derivatives[(i, int(m.group(2)))] = atom + found.append((var, _symbols_of(var), derivatives)) + return found + + +def _tokens_of(variables): + return {_token_of(var): var for var in variables if hasattr(var, "sym")} + + +def field_duals(solver, mu, variables, scratch=None): + r"""``(\partial R/\partial f)^T \mu`` as a dual on each field ``f`` the + solver's residual reads, among ``variables``. + + A field read through its value gives the value part of the load; one + read through its gradient (a Crank–Nicolson step reads the old flux) + gives the gradient part. Both come from :meth:`adjoint_integrand`, the + symbolic derivative of the residual, and are assembled as one load. + """ + scratch = _shared_scratch if scratch is None else scratch + out = {} + for var, symbols, derivatives in _reads_of(solver, solver.u, _tokens_of(variables)): + value = [solver.adjoint_integrand(mu, s) for s in symbols] + g1 = None + if derivatives: + cdim = var.mesh.cdim + g1 = sympy.zeros(len(symbols), cdim) + for (i, k), atom in derivatives.items(): + g1[i, k] = solver.adjoint_integrand(mu, atom) + out[var] = dual_on(var, _as_expression(value), g1, scratch) + return out + + +def gradient(solver, misfit, parameters=(), fields=(), scratch=None, boundary=None): + r"""``dJ/dm`` and the duals on fields, by the adjoint of ONE solve. + + For :math:`J = \int` ``misfit`` over the mesh, evaluated in the state the + solver ended in: the dual of :math:`J` on the unknown is assembled + (:func:`misfit_duals`), the adjoint system :math:`K^T\mu = -\partial J/ + \partial u` is solved (:meth:`adjoint_solve`), and each parameter gets + :math:`\partial J/\partial m + \mu^T\partial R/\partial m` + (:meth:`sensitivity`). Each requested field gets :math:`\partial J/ + \partial f + (\partial R/\partial f)^T\mu` as a dual on its own space + — the derivative through a field the residual reads, an initial + condition or a coefficient field. + + ``boundary`` names a mesh boundary label over which the misfit is + integrated instead of the volume: surface observations. A misfit with + terms on several domains is a dict ``{None: volume integrand, "Top": + surface integrand}``; then ``boundary`` is ignored. + + Returns ``{"J": float, "parameters": {expr: float}, "fields": {var: dual}}``; + the duals are NumPy copies, safe to keep. + """ + scratch = _shared_scratch if scratch is None else scratch + parameters, fields = list(parameters), list(fields) + u = solver.u + mesh = u.mesh + # One misfit, or several terms on different domains: {None: volume + # integrand, "Top": surface integrand, ...}. J and every dual are sums. + terms = dict(misfit) if isinstance(misfit, dict) else {boundary: misfit} + J = 0.0 + duals = {} + grad = {p: 0.0 for p in parameters} + read = [u] + [f for f in fields if f is not u] + for where, term in terms.items(): + J += integral(mesh, term, where) + for var, dual in misfit_duals(term, read, scratch, boundary=where).items(): + if var in duals: + duals[var].array[...] = np.asarray(duals[var].array) + np.asarray(dual.array) + scratch.give(dual) + else: + duals[var] = dual + for p in parameters: + explicit = sympy.diff(_peel_except(term, p), p) + if explicit != 0: + grad[p] += integral(mesh, explicit, where) + # The explicit part on a field the misfit reads directly — never on the + # unknown, whose misfit dual is the adjoint's right-hand side. + out_fields = {var: (np.array(duals[var].array, copy=True) if (var in duals and var is not u) + else np.zeros_like(np.asarray(var.array))) for var in fields} + # A history slot the residual reads (psi_star[0]) holds the tracked field + # at the solve's input, so its dual is the derivative with respect to + # that field: read the slot, route the dual to the field. + route = {} + for history in (getattr(solver, "DuDt", None), getattr(solver, "DFDt", None)): + if history is None or not getattr(history, "psi_star", None): + continue + text = str(history.psi_fn) + for field in fields: + # the unknown itself is a control through its INPUT level, which + # is what the slot holds + if _token_of(field) in text: + route[history.psi_star[0]] = field + read_vars = [f for f in fields if f is not u] + list(route) + if u in duals: + rhs = duals.pop(u) + rhs.array[...] = -np.asarray(rhs.array) + mu = scratch.take(u) + if getattr(solver, "p", None) is not None and hasattr(solver, "_subdict"): + lam = scratch.take(solver.p) + _, reason = solver.adjoint_solve((rhs, None), target=(mu, lam)) + scratch.give(lam) + else: + _, reason = solver.adjoint_solve(rhs, target=mu) + scratch.give(rhs) + if reason <= 0: + raise RuntimeError(f"gradient: the adjoint of {type(solver).__name__}({u.name}) " + f"did not converge ({reason})") + for p in parameters: + grad[p] += solver.sensitivity(mu, p) + for var, dual in field_duals(solver, mu, read_vars, scratch).items(): + target = route.get(var, var) + out_fields[target] = out_fields[target] + np.asarray(dual.array) + scratch.give(dual) + scratch.give(mu) + for dual in duals.values(): + scratch.give(dual) + return {"J": J, "parameters": grad, "fields": out_fields} + + +def minimise(objective, x0, bounds=None, method="lmvm", options=None, max_evaluations=100, + gradient_tolerance=1e-8, callback=None): + r"""Minimise ``objective(x) -> (J, dJ/dx)`` with PETSc TAO. + + The driver for an inversion: the misfit and its gradient come from the + adjoint, the step along the gradient is the line search's, and the + quasi-Newton update is TAO's limited-memory one (``"lmvm"``; ``"blmvm"`` + honours ``bounds``, a pair of arrays). ``x0`` is a NumPy array; a few + scalar controls are replicated on every rank, each rank's TAO doing the + same arithmetic on the same numbers, and the objective's own collective + solves keep the ranks in step. Returns ``(x, info)`` with the iterations, + the converged reason and the history of ``(J, x)`` per evaluation. + """ + from petsc4py import PETSc + x0 = np.asarray(x0, dtype=float).ravel() + x = PETSc.Vec().createSeq(x0.size, comm=PETSc.COMM_SELF) + x.setArray(x0) + history = [] + + def fg(tao, xv, g): + values = np.array(xv.getArray(readonly=True), copy=True) + J, grad = objective(values) + g.setArray(np.asarray(grad, dtype=float).ravel()) + history.append((float(J), values)) + if callback is not None: + callback(J, values) + return float(J) + + tao = PETSc.TAO().create(comm=PETSc.COMM_SELF) + tao.setType(method) + tao.setObjectiveGradient(fg, None) + if bounds is not None: + lo, hi = bounds + lower = x.duplicate(); lower.setArray(np.asarray(lo, dtype=float).ravel()) + upper = x.duplicate(); upper.setArray(np.asarray(hi, dtype=float).ravel()) + tao.setVariableBounds(lower, upper) + tao.setMaximumFunctionEvaluations(int(max_evaluations)) + tao.setTolerances(gatol=gradient_tolerance) + for key, value in (options or {}).items(): + PETSc.Options().setValue(key, value) + tao.setFromOptions() + tao.setSolution(x) + tao.solve() + info = {"iterations": int(tao.getIterationNumber()), + "reason": int(tao.getConvergedReason()), + "history": history} + out = np.array(x.getArray(readonly=True), copy=True) + tao.destroy() + return out, info + + +_n = [0] + + +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] = {} + for var, dual in misfit_duals(misfit, self._tokens().values(), scratch).items(): + self._accumulate(acc, var, dual) + + grad = {p: 0.0 for p in parameters} + for p in parameters: + 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 (:func:`_token_of`). + Matching on the name found the user's fields and silently missed + every history, which cut the chain at the first step.""" + return {_token_of(var): var for var in self.model._variables.values() + if hasattr(var, "sym")} + + def _reads(self, solver, unknown): + return _reads_of(solver, unknown, self._tokens()) + + def _linearise_at(self, step, solves, j): + """Restore the step's snapshot, replay solves 0..j, and put each + 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 + and not getattr(getattr(e, "sym", None), "is_Number", False)] + 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..5c4299b0 100644 --- a/src/underworld3/cython/petsc_generic_snes_solvers.pyx +++ b/src/underworld3/cython/petsc_generic_snes_solvers.pyx @@ -53,6 +53,95 @@ expression = lambda *x, **X: public_expression(*x, _unique_name_generation=True, from underworld3.function.expressions import unwrap_expression as _unwrap_expression +def _transpose_kernels(G0, G1, G2, G3, nc_test, nc_trial, dim): + """The pointwise Jacobian kernels of the TRANSPOSED bilinear form. + + PETSc reads a Jacobian block (test field with ``nc_test`` components, + trial field with ``nc_trial``) from four flat buffers laid out + ``g0[fc, gc]``, ``g1[fc, gc, d]``, ``g2[fc, gc, d]``, ``g3[fc, gc, d, e]`` + — ``fc`` the test component, ``gc`` the trial component, ``d`` the + test-side derivative direction in g2/g3 and the trial-side one in g1, + ``e`` the trial-side direction in g3. The form is + + phi_fc g0 u_gc + phi_fc g1 d_d u_gc + d_d phi_fc g2 u_gc + d_d phi_fc g3 d_e u_gc. + + Swapping trial and test gives the block of K^T with the roles of the + two fields exchanged: h0 is g0 transposed, h1 and h2 are g2 and g1 with + the components exchanged, and h3 is g3 with both index pairs exchanged. + Nothing is differentiated again; the derivatives are the forward ones, + evaluated at the same state. Matrices are returned whose row-major + flattening is PETSc's layout for the transposed block. A ``None`` + kernel is an absent (zero) one. + """ + import sympy + + def flat(G, n): + if G is None: + return [sympy.Integer(0)] * n + values = list(G) + if len(values) != n: + raise ValueError(f"kernel has {len(values)} entries, expected {n}") + return values + + nt, nu = nc_test, nc_trial + g0 = flat(G0, nt * nu) + g1 = flat(G1, nt * nu * dim) + g2 = flat(G2, nt * nu * dim) + g3 = flat(G3, nt * nu * dim * dim) + H0 = sympy.zeros(nu, nt) + H1 = sympy.zeros(nu * nt, dim) + H2 = sympy.zeros(nu * nt, dim) + H3 = sympy.zeros(nu * nt, dim * dim) + for fc in range(nt): + for gc in range(nu): + H0[gc, fc] = g0[fc * nu + gc] + for d in range(dim): + H1[gc * nt + fc, d] = g2[(fc * nu + gc) * dim + d] + H2[gc * nt + fc, d] = g1[(fc * nu + gc) * dim + d] + for e in range(dim): + H3[gc * nt + fc, e * dim + d] = g3[((fc * nu + gc) * dim + d) * dim + e] + return (sympy.ImmutableMatrix(H0), sympy.ImmutableMatrix(H1), + sympy.ImmutableMatrix(H2), sympy.ImmutableMatrix(H3)) + + +def _solve_transposed(ksp, J, P, b, x): + """Solve :math:`J^T x = b` with the solver's OWN KSP and preconditioner. + + ``KSP.solveTranspose`` applies the transpose of the preconditioner, which + PETSc refuses for a multigrid whose smoother is a one-sided SOR sweep — + the velocity block of every Stokes solve with a mesh hierarchy. So the + matrices are transposed explicitly and the same KSP solves the + transposed system forwards: every preconditioner the forward solve can + use, the adjoint can use, set up on :math:`J^T` exactly as it was on + :math:`J`. The forward operators are put back afterwards. A null space + attached to ``J`` (the pressure constant on an enclosed domain) is the + transpose null space of ``J^T``, and is carried across. + """ + # Mat.transpose() with no target transposes IN PLACE; give it a new Mat. + Jt = PETSc.Mat() + J.transpose(Jt) + if P.handle == J.handle: + Pt = Jt + else: + Pt = PETSc.Mat() + P.transpose(Pt) + for source, put in ((J.getNullSpace(), Jt.setTransposeNullSpace), + (J.getTransposeNullSpace(), Jt.setNullSpace), + (J.getNearNullSpace(), Jt.setNearNullSpace)): + if source is not None and source.handle != 0: + put(source) + ksp.setOperators(Jt, Pt) + try: + ksp.solve(b, x) + reason = int(ksp.getConvergedReason()) + finally: + ksp.setOperators(J, P) + if Pt is not Jt: + Pt.destroy() + Jt.destroy() + return reason + + def _jacobian_unwrap(expr): """Expand UWexpressions down to (but NOT including) constant atoms, for use as the input to a Jacobian derivative (``derive_by_array`` / ``diff``). @@ -131,7 +220,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 +483,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 +1227,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 +1490,601 @@ 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. + by_kernels = self._adjoint_by_kernels() + tangent = (self._install_adjoint_kernels() if by_kernels + else self._consistent_tangent_for_adjoint()) + dm = self.dm + + # The Jacobian at the state the forward solve ended in. + 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) + + if by_kernels: + # J IS K^T: assembled from the transposed kernels. Solve forwards. + ksp = self.snes.getKSP() + ksp.setOperators(J, P) + ksp.solve(b, x) + reason = int(ksp.getConvergedReason()) + self._uninstall_adjoint_kernels(tangent) + else: + reason = _solve_transposed(self.snes.getKSP(), J, P, b, x) + self._restore_tangent(tangent) + + if target is not None: + # Homogeneous constraints: the local vector is zeroed before the + # 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 gradient(self, misfit, parameters=(), fields=(), boundary=None): + r"""``dJ/dm`` for each parameter, and the dual on each field, by the + adjoint of this solve. + + ``misfit`` is :math:`J` as an integrand over the mesh in the state + the solve ended in. The three steps — the dual of :math:`J` on the + unknown, the transposed solve, the sensitivities — are + :func:`underworld3.adjoint.gradient`; this is the method form of it. + Returns ``{"J", "parameters": {expr: dJ/dm}, "fields": {var: dual}}``. + """ + from underworld3.adjoint import gradient as _gradient + return _gradient(self, misfit, parameters=parameters, fields=fields, + boundary=boundary) + + def adjoint_kernels(self): + """The pointwise kernels of the adjoint operator, as SymPy matrices. + + ``{"F0_u", "F0_grad_u", "F1_u", "F1_grad_u"}`` per block, in + PETSc's flat layout: the forward Jacobian kernels with trial and + test exchanged (:func:`_transpose_kernels`), which is what + :meth:`adjoint_solve` assembles. Read-only; the solver must have + been built. + """ + if getattr(self, "_G0", None) is None and getattr(self, "_uu_G0", None) is None: + self._build(False, False, None) + dim = self.mesh.cdim + if getattr(self, "_uu_G0", None) is not None: + uu = _transpose_kernels(self._uu_G0, self._uu_G1, self._uu_G2, self._uu_G3, dim, dim, dim) + up = _transpose_kernels(self._pu_G0, self._pu_G1, None, None, 1, dim, dim) + pu = _transpose_kernels(self._up_G0, self._up_G1, self._up_G2, self._up_G3, dim, 1, dim) + names = ("F0_u", "F0_grad_u", "F1_u", "F1_grad_u") + return {"uu": dict(zip(names, uu)), "up": dict(zip(names, up)), "pu": dict(zip(names, pu))} + nc = int(self._G0.shape[0]) + H = _transpose_kernels(self._G0, self._G1, self._G2, self._G3, nc, nc, dim) + return dict(zip(("F0_u", "F0_grad_u", "F1_u", "F1_grad_u"), H)) + + def adjoint_templates(self): + r"""The adjoint problem in the residual template form. + + :math:`K^T\mu = b` is linear in :math:`\mu`, so it is a solver in + the same language as the forward one: + + .. math:: + + f_0^{\rm adj} = g_0^T\mu + g_2^T\!:\!\nabla\mu, \qquad + \mathbf f_1^{\rm adj} = g_1^T\mu + g_3^T\nabla\mu, + + the forward Jacobian kernels with trial and test exchanged, applied + to the adjoint variable. Returns ``(F0_adj, F1_adj)`` as SymPy + expressions in :math:`\mu` (a scalar or a row vector of the + unknown's size) for a single-field solver, and the same per block + for a saddle point, which :meth:`adjoint_view` typesets. + """ + import sympy + kernels = self.adjoint_kernels() + dim = self.mesh.cdim + + def apply(block, nc, mu, grad_mu): + H0, H1, H2, H3 = (block["F0_u"], block["F0_grad_u"], block["F1_u"], block["F1_grad_u"]) + nt = int(H0.shape[0]) # test components of the adjoint block + f0 = sympy.zeros(nt, 1) + f1 = sympy.zeros(nt, dim) + for a in range(nt): + for c in range(nc): + f0[a] += H0[a, c] * mu[c] + for d in range(dim): + f0[a] += H1[a * nc + c, d] * grad_mu[c, d] + f1[a, d] += H2[a * nc + c, d] * mu[c] + for e in range(dim): + f1[a, d] += H3[a * nc + c, d * dim + e] * grad_mu[c, e] + return f0, f1 + + if "uu" in kernels: + mu = sympy.Matrix([sympy.Symbol(f"\\mu_{{{i}}}") for i in range(dim)]) + gmu = sympy.Matrix(dim, dim, lambda i, j: sympy.Symbol(f"\\mu_{{{i},{j}}}")) + lam = sympy.Matrix([sympy.Symbol(r"\lambda")]) + glam = sympy.Matrix(1, dim, lambda i, j: sympy.Symbol(f"\\lambda_{{,{j}}}")) + f0_uu, f1_uu = apply(kernels["uu"], dim, mu, gmu) + f0_up, f1_up = apply(kernels["up"], 1, lam, glam) + f0_pu, f1_pu = apply(kernels["pu"], dim, mu, gmu) + return {"u": (f0_uu + f0_up, f1_uu + f1_up), "p": (f0_pu, f1_pu)} + nc = int(kernels["F0_u"].shape[0]) + if nc == 1: + mu = sympy.Matrix([sympy.Symbol(r"\mu")]) + gmu = sympy.Matrix(1, dim, lambda i, j: sympy.Symbol(f"\\mu_{{,{j}}}")) + else: + mu = sympy.Matrix([sympy.Symbol(f"\\mu_{{{i}}}") for i in range(nc)]) + gmu = sympy.Matrix(nc, dim, lambda i, j: sympy.Symbol(f"\\mu_{{{i},{j}}}")) + return apply(kernels, nc, mu, gmu) + + def adjoint_view(self): + """Typeset the adjoint problem this solver assembles, in a notebook.""" + from IPython.display import Latex, Markdown, display + import sympy + templates = self.adjoint_templates() + display(Markdown("### Adjoint problem, as assembled")) + display(Markdown(r"$\int \phi\, f_0^{\rm adj}(\mu,\nabla\mu) + \nabla\phi\cdot\mathbf f_1^{\rm adj}(\mu,\nabla\mu) = b(\phi)$ with")) + if isinstance(templates, dict): + for block, (f0, f1) in templates.items(): + display(Latex(f"$f_0^{{\\rm adj}}[{block}] = {sympy.latex(f0)}$")) + display(Latex(f"$\\mathbf f_1^{{\\rm adj}}[{block}] = {sympy.latex(f1)}$")) + else: + f0, f1 = templates + display(Latex(f"$f_0^{{\\rm adj}} = {sympy.latex(f0)}$")) + display(Latex(f"$\\mathbf f_1^{{\\rm adj}} = {sympy.latex(f1)}$")) + + def _adjoint_by_kernels(self): + """Whether the adjoint operator is ASSEMBLED from the transposed + kernels rather than obtained by transposing the assembled matrix. + The kernel route covers the volume terms; a boundary Jacobian (a + natural or Nitsche condition with a tangent) is not swapped yet, so + such a solver takes the matrix route.""" + return not (getattr(self, "natural_bcs", None) or []) + + def _install_adjoint_kernels(self): + """Rewire the solver to the kernels of the transposed form. + + The Jacobian the SNES then assembles IS K^T at the state the forward + solve ended in: the consistent tangent's derivatives with trial and + test exchanged (:func:`_transpose_kernels`), compiled and cached like + any other kernel set. Returns what to hand back to + :meth:`_uninstall_adjoint_kernels`. + """ + previous = self._consistent_jacobian + if self.consistent_jacobian is False and not self._residual_is_linear_in_unknown(): + self._consistent_jacobian = True # transpose dR/du, not the Picard kernel + self._adjoint_kernels = True + self._needs_function_rewire = True + self._build(False, False, None) + self.snes.setUp() # a direct computeJacobian needs it + return previous + + def _uninstall_adjoint_kernels(self, previous): + """Mark the forward kernels as wanted again. The next forward solve's + own build rewires; nothing is torn down here, so a second adjoint on + the same forward state reuses the installed set.""" + self._consistent_jacobian = previous + self._adjoint_kernels = False + self._adjoint_kernel_installed = True + self._needs_function_rewire = True + + def _consistent_tangent_for_adjoint(self): + """Make sure the Jacobian kernel the adjoint assembles is dR/du. + + 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. + """ + # A named expression whose value is a bare number is left as it is: + # the JIT passes it as a run-time constant, so the kernel compiled + # for the integrand survives a change of its value. Substituting the + # number recompiled every sensitivity at every evaluation of an + # inversion (four compiles per evaluation on the fault example). + for _ in range(depth): + named = [e for e in uw.function.fn_extract_expressions(expression) + if e is not wrt and e != wrt + and not getattr(getattr(e, "sym", None), "is_Number", False)] + 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. + """ + total = float(uw.maths.Integral(self.mesh, self.adjoint_integrand(mu, wrt)).evaluate()) + # A parameter that enters through a natural condition — a prescribed + # traction or flux — has a facet part: (d bd_F0 / dm) . mu on that boundary. + import sympy + mu_sym = mu.sym + for bc in (getattr(self, "natural_bcs", None) or []): + fn = getattr(bc, "fn_f", None) + if fn is None: + continue + d = sympy.diff(self._peel_except(sympy.Matrix(fn), wrt), wrt) + if d.is_zero_matrix: + continue + values = list(d) + comps = getattr(mu, "num_components", 1) + integrand = sum(values[i] * mu_sym[i] for i in range(min(len(values), comps))) + total += float(uw.maths.BdIntegral(self.mesh, integrand, bc.boundary).evaluate()) + return total + def _constraint_mechanisms(self): """Every way a constraint can have been put on this solver. @@ -2570,7 +3263,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 @@ -4048,6 +4744,12 @@ class SNES_Scalar(SolverBaseClass): self._G1 = sympy.ImmutableMatrix(G1) self._G2 = sympy.ImmutableMatrix(G2) self._G3 = sympy.ImmutableMatrix(G3) + if getattr(self, "_adjoint_kernels", False): + # The adjoint's operator: the same derivatives, trial and test + # exchanged. Compiled as its own kernel set and cached. + _nc = int(self._G0.shape[0]) + self._G0, self._G1, self._G2, self._G3 = _transpose_kernels( + self._G0, self._G1, self._G2, self._G3, _nc, _nc, self.mesh.cdim) ################## @@ -5049,6 +5751,12 @@ class SNES_Vector(SolverBaseClass): self._G1 = sympy.ImmutableMatrix(G1) self._G2 = sympy.ImmutableMatrix(G2) self._G3 = sympy.ImmutableMatrix(G3) + if getattr(self, "_adjoint_kernels", False): + # The adjoint's operator: the same derivatives, trial and test + # exchanged. Compiled as its own kernel set and cached. + _nc = int(self._G0.shape[0]) + self._G0, self._G1, self._G2, self._G3 = _transpose_kernels( + self._G0, self._G1, self._G2, self._G3, _nc, _nc, self.mesh.cdim) ################## @@ -5847,6 +6555,12 @@ class SNES_MultiComponent(SolverBaseClass): self._G1 = sympy.ImmutableMatrix(G1) self._G2 = sympy.ImmutableMatrix(G2) self._G3 = sympy.ImmutableMatrix(G3) + if getattr(self, "_adjoint_kernels", False): + # The adjoint's operator: the same derivatives, trial and test + # exchanged. Compiled as its own kernel set and cached. + _nc = int(self._G0.shape[0]) + self._G0, self._G1, self._G2, self._G3 = _transpose_kernels( + self._G0, self._G1, self._G2, self._G3, _nc, _nc, self.mesh.cdim) fns_jacobian = (self._G0, self._G1, self._G2, self._G3) @@ -8295,6 +9009,25 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): fns_jacobian.append(self._pp_G0) + if getattr(self, "_adjoint_kernels", False): + # K^T block by block: uu transposed in place; the (u,p) block of + # K^T is the transpose of K's (p,u) block and vice versa; pp is a + # scalar. Same derivatives, trial and test exchanged. + forward = [self._uu_G0, self._uu_G1, self._uu_G2, self._uu_G3, + self._up_G0, self._up_G1, self._up_G2, self._up_G3, + self._pu_G0, self._pu_G1] + up_from_pu = _transpose_kernels(self._pu_G0, self._pu_G1, None, None, 1, dim, dim) + pu_from_up = _transpose_kernels(self._up_G0, self._up_G1, self._up_G2, self._up_G3, + dim, 1, dim) + self._uu_G0, self._uu_G1, self._uu_G2, self._uu_G3 = _transpose_kernels( + self._uu_G0, self._uu_G1, self._uu_G2, self._uu_G3, dim, dim, dim) + self._up_G0, self._up_G1, self._up_G2, self._up_G3 = up_from_pu + self._pu_G0, self._pu_G1, self._pu_G2, self._pu_G3 = pu_from_up + fns_jacobian = [f for f in fns_jacobian if not any(f is g for g in forward)] + fns_jacobian += [self._uu_G0, self._uu_G1, self._uu_G2, self._uu_G3, + self._up_G0, self._up_G1, self._up_G2, self._up_G3, + self._pu_G0, self._pu_G1, self._pu_G2, self._pu_G3] + ## Lagrange-multiplier rows (block-constrained Stokes). Guarded: no-op ## for ordinary Stokes. Each multiplier h_k contributes an interior ## screening residual f0 = eps_k * h_k and a diagonal mass Jacobian @@ -8970,10 +9703,17 @@ class SNES_Stokes_SaddlePt(SolverBaseClass): PetscDSSetJacobian( ds.ds, 0, 0, ext.fns_jacobian[i_jac[self._uu_G0]], ext.fns_jacobian[i_jac[self._uu_G1]], ext.fns_jacobian[i_jac[self._uu_G2]], ext.fns_jacobian[i_jac[self._uu_G3]]) PetscDSSetJacobian( ds.ds, 0, 1, ext.fns_jacobian[i_jac[self._up_G0]], ext.fns_jacobian[i_jac[self._up_G1]], ext.fns_jacobian[i_jac[self._up_G2]], ext.fns_jacobian[i_jac[self._up_G3]]) - PetscDSSetJacobian( ds.ds, 1, 0, ext.fns_jacobian[i_jac[self._pu_G0]], ext.fns_jacobian[i_jac[self._pu_G1]], NULL, NULL) + if getattr(self, "_adjoint_kernels", False): + # the transposed (p,u) block carries the flux derivatives of K's (u,p) block + PetscDSSetJacobian( ds.ds, 1, 0, ext.fns_jacobian[i_jac[self._pu_G0]], ext.fns_jacobian[i_jac[self._pu_G1]], ext.fns_jacobian[i_jac[self._pu_G2]], ext.fns_jacobian[i_jac[self._pu_G3]]) + else: + PetscDSSetJacobian( ds.ds, 1, 0, ext.fns_jacobian[i_jac[self._pu_G0]], ext.fns_jacobian[i_jac[self._pu_G1]], NULL, NULL) PetscDSSetJacobianPreconditioner(ds.ds, 0, 0, ext.fns_jacobian[i_jac[self._uu_G0]], ext.fns_jacobian[i_jac[self._uu_G1]], ext.fns_jacobian[i_jac[self._uu_G2]], ext.fns_jacobian[i_jac[self._uu_G3]]) PetscDSSetJacobianPreconditioner(ds.ds, 0, 1, ext.fns_jacobian[i_jac[self._up_G0]], ext.fns_jacobian[i_jac[self._up_G1]], ext.fns_jacobian[i_jac[self._up_G2]], ext.fns_jacobian[i_jac[self._up_G3]]) - PetscDSSetJacobianPreconditioner(ds.ds, 1, 0, ext.fns_jacobian[i_jac[self._pu_G0]], ext.fns_jacobian[i_jac[self._pu_G1]], NULL, NULL) + if getattr(self, "_adjoint_kernels", False): + PetscDSSetJacobianPreconditioner(ds.ds, 1, 0, ext.fns_jacobian[i_jac[self._pu_G0]], ext.fns_jacobian[i_jac[self._pu_G1]], ext.fns_jacobian[i_jac[self._pu_G2]], ext.fns_jacobian[i_jac[self._pu_G3]]) + else: + PetscDSSetJacobianPreconditioner(ds.ds, 1, 0, ext.fns_jacobian[i_jac[self._pu_G0]], ext.fns_jacobian[i_jac[self._pu_G1]], NULL, NULL) PetscDSSetJacobianPreconditioner(ds.ds, 1, 1, ext.fns_jacobian[i_jac[self._pp_G0]], NULL, NULL, NULL) # Lagrange-multiplier rows (block-constrained Stokes). Guarded: no-op @@ -9632,6 +10372,197 @@ 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 + + by_kernels = self._adjoint_by_kernels() # before any DM vector + tangent = (self._install_adjoint_kernels() if by_kernels + else self._consistent_tangent_for_adjoint()) + gvec = self.dm.getGlobalVec() + gvec.setArray(0.0) + self._gather_fields_to_global(gvec) + 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) + + if by_kernels: + # J IS K^T: assembled from the transposed kernels. Solve forwards. + ksp = self.snes.getKSP() + ksp.setOperators(J, P) + ksp.solve(b, x) + reason = int(ksp.getConvergedReason()) + self._uninstall_adjoint_kernels(tangent) + else: + reason = _solve_transposed(self.snes.getKSP(), J, P, b, x) + self._restore_tangent(tangent) + + if target is not None: + u_adj, p_adj = target + 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 +10959,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..1dd66a90 100644 --- a/src/underworld3/model.py +++ b/src/underworld3/model.py @@ -31,6 +31,7 @@ # Import the Pint-native implementation import os +import re import sys sys.path.append(os.path.dirname(__file__)) @@ -243,6 +244,8 @@ def _operator_text(event): return name if kind == "history_shift": return f"shift {name}" + if kind == "adjoint_solve": + return f"adjoint {name}" return f"{kind}:{name}" @@ -259,15 +262,23 @@ def _quantity_parts(value): return None, None +_UNIT_SYMBOLS = { + "second": "s", "minute": "min", "hour": "hr", "day": "d", + "year": "yr", "kiloyear": "kyr", "megayear": "Myr", "gigayear": "Gyr", + "meter": "m", "kilometer": "km", "centimeter": "cm", "millimeter": "mm", + "kelvin": "K", "kilogram": "kg", "pascal": "Pa", "newton": "N", + "joule": "J", "watt": "W", +} + + def _abbreviate_unit(unit): - """A short unit name for a column header. Falls back to the full name.""" + """A short unit name for a column header or a scale: each unit name in a + pint unit string by its symbol, "pascal * second" as "Pa s" and + "millimeter / year" as "mm/yr". Names without a symbol stay as they are.""" if unit is None: return "" - return { - "second": "s", "minute": "min", "hour": "hr", "day": "d", - "year": "yr", "kiloyear": "kyr", "megayear": "Myr", "gigayear": "Gyr", - "meter": "m", "kilometer": "km", "kelvin": "K", "kilogram": "kg", - }.get(str(unit), str(unit)) + text = re.sub(r"[A-Za-z_]+", lambda m: _UNIT_SYMBOLS.get(m.group(0), m.group(0)), str(unit)) + return text.replace(" * ", " ").replace(" / ", "/") def _in_units_of(value, unit): @@ -449,6 +460,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 @@ -1253,12 +1265,31 @@ def _run_header(self): """The record that opens a run in the log, so the file is self-describing.""" from datetime import datetime, timezone - scales = {} + # The record keeps the fundamental scales, and beside them the + # reference quantities as they were declared, in the units they were + # quoted in. The readable header reports the declaration when there is + # one: "domain_depth 10 km", not the fundamental length in metres. + scales, reference = {}, {} try: for name, scale in (self.get_fundamental_scales() or {}).items(): scales[str(name)] = _jsonable_quantity(scale) except Exception: scales = {} + try: + from .scaling import units as ureg + + for name, quantity in (self.get_reference_quantities() or {}).items(): + if not (isinstance(quantity, dict) and "magnitude" in quantity): + continue + magnitude, unit = float(quantity["magnitude"]), str(quantity.get("units")) + # A quantity declared as an expression of others carries the + # raw composite of their units; in SI base units it reads. + if "**" in unit or unit.count("/") > 1: + base = ureg.Quantity(magnitude, unit).to_base_units() + magnitude, unit = float(base.magnitude), str(base.units) + reference[str(name)] = {"magnitude": magnitude, "units": unit} + except Exception: + reference = {} script = None try: entry = sys.argv[0] if sys.argv else "" @@ -1272,6 +1303,7 @@ def _run_header(self): "script": script, "started": datetime.now().astimezone().isoformat(timespec="seconds"), "scales": scales, + "reference": reference, } # ------------------------------------------------------------------ @@ -1284,7 +1316,7 @@ def _render_transcript_text(self, payload): kind = payload.get("kind") if kind == "run": - scales = payload.get("scales") or {} + scales = payload.get("reference") or payload.get("scales") or {} summary = " | ".join( f"{name} {value['magnitude']:.4g} {_abbreviate_unit(value['units'])}" for name, value in scales.items() @@ -1304,6 +1336,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": @@ -1348,7 +1381,7 @@ def _render_transcript_text(self, payload): # older transcript may carry is not an operator and is left out. operators = " > ".join( _operator_text(e) for e in events - if e.get("kind") in ("solve", "history_shift") + if e.get("kind") in ("solve", "history_shift", "adjoint_solve") ) or "(nothing)" notes = [] @@ -1382,6 +1415,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 +1735,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: @@ -2174,12 +2249,12 @@ def set_reference_quantities(self, verbose=False, nondimensional_scaling=True, * # Enable/disable non-dimensionalization based on parameter import underworld3 as uw - if nondimensional_scaling: - uw.use_nondimensional_scaling(True) - uw.pprint("✓ Units system active with automatic non-dimensionalization", proc=0) - else: - uw.use_nondimensional_scaling(False) - uw.pprint("⚠ Expert mode: Units active WITHOUT non-dimensionalization", proc=0) + uw.use_nondimensional_scaling(bool(nondimensional_scaling)) + if verbose: + if nondimensional_scaling: + uw.pprint("Units system active with automatic non-dimensionalization", proc=0) + else: + uw.pprint("Units active without non-dimensionalization", proc=0) uw.pprint(" (Warning: This mode may have numerical conditioning issues)", proc=0) def get_reference_quantities(self): @@ -2489,9 +2564,9 @@ def _solve_available_dimensions(self, matrix, names, fundamental_dims, ureg): import underworld3 as uw # Informational message about missing dimensions (not an error!) - if missing_dims: - uw.pprint(f"ℹ️ Dimensional coverage: {covered_dims}", proc=0) - uw.pprint(f" (Not covered: {missing_dims} - will fail only if needed)", proc=0) + if missing_dims and getattr(self, "_verbose_units", False): + uw.pprint(f"Dimensional coverage: {covered_dims}", proc=0) + uw.pprint(f" (not covered: {missing_dims}; needed only if a quantity uses them)", proc=0) # Extract sub-matrix for covered dimensions only sub_matrix = matrix[:, covered_indices] diff --git a/src/underworld3/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..5647f953 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 --------------------------------------------------------------- @@ -162,6 +163,17 @@ def _magnitude_and_unit(value_text, units): m = re.match(r"^\s*Matrix\(\[\[([-+0-9.eE]+)\]\]\)\s*$", text) if m: return _compact_number(m.group(1)), _plain_unit(units) + # A vector of numbers, as a boundary condition often is: (0, 0), or + # (0.5, free) where a component is left unconstrained. + m = re.match(r"^\s*Matrix\(\[(\[.*\])\]\)\s*$", text) + if m and "(" not in m.group(1): + entries = re.findall(r"\[([^\[\]]*)\]", m.group(1)) + parts = [] + for e in entries: + e = e.strip() + parts.append("free" if e in ("oo", "zoo", "nan") else _compact_number(e)) + if parts and all(re.match(r"^[-+0-9.eE]+$|^free$", q) for q in parts): + return "(" + ", ".join(parts) + ")", _plain_unit(units) if re.match(r"^\s*[-+0-9.eE]+\s*$", text): return _compact_number(text), _plain_unit(units) return text, _plain_unit(units) @@ -184,7 +196,13 @@ def _latex_value(value_latex, value_text, units): if re.match(r"^[-+0-9.eE]+$", number or ""): latex = _number_latex(number) return latex + (rf"\ \mathrm{{{_plain_unit_latex(unit)}}}" if unit else "") - return str(value_latex) if value_latex not in (None, "") else str(value_text) + if number and number.startswith("("): + return number + text = str(value_latex) if value_latex not in (None, "") else str(value_text) + # A quantity whose value is itself an expression of the fields — a + # yield-limited viscosity, say — is not a number to print; the full + # expression is in the record. + return text if len(text) <= 60 else r"\text{(an expression; in the record)}" def _number_latex(number): @@ -413,6 +431,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")] @@ -435,13 +523,19 @@ def _unit(value): def _short_unit(unit): - if unit is None: - return "" - return { - "second": "s", "minute": "min", "hour": "hr", "day": "d", "year": "yr", - "kiloyear": "kyr", "megayear": "Myr", "gigayear": "Gyr", - "meter": "m", "kilometer": "km", "kelvin": "K", "kilogram": "kg", - }.get(str(unit), str(unit)) + """The unit symbol the log uses, so the figure and the log agree.""" + from underworld3.model import _abbreviate_unit + return _abbreviate_unit(unit) + + +def _scales_line(header, sep=" "): + """The run's scales as one line: the reference quantities as they were + declared, in their own units, or the fundamental scales for a run that + declared none. Empty for a nondimensional run.""" + scales = header.get("reference") or header.get("scales") or {} + items = [f"{name} {value['magnitude']:.4g} {_short_unit(value['units'])}" + for name, value in scales.items() if isinstance(value, dict)] + return "scales: " + sep.join(items) if items else "" def _converted(value, unit): @@ -524,18 +618,18 @@ def _signature(step): return tuple( (event["kind"], _short_operator(event["name"])) for event in step.get("events", []) - if event.get("kind") in ("solve", "history_shift") + if event.get("kind") in _OPERATOR_KINDS ) def _describe(signature): return " ".join( - name if kind == "solve" else f"shift {name}" for kind, name in signature + _kind_text(kind, name) for kind, name in signature ) or "(nothing)" def _sequence_text(signature): - parts = [name if kind == "solve" else f"shift {name}" + parts = [_kind_text(kind, name) for kind, name in signature] return " > ".join(parts) or "(nothing)" @@ -782,12 +876,8 @@ def draw_header(y, first): bits.append(f"{len(notes)} backtrack(s)") canvas.text(_MARGIN, y + 8, " · ".join(bits), size=8.5, fill=_MUTED) y += 13 - scales = header.get("scales") or {} - if scales: - canvas.text(_MARGIN, y + 8, "scales: " + " ".join( - f"{name} {value['magnitude']:.4g} {_short_unit(value['units'])}" - for name, value in scales.items() if isinstance(value, dict) - ), size=8, fill=_MUTED) + if _scales_line(header): + canvas.text(_MARGIN, y + 8, _scales_line(header), size=8, fill=_MUTED) y += 12 y += 10 # column captions @@ -1461,6 +1551,33 @@ def transcript_flowchart(source, run=-1, out=None): # The transcript as a chart: parts across the page, steps down it # --------------------------------------------------------------------------- +_OPERATOR_KINDS = ("solve", "history_shift", "adjoint_solve") + + +def _kind_text(kind, name): + if kind == "history_shift": + return f"shift {name}" + if kind == "adjoint_solve": + return f"adjoint {name}" + return name + + +def _part_key(event): + """Which column an event belongs to. An adjoint solve is the same solver + as its forward solve and a different operator, so it gets its own.""" + key = event.get("part") or event.get("name") + return f"{key}/adjoint" if event.get("kind") == "adjoint_solve" else key + + +def _part_label(event): + name = _short_operator(event.get("name", event.get("part", "?"))) + if event.get("kind") == "history_shift": + return name + if event.get("kind") == "adjoint_solve": + return f"adjoint {name}" + return name + + def _parts_of(steps): """The roster, in a stable order, from what actually played. @@ -1473,14 +1590,14 @@ def _parts_of(steps): order, labels = [], {} for step in steps: for event in step.get("events", []): - if event.get("kind") not in ("solve", "history_shift"): + if event.get("kind") not in _OPERATOR_KINDS: continue - key = event.get("part") or event.get("name") + key = _part_key(event) if key not in labels: order.append(key) - labels[key] = _short_operator(event.get("name", key)) + labels[key] = _part_label(event) elif event.get("kind") == "history_shift": - labels[key] = _short_operator(event.get("name", key)) + labels[key] = _part_label(event) # In the order they first ran. A step is drawn as a bar whose events # descend in the order they ran, so with the columns in that same order # the usual step reads as a staircase down and to the right, and any step @@ -1502,7 +1619,7 @@ def _outcome(event): it is neither a clean convergence nor a failure, and reading it as either loses the thing worth seeing. """ - if event.get("kind") != "solve" or "converged" not in event: + if event.get("kind") not in ("solve", "adjoint_solve") or "converged" not in event: return None if not event.get("converged"): return "diverged" @@ -1521,9 +1638,8 @@ def _step_cells(step, parts): """ played = [] for event in step.get("events", []): - if event.get("kind") in ("solve", "history_shift"): - played.append((event.get("part") or event.get("name"), - _outcome(event))) + if event.get("kind") in _OPERATOR_KINDS: + played.append((_part_key(event), _outcome(event))) cells = [] for key, _ in parts: hits = tuple((i + 1, outcome) for i, (k, outcome) in enumerate(played) @@ -1676,7 +1792,11 @@ def _transcript_layout(header, steps, notes, entry, title=None, width=PAGE_W, bits.append("no terminator: still running, or interrupted") bits.append(f"{len(parts)} parts") canvas.text(_MARGIN, y + 8, " · ".join(bits), size=8.5, fill=_MUTED) - y += 22 + y += 13 + if _scales_line(header): + canvas.text(_MARGIN, y + 8, _scales_line(header), size=8, fill=_MUTED) + y += 12 + y += 9 # --- columns --- gutter = _MARGIN + 46.0 @@ -1985,6 +2105,8 @@ def transcript_table(source, run=-1, width=11, collapse=True): out.append(f"transcript · {_run_title(header, fallback='')}".rstrip(" ·")) if header.get("started"): out.append(f"started {header['started']}") + if _scales_line(header, sep=" | "): + out.append(_scales_line(header, sep=" | ")) if ended: out.append(f"complete — {ended.get('steps', len(steps))} step(s)") elif entry.get("live"): 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_0021_adjoint_misfit_through_gradient.py b/tests/test_0021_adjoint_misfit_through_gradient.py new file mode 100644 index 00000000..2e67e852 --- /dev/null +++ b/tests/test_0021_adjoint_misfit_through_gradient.py @@ -0,0 +1,89 @@ +"""A misfit on a strain rate: the dual carries a gradient part. + +The fault-segments example in miniature — a weak plane in a transversely +isotropic viscosity, two segments of unknown weak-plane viscosity, the misfit +on the surface uplift rate and on the shear stress near a point. The stress +reads the velocity through its gradient, so ``misfit_duals`` must assemble +the gradient part of the load; without it the adjoint gradient is wrong by +the whole stress term. +""" +import math + +import numpy as np +import pytest +import sympy + +import underworld3 as uw +from underworld3.adjoint import misfit_duals + + +@pytest.mark.level_2 +@pytest.mark.tier_a +def test_segment_strength_gradient_matches_finite_differences(): + mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0.0, 0.0), maxCoords=(2.0, 1.0), + cellSize=1 / 10, qdegree=3) + x, y = mesh.X + v = uw.discretisation.MeshVariable("v", mesh, 2, degree=2) + p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1, continuous=True) + v_obs = uw.discretisation.MeshVariable("v_obs", mesh, 2, degree=2) + + theta = math.radians(45.0) + n_hat = sympy.Matrix([[-math.sin(theta), math.cos(theta)]]) + d = (x - 0.6) * n_hat[0] + y * n_hat[1] + s = (x - 0.6) * math.cos(theta) + y * math.sin(theta) + band = sympy.exp(-(d / 0.15) ** 2) + strengths = [uw.expression(r"\eta_1", 0.1, "segment 1"), + uw.expression(r"\eta_2", 0.1, "segment 2")] + half = 0.5 / math.sin(theta) + seg = [(1 - sympy.tanh((s - half) / 0.15)) / 2, (1 + sympy.tanh((s - half) / 0.15)) / 2] + eta_1 = 1 - band * sum((1 - strengths[k]) * seg[k] for k in range(2)) + + stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + stokes.constitutive_model = uw.constitutive_models.TransverseIsotropicFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = 1 + stokes.constitutive_model.Parameters.shear_viscosity_1 = eta_1 + stokes.constitutive_model.Parameters.director = n_hat + stokes.tolerance = 1e-9 + stokes.add_essential_bc((0.0, 0.0), "Bottom") + stokes.add_essential_bc((0.5, None), "Left") + stokes.add_essential_bc((-0.5, None), "Right") + + def shear(field): + e = mesh.vector.strain_tensor(field.sym) + return 2 * e[0, 1] + + w_top = sympy.exp(-((1 - y) / 0.15) ** 2) + w_pt = sympy.exp(-((x - 1.3) ** 2 + (y - 0.4) ** 2) / 0.3 ** 2) + misfit = (w_top * (v.sym[1] - v_obs.sym[1]) ** 2 + + w_pt * (shear(v) - shear(v_obs)) ** 2) / 2 + + def set_strengths(a, b): + strengths[0].sym, strengths[1].sym = float(a), float(b) + + def J(): + stokes.solve(zero_init_guess=True) + return float(uw.maths.Integral(mesh, misfit).evaluate()) + + set_strengths(0.02, 0.2) + stokes.solve(zero_init_guess=True) + v_obs.array[...] = np.asarray(v.array) + + set_strengths(0.1, 0.1) + J() + dual = misfit_duals(misfit, [v])[v] + dual.array[...] = -np.asarray(dual.array) + mu = uw.discretisation.MeshVariable("mu", mesh, 2, degree=2) + lam = uw.discretisation.MeshVariable("lam", mesh, 1, degree=1) + _, reason = stokes.adjoint_solve((dual, None), target=(mu, lam)) + assert reason > 0 + adjoint = [stokes.sensitivity(mu, e) for e in strengths] + + h = 1e-4 + for k in range(2): + plus = [0.1, 0.1]; plus[k] += h + minus = [0.1, 0.1]; minus[k] -= h + set_strengths(*plus); Jp = J() + set_strengths(*minus); Jm = J() + fd = (Jp - Jm) / (2 * h) + assert adjoint[k] != 0 + assert abs(fd / adjoint[k] - 1) < 2e-3, (k, fd, adjoint[k]) diff --git a/tests/test_0022_adjoint_operator_from_kernels.py b/tests/test_0022_adjoint_operator_from_kernels.py new file mode 100644 index 00000000..bd4c530d --- /dev/null +++ b/tests/test_0022_adjoint_operator_from_kernels.py @@ -0,0 +1,89 @@ +"""The adjoint operator assembled from the transposed kernels IS K^T. + +``adjoint_solve`` no longer transposes the assembled Jacobian: it rewires the +solver to the pointwise kernels of the transposed bilinear form (g0 and g3 +transposed on their paired indices, g1 and g2 exchanged) and lets the SNES +assemble. This checks that matrix against the explicit transpose to machine +precision, on a non-symmetric operator and on a nonlinear saddle point with +pressure in the viscosity, so every block and every index order is exercised. +""" +import numpy as np +import pytest +import sympy + +import underworld3 as uw +from petsc4py import PETSc + + +def _adjoint_versus_transpose(solver): + def assemble(): + dm = solver.dm + g = dm.getGlobalVec() + if hasattr(solver, "_gather_fields_to_global"): + solver._gather_fields_to_global(g) + else: + dm.localToGlobal(solver.u.vec, g) + solver.mesh.update_lvec() + dm.setAuxiliaryVec(solver.mesh.lvec, None) + J, P = solver.snes.getJacobian()[:2] + solver.snes.computeJacobian(g, J, P) + out = J.copy() + dm.restoreGlobalVec(g) + return out + + K = assemble() + token = solver._install_adjoint_kernels() + K_adj = assemble() + solver._uninstall_adjoint_kernels(token) + Kt = PETSc.Mat() + K.transpose(Kt) # a new matrix; the no-argument form transposes in place + D = K_adj.copy() + D.axpy(-1.0, Kt) + S = K.copy() + S.axpy(-1.0, Kt) + return D.norm() / K.norm(), S.norm() / K.norm() + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_supg_step_adjoint_kernels_give_the_transpose(): + mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0, 0), maxCoords=(1, 1), + cellSize=1 / 8, qdegree=3) + T = uw.discretisation.MeshVariable("T", mesh, 1, degree=2) + adv = uw.systems.AdvDiffusion(mesh, u_Field=T, V_fn=sympy.Matrix([[1.0, 0.3]]), theta=1.0) + adv.constitutive_model = uw.constitutive_models.DiffusionModel + adv.constitutive_model.Parameters.diffusivity = 0.05 + adv.add_essential_bc(1.0, "Bottom") + adv.add_essential_bc(0.0, "Top") + adv.solve(timestep=0.1, zero_init_guess=True) + error, asymmetry = _adjoint_versus_transpose(adv) + assert asymmetry > 1e-3, "the operator must be non-symmetric for this to test anything" + assert error < 1e-12 + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_nonlinear_stokes_adjoint_kernels_give_the_transpose(): + mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0, 0), maxCoords=(1, 1), + cellSize=1 / 6, qdegree=3) + x, y = mesh.X + v = uw.discretisation.MeshVariable("v", mesh, 2, degree=2) + p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1, continuous=True) + stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + # viscosity depending on the strain rate AND the pressure: every block of + # the Jacobian is non-trivial, and (u,p) and (p,u) are not transposes of + # each other + edot = stokes.constitutive_model.Parameters.strainrate_inv_II if hasattr( + stokes.constitutive_model.Parameters, "strainrate_inv_II") else None + E = mesh.vector.strain_tensor(v.sym) + e2 = (E[0, 0] ** 2 + E[1, 1] ** 2 + 2 * E[0, 1] ** 2) / 2 + uw.maths.functions.vanishing + stokes.constitutive_model.Parameters.shear_viscosity_0 = (1 + sympy.Max(p.sym[0], 0)) / (1 + sympy.sqrt(e2)) + stokes.bodyforce = sympy.Matrix([0, -sympy.sin(3 * x)]) + stokes.add_essential_bc((0.0, 0.0), "Bottom") + stokes.add_essential_bc((0.5, None), "Left") + stokes.add_essential_bc((-0.5, None), "Right") + stokes.solve(zero_init_guess=True) + error, asymmetry = _adjoint_versus_transpose(stokes) + assert asymmetry > 1e-3 + assert error < 1e-12 diff --git a/tests/test_0023_solver_gradient_and_adjoint_introspection.py b/tests/test_0023_solver_gradient_and_adjoint_introspection.py new file mode 100644 index 00000000..8e62f048 --- /dev/null +++ b/tests/test_0023_solver_gradient_and_adjoint_introspection.py @@ -0,0 +1,89 @@ +"""One call for the steady adjoint, and the adjoint problem it assembles, in writing. + +``solver.gradient(misfit, parameters=..., fields=...)`` is the three-call +route — dual of the misfit, transposed solve, sensitivities — as one method, +so a steady inversion reads like the time-dependent one. Checked against a +central difference on a non-symmetric SUPG step, in the parameter and in the +initial field. ``adjoint_templates()`` writes the adjoint residual in the +same template language as the forward one; the check here is that it is +the forward Jacobian kernels with trial and test exchanged, by evaluating +both at a point. +""" +import numpy as np +import pytest +import sympy + +import underworld3 as uw + + +def _supg(cell=1 / 8): + mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0, 0), maxCoords=(1, 1), + cellSize=cell, qdegree=3) + x, y = mesh.X + T = uw.discretisation.MeshVariable("T", mesh, 1, degree=2) + T_obs = uw.discretisation.MeshVariable("T_obs", mesh, 1, degree=2) + kappa = uw.expression(r"\kappa", 0.05, "diffusivity") + adv = uw.systems.AdvDiffusion(mesh, u_Field=T, V_fn=sympy.Matrix([[1.0, 0.3]]), theta=1.0) + adv.constitutive_model = uw.constitutive_models.DiffusionModel + adv.constitutive_model.Parameters.diffusivity = kappa + adv.add_essential_bc(1.0, "Bottom") + adv.add_essential_bc(0.0, "Top") + return mesh, T, T_obs, kappa, adv + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_gradient_method_matches_finite_differences_in_parameter_and_initial_field(): + mesh, T, T_obs, kappa, adv = _supg() + x, y = mesh.X + dt = 0.05 + + def initial(c=0.5): + X = np.asarray(T.coords) + return np.exp(-((X[:, 0] - c) ** 2 + (X[:, 1] - 0.5) ** 2) / 0.05) + + def run(k, T0): + kappa.sym = float(k) + T.array[:, 0, 0] = T0 + adv.DuDt.initialise_history() if hasattr(adv.DuDt, "initialise_history") else None + adv.solve(timestep=dt, zero_init_guess=True) + return float(uw.maths.Integral(mesh, misfit).evaluate()) + + misfit = (T.sym[0] - T_obs.sym[0]) ** 2 / 2 + run(0.05, initial(0.55)) + T_obs.array[...] = np.asarray(T.array) + T0 = initial(0.5) + J0 = run(0.05, T0) + adv.DuDt.psi_star[0].array[:, 0, 0] = T0 # the step's input, not its output + out = adv.gradient(misfit, parameters=[kappa], fields=[T]) + assert abs(out["J"] - J0) < 1e-12 + + h = 1e-4 + fd = (run(0.05 + h, T0) - run(0.05 - h, T0)) / (2 * h) + assert abs(fd / out["parameters"][kappa] - 1) < 1e-3, (fd, out["parameters"][kappa]) + + # the initial field enters through the history slot the solve reads + direction = initial(0.45) - initial(0.5) + dual = out["fields"][T] + adjoint = uw.adjoint.inner(T, dual, direction) + fd = (run(0.05, T0 + h * direction) - run(0.05, T0 - h * direction)) / (2 * h) + assert abs(fd / adjoint - 1) < 1e-3, (fd, adjoint) + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_adjoint_templates_are_the_forward_kernels_with_trial_and_test_exchanged(): + mesh, T, T_obs, kappa, adv = _supg(cell=1 / 4) + adv.solve(timestep=0.05, zero_init_guess=True) + f0, f1 = adv.adjoint_templates() + kernels = adv.adjoint_kernels() + # value part: the coefficient of mu in f0 is H0, of mu_{,d} is H1 + mu = sympy.Symbol(r"\mu") + assert sympy.simplify(sympy.diff(f0[0], mu) - kernels["F0_u"][0, 0]) == 0 + for d in range(mesh.cdim): + md = sympy.Symbol(f"\\mu_{{,{d}}}") + assert sympy.simplify(sympy.diff(f0[0], md) - kernels["F0_grad_u"][0, d]) == 0 + assert sympy.simplify(sympy.diff(f1[0, d], mu) - kernels["F1_u"][0, d]) == 0 + # and the forward kernels' transpose: the SUPG term makes g3 non-symmetric only + # through the velocity; g1 <-> g2 is where a wrong swap would show + assert kernels["F0_grad_u"] == sympy.ImmutableMatrix(adv._G2).reshape(1, mesh.cdim) diff --git a/tests/test_0024_adjoint_boundary_misfit_and_boundary_sensitivity.py b/tests/test_0024_adjoint_boundary_misfit_and_boundary_sensitivity.py new file mode 100644 index 00000000..62f76194 --- /dev/null +++ b/tests/test_0024_adjoint_boundary_misfit_and_boundary_sensitivity.py @@ -0,0 +1,85 @@ +"""Surface observations, and a parameter that lives on a boundary. + +Two things an adjoint paper derives by hand and this machinery must get +without derivation. A misfit integrated over a boundary — the uplift rate +along a free surface — whose dual is a facet load. And a parameter that +enters the residual through a natural condition — the traction on a +boundary — whose sensitivity has a facet part. Both against central +differences, on a Stokes problem small enough to run in seconds. +""" +import numpy as np +import pytest +import sympy + +import underworld3 as uw + + +def _box(): + mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0, 0), maxCoords=(1, 1), + cellSize=1 / 6, qdegree=3) + v = uw.discretisation.MeshVariable("v", mesh, 2, degree=2) + p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1, continuous=True) + v_obs = uw.discretisation.MeshVariable("v_obs", mesh, 2, degree=2) + return mesh, v, p, v_obs + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_misfit_on_the_free_surface_matches_finite_differences(): + mesh, v, p, v_obs = _box() + x, y = mesh.X + eta = uw.expression(r"\eta", 1.0, "viscosity") + stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = eta * (1 + 2 * (x - 0.5) ** 2) + stokes.bodyforce = sympy.Matrix([0, -sympy.sin(3 * x)]) + stokes.add_essential_bc((0.0, 0.0), "Bottom") + stokes.add_essential_bc((0.3, None), "Left") + stokes.add_essential_bc((-0.3, None), "Right") # free top + stokes.tolerance = 1e-10 # the gradient is small; the difference must be clean + + misfit = (v.sym[1] - v_obs.sym[1]) ** 2 / 2 # uplift rate along the top + + def J(value): + eta.sym = float(value) + stokes.solve(zero_init_guess=True) + return float(uw.maths.BdIntegral(mesh, misfit, "Top").evaluate()) + + J(2.5); v_obs.array[...] = np.asarray(v.array) + J0 = J(1.0) + out = stokes.gradient(misfit, parameters=[eta], boundary="Top") + assert abs(out["J"] - J0) < 1e-12 + h = 1e-4 + fd = (J(1.0 + h) - J(1.0 - h)) / (2 * h) + assert abs(fd / out["parameters"][eta] - 1) < 1e-3, (fd, out["parameters"][eta]) + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_parameter_in_a_traction_condition_matches_finite_differences(): + mesh, v, p, v_obs = _box() + x, y = mesh.X + tau = uw.expression(r"\tau", 0.5, "traction on the top") + stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p) + stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel + stokes.constitutive_model.Parameters.shear_viscosity_0 = 1 + 2 * (x - 0.5) ** 2 + stokes.add_essential_bc((0.0, 0.0), "Bottom") + stokes.add_essential_bc((0.0, None), "Left") + stokes.add_essential_bc((0.0, None), "Right") + stokes.add_natural_bc(sympy.Matrix([[tau * sympy.sin(sympy.pi * x), 0.0]]), "Top") + stokes.tolerance = 1e-10 + + misfit = (v.sym[0] ** 2 + v.sym[1] ** 2) / 2 # a volume misfit + + def J(value): + tau.sym = float(value) + stokes.solve(zero_init_guess=True) + return float(uw.maths.Integral(mesh, misfit).evaluate()) + + J0 = J(0.5) + out = stokes.gradient(misfit, parameters=[tau]) + assert abs(out["J"] - J0) < 1e-12 + h = 1e-4 + fd = (J(0.5 + h) - J(0.5 - h)) / (2 * h) + assert out["parameters"][tau] != 0.0 + assert abs(fd / out["parameters"][tau] - 1) < 1e-3, (fd, out["parameters"][tau]) diff --git a/tests/test_0025_tao_driver.py b/tests/test_0025_tao_driver.py new file mode 100644 index 00000000..beea9c28 --- /dev/null +++ b/tests/test_0025_tao_driver.py @@ -0,0 +1,29 @@ +"""uw.adjoint.minimise drives an objective with PETSc TAO: a quadratic, bounded and not.""" +import numpy as np +import pytest + +import underworld3 as uw + + +def _quadratic(x): + c = np.array([1.0, -0.5, 2.0]) + return float(np.sum((x - c) ** 2 * [1, 2, 3])), 2 * (x - c) * [1, 2, 3] + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_lmvm_finds_the_minimum(): + x, info = uw.adjoint.minimise(_quadratic, np.zeros(3), gradient_tolerance=1e-12) + assert np.allclose(x, [1.0, -0.5, 2.0], atol=1e-6) + assert info["reason"] > 0 + assert len(info["history"]) == len(info["history"]) and info["history"][0][0] > info["history"][-1][0] + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_blmvm_honours_bounds(): + x, info = uw.adjoint.minimise(_quadratic, np.zeros(3), method="blmvm", + bounds=(np.array([0.0, 0.0, 0.0]), np.array([0.5, 1.0, 1.0])), + gradient_tolerance=1e-12) + assert np.allclose(x, [0.5, 0.0, 1.0], atol=1e-6) + assert info["reason"] > 0 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())