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 -1 alone, 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 == 0 and no well is on BHP control, $ p^{n+1} $ does not depend on $ p^n $ at all, so the gradient with respect to P0 is then 0.
  • The BHP controls, $ p_\mathrm{bh} $ of the BHP-controlled completions at each step (whatever well_controls returned, so, by default, the spec Wells.bhp at 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 (TPFA adds $ \sum K_{00} $ to the first diagonal entry) fixes $ p_0 = 0 $ whatever the value, so the derivative is exactly zero; and the well index Wells.WI is a stored parameter (minires.wells.peaceman_WI evaluates it once, from the K of that moment, and does not track K thereafter), 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 in sim. (The rates would be as easy as the BHPs: they enter Q directly, so their adjoint is adQ gathered at the rate-controlled completions. But with ct == 0 they must sum to zero, which the gradient does not know.)
  • Not the controls' dependence on the state: well_controls is 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, ref minires.ResSim.estimate_1CFL), the upwind directions and the signs of the well fluxes (the clips of upwind_diff). They are piecewise constant, so this is the derivative almost everywhere; at a switch (a face with exactly zero flux, nT on 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)
class Gradient(typing.NamedTuple):
196class Gradient(NamedTuple):
197    """The gradient of an objective, as returned by `adjoint`."""
198
199    S0: np.ndarray
200    """W.r.t. the initial saturation, `(Nxy,)`."""
201    P0: np.ndarray
202    """W.r.t. the initial pressure, `(Nxy,)`. Zero unless `ct > 0` or a well is on BHP."""
203    logK: np.ndarray
204    """W.r.t. $ \\log K $, shaped like `K`: `(2, Nx, Ny)`. Sum over axis `0` if isotropic."""
205    bhp: np.ndarray
206    """W.r.t. the BHP controls, `(nComp, nSteps)`: entry `[i, k]` is for
207    completion `i`'s control at step `k` (`0` where rate-controlled). Sum over
208    axis `1` if the control is constant in time."""

The gradient of an objective, as returned by adjoint.

Gradient( S0: numpy.ndarray, P0: numpy.ndarray, logK: numpy.ndarray, bhp: numpy.ndarray)

Create new instance of Gradient(S0, P0, logK, bhp)

S0: numpy.ndarray

W.r.t. the initial saturation, (Nxy,).

P0: numpy.ndarray

W.r.t. the initial pressure, (Nxy,). Zero unless ct > 0 or a well is on BHP.

logK: numpy.ndarray

W.r.t. $ \log K $, shaped like K: (2, Nx, Ny). Sum over axis 0 if isotropic.

bhp: numpy.ndarray

W.r.t. the BHP controls, (nComp, nSteps): entry [i, k] is for completion i's control at step k (0 where rate-controlled). Sum over axis 1 if the control is constant in time.

def face_operators(model: minires.ResSim) -> tuple:
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 (lo has the smaller index), shape (nF,). The x-faces come first ((Nx-1) * Ny of 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], as TPFA computes the fluxes; Grad.T is 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 that T = g / (Sum @ (1/KM)).
>>> from minires import ResSim
>>> lo, hi, Grad, Sum, g = face_operators(ResSim(Nx=3, Ny=2))
>>> lo, hi   # the 4 x-faces, then the 3 y-faces
(array([0, 1, 2, 3, 0, 2, 4]), array([2, 3, 4, 5, 1, 3, 5]))
>>> (Grad.T @ np.ones(7)).astype(int)   # #high faces - #low faces, per cell
array([ 2,  0,  1, -1,  0, -2])
@dataclass
class Tape(minires._repr.AlignedRepr):
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.

Tape( model: minires.ResSim, dt: float, k: int, S: numpy.ndarray, P: numpy.ndarray, S1: numpy.ndarray, P1: numpy.ndarray, Grad: Any, dMt_dS: numpy.ndarray, dT_dMt: Any, dT_dlogK: Any, T: numpy.ndarray, gradP: numpy.ndarray, V: numpy.ndarray, accum: numpy.ndarray, solve: Callable, Gb: Any, WI_b: numpy.ndarray, WI_lam_b: numpy.ndarray, is_bhp: numpy.ndarray, p_bh_b: numpy.ndarray, bhp_diag: numpy.ndarray, Q: numpy.ndarray, st: numpy.ndarray, dtx: numpy.ndarray, Up: Any, Ssub: numpy.ndarray)

The model whose step this linearizes (for fluid, ct, ...).

dt: float

The time step.

k: int

The time index (of the well controls).

S: numpy.ndarray

Saturation at the start of the step, (Nxy,).

P: numpy.ndarray

Pressure at the start of the step, (Nxy,).

S1: numpy.ndarray

Saturation at the end of the step -- the recomputed forward result.

P1: numpy.ndarray

Pressure at the end of the step -- the recomputed forward result.

Grad: Any

(nF, Nxy) face difference operator, ref face_operators. Grad.T is the divergence.

dMt_dS: numpy.ndarray

(Nxy,) derivative of the total mobility $ λ_t $ wrt. S.

dT_dMt: Any

(nF, Nxy) Jacobian of the transmissibilities wrt. the cell mobilities: $ ∂T_f / ∂λ_i = T_f^2 / (g_f \, K_i \, λ_i^2) $ for each of the face's two cells (harmonic averaging), 0 otherwise.

dT_dlogK: Any

(nF, 2*Nxy) Jacobian of the transmissibilities wrt. $ \log K $ (flattened like K): $ ∂T_f / ∂\log K_c = T_f^2 / (g_f \, K_c \, λ_c) $ for each of the face's two cells, in the face's direction, 0 otherwise.

T: numpy.ndarray

(nF,) transmissibilities, $ T = g / Σ_\mathrm{cells} 1/(K λ_t) $.

gradP: numpy.ndarray

(nF,) pressure differences, Grad @ P1; the fluxes are V = T * gradP.

V: numpy.ndarray

(nF,) fluxes through the interior faces (positive from lo to hi).

accum: numpy.ndarray

(Nxy,) accumulation coefficient, $ φ c_t h^2 / Δt $ (0 if ct == 0).

solve: Callable

x ↦ A⁻¹ x, by the LU factorization of the (symmetric) pressure matrix.

Gb: Any

(nBHP, Nxy) gathers cell values at the BHP-controlled completions; Gb.T scatters.

WI_b: numpy.ndarray

(nBHP,) their well indices.

WI_lam_b: numpy.ndarray

(nBHP,) their $ WI λ_t $, the coefficient of $ p_\mathrm{bh} $ in the right-hand side.

is_bhp: numpy.ndarray

(nComp,) boolean mask of the BHP-controlled completions.

p_bh_b: numpy.ndarray

(nBHP,) their bottom-hole pressures.

bhp_diag: numpy.ndarray

(Nxy,) their $ WI λ_t $, scattered onto the cells (0 elsewhere).

Q: numpy.ndarray

(Nxy,) total well flux per cell, signed (realize_bhp included).

st: numpy.ndarray

(Nxy,) storage rate, Q - Grad.T @ V (0 if ct == 0).

dtx: numpy.ndarray

(Nxy,) sub-step over pore volume, dt / nT / pv.

Up: Any

(nF, Nxy) upwind selector: (Up @ f)[face] is f of the face's upwind cell.

Ssub: numpy.ndarray

(nT, Nxy) saturation at the start of each transport sub-step.

nT: int
334    @property
335    def nT(self) -> int:
336        """Number of transport sub-steps."""
337        return len(self.Ssub)

Number of transport sub-steps.

def linearize( model: minires.ResSim, dt: float, S: numpy.ndarray, P: numpy.ndarray | None, k: int = 0) -> Tape:
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.

def adj_step( tape: Tape, aS1: numpy.ndarray, aP1: numpy.ndarray) -> tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray, numpy.ndarray]:
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.

def adjoint( model: minires.ResSim, dt: float, SS: numpy.ndarray, PP: numpy.ndarray, dJ_dSS: numpy.ndarray, dJ_dPP: numpy.ndarray | None = None) -> Gradient:
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.