minires.tlm
The adjoint of minires.ResSim.time_stepper, derived by hand.
I.e. the transpose of the Jacobian of one time step,
$ (s^n, p^n) ↦ (s^{n+1}, p^{n+1}) $, with respect to the state -- and to the
parameters: $ \log K $, and the BHP controls of the step -- applied to a
sensitivity (adj_step), rather than formed (it is dense, through $A^{-1}$).
Chained backwards along a trajectory (adjoint), it yields the gradient of an
objective with respect to the initial state, S0 and P0 of
minires.ResSim.sim, to $ \log K $, and to the BHP schedule,
minires.wells.Wells.bhp, at the cost of about one more simulation --
whatever the number of parameters.
linearize recomputes one forward step -- from the state it is given, so the
trajectory that sim returns is all the record it needs (checkpointing) -- and
returns a Tape of the intermediates, about which adj_step transposes. The
following checks the gradient of a final-time quantity -- the water saturation
at the producer, $ J = S_N[i_\mathrm{prd}] $ -- against a finite difference:
>>> from minires import ResSim
>>> model = ResSim(Lx=1, Ly=1, Nx=8, Ny=8, ct=.1, cached_precond=False, wells=[
... dict(xy=[0, 0], rate=+1),
... dict(xy=[1, 1], rate=-.5),
... ])
>>> rng = np.random.default_rng(3)
>>> S0 = .2 + .6 * rng.random(model.Nxy)
>>> P0 = rng.random(model.Nxy)
>>> dt, nSteps = .0337, 3
>>> SS, PP = model.sim(dt, nSteps, S0, P0, pbar=False)
>>> prd = model.xy2ind(1, 1)
>>> dJ_dSS = np.zeros_like(SS)
>>> dJ_dSS[-1, prd] = 1
>>> grad = adjoint(model, dt, SS, PP, dJ_dSS)
Its directional derivative along a random $ (δS_0, δP_0, δ\log K) $ ...
>>> dS0, dP0 = rng.standard_normal((2, model.Nxy))
>>> dlogK = rng.standard_normal(model.K.shape)
>>> directional = grad.S0 @ dS0 + grad.P0 @ dP0 + (grad.logK * dlogK).sum()
... is that of the simulation:
>>> eps, K = 1e-6, model.K.copy()
>>> model.K = K * np.exp(+eps*dlogK)
>>> SSp, _ = model.sim(dt, nSteps, S0 + eps*dS0, P0 + eps*dP0, pbar=False)
>>> model.K = K * np.exp(-eps*dlogK)
>>> SSm, _ = model.sim(dt, nSteps, S0 - eps*dS0, P0 - eps*dP0, pbar=False)
>>> model.K = K
>>> fd = (SSp[-1, prd] - SSm[-1, prd]) / (2*eps)
>>> bool(abs(fd - directional) < 1e-6 * abs(directional))
True
(cached_precond=False merely spares the finite difference the noise of the
iterative solver's tolerance, $10^{-10}$, which eps would amplify a
million-fold; the adjoint itself does not care.)
The same check, on several configurations (incompressible, compressible,
BHP-controlled, 1D, inactive cells), is tests/test_tlm.py. See
examples.water_cut_gradient for the gradient of one producer's water cut with
respect to the $ \log K $ field and to the producers' BHP schedule, and
examples.history_match_gradient for that of a production-history misfit, put
to use in a few descent steps.
Seeding with an objective
adjoint sweeps backwards along the trajectory, from the partial derivatives
of a scalar objective, $ J(S, P) $, with respect to the stored states --
dJ_dSS[k] $ = ∂J/∂S_k $ and dJ_dPP[k] $ = ∂J/∂P_k $, shaped like SS and
PP -- and returns $ ∂J/∂S_0 $, $ ∂J/∂P_0 $, $ ∂J/∂\log K $ and
$ ∂J/∂p_\mathrm{bh} $ (a Gradient).
The seeds are simply what the objective says they are:
- A quantity of the final state seeds index
-1alone, as above. - A data misfit, $ J = \frac{1}{2} \sum_k \| H_k S_k - y_k \|^2 / σ^2 $ (the
history-matching case), seeds every observed time with its weighted
residual, scattered back onto the observed cells:
dJ_dSS[k] = H_k.T @ (H_k @ SS[k] - y_k) / σ**2. - Objectives on the well reports (
actual_rates,actual_bhp) are not seedable directly: those are functions of $ (S_k, P_{k+1}) $ through the well model (minires.ResSim.bhp,minires.ResSim.realize_bhp), which must be differentiated by hand into the seeds. Not built in.
The gradient with respect to $ \log K $ has the shape of K, (2, Nx, Ny):
both permeability components. For an isotropic field (a scalar or
(Nx, Ny)-shaped K, which ResSim.__setattr__ broadcasts to both), the
gradient with respect to the single $ \log k $ field is grad.logK.sum(0).
Likewise, the gradient with respect to the BHP controls is per step,
(nComp, nSteps), since the control may be scheduled (Wells.bhp is
(nComp, nTime)); for a constant-in-time control (a (nComp, 1) spec), it is
grad.bhp.sum(1). Rate-controlled completions have no BHP to differentiate
with respect to, and get 0.
How it is derived
The step is first expressed on the interior faces (the boundary ones carry
no flux), numbered x-faces first, in C-order, as TX[1:-1, :] and
TY[:, 1:-1] of TPFA flatten. With Grad the $ (n_F × N) $ difference
operator, $ (∇p)_f = p_\mathrm{lo} - p_\mathrm{hi} $, whose transpose is the
divergence (cell $i$ gets its high faces minus its low ones), the whole step
reads
$$ A = ∇^T \, \mathrm{diag}(T) \, ∇ + \mathrm{diag}(a + w) \,, \qquad p^{n+1} = A^{-1} (Q + a \, p^n + r) \,, \qquad v = T ⊙ ∇ p^{n+1} \,, $$
with $T$ the transmissibilities (a function of the cell mobilities
$λ_t(s^n)$ and of $K$, ref Tape.dT_dMt, Tape.dT_dlogK), $a$ the
accumulation term (Tape.accum) and $w$, $r$ the well model of the
BHP-controlled wells (Tape.bhp_diag, and the matching right-hand side, both
proportional to $λ_t$ in the well cells); then, for each of the nT sub-steps,
$$ s ← s + \frac{Δt}{n_T \, |Ω|} \Big( -∇^T \big( v ⊙ \mathrm{Up}(v) \, f(s) \big) + Q^- f(s) + Q^+ - s \, \mathrm{st} \Big) \,, $$
with Up the $ (n_F × N) $ upwind selector (1 in the column of the face's
upwind cell) and st the storage rate (0 if incompressible). This recasts the
5-diagonal assemblies of minires.ResSim.TPFA and
minires.ResSim.upwind_diff (which tests/test_tlm.py checks) on
operators whose transposes are just .T -- so that the tangent of the step
is a straight line of statements of three kinds, each linear in the
perturbations: y = M @ x (a sparse M on the tape), y = a * x (a vector
a), and y = tape.solve(x), the pressure solve. adj_step is that tangent
in reverse order, each statement transposed -- x̄ += M.T @ ȳ, x̄ += a * ȳ
-- with the solve its own transpose, $ A $ being symmetric (that is what makes
the TPFA system SPD, ref minires.ResSim.cached_precond), so Tape.solve
serves as it is. Each line's comment names the tangent statement it
transposes. (The tangent model itself, and the dot-product test binding the
two to round-off, are in the git history, should a statement need
re-deriving.)
What is (and is not) differentiated
- The state, $ (s^n, p^n) $, and everything downstream of it: the
mobilities, the transmissibilities, the pressure and the fluxes, the well
model of a BHP-controlled well (its $ WI λ_t $ and the rate it realizes),
the storage rate and the transport. If
ct == 0and no well is on BHP control, $ p^{n+1} $ does not depend on $ p^n $ at all, so the gradient with respect toP0is then0. - The BHP controls, $ p_\mathrm{bh} $ of the BHP-controlled completions
at each step (whatever
well_controlsreturned, so, by default, the specWells.bhpat that time), which enter the step at one point only: the right-hand side of the pressure system, as $ WI λ_t \, p_\mathrm{bh} $. The natural optimization (rather than history-matching) parameter, and the more so as the physics makes the rates derived quantities under BHP control. Note the control at step $k$ cannot affect the states before $ k+1 $, so the gradient of an objective of the state at time $k$ vanishes for the controls from $k$ on. - The permeability, as $ \log K $ (positivity built in, and the natural
history-matching parameter), through the transmissibilities alone. Its two
other appearances are not differentiated, and rightly so: the pin of the
incompressible pressure system (
TPFAadds $ \sum K_{00} $ to the first diagonal entry) fixes $ p_0 = 0 $ whatever the value, so the derivative is exactly zero; and the well indexWells.WIis a stored parameter (minires.wells.peaceman_WIevaluates it once, from theKof that moment, and does not trackKthereafter), so it is held fixed, like the other well parameters. - Not the other parameters:
por,ct, the viscosities, the well positions, indices and rate specifications. These stay fixed, as they are insim. (The rates would be as easy as the BHPs: they enterQdirectly, so their adjoint isadQgathered at the rate-controlled completions. But withct == 0they must sum to zero, which the gradient does not know.) - Not the controls' dependence on the state:
well_controlsis assumed open-loop (the default). An override that feeds the state back is not seen -- the controls enter as constants. - The discrete decisions are frozen at the linearization point: the
sub-step count
nT(a ceiling, refminires.ResSim.estimate_1CFL), the upwind directions and the signs of the well fluxes (theclips ofupwind_diff). They are piecewise constant, so this is the derivative almost everywhere; at a switch (a face with exactly zero flux,nTon an integer) the forward map has a kink, and a finite difference across it will disagree. - Only the explicit transport scheme (
saturation_step_upwind, the default). The implicit one has no adjoint here (ref the "How to solve" section of the docs for why it is not used).
linearize mutates the model, as a forward step does.
Ref its docstring. In short: it leaves _Q, _wells_now and the
preconditioner cache _pLU as of the step it linearized -- after a
reverse sweep, those of the first step -- but not the reports.
1"""The adjoint of `minires.ResSim.time_stepper`, derived by hand. 2 3I.e. the transpose of the Jacobian of one time step, 4$ (s^n, p^n) ↦ (s^{n+1}, p^{n+1}) $, with respect to the *state* -- and to the 5parameters: $ \\log K $, and the BHP controls of the step -- applied to a 6sensitivity (`adj_step`), rather than formed (it is dense, through $A^{-1}$). 7Chained backwards along a trajectory (`adjoint`), it yields the gradient of an 8objective with respect to the initial state, `S0` and `P0` of 9`minires.ResSim.sim`, to $ \\log K $, and to the BHP *schedule*, 10`minires.wells.Wells.bhp`, at the cost of about one more simulation -- 11whatever the number of parameters. 12 13`linearize` recomputes one forward step -- from the state it is given, so the 14trajectory that `sim` returns is all the record it needs (checkpointing) -- and 15returns a `Tape` of the intermediates, about which `adj_step` transposes. The 16following checks the gradient of a final-time quantity -- the water saturation 17at the producer, $ J = S_N[i_\\mathrm{prd}] $ -- against a finite difference: 18 19>>> from minires import ResSim 20>>> model = ResSim(Lx=1, Ly=1, Nx=8, Ny=8, ct=.1, cached_precond=False, wells=[ 21... dict(xy=[0, 0], rate=+1), 22... dict(xy=[1, 1], rate=-.5), 23... ]) 24>>> rng = np.random.default_rng(3) 25>>> S0 = .2 + .6 * rng.random(model.Nxy) 26>>> P0 = rng.random(model.Nxy) 27>>> dt, nSteps = .0337, 3 28>>> SS, PP = model.sim(dt, nSteps, S0, P0, pbar=False) 29>>> prd = model.xy2ind(1, 1) 30>>> dJ_dSS = np.zeros_like(SS) 31>>> dJ_dSS[-1, prd] = 1 32>>> grad = adjoint(model, dt, SS, PP, dJ_dSS) 33 34Its directional derivative along a random $ (δS_0, δP_0, δ\\log K) $ ... 35 36>>> dS0, dP0 = rng.standard_normal((2, model.Nxy)) 37>>> dlogK = rng.standard_normal(model.K.shape) 38>>> directional = grad.S0 @ dS0 + grad.P0 @ dP0 + (grad.logK * dlogK).sum() 39 40... is that of the simulation: 41 42>>> eps, K = 1e-6, model.K.copy() 43>>> model.K = K * np.exp(+eps*dlogK) 44>>> SSp, _ = model.sim(dt, nSteps, S0 + eps*dS0, P0 + eps*dP0, pbar=False) 45>>> model.K = K * np.exp(-eps*dlogK) 46>>> SSm, _ = model.sim(dt, nSteps, S0 - eps*dS0, P0 - eps*dP0, pbar=False) 47>>> model.K = K 48>>> fd = (SSp[-1, prd] - SSm[-1, prd]) / (2*eps) 49>>> bool(abs(fd - directional) < 1e-6 * abs(directional)) 50True 51 52(`cached_precond=False` merely spares the finite difference the noise of the 53iterative solver's tolerance, $10^{-10}$, which `eps` would amplify a 54million-fold; the adjoint itself does not care.) 55 56The same check, on several configurations (incompressible, compressible, 57BHP-controlled, 1D, inactive cells), is `tests/test_tlm.py`. See 58`examples.water_cut_gradient` for the gradient of one producer's water cut with 59respect to the $ \\log K $ field and to the producers' BHP schedule, and 60`examples.history_match_gradient` for that of a production-history misfit, put 61to use in a few descent steps. 62 63## Seeding with an objective 64 65`adjoint` sweeps backwards along the trajectory, from the partial derivatives 66of a scalar objective, $ J(S, P) $, with respect to the *stored* states -- 67`dJ_dSS[k]` $ = ∂J/∂S_k $ and `dJ_dPP[k]` $ = ∂J/∂P_k $, shaped like `SS` and 68`PP` -- and returns $ ∂J/∂S_0 $, $ ∂J/∂P_0 $, $ ∂J/∂\\log K $ and 69$ ∂J/∂p_\\mathrm{bh} $ (a `Gradient`). 70The seeds are simply what the objective says they are: 71 72- A quantity of the *final* state seeds index `-1` alone, as above. 73- A data misfit, $ J = \\frac{1}{2} \\sum_k \\| H_k S_k - y_k \\|^2 / σ^2 $ (the 74 history-matching case), seeds every observed time with its weighted 75 residual, scattered back onto the observed cells: 76 `dJ_dSS[k] = H_k.T @ (H_k @ SS[k] - y_k) / σ**2`. 77- Objectives on the well *reports* (`actual_rates`, `actual_bhp`) are not 78 seedable directly: those are functions of $ (S_k, P_{k+1}) $ through the 79 well model (`minires.ResSim.bhp`, `minires.ResSim.realize_bhp`), 80 which must be differentiated by hand into the seeds. Not built in. 81 82The gradient with respect to $ \\log K $ has the shape of `K`, `(2, Nx, Ny)`: 83both permeability components. For an *isotropic* field (a scalar or 84`(Nx, Ny)`-shaped `K`, which `ResSim.__setattr__` broadcasts to both), the 85gradient with respect to the single $ \\log k $ field is `grad.logK.sum(0)`. 86 87Likewise, the gradient with respect to the BHP controls is per *step*, 88`(nComp, nSteps)`, since the control may be scheduled (`Wells.bhp` is 89`(nComp, nTime)`); for a constant-in-time control (a `(nComp, 1)` spec), it is 90`grad.bhp.sum(1)`. Rate-controlled completions have no BHP to differentiate 91with respect to, and get `0`. 92 93## How it is derived 94 95The step is first expressed on the **interior faces** (the boundary ones carry 96no flux), numbered x-faces first, in C-order, as `TX[1:-1, :]` and 97`TY[:, 1:-1]` of `TPFA` flatten. With `Grad` the $ (n_F × N) $ difference 98operator, $ (∇p)_f = p_\\mathrm{lo} - p_\\mathrm{hi} $, whose transpose is the 99divergence (cell $i$ gets its high faces minus its low ones), the whole step 100reads 101 102$$ A = ∇^T \\, \\mathrm{diag}(T) \\, ∇ + \\mathrm{diag}(a + w) \\,, \\qquad 103 p^{n+1} = A^{-1} (Q + a \\, p^n + r) \\,, \\qquad 104 v = T ⊙ ∇ p^{n+1} \\,, $$ 105 106with $T$ the transmissibilities (a function of the cell mobilities 107$λ_t(s^n)$ and of $K$, ref `Tape.dT_dMt`, `Tape.dT_dlogK`), $a$ the 108accumulation term (`Tape.accum`) and $w$, $r$ the well model of the 109BHP-controlled wells (`Tape.bhp_diag`, and the matching right-hand side, both 110proportional to $λ_t$ in the well cells); then, for each of the `nT` sub-steps, 111 112$$ s ← s + \\frac{Δt}{n_T \\, |Ω|} 113 \\Big( -∇^T \\big( v ⊙ \\mathrm{Up}(v) \\, f(s) \\big) + 114 Q^- f(s) + Q^+ - s \\, \\mathrm{st} \\Big) \\,, $$ 115 116with `Up` the $ (n_F × N) $ upwind *selector* (1 in the column of the face's 117upwind cell) and `st` the storage rate (0 if incompressible). This recasts the 1185-diagonal assemblies of `minires.ResSim.TPFA` and 119`minires.ResSim.upwind_diff` (which `tests/test_tlm.py` checks) on 120operators whose transposes are just `.T` -- so that the *tangent* of the step 121is a straight line of statements of three kinds, each linear in the 122perturbations: `y = M @ x` (a sparse `M` on the tape), `y = a * x` (a vector 123`a`), and `y = tape.solve(x)`, the pressure solve. `adj_step` is that tangent 124in *reverse* order, each statement transposed -- `x̄ += M.T @ ȳ`, `x̄ += a * ȳ` 125-- with the solve its own transpose, $ A $ being symmetric (that is what makes 126the TPFA system SPD, ref `minires.ResSim.cached_precond`), so `Tape.solve` 127serves as it is. Each line's comment names the tangent statement it 128transposes. (The tangent model itself, and the dot-product test binding the 129two to round-off, are in the git history, should a statement need 130re-deriving.) 131 132## What is (and is not) differentiated 133 134- The **state**, $ (s^n, p^n) $, and everything downstream of it: the 135 mobilities, the transmissibilities, the pressure and the fluxes, the well 136 model of a BHP-controlled well (its $ WI λ_t $ and the rate it realizes), 137 the storage rate and the transport. If `ct == 0` and no well is on BHP 138 control, $ p^{n+1} $ does not depend on $ p^n $ at all, so the gradient 139 with respect to `P0` is then `0`. 140- The **BHP controls**, $ p_\\mathrm{bh} $ of the BHP-controlled completions 141 at each step (whatever `well_controls` returned, so, by default, the spec 142 `Wells.bhp` at that time), which enter the step at one point only: the 143 right-hand side of the pressure system, as $ WI λ_t \\, p_\\mathrm{bh} $. 144 The natural *optimization* (rather than history-matching) parameter, and 145 the more so as the physics makes the rates derived quantities under BHP 146 control. Note the control at step $k$ cannot affect the states before 147 $ k+1 $, so the gradient of an objective of the state at time $k$ vanishes 148 for the controls from $k$ on. 149- The **permeability**, as $ \\log K $ (positivity built in, and the natural 150 history-matching parameter), through the transmissibilities alone. Its two 151 other appearances are *not* differentiated, and rightly so: the pin of the 152 incompressible pressure system (`TPFA` adds $ \\sum K_{00} $ to the first 153 diagonal entry) fixes $ p_0 = 0 $ whatever the value, so the derivative is 154 exactly zero; and the well index `Wells.WI` is a *stored* parameter 155 (`minires.wells.peaceman_WI` evaluates it once, from the `K` of that 156 moment, and does not track `K` thereafter), so it is held fixed, like the 157 other well parameters. 158- **Not** the other parameters: `por`, `ct`, the viscosities, the well 159 positions, indices and *rate* specifications. These stay fixed, as they are 160 in `sim`. (The rates would be as easy as the BHPs: they enter `Q` directly, 161 so their adjoint is `adQ` gathered at the rate-controlled completions. But 162 with `ct == 0` they must sum to zero, which the gradient does not know.) 163- **Not** the controls' dependence on the state: `well_controls` is assumed 164 *open-loop* (the default). An override that feeds the state back is not 165 seen -- the controls enter as constants. 166- The **discrete decisions** are frozen at the linearization point: the 167 sub-step count `nT` (a ceiling, ref `minires.ResSim.estimate_1CFL`), 168 the upwind directions and the signs of the well fluxes (the `clip`s of 169 `upwind_diff`). They are piecewise constant, so this is the derivative 170 almost everywhere; at a switch (a face with exactly zero flux, `nT` on an 171 integer) the forward map has a kink, and a finite difference across it 172 will disagree. 173- Only the **explicit** transport scheme (`saturation_step_upwind`, the 174 default). The implicit one has no adjoint here (ref the "How to solve" 175 section of the docs for why it is not used). 176 177.. warning:: `linearize` mutates the model, as a forward step does. 178 179 Ref its docstring. In short: it leaves `_Q`, `_wells_now` and the 180 preconditioner cache `_pLU` as of the step it linearized -- after a 181 reverse sweep, those of the *first* step -- but not the reports. 182""" 183 184from dataclasses import dataclass 185from typing import Any, Callable, NamedTuple 186 187import numpy as np 188from scipy import sparse 189from scipy.sparse.linalg import splu 190from struct_tools import AlignedRepr 191 192from minires.core import ResSim 193 194 195class Gradient(NamedTuple): 196 """The gradient of an objective, as returned by `adjoint`.""" 197 198 S0: np.ndarray 199 """W.r.t. the initial saturation, `(Nxy,)`.""" 200 P0: np.ndarray 201 """W.r.t. the initial pressure, `(Nxy,)`. Zero unless `ct > 0` or a well is on BHP.""" 202 logK: np.ndarray 203 """W.r.t. $ \\log K $, shaped like `K`: `(2, Nx, Ny)`. Sum over axis `0` if isotropic.""" 204 bhp: np.ndarray 205 """W.r.t. the BHP controls, `(nComp, nSteps)`: entry `[i, k]` is for 206 completion `i`'s control at step `k` (`0` where rate-controlled). Sum over 207 axis `1` if the control is constant in time.""" 208 209 210def adjoint( 211 model: ResSim, 212 dt: float, 213 SS: np.ndarray, 214 PP: np.ndarray, 215 dJ_dSS: np.ndarray, 216 dJ_dPP: np.ndarray | None = None, 217) -> Gradient: 218 """The gradient of $ J(S, P) $ wrt. `S0`, `P0`, $ \\log K $ and the BHP controls, by the adjoint sweep. 219 220 Seeded by the partials of the objective wrt. the *stored* trajectory, 221 `dJ_dSS[k]` $ = ∂J/∂S_k $, `dJ_dPP[k]` $ = ∂J/∂P_k $ (shaped like `SS`, 222 `PP`; the latter defaults to `0`), ref the module docstring. Sweeps 223 backwards from the final time, re-linearizing each step (`linearize`) on 224 the way, so the trajectory `(SS, PP)` of `sim(dt, ...)` is all it needs. 225 226 The cost is about that of a `sim` (one forward step and one factorization 227 per step, ref `linearize`), independently of the number of parameters -- 228 which is the point of an adjoint. 229 """ 230 nSteps = len(SS) - 1 231 aS = np.array(dJ_dSS[-1], float) 232 aP = np.zeros(model.Nxy) if dJ_dPP is None else np.array(dJ_dPP[-1], float) 233 alogK = np.zeros(model.K.shape) 234 abhp = np.zeros((model.nComp, nSteps)) 235 for k in reversed(range(nSteps)): 236 tape = linearize(model, dt, SS[k], PP[k], k) 237 aS, aP, aK, abhp[:, k] = adj_step(tape, aS, aP) 238 alogK += aK 239 aS += dJ_dSS[k] 240 if dJ_dPP is not None: 241 aP += dJ_dPP[k] 242 return Gradient(aS, aP, alogK, abhp) 243 244 245 246def face_operators(model: ResSim) -> tuple: 247 """The interior faces of the grid, and the sparse operators on them. 248 249 Returns `(lo, hi, Grad, Sum, g)`: 250 251 - `lo`, `hi`: the flat indices of the two cells each face separates 252 (`lo` has the smaller index), shape `(nF,)`. The x-faces come first 253 (`(Nx-1) * Ny` of them, in C-order), then the y-faces (`Nx * (Ny-1)`) 254 -- less those of inactive cells (`ResSim.active`), which are omitted. 255 - `Grad`: `(nF, Nxy)`, `(Grad @ p)[f] = p[lo] - p[hi]`, as `TPFA` computes 256 the fluxes; `Grad.T` is the divergence (high faces minus low faces). 257 - `Sum`: `(nF, 2*Nxy)`, summing over the face's two cells the *directional* 258 component of a `(2, Nx, Ny)`-shaped field (flattened): the x-component 259 for x-faces, the y-component for y-faces. This is how the transmissibility 260 harmonically averages the permeabilities. 261 - `g`: `(nF,)`, the geometric factor of the transmissibilities, 262 $ 2 C h_y / h_x $ resp. $ 2 C h_x / h_y $, such that `T = g / (Sum @ (1/KM))`. 263 264 >>> from minires import ResSim 265 >>> lo, hi, Grad, Sum, g = face_operators(ResSim(Nx=3, Ny=2)) 266 >>> lo, hi # the 4 x-faces, then the 3 y-faces 267 (array([0, 1, 2, 3, 0, 2, 4]), array([2, 3, 4, 5, 1, 3, 5])) 268 >>> (Grad.T @ np.ones(7)).astype(int) # #high faces - #low faces, per cell 269 array([ 2, 0, 1, -1, 0, -2]) 270 """ 271 N = model.Nxy 272 idx = np.arange(N).reshape(model.shape) 273 act = model.active 274 fx = (act[:-1, :] & act[1:, :]).ravel() # faces between active cells 275 fy = (act[:, :-1] & act[:, 1:]).ravel() 276 lo = np.concatenate([idx[:-1, :].ravel()[fx], idx[:, :-1].ravel()[fy]]) 277 hi = np.concatenate([idx[1:, :].ravel()[fx], idx[:, 1:].ravel()[fy]]) 278 nF = len(lo) 279 nFx = int(fx.sum()) 280 ff = np.r_[np.arange(nF), np.arange(nF)] 281 ones = np.ones(nF) 282 Grad = sparse.csr_matrix((np.r_[ones, -ones], (ff, np.r_[lo, hi])), shape=(nF, N)) 283 comp = np.r_[np.zeros(nFx, int), np.full(nF - nFx, N)] # offset into 2nd component 284 Sum = sparse.csr_matrix( 285 (np.r_[ones, ones], (ff, np.r_[lo + comp, hi + comp])), shape=(nF, 2 * N) 286 ) 287 C = model.cdarcy 288 g = 2 * C * np.r_[np.full(nFx, model.hy / model.hx), np.full(nF - nFx, model.hx / model.hy)] 289 return lo, hi, Grad, Sum, g 290 291 292@dataclass(repr=False) 293class Tape(AlignedRepr): 294 """The linearization of one step of `minires.ResSim.time_stepper`. 295 296 Produced by `linearize`, consumed by `adj_step`. Holds the state it was 297 taken about, the step's result, and the coefficients -- sparse matrices 298 and vectors -- of the linear statements of the step's tangent, evaluated 299 at that state. 300 """ 301 302 model: ResSim 303 """The model whose step this linearizes (for `fluid`, `ct`, ...).""" 304 dt: float 305 """The time step.""" 306 k: int 307 """The time index (of the well controls).""" 308 S: np.ndarray 309 """Saturation at the start of the step, `(Nxy,)`.""" 310 P: np.ndarray 311 """Pressure at the start of the step, `(Nxy,)`.""" 312 S1: np.ndarray 313 """Saturation at the end of the step -- the recomputed forward result.""" 314 P1: np.ndarray 315 """Pressure at the end of the step -- the recomputed forward result.""" 316 317 # Pressure equation 318 Grad: Any 319 """`(nF, Nxy)` face difference operator, ref `face_operators`. `Grad.T` is the divergence.""" 320 dMt_dS: np.ndarray 321 """`(Nxy,)` derivative of the total mobility $ λ_t $ wrt. `S`.""" 322 dT_dMt: Any 323 """`(nF, Nxy)` Jacobian of the transmissibilities wrt. the cell mobilities: 324 $ ∂T_f / ∂λ_i = T_f^2 / (g_f \\, K_i \\, λ_i^2) $ for each of the face's 325 two cells (harmonic averaging), `0` otherwise.""" 326 dT_dlogK: Any 327 """`(nF, 2*Nxy)` Jacobian of the transmissibilities wrt. $ \\log K $ (flattened 328 like `K`): $ ∂T_f / ∂\\log K_c = T_f^2 / (g_f \\, K_c \\, λ_c) $ for each of the 329 face's two cells, in the face's direction, `0` otherwise.""" 330 T: np.ndarray 331 """`(nF,)` transmissibilities, $ T = g / Σ_\\mathrm{cells} 1/(K λ_t) $.""" 332 gradP: np.ndarray 333 """`(nF,)` pressure differences, `Grad @ P1`; the fluxes are `V = T * gradP`.""" 334 V: np.ndarray 335 """`(nF,)` fluxes through the interior faces (positive from `lo` to `hi`).""" 336 accum: np.ndarray 337 """`(Nxy,)` accumulation coefficient, $ φ c_t h^2 / Δt $ (`0` if `ct == 0`).""" 338 solve: Callable 339 """`x ↦ A⁻¹ x`, by the LU factorization of the (symmetric) pressure matrix.""" 340 341 # Well model of the BHP-controlled completions (empty if none) 342 Gb: Any 343 """`(nBHP, Nxy)` gathers cell values at the BHP-controlled completions; `Gb.T` scatters.""" 344 WI_b: np.ndarray 345 """`(nBHP,)` their well indices.""" 346 WI_lam_b: np.ndarray 347 """`(nBHP,)` their $ WI λ_t $, the coefficient of $ p_\\mathrm{bh} $ in the right-hand side.""" 348 is_bhp: np.ndarray 349 """`(nComp,)` boolean mask of the BHP-controlled completions.""" 350 p_bh_b: np.ndarray 351 """`(nBHP,)` their bottom-hole pressures.""" 352 bhp_diag: np.ndarray 353 """`(Nxy,)` their $ WI λ_t $, scattered onto the cells (`0` elsewhere).""" 354 355 # Transport equation 356 Q: np.ndarray 357 """`(Nxy,)` total well flux per cell, signed (`realize_bhp` included).""" 358 st: np.ndarray 359 """`(Nxy,)` storage rate, `Q - Grad.T @ V` (`0` if `ct == 0`).""" 360 dtx: np.ndarray 361 """`(Nxy,)` sub-step over pore volume, `dt / nT / pv`.""" 362 Up: Any 363 """`(nF, Nxy)` upwind selector: `(Up @ f)[face]` is `f` of the face's upwind cell.""" 364 Ssub: np.ndarray 365 """`(nT, Nxy)` saturation at the start of each transport sub-step.""" 366 367 @property 368 def nT(self) -> int: 369 """Number of transport sub-steps.""" 370 return len(self.Ssub) 371 372 373def linearize( 374 model: ResSim, 375 dt: float, 376 S: np.ndarray, 377 P: np.ndarray | None, 378 k: int = 0, 379) -> Tape: 380 """Recompute the step of `time_stepper` from `(S, P)` at time `k`; return its `Tape`. 381 382 The recomputation calls the very methods of the forward model (so the 383 result, `Tape.S1`/`Tape.P1`, is that of `sim`, to the solver tolerance), 384 except that it records the transport sub-steps, and does not write the 385 `actual_rates`/`actual_bhp` reports. Then the coefficients of the 386 linearization are evaluated at that state. 387 388 The cost is about that of a forward step plus a factorization of the 389 pressure matrix (held in `Tape.solve`, for the adjoint solve). 390 391 .. warning:: This mutates the model, exactly as a step of `sim` does. 392 393 Because it *is* one: `assemble_wells` (hence `well_controls`), 394 `pressure_step` and `realize_bhp` are called, and they write the 395 source field `_Q`, the bundle `_wells_now` and -- if `cached_precond` 396 -- the factorization cache `_pLU`, all as of the step linearized. So 397 after `sim`, then `adjoint` (which linearizes the steps in *reverse*), 398 these hold the values of the first step rather than the last. None of 399 it is consequential: the next step (or `linearize`) overwrites them, 400 and `_pLU` is a mere preconditioner (ref `_solve_pressure`). The 401 well reports, `actual_rates`/`actual_bhp`, are *not* written, so they 402 remain those of the `sim` that produced the trajectory. Nothing else 403 is touched: `K`, `por`, the well specifications are read only. 404 """ 405 N = model.Nxy 406 S = np.asarray(S, float).ravel() 407 P = np.zeros(N) if P is None else np.asarray(P, float).ravel() 408 lo, hi, Grad, Sum, g = face_operators(model) 409 nF = len(lo) 410 411 # Forward step, recomputed. Mirrors `time_stepper` (minus the reporting) 412 model.assemble_wells(S, P, k) 413 model._validate() 414 P1, VV = model.pressure_step(S, P, dt) 415 model.realize_bhp(P1) 416 wls = model._wells_now 417 Q = model._Q # total well flux, the BHP wells' rates now realized 418 # ... and `saturation_step_upwind`, recording the sub-steps 419 A_up = model.upwind_diff(VV) 420 pv = model.pore_volume() 421 fi = Q.clip(min=0) 422 st = model.storage_rate(VV) 423 nT = max(1, int(np.ceil(dt * model.estimate_1CFL(pv, VV, fi)))) 424 dtx = dt / nT / pv 425 B = model._spdiags(dtx, 0) @ A_up 426 Ssub = np.zeros((nT, N)) 427 Sj = S 428 for j in range(nT): 429 Ssub[j] = Sj 430 fw = model.fluid.fractional_flow(Sj) 431 Sj = Sj + (B @ fw + (fi - Sj * st) * dtx) 432 S1 = Sj 433 434 # Coefficients of the linearization, at that state 435 # -- mobilities and transmissibilities 436 Mw, Mo = model.fluid.RelPerm(S) 437 dMw, dMo = model.fluid.dRelPerm(S) 438 Mt = Mw + Mo 439 Kflat = model.K.reshape(-1) 440 KM = model.K.reshape(2, N) * Mt # K λ_t, both components, flattened 441 KM = KM.reshape(-1) 442 T = g / (Sum @ (1 / KM)) 443 Dup = sparse.vstack([sparse.eye(N), sparse.eye(N)]) # cell field → both components 444 dT_dKM = sparse.diags(T**2 / g) @ Sum # ... @ diag(1/KM²), below 445 dT_dMt = dT_dKM @ sparse.diags(Kflat / KM**2) @ Dup 446 dT_dlogK = dT_dKM @ sparse.diags(1 / KM) # since ∂(Kλ)/∂log K = Kλ 447 gradP = Grad @ P1 448 V = T * gradP 449 # -- the BHP wells 450 is_bhp = np.isfinite(wls["WI_lam"]) 451 nB = int(is_bhp.sum()) 452 inds_b = wls["inds"][is_bhp] 453 Gb = sparse.csr_matrix((np.ones(nB), (np.arange(nB), inds_b)), shape=(nB, N)) 454 WI = model.wells.WI 455 WI_b = WI[is_bhp] if WI is not None else np.zeros(0) 456 WI_lam_b = wls["WI_lam"][is_bhp] 457 p_bh_b = wls["p_bh"][is_bhp] 458 bhp_diag = wls["bhp_diag"] 459 # -- the pressure system (as `TPFA` assembles it, incl. its pin and the 460 # identity rows of the inactive cells, whose `accum` is thereby `1`) 461 accum = pv * model.ct / dt if model.ct > 0 else np.zeros(N) 462 accum = np.where(model.active.ravel(), accum, 1.0) 463 diag = accum + bhp_diag 464 if model.ct == 0 and not bhp_diag.any(): 465 diag[model._pin] += np.sum(model.K.reshape(2, -1)[:, model._pin]) 466 A = Grad.T @ sparse.diags(T) @ Grad + sparse.diags(diag) 467 solve = splu(A.tocsc(), permc_spec="MMD_AT_PLUS_A").solve 468 # -- the upwind directions 469 Up = sparse.csr_matrix( 470 (np.ones(nF), (np.arange(nF), np.where(V >= 0, lo, hi))), shape=(nF, N) 471 ) 472 473 return Tape( 474 model=model, dt=dt, k=k, S=S, P=P, S1=S1, P1=P1, 475 Grad=Grad, dMt_dS=dMw + dMo, dT_dMt=dT_dMt, dT_dlogK=dT_dlogK, 476 T=T, gradP=gradP, V=V, accum=accum, solve=solve, 477 Gb=Gb, WI_b=WI_b, WI_lam_b=WI_lam_b, is_bhp=is_bhp, p_bh_b=p_bh_b, 478 bhp_diag=bhp_diag, 479 Q=Q, st=st, dtx=dtx, Up=Up, Ssub=Ssub, 480 ) # fmt: skip 481 482 483def adj_step( 484 tape: Tape, aS1: np.ndarray, aP1: np.ndarray 485) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: 486 """Propagate the sensitivity `(aS1, aP1)` back through the step of `tape`. 487 488 The transpose of the step's tangent: given $ ∂J/∂s^{n+1} $ and 489 $ ∂J/∂p^{n+1} $, returns `(aS, aP, alogK, abhp)` 490 $ = (∂J/∂s^n, ∂J/∂p^n, ∂J/∂\\log K, ∂J/∂p_\\mathrm{bh}) $ -- the third 491 shaped like `K` and being this step's *contribution* (to be summed over 492 the steps, as `adjoint` does), the fourth `(nComp,)`, wrt. this step's 493 controls (`0` at the rate-controlled completions). Each statement of the tangent appears here, in reverse order, 494 transposed: `y = M @ x` as `x̄ += M.T @ ȳ`, `y = a * x` as `x̄ += a * ȳ`, 495 the (symmetric) solve as itself. The comments name the tangent statement 496 being transposed (ref the module docstring). 497 498 Does not modify `aS1`, `aP1`. 499 """ 500 t = tape 501 aS = np.array(aS1, float) # the running adjoint of `dS`, which the loop updates 502 aP1 = np.array(aP1, float) # `dP1` is the output, and feeds `dV` and `dQ` 503 zeros = np.zeros(len(aS)) 504 adfi, adfp, adst, adVm = zeros, zeros, zeros, np.zeros_like(t.V) 505 fp = t.Q.clip(max=0) 506 # fmt: off 507 # Transport equation, sub-steps in reverse 508 for Sj in t.Ssub[::-1]: 509 fw = t.model.fluid.fractional_flow(Sj) 510 dfw_dS = t.model.fluid.dfractional_flow(Sj) 511 adrhs = t.dtx * aS # dS = dS + dtx*drhs 512 adF = -(t.Grad @ adrhs) # drhs = -Grad.T @ dF ... 513 adfp = adfp + fw * adrhs # + dfp*fw 514 adfw = fp * adrhs # + fp*dfw 515 adfi = adfi + adrhs # + dfi 516 aS = aS - t.st * adrhs # - dS*st 517 adst = adst - Sj * adrhs # - Sj*dst 518 adVm = adVm + (t.Up @ fw) * adF # dF = dVm*(Up@fw) ... 519 adfw = adfw + t.Up.T @ (t.V * adF) # + V*(Up@dfw) 520 aS = aS + dfw_dS * adfw # dfw = dfw_dS * dS 521 adV = (t.V != 0) * adVm # dVm = dV * (V != 0) 522 if t.model.ct > 0: # dst = dQ - Grad.T @ dV 523 adQ = adst 524 adV = adV - t.Grad @ adst 525 else: 526 adQ = zeros 527 adQ = adQ + (t.Q < 0) * adfp + (t.Q > 0) * adfi # dfp, dfi = dQ*(Q<0), dQ*(Q>0) 528 529 # Pressure equation 530 adr = adQ # dQ = dr - dw*P1 - bhp_diag*dP1 531 adw = -t.P1 * adQ 532 aP1 = aP1 - t.bhp_diag * adQ 533 aP1 = aP1 + t.Grad.T @ (t.T * adV) # dV = T*(Grad@dP1) + gradP*dT 534 adT = t.gradP * adV 535 adq = t.solve(aP1) # dP1 = solve(dq - dAP) 536 adAP = -adq 537 aP = t.accum * adq # dq = accum*dP + dr 538 adr = adr + adq 539 adT = adT + t.gradP * (t.Grad @ adAP) # dAP = Grad.T@(gradP*dT) + dw*P1 540 adw = adw + t.P1 * adAP 541 adr_b = t.Gb @ adr # dr = Gb.T @ (p_bh_b * dWI_lam ... 542 adWI_lam = t.p_bh_b * adr_b 543 abhp = np.zeros(len(t.is_bhp)) # + WI_lam_b * dp_bh_b) 544 abhp[t.is_bhp] = t.WI_lam_b * adr_b 545 adWI_lam = adWI_lam + t.Gb @ adw # dw = Gb.T @ dWI_lam 546 adMt = t.Gb.T @ (t.WI_b * adWI_lam) # dWI_lam = WI_b * (Gb @ dMt) 547 adMt = adMt + t.dT_dMt.T @ adT # dT = dT_dMt@dMt + dT_dlogK@dlogK 548 alogK = t.dT_dlogK.T @ adT 549 aS = aS + t.dMt_dS * adMt # dMt = dMt_dS * dS 550 # fmt: on 551 return aS, aP, alogK.reshape(t.model.K.shape), abhp
class
Gradient(typing.NamedTuple):
196class Gradient(NamedTuple): 197 """The gradient of an objective, as returned by `adjoint`.""" 198 199 S0: np.ndarray 200 """W.r.t. the initial saturation, `(Nxy,)`.""" 201 P0: np.ndarray 202 """W.r.t. the initial pressure, `(Nxy,)`. Zero unless `ct > 0` or a well is on BHP.""" 203 logK: np.ndarray 204 """W.r.t. $ \\log K $, shaped like `K`: `(2, Nx, Ny)`. Sum over axis `0` if isotropic.""" 205 bhp: np.ndarray 206 """W.r.t. the BHP controls, `(nComp, nSteps)`: entry `[i, k]` is for 207 completion `i`'s control at step `k` (`0` where rate-controlled). Sum over 208 axis `1` if the control is constant in time."""
The gradient of an objective, as returned by adjoint.
def
adjoint( model: minires.ResSim, dt: float, SS: numpy.ndarray, PP: numpy.ndarray, dJ_dSS: numpy.ndarray, dJ_dPP: numpy.ndarray | None = None) -> Gradient:
211def adjoint( 212 model: ResSim, 213 dt: float, 214 SS: np.ndarray, 215 PP: np.ndarray, 216 dJ_dSS: np.ndarray, 217 dJ_dPP: np.ndarray | None = None, 218) -> Gradient: 219 """The gradient of $ J(S, P) $ wrt. `S0`, `P0`, $ \\log K $ and the BHP controls, by the adjoint sweep. 220 221 Seeded by the partials of the objective wrt. the *stored* trajectory, 222 `dJ_dSS[k]` $ = ∂J/∂S_k $, `dJ_dPP[k]` $ = ∂J/∂P_k $ (shaped like `SS`, 223 `PP`; the latter defaults to `0`), ref the module docstring. Sweeps 224 backwards from the final time, re-linearizing each step (`linearize`) on 225 the way, so the trajectory `(SS, PP)` of `sim(dt, ...)` is all it needs. 226 227 The cost is about that of a `sim` (one forward step and one factorization 228 per step, ref `linearize`), independently of the number of parameters -- 229 which is the point of an adjoint. 230 """ 231 nSteps = len(SS) - 1 232 aS = np.array(dJ_dSS[-1], float) 233 aP = np.zeros(model.Nxy) if dJ_dPP is None else np.array(dJ_dPP[-1], float) 234 alogK = np.zeros(model.K.shape) 235 abhp = np.zeros((model.nComp, nSteps)) 236 for k in reversed(range(nSteps)): 237 tape = linearize(model, dt, SS[k], PP[k], k) 238 aS, aP, aK, abhp[:, k] = adj_step(tape, aS, aP) 239 alogK += aK 240 aS += dJ_dSS[k] 241 if dJ_dPP is not None: 242 aP += dJ_dPP[k] 243 return Gradient(aS, aP, alogK, abhp)
The gradient of $ J(S, P) $ wrt. S0, P0, $ \log K $ and the BHP controls, by the adjoint sweep.
Seeded by the partials of the objective wrt. the stored trajectory,
dJ_dSS[k] $ = ∂J/∂S_k $, dJ_dPP[k] $ = ∂J/∂P_k $ (shaped like SS,
PP; the latter defaults to 0), ref the module docstring. Sweeps
backwards from the final time, re-linearizing each step (linearize) on
the way, so the trajectory (SS, PP) of sim(dt, ...) is all it needs.
The cost is about that of a sim (one forward step and one factorization
per step, ref linearize), independently of the number of parameters --
which is the point of an adjoint.
247def face_operators(model: ResSim) -> tuple: 248 """The interior faces of the grid, and the sparse operators on them. 249 250 Returns `(lo, hi, Grad, Sum, g)`: 251 252 - `lo`, `hi`: the flat indices of the two cells each face separates 253 (`lo` has the smaller index), shape `(nF,)`. The x-faces come first 254 (`(Nx-1) * Ny` of them, in C-order), then the y-faces (`Nx * (Ny-1)`) 255 -- less those of inactive cells (`ResSim.active`), which are omitted. 256 - `Grad`: `(nF, Nxy)`, `(Grad @ p)[f] = p[lo] - p[hi]`, as `TPFA` computes 257 the fluxes; `Grad.T` is the divergence (high faces minus low faces). 258 - `Sum`: `(nF, 2*Nxy)`, summing over the face's two cells the *directional* 259 component of a `(2, Nx, Ny)`-shaped field (flattened): the x-component 260 for x-faces, the y-component for y-faces. This is how the transmissibility 261 harmonically averages the permeabilities. 262 - `g`: `(nF,)`, the geometric factor of the transmissibilities, 263 $ 2 C h_y / h_x $ resp. $ 2 C h_x / h_y $, such that `T = g / (Sum @ (1/KM))`. 264 265 >>> from minires import ResSim 266 >>> lo, hi, Grad, Sum, g = face_operators(ResSim(Nx=3, Ny=2)) 267 >>> lo, hi # the 4 x-faces, then the 3 y-faces 268 (array([0, 1, 2, 3, 0, 2, 4]), array([2, 3, 4, 5, 1, 3, 5])) 269 >>> (Grad.T @ np.ones(7)).astype(int) # #high faces - #low faces, per cell 270 array([ 2, 0, 1, -1, 0, -2]) 271 """ 272 N = model.Nxy 273 idx = np.arange(N).reshape(model.shape) 274 act = model.active 275 fx = (act[:-1, :] & act[1:, :]).ravel() # faces between active cells 276 fy = (act[:, :-1] & act[:, 1:]).ravel() 277 lo = np.concatenate([idx[:-1, :].ravel()[fx], idx[:, :-1].ravel()[fy]]) 278 hi = np.concatenate([idx[1:, :].ravel()[fx], idx[:, 1:].ravel()[fy]]) 279 nF = len(lo) 280 nFx = int(fx.sum()) 281 ff = np.r_[np.arange(nF), np.arange(nF)] 282 ones = np.ones(nF) 283 Grad = sparse.csr_matrix((np.r_[ones, -ones], (ff, np.r_[lo, hi])), shape=(nF, N)) 284 comp = np.r_[np.zeros(nFx, int), np.full(nF - nFx, N)] # offset into 2nd component 285 Sum = sparse.csr_matrix( 286 (np.r_[ones, ones], (ff, np.r_[lo + comp, hi + comp])), shape=(nF, 2 * N) 287 ) 288 C = model.cdarcy 289 g = 2 * C * np.r_[np.full(nFx, model.hy / model.hx), np.full(nF - nFx, model.hx / model.hy)] 290 return lo, hi, Grad, Sum, g
The interior faces of the grid, and the sparse operators on them.
Returns (lo, hi, Grad, Sum, g):
lo,hi: the flat indices of the two cells each face separates (lohas the smaller index), shape(nF,). The x-faces come first ((Nx-1) * Nyof them, in C-order), then the y-faces (Nx * (Ny-1)) -- less those of inactive cells (ResSim.active), which are omitted.Grad:(nF, Nxy),(Grad @ p)[f] = p[lo] - p[hi], asTPFAcomputes the fluxes;Grad.Tis the divergence (high faces minus low faces).Sum:(nF, 2*Nxy), summing over the face's two cells the directional component of a(2, Nx, Ny)-shaped field (flattened): the x-component for x-faces, the y-component for y-faces. This is how the transmissibility harmonically averages the permeabilities.g:(nF,), the geometric factor of the transmissibilities, $ 2 C h_y / h_x $ resp. $ 2 C h_x / h_y $, such thatT = g / (Sum @ (1/KM)).
>>> from minires import ResSim
>>> lo, hi, Grad, Sum, g = face_operators(ResSim(Nx=3, Ny=2))
>>> lo, hi # the 4 x-faces, then the 3 y-faces
(array([0, 1, 2, 3, 0, 2, 4]), array([2, 3, 4, 5, 1, 3, 5]))
>>> (Grad.T @ np.ones(7)).astype(int) # #high faces - #low faces, per cell
array([ 2, 0, 1, -1, 0, -2])
@dataclass(repr=False)
class
Tape293@dataclass(repr=False) 294class Tape(AlignedRepr): 295 """The linearization of one step of `minires.ResSim.time_stepper`. 296 297 Produced by `linearize`, consumed by `adj_step`. Holds the state it was 298 taken about, the step's result, and the coefficients -- sparse matrices 299 and vectors -- of the linear statements of the step's tangent, evaluated 300 at that state. 301 """ 302 303 model: ResSim 304 """The model whose step this linearizes (for `fluid`, `ct`, ...).""" 305 dt: float 306 """The time step.""" 307 k: int 308 """The time index (of the well controls).""" 309 S: np.ndarray 310 """Saturation at the start of the step, `(Nxy,)`.""" 311 P: np.ndarray 312 """Pressure at the start of the step, `(Nxy,)`.""" 313 S1: np.ndarray 314 """Saturation at the end of the step -- the recomputed forward result.""" 315 P1: np.ndarray 316 """Pressure at the end of the step -- the recomputed forward result.""" 317 318 # Pressure equation 319 Grad: Any 320 """`(nF, Nxy)` face difference operator, ref `face_operators`. `Grad.T` is the divergence.""" 321 dMt_dS: np.ndarray 322 """`(Nxy,)` derivative of the total mobility $ λ_t $ wrt. `S`.""" 323 dT_dMt: Any 324 """`(nF, Nxy)` Jacobian of the transmissibilities wrt. the cell mobilities: 325 $ ∂T_f / ∂λ_i = T_f^2 / (g_f \\, K_i \\, λ_i^2) $ for each of the face's 326 two cells (harmonic averaging), `0` otherwise.""" 327 dT_dlogK: Any 328 """`(nF, 2*Nxy)` Jacobian of the transmissibilities wrt. $ \\log K $ (flattened 329 like `K`): $ ∂T_f / ∂\\log K_c = T_f^2 / (g_f \\, K_c \\, λ_c) $ for each of the 330 face's two cells, in the face's direction, `0` otherwise.""" 331 T: np.ndarray 332 """`(nF,)` transmissibilities, $ T = g / Σ_\\mathrm{cells} 1/(K λ_t) $.""" 333 gradP: np.ndarray 334 """`(nF,)` pressure differences, `Grad @ P1`; the fluxes are `V = T * gradP`.""" 335 V: np.ndarray 336 """`(nF,)` fluxes through the interior faces (positive from `lo` to `hi`).""" 337 accum: np.ndarray 338 """`(Nxy,)` accumulation coefficient, $ φ c_t h^2 / Δt $ (`0` if `ct == 0`).""" 339 solve: Callable 340 """`x ↦ A⁻¹ x`, by the LU factorization of the (symmetric) pressure matrix.""" 341 342 # Well model of the BHP-controlled completions (empty if none) 343 Gb: Any 344 """`(nBHP, Nxy)` gathers cell values at the BHP-controlled completions; `Gb.T` scatters.""" 345 WI_b: np.ndarray 346 """`(nBHP,)` their well indices.""" 347 WI_lam_b: np.ndarray 348 """`(nBHP,)` their $ WI λ_t $, the coefficient of $ p_\\mathrm{bh} $ in the right-hand side.""" 349 is_bhp: np.ndarray 350 """`(nComp,)` boolean mask of the BHP-controlled completions.""" 351 p_bh_b: np.ndarray 352 """`(nBHP,)` their bottom-hole pressures.""" 353 bhp_diag: np.ndarray 354 """`(Nxy,)` their $ WI λ_t $, scattered onto the cells (`0` elsewhere).""" 355 356 # Transport equation 357 Q: np.ndarray 358 """`(Nxy,)` total well flux per cell, signed (`realize_bhp` included).""" 359 st: np.ndarray 360 """`(Nxy,)` storage rate, `Q - Grad.T @ V` (`0` if `ct == 0`).""" 361 dtx: np.ndarray 362 """`(Nxy,)` sub-step over pore volume, `dt / nT / pv`.""" 363 Up: Any 364 """`(nF, Nxy)` upwind selector: `(Up @ f)[face]` is `f` of the face's upwind cell.""" 365 Ssub: np.ndarray 366 """`(nT, Nxy)` saturation at the start of each transport sub-step.""" 367 368 @property 369 def nT(self) -> int: 370 """Number of transport sub-steps.""" 371 return len(self.Ssub)
The linearization of one step of minires.ResSim.time_stepper.
Produced by linearize, consumed by adj_step. Holds the state it was
taken about, the step's result, and the coefficients -- sparse matrices
and vectors -- of the linear statements of the step's tangent, evaluated
at that state.
Tape( model: minires.ResSim, dt: float, k: int, S: numpy.ndarray, P: numpy.ndarray, S1: numpy.ndarray, P1: numpy.ndarray, Grad: Any, dMt_dS: numpy.ndarray, dT_dMt: Any, dT_dlogK: Any, T: numpy.ndarray, gradP: numpy.ndarray, V: numpy.ndarray, accum: numpy.ndarray, solve: Callable, Gb: Any, WI_b: numpy.ndarray, WI_lam_b: numpy.ndarray, is_bhp: numpy.ndarray, p_bh_b: numpy.ndarray, bhp_diag: numpy.ndarray, Q: numpy.ndarray, st: numpy.ndarray, dtx: numpy.ndarray, Up: Any, Ssub: numpy.ndarray)
dT_dMt: Any
(nF, Nxy) Jacobian of the transmissibilities wrt. the cell mobilities:
$ ∂T_f / ∂λ_i = T_f^2 / (g_f \, K_i \, λ_i^2) $ for each of the face's
two cells (harmonic averaging), 0 otherwise.
dT_dlogK: Any
(nF, 2*Nxy) Jacobian of the transmissibilities wrt. $ \log K $ (flattened
like K): $ ∂T_f / ∂\log K_c = T_f^2 / (g_f \, K_c \, λ_c) $ for each of the
face's two cells, in the face's direction, 0 otherwise.
def
linearize( model: minires.ResSim, dt: float, S: numpy.ndarray, P: numpy.ndarray | None, k: int = 0) -> Tape:
374def linearize( 375 model: ResSim, 376 dt: float, 377 S: np.ndarray, 378 P: np.ndarray | None, 379 k: int = 0, 380) -> Tape: 381 """Recompute the step of `time_stepper` from `(S, P)` at time `k`; return its `Tape`. 382 383 The recomputation calls the very methods of the forward model (so the 384 result, `Tape.S1`/`Tape.P1`, is that of `sim`, to the solver tolerance), 385 except that it records the transport sub-steps, and does not write the 386 `actual_rates`/`actual_bhp` reports. Then the coefficients of the 387 linearization are evaluated at that state. 388 389 The cost is about that of a forward step plus a factorization of the 390 pressure matrix (held in `Tape.solve`, for the adjoint solve). 391 392 .. warning:: This mutates the model, exactly as a step of `sim` does. 393 394 Because it *is* one: `assemble_wells` (hence `well_controls`), 395 `pressure_step` and `realize_bhp` are called, and they write the 396 source field `_Q`, the bundle `_wells_now` and -- if `cached_precond` 397 -- the factorization cache `_pLU`, all as of the step linearized. So 398 after `sim`, then `adjoint` (which linearizes the steps in *reverse*), 399 these hold the values of the first step rather than the last. None of 400 it is consequential: the next step (or `linearize`) overwrites them, 401 and `_pLU` is a mere preconditioner (ref `_solve_pressure`). The 402 well reports, `actual_rates`/`actual_bhp`, are *not* written, so they 403 remain those of the `sim` that produced the trajectory. Nothing else 404 is touched: `K`, `por`, the well specifications are read only. 405 """ 406 N = model.Nxy 407 S = np.asarray(S, float).ravel() 408 P = np.zeros(N) if P is None else np.asarray(P, float).ravel() 409 lo, hi, Grad, Sum, g = face_operators(model) 410 nF = len(lo) 411 412 # Forward step, recomputed. Mirrors `time_stepper` (minus the reporting) 413 model.assemble_wells(S, P, k) 414 model._validate() 415 P1, VV = model.pressure_step(S, P, dt) 416 model.realize_bhp(P1) 417 wls = model._wells_now 418 Q = model._Q # total well flux, the BHP wells' rates now realized 419 # ... and `saturation_step_upwind`, recording the sub-steps 420 A_up = model.upwind_diff(VV) 421 pv = model.pore_volume() 422 fi = Q.clip(min=0) 423 st = model.storage_rate(VV) 424 nT = max(1, int(np.ceil(dt * model.estimate_1CFL(pv, VV, fi)))) 425 dtx = dt / nT / pv 426 B = model._spdiags(dtx, 0) @ A_up 427 Ssub = np.zeros((nT, N)) 428 Sj = S 429 for j in range(nT): 430 Ssub[j] = Sj 431 fw = model.fluid.fractional_flow(Sj) 432 Sj = Sj + (B @ fw + (fi - Sj * st) * dtx) 433 S1 = Sj 434 435 # Coefficients of the linearization, at that state 436 # -- mobilities and transmissibilities 437 Mw, Mo = model.fluid.RelPerm(S) 438 dMw, dMo = model.fluid.dRelPerm(S) 439 Mt = Mw + Mo 440 Kflat = model.K.reshape(-1) 441 KM = model.K.reshape(2, N) * Mt # K λ_t, both components, flattened 442 KM = KM.reshape(-1) 443 T = g / (Sum @ (1 / KM)) 444 Dup = sparse.vstack([sparse.eye(N), sparse.eye(N)]) # cell field → both components 445 dT_dKM = sparse.diags(T**2 / g) @ Sum # ... @ diag(1/KM²), below 446 dT_dMt = dT_dKM @ sparse.diags(Kflat / KM**2) @ Dup 447 dT_dlogK = dT_dKM @ sparse.diags(1 / KM) # since ∂(Kλ)/∂log K = Kλ 448 gradP = Grad @ P1 449 V = T * gradP 450 # -- the BHP wells 451 is_bhp = np.isfinite(wls["WI_lam"]) 452 nB = int(is_bhp.sum()) 453 inds_b = wls["inds"][is_bhp] 454 Gb = sparse.csr_matrix((np.ones(nB), (np.arange(nB), inds_b)), shape=(nB, N)) 455 WI = model.wells.WI 456 WI_b = WI[is_bhp] if WI is not None else np.zeros(0) 457 WI_lam_b = wls["WI_lam"][is_bhp] 458 p_bh_b = wls["p_bh"][is_bhp] 459 bhp_diag = wls["bhp_diag"] 460 # -- the pressure system (as `TPFA` assembles it, incl. its pin and the 461 # identity rows of the inactive cells, whose `accum` is thereby `1`) 462 accum = pv * model.ct / dt if model.ct > 0 else np.zeros(N) 463 accum = np.where(model.active.ravel(), accum, 1.0) 464 diag = accum + bhp_diag 465 if model.ct == 0 and not bhp_diag.any(): 466 diag[model._pin] += np.sum(model.K.reshape(2, -1)[:, model._pin]) 467 A = Grad.T @ sparse.diags(T) @ Grad + sparse.diags(diag) 468 solve = splu(A.tocsc(), permc_spec="MMD_AT_PLUS_A").solve 469 # -- the upwind directions 470 Up = sparse.csr_matrix( 471 (np.ones(nF), (np.arange(nF), np.where(V >= 0, lo, hi))), shape=(nF, N) 472 ) 473 474 return Tape( 475 model=model, dt=dt, k=k, S=S, P=P, S1=S1, P1=P1, 476 Grad=Grad, dMt_dS=dMw + dMo, dT_dMt=dT_dMt, dT_dlogK=dT_dlogK, 477 T=T, gradP=gradP, V=V, accum=accum, solve=solve, 478 Gb=Gb, WI_b=WI_b, WI_lam_b=WI_lam_b, is_bhp=is_bhp, p_bh_b=p_bh_b, 479 bhp_diag=bhp_diag, 480 Q=Q, st=st, dtx=dtx, Up=Up, Ssub=Ssub, 481 ) # fmt: skip
Recompute the step of time_stepper from (S, P) at time k; return its Tape.
The recomputation calls the very methods of the forward model (so the
result, Tape.S1/Tape.P1, is that of sim, to the solver tolerance),
except that it records the transport sub-steps, and does not write the
actual_rates/actual_bhp reports. Then the coefficients of the
linearization are evaluated at that state.
The cost is about that of a forward step plus a factorization of the
pressure matrix (held in Tape.solve, for the adjoint solve).
This mutates the model, exactly as a step of sim does.
Because it is one: assemble_wells (hence well_controls),
pressure_step and realize_bhp are called, and they write the
source field _Q, the bundle _wells_now and -- if cached_precond
-- the factorization cache _pLU, all as of the step linearized. So
after sim, then adjoint (which linearizes the steps in reverse),
these hold the values of the first step rather than the last. None of
it is consequential: the next step (or linearize) overwrites them,
and _pLU is a mere preconditioner (ref _solve_pressure). The
well reports, actual_rates/actual_bhp, are not written, so they
remain those of the sim that produced the trajectory. Nothing else
is touched: K, por, the well specifications are read only.
def
adj_step( tape: Tape, aS1: numpy.ndarray, aP1: numpy.ndarray) -> tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray, numpy.ndarray]:
484def adj_step( 485 tape: Tape, aS1: np.ndarray, aP1: np.ndarray 486) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: 487 """Propagate the sensitivity `(aS1, aP1)` back through the step of `tape`. 488 489 The transpose of the step's tangent: given $ ∂J/∂s^{n+1} $ and 490 $ ∂J/∂p^{n+1} $, returns `(aS, aP, alogK, abhp)` 491 $ = (∂J/∂s^n, ∂J/∂p^n, ∂J/∂\\log K, ∂J/∂p_\\mathrm{bh}) $ -- the third 492 shaped like `K` and being this step's *contribution* (to be summed over 493 the steps, as `adjoint` does), the fourth `(nComp,)`, wrt. this step's 494 controls (`0` at the rate-controlled completions). Each statement of the tangent appears here, in reverse order, 495 transposed: `y = M @ x` as `x̄ += M.T @ ȳ`, `y = a * x` as `x̄ += a * ȳ`, 496 the (symmetric) solve as itself. The comments name the tangent statement 497 being transposed (ref the module docstring). 498 499 Does not modify `aS1`, `aP1`. 500 """ 501 t = tape 502 aS = np.array(aS1, float) # the running adjoint of `dS`, which the loop updates 503 aP1 = np.array(aP1, float) # `dP1` is the output, and feeds `dV` and `dQ` 504 zeros = np.zeros(len(aS)) 505 adfi, adfp, adst, adVm = zeros, zeros, zeros, np.zeros_like(t.V) 506 fp = t.Q.clip(max=0) 507 # fmt: off 508 # Transport equation, sub-steps in reverse 509 for Sj in t.Ssub[::-1]: 510 fw = t.model.fluid.fractional_flow(Sj) 511 dfw_dS = t.model.fluid.dfractional_flow(Sj) 512 adrhs = t.dtx * aS # dS = dS + dtx*drhs 513 adF = -(t.Grad @ adrhs) # drhs = -Grad.T @ dF ... 514 adfp = adfp + fw * adrhs # + dfp*fw 515 adfw = fp * adrhs # + fp*dfw 516 adfi = adfi + adrhs # + dfi 517 aS = aS - t.st * adrhs # - dS*st 518 adst = adst - Sj * adrhs # - Sj*dst 519 adVm = adVm + (t.Up @ fw) * adF # dF = dVm*(Up@fw) ... 520 adfw = adfw + t.Up.T @ (t.V * adF) # + V*(Up@dfw) 521 aS = aS + dfw_dS * adfw # dfw = dfw_dS * dS 522 adV = (t.V != 0) * adVm # dVm = dV * (V != 0) 523 if t.model.ct > 0: # dst = dQ - Grad.T @ dV 524 adQ = adst 525 adV = adV - t.Grad @ adst 526 else: 527 adQ = zeros 528 adQ = adQ + (t.Q < 0) * adfp + (t.Q > 0) * adfi # dfp, dfi = dQ*(Q<0), dQ*(Q>0) 529 530 # Pressure equation 531 adr = adQ # dQ = dr - dw*P1 - bhp_diag*dP1 532 adw = -t.P1 * adQ 533 aP1 = aP1 - t.bhp_diag * adQ 534 aP1 = aP1 + t.Grad.T @ (t.T * adV) # dV = T*(Grad@dP1) + gradP*dT 535 adT = t.gradP * adV 536 adq = t.solve(aP1) # dP1 = solve(dq - dAP) 537 adAP = -adq 538 aP = t.accum * adq # dq = accum*dP + dr 539 adr = adr + adq 540 adT = adT + t.gradP * (t.Grad @ adAP) # dAP = Grad.T@(gradP*dT) + dw*P1 541 adw = adw + t.P1 * adAP 542 adr_b = t.Gb @ adr # dr = Gb.T @ (p_bh_b * dWI_lam ... 543 adWI_lam = t.p_bh_b * adr_b 544 abhp = np.zeros(len(t.is_bhp)) # + WI_lam_b * dp_bh_b) 545 abhp[t.is_bhp] = t.WI_lam_b * adr_b 546 adWI_lam = adWI_lam + t.Gb @ adw # dw = Gb.T @ dWI_lam 547 adMt = t.Gb.T @ (t.WI_b * adWI_lam) # dWI_lam = WI_b * (Gb @ dMt) 548 adMt = adMt + t.dT_dMt.T @ adT # dT = dT_dMt@dMt + dT_dlogK@dlogK 549 alogK = t.dT_dlogK.T @ adT 550 aS = aS + t.dMt_dS * adMt # dMt = dMt_dS * dS 551 # fmt: on 552 return aS, aP, alogK.reshape(t.model.K.shape), abhp
Propagate the sensitivity (aS1, aP1) back through the step of tape.
The transpose of the step's tangent: given $ ∂J/∂s^{n+1} $ and
$ ∂J/∂p^{n+1} $, returns (aS, aP, alogK, abhp)
$ = (∂J/∂s^n, ∂J/∂p^n, ∂J/∂\log K, ∂J/∂p_\mathrm{bh}) $ -- the third
shaped like K and being this step's contribution (to be summed over
the steps, as adjoint does), the fourth (nComp,), wrt. this step's
controls (0 at the rate-controlled completions). Each statement of the tangent appears here, in reverse order,
transposed: y = M @ x as x̄ += M.T @ ȳ, y = a * x as x̄ += a * ȳ,
the (symmetric) solve as itself. The comments name the tangent statement
being transposed (ref the module docstring).
Does not modify aS1, aP1.