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