Coverage for src/gwtransport/_radial_asr_reuse.py: 100%
157 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-07 21:17 +0000
1r"""Reused-propagator-matrix acceleration of the multi-cycle radial advection-dispersion engine.
3For a signed-flow schedule with ``K`` flow reversals, the per-reversal grid-free composition performs one
4de Hoog inversion of the interior resolvent per intermediate reversal to hand the resident field across
5the reversal -- the ``O(n_quad^2)`` per-reversal
6cost. That hand-off is a *linear* operator on the resident field: the propagated field is
7``f_out(r_i) = L^{-1}_s[ sum_j Ghat(r_i, r_j; s) w(r_j) f(r_j) ](tau)``, i.e. ``f_out = P @ f`` with the
8propagator matrix
10``P_{dir,tau}[i, j] = L^{-1}_s[ Ghat(r_i, r_j; s) w_dir(r_j) ](tau)``.
12For a schedule with repeated phase volumes (periodic ASR / SWIW) the same ``(direction, tau)`` recurs every
13cycle, so ``P`` is identical across cycles. This module builds each distinct ``P`` *once* (one batched de
14Hoog inversion over all ``n_quad^2`` source/output entries, reusing the exact Airy interior Green's
15function) and reuses it as a bounded matrix-multiply at every reversal -- the special-function + de Hoog
16cost becomes ``O(number of distinct phase volumes)`` instead of ``O(K)``.
18The matrix is assembled on the Bromwich contour (``Re s > 0``), where the Airy exponent is growing-real
19and tames the divergent injection gauge ``e^{+r/alpha_L}`` (carried in log space by
20:func:`gwtransport._radial_asr_kernels.assemble_airy_resolvent`), so ``P`` is well-conditioned at any
21Peclet (``|P| ~ O(1)``, the same bounded physical operator the per-reversal engine applies). ``P @ f``
22equals the per-reversal engine's ``_propagate(f)`` to the de Hoog floor; the only difference is the de Hoog
23Pade acceleration's mild nonlinearity (``invert-then-sum`` vs ``sum-then-invert``), which vanishes as the
24de Hoog ``n_terms``/``tol`` tighten.
26This file is part of gwtransport which is released under AGPL-3.0 license.
27See the ./LICENSE file or go to https://github.com/gwtransport/gwtransport/blob/main/LICENSE for full license details.
28"""
30import numpy as np
31import numpy.typing as npt
32from scipy.special import ive, kve
34from gwtransport._radial_asr_compose import _DEHOOG_TERMS, _SCALE_MARGIN, _fr_step_response
35from gwtransport._radial_asr_dehoog import dehoog_inverse
36from gwtransport._radial_asr_kernels import (
37 _integrate_logderiv,
38 _resolvent_airy_pieces,
39 assemble_airy_resolvent,
40)
42_INJECTION = "injection"
45def _phase_slices(flow: npt.NDArray[np.floating]) -> list[tuple[int, slice]]:
46 """Group the schedule into maximal one-signed phases.
48 Returns
49 -------
50 list of (sign, slice)
51 ``sign`` in ``{+1, -1, 0}`` (injection / extraction / rest) and the contiguous bin slice.
52 """
53 signs = np.sign(flow).astype(int)
54 edges = np.flatnonzero(np.diff(signs) != 0) + 1
55 starts = np.concatenate(([0], edges))
56 stops = np.concatenate((edges, [len(flow)]))
57 return [(int(signs[a]), slice(a, b)) for a, b in zip(starts, stops, strict=True)]
60def _field_grid(
61 flow: npt.NDArray[np.floating],
62 dt_days: npt.NDArray[np.floating],
63 c_geo: float,
64 r_w: float,
65 alpha_l: float,
66 molecular_diffusivity: float,
67 n_quad: int,
68) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating], npt.NDArray[np.floating]]:
69 """Radial Gauss-Legendre quadrature grid spanning the plume front plus its dispersive tail.
71 The grid is kept **tight**: it covers the advective front radius ``r_front`` (where the plume
72 volume ``V(r) = peak net injected volume``) plus a margin of breakthrough widths (radial std
73 ``~ sqrt(alpha_L r_front)``) and a molecular reach, and no further. An over-extended grid would
74 push the Peclet so high that the divergent (injection) interior Green's function, which grows as
75 ``e^{+r/alpha_L}`` before its ``e^{-r/alpha_L}`` Sturm-Liouville weight tames the product, overflows
76 double precision; the resident profile is ``~0`` beyond the tail anyway (verified by mass
77 conservation). Retardation rescales the clock, not the radius, so it does not enter ``r_max``.
79 Returns
80 -------
81 r_nodes : ndarray
82 Radial nodes (m), shape ``(n_quad,)``.
83 v_nodes : ndarray
84 Volume coordinate ``V(r) = c_geo (r^2 - r_w^2)`` at the nodes.
85 dr_weights : ndarray
86 Gauss-Legendre weights in ``r`` (so ``int g dr ~ sum g(r_k) dr_weights_k``).
87 """
88 net_volume = np.concatenate(([0.0], np.cumsum(flow * dt_days)))
89 peak_volume = max(float(net_volume.max()), 0.0)
90 r_front = np.sqrt(r_w**2 + peak_volume / c_geo) # advective plume-front radius
91 total_time = float(np.sum(dt_days))
92 reach = 12.0 * np.sqrt(alpha_l * r_front + alpha_l**2) + 6.0 * np.sqrt(molecular_diffusivity * total_time)
93 r_max = r_front + reach + r_w
94 nodes, weights = np.polynomial.legendre.leggauss(n_quad)
95 r_nodes = 0.5 * (r_max - r_w) * (nodes + 1.0) + r_w
96 dr_weights = 0.5 * (r_max - r_w) * weights
97 v_nodes = c_geo * (r_nodes**2 - r_w**2)
98 return r_nodes, v_nodes, dr_weights
101def _airy_propagator_matrix(
102 direction: str,
103 tau: float,
104 r_nodes: npt.NDArray[np.floating],
105 dr_weights: npt.NDArray[np.floating],
106 *,
107 r_w: float,
108 alpha_l: float,
109 c_geo: float,
110 flow_scale: float,
111 n_terms: int,
112 tol: float,
113) -> npt.NDArray[np.floating]:
114 r"""Field-propagator matrix ``P[i, j]`` for one constant-Q phase (``D_m = 0`` Airy branch).
116 ``f_out = P @ f`` with ``P[i, j] = L^{-1}_s[ Ghat(r_i, r_j; s) w(r_j) ](tau)`` -- the same Airy interior
117 resolvent and flushed-volume clock as the per-reversal ``_propagate``, assembled for
118 every output/source pair and inverted in a single batched de Hoog pass instead of contracting against a
119 specific field. The Airy pieces are evaluated once on the grid per de Hoog node (as in ``_propagate``)
120 and every output row is assembled by prefix selection (``O(n_quad)`` Airy evaluations).
122 The Sturm-Liouville source weight is ``w(r') = (2 c_geo r'/alpha_L) e^{-gauge_sign r'/alpha_L} dr'``
123 (``gauge_sign = +1`` injection / ``-1`` extraction); the flushed-volume clock is set by ``s =
124 flow_scale * p`` with ``a0 = flow_scale/(2 c_geo)`` so ``beta = 2 c_geo p/alpha_L`` is flow-magnitude
125 independent (the autonomy of the S-clock).
127 Returns
128 -------
129 ndarray
130 Propagator matrix ``P``, shape ``(n_quad, n_quad)``.
131 """
132 a0 = flow_scale / (2.0 * c_geo)
133 gauge_sign = 1.0 if direction == _INJECTION else -1.0
134 n = r_nodes.size
135 # Split the Sturm-Liouville source weight (2 c_geo r'/alpha_L) e^{-gauge_sign r'/alpha_L} dr' into its
136 # non-exponential prefactor and its LOG exponent: the exponent is folded into assemble_airy_resolvent
137 # so the divergent gauge cancels before np.exp (overflow-safe at any Peclet). Applying the exponential
138 # weight afterwards would overflow -- the injection assemble exponent and the extraction weight each blow
139 # up past r/alpha_L ~ 700 and meet as Inf * 0 = NaN.
140 sl_prefactor = (2.0 * c_geo * r_nodes / alpha_l) * dr_weights
141 sl_log = (-gauge_sign * r_nodes / alpha_l)[None, :]
142 mu = c_geo * ((float(r_nodes.max()) + alpha_l) ** 2 + alpha_l**2 - r_w**2)
143 scaling = _SCALE_MARGIN * max(mu, tau)
144 below = r_nodes[None, :] < r_nodes[:, None] # below[i, j] = r_j < r_i
146 def f_hat(p: npt.NDArray[np.complexfloating]) -> npt.NDArray[np.complexfloating]:
147 s = (flow_scale * p).reshape(-1, 1)
148 grid_p = _resolvent_airy_pieces(s, r_nodes.reshape(1, -1), alpha_l, a0, gauge_sign)
149 piece_w = _resolvent_airy_pieces(s, np.array([[r_w]]), alpha_l, a0, gauge_sign)
150 ghat = np.empty((p.size, n, n), dtype=complex)
151 for i in range(n):
152 # r_< / r_> selection at output i (the SL kernel is symmetric; radii enter only via r_i + r_j).
153 mask = below[i][None, :]
154 piece_a = {f: np.where(mask, grid_p[f], grid_p[f][:, i : i + 1]) for f in grid_p}
155 piece_b = {f: np.where(mask, grid_p[f][:, i : i + 1], grid_p[f]) for f in grid_p}
156 r_sum = (r_nodes[i] + r_nodes)[None, :]
157 ghat[:, i, :] = assemble_airy_resolvent(
158 piece_a, piece_b, piece_w, r_sum, alpha_l, gauge_sign, source_log_weight=sl_log
159 )
160 ghat *= sl_prefactor[None, None, :] # non-exponential SL prefactor (the gauge is already folded in)
161 return ghat
163 return dehoog_inverse(f_hat=f_hat, t=tau, n_terms=n_terms, scaling=scaling, tol=tol)
166def _riccati_propagator_matrix(
167 direction: str,
168 tau: float,
169 r_nodes: npt.NDArray[np.floating],
170 dr_weights: npt.NDArray[np.floating],
171 *,
172 r_w: float,
173 alpha_l: float,
174 c_geo: float,
175 flow_scale: float,
176 molecular_diffusivity: float,
177 retardation_factor: float,
178 n_terms: int,
179 tol: float,
180) -> npt.NDArray[np.floating]:
181 r"""Field-propagator matrix ``P[i, j]`` for one constant-Q phase (``D_m > 0`` Riccati branch).
183 The matrix form of the per-reversal ``_propagate_diffusive``: instead of contracting
184 the interior resolvent against a specific field by the prefix/suffix cumsums of
185 :func:`gwtransport._radial_asr_kernels.resolvent_riccati`, the full Sturm-Liouville Green's function is
186 assembled as the triangular outer product of the field-independent log-derivative integrals,
188 ``Ghat(r_i, r_j) w_j = -[exp(J_-(r_i) + J_+(r_j) + LG_j) if j <= i else exp(J_+(r_i) + J_-(r_j) + LG_j)]
189 / (L_-(r_w) - L_+(r_w)) * r_j dr_j``,
191 with ``J_-`` / ``J_+`` the inward (decaying) and outward (well-regular) running integrals of the
192 log-derivative (:func:`_integrate_logderiv`) and ``LG_j = b ln(G_j/G_w) - ln G_j`` the divergent gauge
193 carried in log space (``b = 1 - sigma_A A_0/D_m``). Wall-clock time ``tau`` (the volume clock is not
194 autonomous when ``D_m > 0``); retardation rescales ``A_0 -> A_0/R``, ``D_m -> D_m/R``.
196 Returns
197 -------
198 ndarray
199 Propagator matrix ``P``, shape ``(n_quad, n_quad)``.
200 """
201 a0_eff = (flow_scale / (2.0 * c_geo)) / retardation_factor
202 d_m_eff = molecular_diffusivity / retardation_factor
203 sigma_a = 1 if direction == _INJECTION else -1
204 n = r_nodes.size
205 mu_t = c_geo * ((float(r_nodes.max()) + alpha_l) ** 2 + alpha_l**2 - r_w**2) / flow_scale
206 scaling = _SCALE_MARGIN * max(mu_t, tau)
207 b = 1.0 - sigma_a * a0_eff / d_m_eff
208 rad = np.concatenate(([r_w], r_nodes))
209 g_nodes = alpha_l * a0_eff + d_m_eff * r_nodes
210 g_w = alpha_l * a0_eff + d_m_eff * r_w
211 lg = b * np.log(g_nodes / g_w) - np.log(g_nodes) # bounded divergent gauge in log space
212 lp_w = a0_eff / (alpha_l * a0_eff + d_m_eff * r_w) if sigma_a > 0 else 0.0
213 lower = np.tril(np.ones((n, n)))[None] # lower[i, j] = (j <= i)
214 weight = (r_nodes * dr_weights)[None, None, :]
216 def f_hat(s: npt.NDArray[np.complexfloating]) -> npt.NDArray[np.complexfloating]:
217 ld_m, jj_m = _integrate_logderiv(s, rad, r_w, alpha_l, a0_eff, d_m_eff, sigma_a, "decaying")
218 lm_w = ld_m[:, 0] # L_-(r_w)
219 jm = jj_m[:, 1:] # int_{r_w}^{r_i} L_- (n_s, n)
220 _, jp = _integrate_logderiv(s, r_nodes, r_w, alpha_l, a0_eff, d_m_eff, sigma_a, "regular") # int L_+
221 denom = (lm_w - lp_w)[:, None, None]
222 em_i = np.exp(jm)[:, :, None]
223 ep_jl = np.exp(jp + lg[None, :])[:, None, :]
224 ep_i = np.exp(jp)[:, :, None]
225 em_jl = np.exp(jm + lg[None, :])[:, None, :]
226 return -np.where(lower > 0, em_i * ep_jl, ep_i * em_jl) / denom * weight
228 return dehoog_inverse(f_hat=f_hat, t=tau, n_terms=n_terms, scaling=scaling, tol=tol)
231def _rest_propagator_matrix(
232 tau: float,
233 r_nodes: npt.NDArray[np.floating],
234 dr_weights: npt.NDArray[np.floating],
235 *,
236 r_w: float,
237 d_m_eff: float,
238 n_terms: int,
239 tol: float,
240) -> npt.NDArray[np.floating]:
241 r"""Field-propagator matrix ``P[i, j]`` for a rest (``Q = 0``) phase -- pure molecular diffusion.
243 The matrix form of the per-reversal ``_propagate_rest``: the order-0 modified Bessel
244 interior resolvent (:func:`gwtransport._radial_asr_kernels.rest_resolvent`) with the Sturm-Liouville rest
245 measure ``w(r') dr' = (r'/D_m) dr'``, on the wall-clock clock. Retardation enters as ``d_m_eff = D_m/R``.
247 **Resolution limit.** Molecular diffusion over the rest spreads mass by the diffusion length
248 ``sqrt(D_m tau)``. Once that length drops below half the coarsest node gap the Gauss-Legendre grid can no
249 longer resolve the near-delta diffusive Green's function: the quadratured resolvent stops conserving mass
250 (``dv^T P`` amplifies above ``1`` -- a maximum-principle violation, ``cout > cin``). In that under-resolved
251 regime the physical propagator *is* the identity to grid accuracy (diffusion has not crossed a cell), so
252 the identity is returned -- mass-exact and equal to the ``D_m = 0`` echo, the correct small-``D_m`` limit.
254 Returns
255 -------
256 ndarray
257 Propagator matrix ``P``, shape ``(n_quad, n_quad)``.
258 """
259 if np.sqrt(d_m_eff * tau) < 0.5 * float(np.diff(r_nodes).max()):
260 return np.eye(r_nodes.size)
261 scaling = _SCALE_MARGIN * tau
262 weight = (r_nodes / d_m_eff * dr_weights)[None, None, :]
263 # r_nodes ascend, so r_< = r_nodes[min(i, j)], r_> = r_nodes[max(i, j)]: the order-0 scaled Bessels are
264 # evaluated once on the grid per s-node and gathered by these index maps instead of O(n_quad) times.
265 idx = np.arange(r_nodes.size)
266 lt = np.minimum(idx[:, None], idx[None, :])
267 gt = np.maximum(idx[:, None], idx[None, :])
269 def f_hat(s: npt.NDArray[np.complexfloating]) -> npt.NDArray[np.complexfloating]:
270 kappa = np.sqrt(np.asarray(s, dtype=complex).reshape(-1, 1) / d_m_eff) # (n_s, 1), Re(kappa) > 0
271 zr = kappa * r_nodes[None, :] # (n_s, n): kappa r' on the grid, Bessels evaluated once here
272 iv0, kv0 = ive(0, zr), kve(0, zr)
273 z_w = kappa * r_w
274 ratio_w = ive(1, z_w) / kve(1, z_w) # (n_s, 1)
275 z_lt, z_gt = zr[:, lt], zr[:, gt] # (n_s, n, n)
276 # Ghat = [I_0(z_<) + (I_1(z_w)/K_1(z_w)) K_0(z_<)] K_0(z_>) with scaled Bessels (see rest_resolvent):
277 # the scaling exponents difference to <= 0 so the growing I_0 never overflows at high kappa r.
278 ghat = iv0[:, lt] * kv0[:, gt] * np.exp(np.abs(z_lt.real) - z_gt)
279 outer = ratio_w[:, :, None] * kv0[:, lt]
280 outer *= kv0[:, gt]
281 outer *= np.exp((np.abs(z_w.real) + z_w)[:, :, None] - z_lt - z_gt)
282 ghat += outer
283 ghat *= weight
284 return ghat
286 # The de Hoog underflow guard zeros far-field cells whose transform decayed to the double-precision floor
287 # (physical ~0); a cell whose transform genuinely OVERFLOWED stays NaN so a real breakdown cannot read as
288 # a physical zero. No extra masking here (that would zero overflowed cells too, reversing that contract).
289 return dehoog_inverse(f_hat=f_hat, t=tau, n_terms=n_terms, scaling=scaling, tol=tol)
292def _fr_source_matrix(
293 v_nodes: npt.NDArray[np.floating], edges: npt.NDArray[np.floating], **readout: float
294) -> npt.NDArray[np.floating]:
295 r"""Injection-source operator ``M[k, j]`` such that ``field_k += sum_j M[k, j] cin_dev_j`` over one phase.
297 Matrix form of the per-reversal ``_fr_profile``:
298 ``M[k, j] = G1_FR(S_inj - sigma_j; V'_k) - G1_FR(S_inj - sigma_{j+1}; V'_k)`` (the same FR step response,
299 KB Sec. 10a), built once per distinct injection phase and reused across recurring cycles.
301 Returns
302 -------
303 ndarray
304 Source operator, shape ``(n_quad, n_inj_bins)``.
305 """
306 # edges are phase-relative (edges[0] == 0 by caller construction); corners descend to 0 at the phase end
307 inj_corners = edges[-1] - edges # descending within-phase cumulative volume corners, last is 0
308 g = np.array([_fr_step_response(v, inj_corners, **readout) for v in v_nodes]) # (n_quad, n_bins + 1)
309 return g[:, :-1] - g[:, 1:]
312def _cout_readout_matrix(
313 v_nodes: npt.NDArray[np.floating],
314 dv_weights: npt.NDArray[np.floating],
315 edges: npt.NDArray[np.floating],
316 **readout: float,
317) -> npt.NDArray[np.floating]:
318 r"""Cout-readout operator ``M[i, k]`` such that ``cout_i = sum_k M[i, k] field_k`` over one phase.
320 Matrix form of the per-reversal ``_cout_phase``:
321 ``M[i, k] = R dv_k [G1_FR(T_{i+1}; V'_k) - G1_FR(T_i; V'_k)] / dT_i`` (the KB Sec. 7 duality arrival
322 kernel, same FR step response; the ``R`` amplitude mobilises the sorbed companion), built once per
323 distinct extraction phase and reused.
325 Returns
326 -------
327 ndarray
328 Readout operator, shape ``(n_ext_bins, n_quad)``.
329 """
330 # edges are phase-relative (edges[0] == 0 by caller construction)
331 dt = np.diff(edges)
332 g = np.array([_fr_step_response(v, edges, **readout) for v in v_nodes]) # (n_quad, n_bins + 1)
333 diffs = (g[:, 1:] - g[:, :-1]) / dt[None, :]
334 return (readout["retardation_factor"] * dv_weights[:, None] * diffs).T
337def cout_deviation(
338 *,
339 cin_deviation: npt.NDArray[np.floating],
340 flow: npt.NDArray[np.floating],
341 dt_days: npt.NDArray[np.floating],
342 c_geo: float,
343 r_w: float,
344 alpha_l: float,
345 molecular_diffusivity: float = 0.0,
346 retardation_factor: float = 1.0,
347 n_quad: int = 240,
348 n_terms: int = _DEHOOG_TERMS,
349 tol: float = 1e-9,
350) -> npt.NDArray[np.floating]:
351 r"""Multi-cycle extracted-flux deviation via reused per-phase propagator matrices (any ``D_m``).
353 Drop-in for the per-reversal ``gridfree_cout_deviation`` reference: every per-phase linear
354 operator -- the across-reversal field hand-off (``field = P @ field``), the injection source
355 (:func:`_fr_source_matrix`, ``field += M_fr @ cin``), and the cout readout (:func:`_cout_readout_matrix`,
356 ``cout = M_cout @ field``) -- is built once per distinct phase and reused as a matrix-multiply whenever
357 that phase recurs, instead of recomputing the de Hoog inversion / FR step response each cycle. The
358 propagator hand-off uses the Airy branch (:func:`_airy_propagator_matrix`,
359 flushed-volume clock) for ``D_m = 0``, the Riccati branch (:func:`_riccati_propagator_matrix`,
360 wall-clock) for ``D_m > 0``, and the Bessel rest branch (:func:`_rest_propagator_matrix`) for ``Q = 0``
361 phases when ``D_m > 0`` -- so the special-function + de Hoog cost is ``O(number of distinct phases)``
362 rather than ``O(K)`` reversals. Equal to the per-reversal grid-free engine to the de Hoog floor at
363 matched ``n_terms``/``tol``; tighten both for machine precision (the de Hoog Pade nonlinearity, dominant
364 at low Peclet / high retardation).
366 Parameters
367 ----------
368 cin_deviation : ndarray, shape (n,) or (n, k)
369 Injected concentration deviation per bin (used on injection bins, ``flow > 0``). A trailing column
370 axis is transported through one engine pass (the per-phase matrices are cin-independent), used by
371 the reverse operator build to apply all unit-pulse columns at once.
372 flow : ndarray, shape (n,)
373 Signed flow per bin [m^3/day]: ``> 0`` injection, ``< 0`` extraction, ``0`` rest.
374 dt_days : ndarray, shape (n,)
375 Bin widths [day].
376 c_geo : float
377 Geometry constant ``pi b n`` (``V = c_geo (r^2 - r_w^2)``).
378 r_w : float
379 Well radius [m].
380 alpha_l : float
381 Longitudinal dispersivity [m].
382 molecular_diffusivity : float, optional
383 Molecular diffusivity [m^2/day]. ``0`` selects the Airy branch (flushed-volume clock); ``> 0`` the
384 Riccati log-derivative pumping branch and the order-0 Bessel rest branch (both wall-clock). Default 0.
385 retardation_factor : float, optional
386 Linear retardation ``R >= 1`` (rescales the flushed-volume clock). Default 1.
387 n_quad : int, optional
388 Number of radial Gauss-Legendre nodes. Default 240.
389 n_terms : int, optional
390 de Hoog series length for the matrix build. Default ``_DEHOOG_TERMS``.
391 tol : float, optional
392 de Hoog target accuracy for the matrix build. Default ``1e-9``. Tighten (e.g.
393 ``tol=1e-13, n_terms=64``) for machine precision at low Peclet / high retardation.
395 Returns
396 -------
397 ndarray, shape (n,) or (n, k)
398 Extracted-flux deviation per bin (matching ``cin_deviation``); ``NaN`` on injection / rest bins.
399 """
400 flow = np.asarray(flow, dtype=float)
401 cin_deviation = np.asarray(cin_deviation, dtype=float)
402 batch = cin_deviation.shape[1:] # () for a single series, (k,) for a column batch (reverse operator build)
403 flushed = np.abs(flow) * dt_days
404 # D_m = 0 keeps the advective grid; D_m > 0 widens it to the molecular reach (the diffusive pumping and
405 # rest kernels both spread beyond the advective front), matching gridfree_cout_deviation.
406 r_nodes, v_nodes, dr_weights = _field_grid(flow, dt_days, c_geo, r_w, alpha_l, molecular_diffusivity, n_quad)
407 dv_weights = 2.0 * c_geo * r_nodes * dr_weights
409 phases = _phase_slices(flow)
410 matrices: dict[tuple, npt.NDArray[np.floating]] = {}
412 def propagate(field, direction, flow_scale, phase_volume, phase_time):
413 if molecular_diffusivity == 0.0: # Airy: flushed-volume clock, retardation rescales the clock
414 tau = phase_volume / retardation_factor
415 key = ("airy", direction, round(tau, 9), round(flow_scale, 9))
416 if key not in matrices:
417 matrices[key] = _airy_propagator_matrix(
418 direction,
419 tau,
420 r_nodes,
421 dr_weights,
422 r_w=r_w,
423 alpha_l=alpha_l,
424 c_geo=c_geo,
425 flow_scale=flow_scale,
426 n_terms=n_terms,
427 tol=tol,
428 )
429 return matrices[key] @ field
430 key = ("riccati", direction, round(phase_time, 9), round(flow_scale, 9)) # D_m > 0: wall-clock time
431 if key not in matrices:
432 matrices[key] = _riccati_propagator_matrix(
433 direction,
434 phase_time,
435 r_nodes,
436 dr_weights,
437 r_w=r_w,
438 alpha_l=alpha_l,
439 c_geo=c_geo,
440 flow_scale=flow_scale,
441 molecular_diffusivity=molecular_diffusivity,
442 retardation_factor=retardation_factor,
443 n_terms=n_terms,
444 tol=tol,
445 )
446 return matrices[key] @ field
448 field = np.zeros((n_quad, *batch))
449 cout = np.full((len(flow), *batch), np.nan)
450 for idx, (sign, sl) in enumerate(phases):
451 phase_volume = float(flushed[sl].sum())
452 if phase_volume == 0.0: # rest: pure molecular diffusion on the wall-clock clock (D_m = 0 -> identity)
453 if molecular_diffusivity > 0.0:
454 tau = float(np.sum(dt_days[sl]))
455 key = ("rest", round(tau, 9))
456 if key not in matrices:
457 matrices[key] = _rest_propagator_matrix(
458 tau,
459 r_nodes,
460 dr_weights,
461 r_w=r_w,
462 d_m_eff=molecular_diffusivity / retardation_factor,
463 n_terms=n_terms,
464 tol=tol,
465 )
466 field = matrices[key] @ field
467 continue
468 flow_scale = float(np.mean(np.abs(flow[sl])))
469 phase_time = float(np.sum(dt_days[sl]))
470 edges = np.concatenate(([0.0], np.cumsum(flushed[sl])))
471 readout = {
472 "c_geo": c_geo,
473 "r_w": r_w,
474 "alpha_l": alpha_l,
475 "retardation_factor": retardation_factor,
476 "flow_scale": flow_scale,
477 "molecular_diffusivity": molecular_diffusivity,
478 }
479 # The injection-source and cout-readout operators are also linear and phase-structure-determined, so
480 # they are matrix-cached and reused across recurring cycles too (retardation R and D_m are fixed for
481 # the whole call, so the cache key is the within-phase volume edges + the flow scale).
482 ekey = (tuple(np.round(edges, 9)), round(flow_scale, 9))
483 if sign > 0: # injection: propagate the buffer, then add the new injected resident profile
484 field = propagate(field, _INJECTION, flow_scale, phase_volume, phase_time)
485 mkey = ("fr", *ekey)
486 if mkey not in matrices:
487 matrices[mkey] = _fr_source_matrix(v_nodes, edges, **readout)
488 field += matrices[mkey] @ cin_deviation[sl]
489 else: # extraction: read cout, then propagate the residual if more pumping follows
490 mkey = ("cout", *ekey)
491 if mkey not in matrices:
492 matrices[mkey] = _cout_readout_matrix(v_nodes, dv_weights, edges, **readout)
493 cout[sl] = matrices[mkey] @ field
494 if idx != len(phases) - 1:
495 field = propagate(field, "extraction", flow_scale, phase_volume, phase_time)
496 return cout