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 190 191from minires._repr import AlignedRepr 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 face_operators(model: ResSim) -> tuple: 211 """The interior faces of the grid, and the sparse operators on them. 212 213 Returns `(lo, hi, Grad, Sum, g)`: 214 215 - `lo`, `hi`: the flat indices of the two cells each face separates 216 (`lo` has the smaller index), shape `(nF,)`. The x-faces come first 217 (`(Nx-1) * Ny` of them, in C-order), then the y-faces (`Nx * (Ny-1)`) 218 -- less those of inactive cells (`ResSim.active`), which are omitted. 219 - `Grad`: `(nF, Nxy)`, `(Grad @ p)[f] = p[lo] - p[hi]`, as `TPFA` computes 220 the fluxes; `Grad.T` is the divergence (high faces minus low faces). 221 - `Sum`: `(nF, 2*Nxy)`, summing over the face's two cells the *directional* 222 component of a `(2, Nx, Ny)`-shaped field (flattened): the x-component 223 for x-faces, the y-component for y-faces. This is how the transmissibility 224 harmonically averages the permeabilities. 225 - `g`: `(nF,)`, the geometric factor of the transmissibilities, 226 $ 2 C h_y / h_x $ resp. $ 2 C h_x / h_y $, such that `T = g / (Sum @ (1/KM))`. 227 228 >>> from minires import ResSim 229 >>> lo, hi, Grad, Sum, g = face_operators(ResSim(Nx=3, Ny=2)) 230 >>> lo, hi # the 4 x-faces, then the 3 y-faces 231 (array([0, 1, 2, 3, 0, 2, 4]), array([2, 3, 4, 5, 1, 3, 5])) 232 >>> (Grad.T @ np.ones(7)).astype(int) # #high faces - #low faces, per cell 233 array([ 2, 0, 1, -1, 0, -2]) 234 """ 235 N = model.Nxy 236 idx = np.arange(N).reshape(model.shape) 237 act = model.active 238 fx = (act[:-1, :] & act[1:, :]).ravel() # faces between active cells 239 fy = (act[:, :-1] & act[:, 1:]).ravel() 240 lo = np.concatenate([idx[:-1, :].ravel()[fx], idx[:, :-1].ravel()[fy]]) 241 hi = np.concatenate([idx[1:, :].ravel()[fx], idx[:, 1:].ravel()[fy]]) 242 nF = len(lo) 243 nFx = int(fx.sum()) 244 ff = np.r_[np.arange(nF), np.arange(nF)] 245 ones = np.ones(nF) 246 Grad = sparse.csr_matrix((np.r_[ones, -ones], (ff, np.r_[lo, hi])), shape=(nF, N)) 247 comp = np.r_[np.zeros(nFx, int), np.full(nF - nFx, N)] # offset into 2nd component 248 Sum = sparse.csr_matrix( 249 (np.r_[ones, ones], (ff, np.r_[lo + comp, hi + comp])), shape=(nF, 2 * N) 250 ) 251 C = model.cdarcy 252 g = 2 * C * np.r_[np.full(nFx, model.hy / model.hx), np.full(nF - nFx, model.hx / model.hy)] 253 return lo, hi, Grad, Sum, g 254 255 256@dataclass 257class Tape(AlignedRepr): 258 """The linearization of one step of `minires.ResSim.time_stepper`. 259 260 Produced by `linearize`, consumed by `adj_step`. Holds the state it was 261 taken about, the step's result, and the coefficients -- sparse matrices 262 and vectors -- of the linear statements of the step's tangent, evaluated 263 at that state. 264 """ 265 266 __repr__ = AlignedRepr.__repr__ 267 268 model: ResSim 269 """The model whose step this linearizes (for `fluid`, `ct`, ...).""" 270 dt: float 271 """The time step.""" 272 k: int 273 """The time index (of the well controls).""" 274 S: np.ndarray 275 """Saturation at the start of the step, `(Nxy,)`.""" 276 P: np.ndarray 277 """Pressure at the start of the step, `(Nxy,)`.""" 278 S1: np.ndarray 279 """Saturation at the end of the step -- the recomputed forward result.""" 280 P1: np.ndarray 281 """Pressure at the end of the step -- the recomputed forward result.""" 282 283 # Pressure equation 284 Grad: Any 285 """`(nF, Nxy)` face difference operator, ref `face_operators`. `Grad.T` is the divergence.""" 286 dMt_dS: np.ndarray 287 """`(Nxy,)` derivative of the total mobility $ λ_t $ wrt. `S`.""" 288 dT_dMt: Any 289 """`(nF, Nxy)` Jacobian of the transmissibilities wrt. the cell mobilities: 290 $ ∂T_f / ∂λ_i = T_f^2 / (g_f \\, K_i \\, λ_i^2) $ for each of the face's 291 two cells (harmonic averaging), `0` otherwise.""" 292 dT_dlogK: Any 293 """`(nF, 2*Nxy)` Jacobian of the transmissibilities wrt. $ \\log K $ (flattened 294 like `K`): $ ∂T_f / ∂\\log K_c = T_f^2 / (g_f \\, K_c \\, λ_c) $ for each of the 295 face's two cells, in the face's direction, `0` otherwise.""" 296 T: np.ndarray 297 """`(nF,)` transmissibilities, $ T = g / Σ_\\mathrm{cells} 1/(K λ_t) $.""" 298 gradP: np.ndarray 299 """`(nF,)` pressure differences, `Grad @ P1`; the fluxes are `V = T * gradP`.""" 300 V: np.ndarray 301 """`(nF,)` fluxes through the interior faces (positive from `lo` to `hi`).""" 302 accum: np.ndarray 303 """`(Nxy,)` accumulation coefficient, $ φ c_t h^2 / Δt $ (`0` if `ct == 0`).""" 304 solve: Callable 305 """`x ↦ A⁻¹ x`, by the LU factorization of the (symmetric) pressure matrix.""" 306 307 # Well model of the BHP-controlled completions (empty if none) 308 Gb: Any 309 """`(nBHP, Nxy)` gathers cell values at the BHP-controlled completions; `Gb.T` scatters.""" 310 WI_b: np.ndarray 311 """`(nBHP,)` their well indices.""" 312 WI_lam_b: np.ndarray 313 """`(nBHP,)` their $ WI λ_t $, the coefficient of $ p_\\mathrm{bh} $ in the right-hand side.""" 314 is_bhp: np.ndarray 315 """`(nComp,)` boolean mask of the BHP-controlled completions.""" 316 p_bh_b: np.ndarray 317 """`(nBHP,)` their bottom-hole pressures.""" 318 bhp_diag: np.ndarray 319 """`(Nxy,)` their $ WI λ_t $, scattered onto the cells (`0` elsewhere).""" 320 321 # Transport equation 322 Q: np.ndarray 323 """`(Nxy,)` total well flux per cell, signed (`realize_bhp` included).""" 324 st: np.ndarray 325 """`(Nxy,)` storage rate, `Q - Grad.T @ V` (`0` if `ct == 0`).""" 326 dtx: np.ndarray 327 """`(Nxy,)` sub-step over pore volume, `dt / nT / pv`.""" 328 Up: Any 329 """`(nF, Nxy)` upwind selector: `(Up @ f)[face]` is `f` of the face's upwind cell.""" 330 Ssub: np.ndarray 331 """`(nT, Nxy)` saturation at the start of each transport sub-step.""" 332 333 @property 334 def nT(self) -> int: 335 """Number of transport sub-steps.""" 336 return len(self.Ssub) 337 338 339def linearize( 340 model: ResSim, 341 dt: float, 342 S: np.ndarray, 343 P: np.ndarray | None, 344 k: int = 0, 345) -> Tape: 346 """Recompute the step of `time_stepper` from `(S, P)` at time `k`; return its `Tape`. 347 348 The recomputation calls the very methods of the forward model (so the 349 result, `Tape.S1`/`Tape.P1`, is that of `sim`, to the solver tolerance), 350 except that it records the transport sub-steps, and does not write the 351 `actual_rates`/`actual_bhp` reports. Then the coefficients of the 352 linearization are evaluated at that state. 353 354 The cost is about that of a forward step plus a factorization of the 355 pressure matrix (held in `Tape.solve`, for the adjoint solve). 356 357 .. warning:: This mutates the model, exactly as a step of `sim` does. 358 359 Because it *is* one: `assemble_wells` (hence `well_controls`), 360 `pressure_step` and `realize_bhp` are called, and they write the 361 source field `_Q`, the bundle `_wells_now` and -- if `cached_precond` 362 -- the factorization cache `_pLU`, all as of the step linearized. So 363 after `sim`, then `adjoint` (which linearizes the steps in *reverse*), 364 these hold the values of the first step rather than the last. None of 365 it is consequential: the next step (or `linearize`) overwrites them, 366 and `_pLU` is a mere preconditioner (ref `_solve_pressure`). The 367 well reports, `actual_rates`/`actual_bhp`, are *not* written, so they 368 remain those of the `sim` that produced the trajectory. Nothing else 369 is touched: `K`, `por`, the well specifications are read only. 370 """ 371 N = model.Nxy 372 S = np.asarray(S, float).ravel() 373 P = np.zeros(N) if P is None else np.asarray(P, float).ravel() 374 lo, hi, Grad, Sum, g = face_operators(model) 375 nF = len(lo) 376 377 # Forward step, recomputed. Mirrors `time_stepper` (minus the reporting) 378 model.assemble_wells(S, P, k) 379 model._validate() 380 P1, VV = model.pressure_step(S, P, dt) 381 model.realize_bhp(P1) 382 wls = model._wells_now 383 Q = model._Q # total well flux, the BHP wells' rates now realized 384 # ... and `saturation_step_upwind`, recording the sub-steps 385 A_up = model.upwind_diff(VV) 386 pv = model.pore_volume() 387 fi = Q.clip(min=0) 388 st = model.storage_rate(VV) 389 nT = max(1, int(np.ceil(dt * model.estimate_1CFL(pv, VV, fi)))) 390 dtx = dt / nT / pv 391 B = model._spdiags(dtx, 0) @ A_up 392 Ssub = np.zeros((nT, N)) 393 Sj = S 394 for j in range(nT): 395 Ssub[j] = Sj 396 fw = model.fluid.fractional_flow(Sj) 397 Sj = Sj + (B @ fw + (fi - Sj * st) * dtx) 398 S1 = Sj 399 400 # Coefficients of the linearization, at that state 401 # -- mobilities and transmissibilities 402 Mw, Mo = model.fluid.RelPerm(S) 403 dMw, dMo = model.fluid.dRelPerm(S) 404 Mt = Mw + Mo 405 Kflat = model.K.reshape(-1) 406 KM = model.K.reshape(2, N) * Mt # K λ_t, both components, flattened 407 KM = KM.reshape(-1) 408 T = g / (Sum @ (1 / KM)) 409 Dup = sparse.vstack([sparse.eye(N), sparse.eye(N)]) # cell field → both components 410 dT_dKM = sparse.diags(T**2 / g) @ Sum # ... @ diag(1/KM²), below 411 dT_dMt = dT_dKM @ sparse.diags(Kflat / KM**2) @ Dup 412 dT_dlogK = dT_dKM @ sparse.diags(1 / KM) # since ∂(Kλ)/∂log K = Kλ 413 gradP = Grad @ P1 414 V = T * gradP 415 # -- the BHP wells 416 is_bhp = np.isfinite(wls["WI_lam"]) 417 nB = int(is_bhp.sum()) 418 inds_b = wls["inds"][is_bhp] 419 Gb = sparse.csr_matrix((np.ones(nB), (np.arange(nB), inds_b)), shape=(nB, N)) 420 WI = model.wells.WI 421 WI_b = WI[is_bhp] if WI is not None else np.zeros(0) 422 WI_lam_b = wls["WI_lam"][is_bhp] 423 p_bh_b = wls["p_bh"][is_bhp] 424 bhp_diag = wls["bhp_diag"] 425 # -- the pressure system (as `TPFA` assembles it, incl. its pin and the 426 # identity rows of the inactive cells, whose `accum` is thereby `1`) 427 accum = pv * model.ct / dt if model.ct > 0 else np.zeros(N) 428 accum = np.where(model.active.ravel(), accum, 1.0) 429 diag = accum + bhp_diag 430 if model.ct == 0 and not bhp_diag.any(): 431 diag[model._pin] += np.sum(model.K.reshape(2, -1)[:, model._pin]) 432 A = Grad.T @ sparse.diags(T) @ Grad + sparse.diags(diag) 433 solve = splu(A.tocsc(), permc_spec="MMD_AT_PLUS_A").solve 434 # -- the upwind directions 435 Up = sparse.csr_matrix( 436 (np.ones(nF), (np.arange(nF), np.where(V >= 0, lo, hi))), shape=(nF, N) 437 ) 438 439 return Tape( 440 model=model, dt=dt, k=k, S=S, P=P, S1=S1, P1=P1, 441 Grad=Grad, dMt_dS=dMw + dMo, dT_dMt=dT_dMt, dT_dlogK=dT_dlogK, 442 T=T, gradP=gradP, V=V, accum=accum, solve=solve, 443 Gb=Gb, WI_b=WI_b, WI_lam_b=WI_lam_b, is_bhp=is_bhp, p_bh_b=p_bh_b, 444 bhp_diag=bhp_diag, 445 Q=Q, st=st, dtx=dtx, Up=Up, Ssub=Ssub, 446 ) # fmt: skip 447 448 449def adj_step( 450 tape: Tape, aS1: np.ndarray, aP1: np.ndarray 451) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: 452 """Propagate the sensitivity `(aS1, aP1)` back through the step of `tape`. 453 454 The transpose of the step's tangent: given $ ∂J/∂s^{n+1} $ and 455 $ ∂J/∂p^{n+1} $, returns `(aS, aP, alogK, abhp)` 456 $ = (∂J/∂s^n, ∂J/∂p^n, ∂J/∂\\log K, ∂J/∂p_\\mathrm{bh}) $ -- the third 457 shaped like `K` and being this step's *contribution* (to be summed over 458 the steps, as `adjoint` does), the fourth `(nComp,)`, wrt. this step's 459 controls (`0` at the rate-controlled completions). Each statement of the tangent appears here, in reverse order, 460 transposed: `y = M @ x` as `x̄ += M.T @ ȳ`, `y = a * x` as `x̄ += a * ȳ`, 461 the (symmetric) solve as itself. The comments name the tangent statement 462 being transposed (ref the module docstring). 463 464 Does not modify `aS1`, `aP1`. 465 """ 466 t = tape 467 aS = np.array(aS1, float) # the running adjoint of `dS`, which the loop updates 468 aP1 = np.array(aP1, float) # `dP1` is the output, and feeds `dV` and `dQ` 469 zeros = np.zeros(len(aS)) 470 adfi, adfp, adst, adVm = zeros, zeros, zeros, np.zeros_like(t.V) 471 fp = t.Q.clip(max=0) 472 # fmt: off 473 # Transport equation, sub-steps in reverse 474 for Sj in t.Ssub[::-1]: 475 fw = t.model.fluid.fractional_flow(Sj) 476 dfw_dS = t.model.fluid.dfractional_flow(Sj) 477 adrhs = t.dtx * aS # dS = dS + dtx*drhs 478 adF = -(t.Grad @ adrhs) # drhs = -Grad.T @ dF ... 479 adfp = adfp + fw * adrhs # + dfp*fw 480 adfw = fp * adrhs # + fp*dfw 481 adfi = adfi + adrhs # + dfi 482 aS = aS - t.st * adrhs # - dS*st 483 adst = adst - Sj * adrhs # - Sj*dst 484 adVm = adVm + (t.Up @ fw) * adF # dF = dVm*(Up@fw) ... 485 adfw = adfw + t.Up.T @ (t.V * adF) # + V*(Up@dfw) 486 aS = aS + dfw_dS * adfw # dfw = dfw_dS * dS 487 adV = (t.V != 0) * adVm # dVm = dV * (V != 0) 488 if t.model.ct > 0: # dst = dQ - Grad.T @ dV 489 adQ = adst 490 adV = adV - t.Grad @ adst 491 else: 492 adQ = zeros 493 adQ = adQ + (t.Q < 0) * adfp + (t.Q > 0) * adfi # dfp, dfi = dQ*(Q<0), dQ*(Q>0) 494 495 # Pressure equation 496 adr = adQ # dQ = dr - dw*P1 - bhp_diag*dP1 497 adw = -t.P1 * adQ 498 aP1 = aP1 - t.bhp_diag * adQ 499 aP1 = aP1 + t.Grad.T @ (t.T * adV) # dV = T*(Grad@dP1) + gradP*dT 500 adT = t.gradP * adV 501 adq = t.solve(aP1) # dP1 = solve(dq - dAP) 502 adAP = -adq 503 aP = t.accum * adq # dq = accum*dP + dr 504 adr = adr + adq 505 adT = adT + t.gradP * (t.Grad @ adAP) # dAP = Grad.T@(gradP*dT) + dw*P1 506 adw = adw + t.P1 * adAP 507 adr_b = t.Gb @ adr # dr = Gb.T @ (p_bh_b * dWI_lam ... 508 adWI_lam = t.p_bh_b * adr_b 509 abhp = np.zeros(len(t.is_bhp)) # + WI_lam_b * dp_bh_b) 510 abhp[t.is_bhp] = t.WI_lam_b * adr_b 511 adWI_lam = adWI_lam + t.Gb @ adw # dw = Gb.T @ dWI_lam 512 adMt = t.Gb.T @ (t.WI_b * adWI_lam) # dWI_lam = WI_b * (Gb @ dMt) 513 adMt = adMt + t.dT_dMt.T @ adT # dT = dT_dMt@dMt + dT_dlogK@dlogK 514 alogK = t.dT_dlogK.T @ adT 515 aS = aS + t.dMt_dS * adMt # dMt = dMt_dS * dS 516 # fmt: on 517 return aS, aP, alogK.reshape(t.model.K.shape), abhp 518 519 520def adjoint( 521 model: ResSim, 522 dt: float, 523 SS: np.ndarray, 524 PP: np.ndarray, 525 dJ_dSS: np.ndarray, 526 dJ_dPP: np.ndarray | None = None, 527) -> Gradient: 528 """The gradient of $ J(S, P) $ wrt. `S0`, `P0`, $ \\log K $ and the BHP controls, by the adjoint sweep. 529 530 Seeded by the partials of the objective wrt. the *stored* trajectory, 531 `dJ_dSS[k]` $ = ∂J/∂S_k $, `dJ_dPP[k]` $ = ∂J/∂P_k $ (shaped like `SS`, 532 `PP`; the latter defaults to `0`), ref the module docstring. Sweeps 533 backwards from the final time, re-linearizing each step (`linearize`) on 534 the way, so the trajectory `(SS, PP)` of `sim(dt, ...)` is all it needs. 535 536 The cost is about that of a `sim` (one forward step and one factorization 537 per step, ref `linearize`), independently of the number of parameters -- 538 which is the point of an adjoint. 539 """ 540 nSteps = len(SS) - 1 541 aS = np.array(dJ_dSS[-1], float) 542 aP = np.zeros(model.Nxy) if dJ_dPP is None else np.array(dJ_dPP[-1], float) 543 alogK = np.zeros(model.K.shape) 544 abhp = np.zeros((model.nComp, nSteps)) 545 for k in reversed(range(nSteps)): 546 tape = linearize(model, dt, SS[k], PP[k], k) 547 aS, aP, aK, abhp[:, k] = adj_step(tape, aS, aP) 548 alogK += aK 549 aS += dJ_dSS[k] 550 if dJ_dPP is not None: 551 aP += dJ_dPP[k] 552 return Gradient(aS, aP, alogK, abhp)
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.
211def face_operators(model: ResSim) -> tuple: 212 """The interior faces of the grid, and the sparse operators on them. 213 214 Returns `(lo, hi, Grad, Sum, g)`: 215 216 - `lo`, `hi`: the flat indices of the two cells each face separates 217 (`lo` has the smaller index), shape `(nF,)`. The x-faces come first 218 (`(Nx-1) * Ny` of them, in C-order), then the y-faces (`Nx * (Ny-1)`) 219 -- less those of inactive cells (`ResSim.active`), which are omitted. 220 - `Grad`: `(nF, Nxy)`, `(Grad @ p)[f] = p[lo] - p[hi]`, as `TPFA` computes 221 the fluxes; `Grad.T` is the divergence (high faces minus low faces). 222 - `Sum`: `(nF, 2*Nxy)`, summing over the face's two cells the *directional* 223 component of a `(2, Nx, Ny)`-shaped field (flattened): the x-component 224 for x-faces, the y-component for y-faces. This is how the transmissibility 225 harmonically averages the permeabilities. 226 - `g`: `(nF,)`, the geometric factor of the transmissibilities, 227 $ 2 C h_y / h_x $ resp. $ 2 C h_x / h_y $, such that `T = g / (Sum @ (1/KM))`. 228 229 >>> from minires import ResSim 230 >>> lo, hi, Grad, Sum, g = face_operators(ResSim(Nx=3, Ny=2)) 231 >>> lo, hi # the 4 x-faces, then the 3 y-faces 232 (array([0, 1, 2, 3, 0, 2, 4]), array([2, 3, 4, 5, 1, 3, 5])) 233 >>> (Grad.T @ np.ones(7)).astype(int) # #high faces - #low faces, per cell 234 array([ 2, 0, 1, -1, 0, -2]) 235 """ 236 N = model.Nxy 237 idx = np.arange(N).reshape(model.shape) 238 act = model.active 239 fx = (act[:-1, :] & act[1:, :]).ravel() # faces between active cells 240 fy = (act[:, :-1] & act[:, 1:]).ravel() 241 lo = np.concatenate([idx[:-1, :].ravel()[fx], idx[:, :-1].ravel()[fy]]) 242 hi = np.concatenate([idx[1:, :].ravel()[fx], idx[:, 1:].ravel()[fy]]) 243 nF = len(lo) 244 nFx = int(fx.sum()) 245 ff = np.r_[np.arange(nF), np.arange(nF)] 246 ones = np.ones(nF) 247 Grad = sparse.csr_matrix((np.r_[ones, -ones], (ff, np.r_[lo, hi])), shape=(nF, N)) 248 comp = np.r_[np.zeros(nFx, int), np.full(nF - nFx, N)] # offset into 2nd component 249 Sum = sparse.csr_matrix( 250 (np.r_[ones, ones], (ff, np.r_[lo + comp, hi + comp])), shape=(nF, 2 * N) 251 ) 252 C = model.cdarcy 253 g = 2 * C * np.r_[np.full(nFx, model.hy / model.hx), np.full(nF - nFx, model.hx / model.hy)] 254 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])
257@dataclass 258class Tape(AlignedRepr): 259 """The linearization of one step of `minires.ResSim.time_stepper`. 260 261 Produced by `linearize`, consumed by `adj_step`. Holds the state it was 262 taken about, the step's result, and the coefficients -- sparse matrices 263 and vectors -- of the linear statements of the step's tangent, evaluated 264 at that state. 265 """ 266 267 __repr__ = AlignedRepr.__repr__ 268 269 model: ResSim 270 """The model whose step this linearizes (for `fluid`, `ct`, ...).""" 271 dt: float 272 """The time step.""" 273 k: int 274 """The time index (of the well controls).""" 275 S: np.ndarray 276 """Saturation at the start of the step, `(Nxy,)`.""" 277 P: np.ndarray 278 """Pressure at the start of the step, `(Nxy,)`.""" 279 S1: np.ndarray 280 """Saturation at the end of the step -- the recomputed forward result.""" 281 P1: np.ndarray 282 """Pressure at the end of the step -- the recomputed forward result.""" 283 284 # Pressure equation 285 Grad: Any 286 """`(nF, Nxy)` face difference operator, ref `face_operators`. `Grad.T` is the divergence.""" 287 dMt_dS: np.ndarray 288 """`(Nxy,)` derivative of the total mobility $ λ_t $ wrt. `S`.""" 289 dT_dMt: Any 290 """`(nF, Nxy)` Jacobian of the transmissibilities wrt. the cell mobilities: 291 $ ∂T_f / ∂λ_i = T_f^2 / (g_f \\, K_i \\, λ_i^2) $ for each of the face's 292 two cells (harmonic averaging), `0` otherwise.""" 293 dT_dlogK: Any 294 """`(nF, 2*Nxy)` Jacobian of the transmissibilities wrt. $ \\log K $ (flattened 295 like `K`): $ ∂T_f / ∂\\log K_c = T_f^2 / (g_f \\, K_c \\, λ_c) $ for each of the 296 face's two cells, in the face's direction, `0` otherwise.""" 297 T: np.ndarray 298 """`(nF,)` transmissibilities, $ T = g / Σ_\\mathrm{cells} 1/(K λ_t) $.""" 299 gradP: np.ndarray 300 """`(nF,)` pressure differences, `Grad @ P1`; the fluxes are `V = T * gradP`.""" 301 V: np.ndarray 302 """`(nF,)` fluxes through the interior faces (positive from `lo` to `hi`).""" 303 accum: np.ndarray 304 """`(Nxy,)` accumulation coefficient, $ φ c_t h^2 / Δt $ (`0` if `ct == 0`).""" 305 solve: Callable 306 """`x ↦ A⁻¹ x`, by the LU factorization of the (symmetric) pressure matrix.""" 307 308 # Well model of the BHP-controlled completions (empty if none) 309 Gb: Any 310 """`(nBHP, Nxy)` gathers cell values at the BHP-controlled completions; `Gb.T` scatters.""" 311 WI_b: np.ndarray 312 """`(nBHP,)` their well indices.""" 313 WI_lam_b: np.ndarray 314 """`(nBHP,)` their $ WI λ_t $, the coefficient of $ p_\\mathrm{bh} $ in the right-hand side.""" 315 is_bhp: np.ndarray 316 """`(nComp,)` boolean mask of the BHP-controlled completions.""" 317 p_bh_b: np.ndarray 318 """`(nBHP,)` their bottom-hole pressures.""" 319 bhp_diag: np.ndarray 320 """`(Nxy,)` their $ WI λ_t $, scattered onto the cells (`0` elsewhere).""" 321 322 # Transport equation 323 Q: np.ndarray 324 """`(Nxy,)` total well flux per cell, signed (`realize_bhp` included).""" 325 st: np.ndarray 326 """`(Nxy,)` storage rate, `Q - Grad.T @ V` (`0` if `ct == 0`).""" 327 dtx: np.ndarray 328 """`(Nxy,)` sub-step over pore volume, `dt / nT / pv`.""" 329 Up: Any 330 """`(nF, Nxy)` upwind selector: `(Up @ f)[face]` is `f` of the face's upwind cell.""" 331 Ssub: np.ndarray 332 """`(nT, Nxy)` saturation at the start of each transport sub-step.""" 333 334 @property 335 def nT(self) -> int: 336 """Number of transport sub-steps.""" 337 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.
(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.
(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.
340def linearize( 341 model: ResSim, 342 dt: float, 343 S: np.ndarray, 344 P: np.ndarray | None, 345 k: int = 0, 346) -> Tape: 347 """Recompute the step of `time_stepper` from `(S, P)` at time `k`; return its `Tape`. 348 349 The recomputation calls the very methods of the forward model (so the 350 result, `Tape.S1`/`Tape.P1`, is that of `sim`, to the solver tolerance), 351 except that it records the transport sub-steps, and does not write the 352 `actual_rates`/`actual_bhp` reports. Then the coefficients of the 353 linearization are evaluated at that state. 354 355 The cost is about that of a forward step plus a factorization of the 356 pressure matrix (held in `Tape.solve`, for the adjoint solve). 357 358 .. warning:: This mutates the model, exactly as a step of `sim` does. 359 360 Because it *is* one: `assemble_wells` (hence `well_controls`), 361 `pressure_step` and `realize_bhp` are called, and they write the 362 source field `_Q`, the bundle `_wells_now` and -- if `cached_precond` 363 -- the factorization cache `_pLU`, all as of the step linearized. So 364 after `sim`, then `adjoint` (which linearizes the steps in *reverse*), 365 these hold the values of the first step rather than the last. None of 366 it is consequential: the next step (or `linearize`) overwrites them, 367 and `_pLU` is a mere preconditioner (ref `_solve_pressure`). The 368 well reports, `actual_rates`/`actual_bhp`, are *not* written, so they 369 remain those of the `sim` that produced the trajectory. Nothing else 370 is touched: `K`, `por`, the well specifications are read only. 371 """ 372 N = model.Nxy 373 S = np.asarray(S, float).ravel() 374 P = np.zeros(N) if P is None else np.asarray(P, float).ravel() 375 lo, hi, Grad, Sum, g = face_operators(model) 376 nF = len(lo) 377 378 # Forward step, recomputed. Mirrors `time_stepper` (minus the reporting) 379 model.assemble_wells(S, P, k) 380 model._validate() 381 P1, VV = model.pressure_step(S, P, dt) 382 model.realize_bhp(P1) 383 wls = model._wells_now 384 Q = model._Q # total well flux, the BHP wells' rates now realized 385 # ... and `saturation_step_upwind`, recording the sub-steps 386 A_up = model.upwind_diff(VV) 387 pv = model.pore_volume() 388 fi = Q.clip(min=0) 389 st = model.storage_rate(VV) 390 nT = max(1, int(np.ceil(dt * model.estimate_1CFL(pv, VV, fi)))) 391 dtx = dt / nT / pv 392 B = model._spdiags(dtx, 0) @ A_up 393 Ssub = np.zeros((nT, N)) 394 Sj = S 395 for j in range(nT): 396 Ssub[j] = Sj 397 fw = model.fluid.fractional_flow(Sj) 398 Sj = Sj + (B @ fw + (fi - Sj * st) * dtx) 399 S1 = Sj 400 401 # Coefficients of the linearization, at that state 402 # -- mobilities and transmissibilities 403 Mw, Mo = model.fluid.RelPerm(S) 404 dMw, dMo = model.fluid.dRelPerm(S) 405 Mt = Mw + Mo 406 Kflat = model.K.reshape(-1) 407 KM = model.K.reshape(2, N) * Mt # K λ_t, both components, flattened 408 KM = KM.reshape(-1) 409 T = g / (Sum @ (1 / KM)) 410 Dup = sparse.vstack([sparse.eye(N), sparse.eye(N)]) # cell field → both components 411 dT_dKM = sparse.diags(T**2 / g) @ Sum # ... @ diag(1/KM²), below 412 dT_dMt = dT_dKM @ sparse.diags(Kflat / KM**2) @ Dup 413 dT_dlogK = dT_dKM @ sparse.diags(1 / KM) # since ∂(Kλ)/∂log K = Kλ 414 gradP = Grad @ P1 415 V = T * gradP 416 # -- the BHP wells 417 is_bhp = np.isfinite(wls["WI_lam"]) 418 nB = int(is_bhp.sum()) 419 inds_b = wls["inds"][is_bhp] 420 Gb = sparse.csr_matrix((np.ones(nB), (np.arange(nB), inds_b)), shape=(nB, N)) 421 WI = model.wells.WI 422 WI_b = WI[is_bhp] if WI is not None else np.zeros(0) 423 WI_lam_b = wls["WI_lam"][is_bhp] 424 p_bh_b = wls["p_bh"][is_bhp] 425 bhp_diag = wls["bhp_diag"] 426 # -- the pressure system (as `TPFA` assembles it, incl. its pin and the 427 # identity rows of the inactive cells, whose `accum` is thereby `1`) 428 accum = pv * model.ct / dt if model.ct > 0 else np.zeros(N) 429 accum = np.where(model.active.ravel(), accum, 1.0) 430 diag = accum + bhp_diag 431 if model.ct == 0 and not bhp_diag.any(): 432 diag[model._pin] += np.sum(model.K.reshape(2, -1)[:, model._pin]) 433 A = Grad.T @ sparse.diags(T) @ Grad + sparse.diags(diag) 434 solve = splu(A.tocsc(), permc_spec="MMD_AT_PLUS_A").solve 435 # -- the upwind directions 436 Up = sparse.csr_matrix( 437 (np.ones(nF), (np.arange(nF), np.where(V >= 0, lo, hi))), shape=(nF, N) 438 ) 439 440 return Tape( 441 model=model, dt=dt, k=k, S=S, P=P, S1=S1, P1=P1, 442 Grad=Grad, dMt_dS=dMw + dMo, dT_dMt=dT_dMt, dT_dlogK=dT_dlogK, 443 T=T, gradP=gradP, V=V, accum=accum, solve=solve, 444 Gb=Gb, WI_b=WI_b, WI_lam_b=WI_lam_b, is_bhp=is_bhp, p_bh_b=p_bh_b, 445 bhp_diag=bhp_diag, 446 Q=Q, st=st, dtx=dtx, Up=Up, Ssub=Ssub, 447 ) # 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.
450def adj_step( 451 tape: Tape, aS1: np.ndarray, aP1: np.ndarray 452) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: 453 """Propagate the sensitivity `(aS1, aP1)` back through the step of `tape`. 454 455 The transpose of the step's tangent: given $ ∂J/∂s^{n+1} $ and 456 $ ∂J/∂p^{n+1} $, returns `(aS, aP, alogK, abhp)` 457 $ = (∂J/∂s^n, ∂J/∂p^n, ∂J/∂\\log K, ∂J/∂p_\\mathrm{bh}) $ -- the third 458 shaped like `K` and being this step's *contribution* (to be summed over 459 the steps, as `adjoint` does), the fourth `(nComp,)`, wrt. this step's 460 controls (`0` at the rate-controlled completions). Each statement of the tangent appears here, in reverse order, 461 transposed: `y = M @ x` as `x̄ += M.T @ ȳ`, `y = a * x` as `x̄ += a * ȳ`, 462 the (symmetric) solve as itself. The comments name the tangent statement 463 being transposed (ref the module docstring). 464 465 Does not modify `aS1`, `aP1`. 466 """ 467 t = tape 468 aS = np.array(aS1, float) # the running adjoint of `dS`, which the loop updates 469 aP1 = np.array(aP1, float) # `dP1` is the output, and feeds `dV` and `dQ` 470 zeros = np.zeros(len(aS)) 471 adfi, adfp, adst, adVm = zeros, zeros, zeros, np.zeros_like(t.V) 472 fp = t.Q.clip(max=0) 473 # fmt: off 474 # Transport equation, sub-steps in reverse 475 for Sj in t.Ssub[::-1]: 476 fw = t.model.fluid.fractional_flow(Sj) 477 dfw_dS = t.model.fluid.dfractional_flow(Sj) 478 adrhs = t.dtx * aS # dS = dS + dtx*drhs 479 adF = -(t.Grad @ adrhs) # drhs = -Grad.T @ dF ... 480 adfp = adfp + fw * adrhs # + dfp*fw 481 adfw = fp * adrhs # + fp*dfw 482 adfi = adfi + adrhs # + dfi 483 aS = aS - t.st * adrhs # - dS*st 484 adst = adst - Sj * adrhs # - Sj*dst 485 adVm = adVm + (t.Up @ fw) * adF # dF = dVm*(Up@fw) ... 486 adfw = adfw + t.Up.T @ (t.V * adF) # + V*(Up@dfw) 487 aS = aS + dfw_dS * adfw # dfw = dfw_dS * dS 488 adV = (t.V != 0) * adVm # dVm = dV * (V != 0) 489 if t.model.ct > 0: # dst = dQ - Grad.T @ dV 490 adQ = adst 491 adV = adV - t.Grad @ adst 492 else: 493 adQ = zeros 494 adQ = adQ + (t.Q < 0) * adfp + (t.Q > 0) * adfi # dfp, dfi = dQ*(Q<0), dQ*(Q>0) 495 496 # Pressure equation 497 adr = adQ # dQ = dr - dw*P1 - bhp_diag*dP1 498 adw = -t.P1 * adQ 499 aP1 = aP1 - t.bhp_diag * adQ 500 aP1 = aP1 + t.Grad.T @ (t.T * adV) # dV = T*(Grad@dP1) + gradP*dT 501 adT = t.gradP * adV 502 adq = t.solve(aP1) # dP1 = solve(dq - dAP) 503 adAP = -adq 504 aP = t.accum * adq # dq = accum*dP + dr 505 adr = adr + adq 506 adT = adT + t.gradP * (t.Grad @ adAP) # dAP = Grad.T@(gradP*dT) + dw*P1 507 adw = adw + t.P1 * adAP 508 adr_b = t.Gb @ adr # dr = Gb.T @ (p_bh_b * dWI_lam ... 509 adWI_lam = t.p_bh_b * adr_b 510 abhp = np.zeros(len(t.is_bhp)) # + WI_lam_b * dp_bh_b) 511 abhp[t.is_bhp] = t.WI_lam_b * adr_b 512 adWI_lam = adWI_lam + t.Gb @ adw # dw = Gb.T @ dWI_lam 513 adMt = t.Gb.T @ (t.WI_b * adWI_lam) # dWI_lam = WI_b * (Gb @ dMt) 514 adMt = adMt + t.dT_dMt.T @ adT # dT = dT_dMt@dMt + dT_dlogK@dlogK 515 alogK = t.dT_dlogK.T @ adT 516 aS = aS + t.dMt_dS * adMt # dMt = dMt_dS * dS 517 # fmt: on 518 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.
521def adjoint( 522 model: ResSim, 523 dt: float, 524 SS: np.ndarray, 525 PP: np.ndarray, 526 dJ_dSS: np.ndarray, 527 dJ_dPP: np.ndarray | None = None, 528) -> Gradient: 529 """The gradient of $ J(S, P) $ wrt. `S0`, `P0`, $ \\log K $ and the BHP controls, by the adjoint sweep. 530 531 Seeded by the partials of the objective wrt. the *stored* trajectory, 532 `dJ_dSS[k]` $ = ∂J/∂S_k $, `dJ_dPP[k]` $ = ∂J/∂P_k $ (shaped like `SS`, 533 `PP`; the latter defaults to `0`), ref the module docstring. Sweeps 534 backwards from the final time, re-linearizing each step (`linearize`) on 535 the way, so the trajectory `(SS, PP)` of `sim(dt, ...)` is all it needs. 536 537 The cost is about that of a `sim` (one forward step and one factorization 538 per step, ref `linearize`), independently of the number of parameters -- 539 which is the point of an adjoint. 540 """ 541 nSteps = len(SS) - 1 542 aS = np.array(dJ_dSS[-1], float) 543 aP = np.zeros(model.Nxy) if dJ_dPP is None else np.array(dJ_dPP[-1], float) 544 alogK = np.zeros(model.K.shape) 545 abhp = np.zeros((model.nComp, nSteps)) 546 for k in reversed(range(nSteps)): 547 tape = linearize(model, dt, SS[k], PP[k], k) 548 aS, aP, aK, abhp[:, k] = adj_step(tape, aS, aP) 549 alogK += aK 550 aS += dJ_dSS[k] 551 if dJ_dPP is not None: 552 aP += dJ_dPP[k] 553 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.