Coverage for src/gwtransport/_radial_asr_drift_kernels.py: 97%

1r"""Block (azimuthal-mode) Laplace kernels for single-well ASR with steady regional drift. 

2 

3Generalizes the radial per-phase kernels (:mod:`gwtransport._radial_asr_kernels`) to a steady uniform 

4regional drift ``v_d = U/n`` superimposed on the radial well field ``A_0/r``. The 2-D 

5advection-dispersion generator ``O[c] = v.grad c - div(D grad c)`` (rank-1 Scheidegger tensor 

6``D = alpha_L (v outer v)/|v| + D_m I``, transverse mechanical dispersion neglected) is expanded in 

7azimuthal Fourier modes ``c(r, theta) = sum_m c_m(r) e^{i m theta}``. The radial well flow acts within 

8each mode (``m = 0`` is the radial engine); the drift couples mode ``m`` to its neighbours, giving a 

9coupled second-order block ODE 

10 

11 A(r) c'' + B(r) c' + (S0(r) + R s I) c = 0, c = (c_{-M}, ..., c_M)^T, 

12 

13with the ``s``-independent coefficient matrices (``M[f]`` the azimuthal Toeplitz coupling matrix whose 

14``(m, m')`` entry is the ``(m - m')``-th Fourier coefficient of ``f(r, .)``; ``Dm = diag(modes)``) 

15 

16 A = -M[D_rr] 

17 B = M[v_r - D_rr/r - d_r D_rr - (1/r) d_theta D_rtheta] + M[-2 D_rtheta/r] (i Dm) 

18 S0 = M[v_theta/r - (d_r D_rtheta)/r - (1/r^2) d_theta D_thth] (i Dm) + M[-D_thth/r^2] (-Dm^2), 

19 

20derived by collecting ``c''``, ``c'`` and ``c`` (the ``d_r d_theta c`` cross term folds into ``B`` via 

21``d_theta -> i m``; the ``d_theta^2`` term into ``S0`` via ``-m^2``). At ``v_d = 0`` every ``M[.]`` is 

22diagonal and the system decouples into the radial per-mode operator, so the ``m = 0`` block is exactly 

23the radial engine's ODE ``G c'' + (D_m - A_0) c' - R s r c = 0`` (``G = alpha_L A_0 + D_m r``). 

24 

25This file is part of gwtransport which is released under AGPL-3.0 license. 

26See the ./LICENSE file or go to https://github.com/gwtransport/gwtransport/blob/main/LICENSE for full license details. 

27""" 

28 

29from collections.abc import Callable 

30 

31import numpy as np 

32import numpy.typing as npt 

33from scipy.integrate import solve_ivp 

34from scipy.interpolate import BarycentricInterpolator 

35from scipy.special import airye 

36 

37from gwtransport._radial_asr_dehoog import dehoog_inverse 

38from gwtransport._radial_asr_reuse import _phase_slices # generic signed-schedule grouping (shared) 

39 

40# Per-phase orientation. Injection (Q > 0) is the divergent operator (Robin/flux well BC); extraction 

41# (Q < 0) the convergent operator (Danckwerts/Neumann well BC). The signed well strength A_0 carries the 

42# sign, so the azimuthal drift coupling uses the signed eps = v_d r / A_0 per phase automatically. 

43_INJECTION = "injection" 

44 

45# Riccati integration tolerances (matched to the scalar log-derivative kernel) and the drift-specific 

46# outer-boundary policy: the recessive initial condition is set at r_far and washed inward; r_far is 

47# floored well past the field but hard-capped below the stagnation radius r_s = |A_0|/v_d, where the 

48# coordinate finite-escape of the matrix Riccati would otherwise blow up (the plan's r_far <= 0.6 r_s). 

49# The slow-drift envelope is honest: the SIGNIFICANT plume (front + _PLUME_WIDTHS breakthrough widths, 

50# where the resident field has decayed to ~1% of peak) must fit below that cap, else the engine raises 

51# rather than silently truncate it. Only the negligible far tail beyond r_far is dropped (the recessive 

52# IC washes it out anyway). At v_d = 0 (r_s = inf) the cap is inactive and this is the scalar grid. 

53_RICCATI_RTOL = 1e-9 

54_RICCATI_ATOL = 1e-10 

55_RFAR_FIELD_MULT = 8.0 

56_RFAR_DECAY = 44.0 

57_RS_FRAC = 0.6 

58_GRID_WIDTHS = 12.0 

59_PLUME_WIDTHS = 3.0 

60# Rest-with-drift kernel: Gauss-Hermite quadrature size for the Gaussian spread, and the honest azimuthal 

61# truncation guard -- the relative spectral tail (harmonics M < |m| <= 2M after reprojection) above which 

62# the translated plume is declared too eccentric for the kept modes (raise, never silently fold the tail). 

63_REST_HERMITE = 20 

64_REST_TAIL_MAX = 1e-2 

65# Per-call cap on cached per-phase kernel solutions (each entry is O(n_quad n_s (2M+1)^2) complex, ~70 MB 

66# at the defaults and ~6x that at the auto-sizing ceiling n_modes=8 -- up to ~2.4 GB at the cap; a periodic 

67# schedule needs one entry per pumping direction, so the cap only sheds entries on long aperiodic 

68# schedules, where nothing recurs anyway). 

69_SOLUTIONS_CACHE_MAX = 8 

70 

71 

72def _toeplitz_from_theta(f_theta: npt.NDArray[np.floating], n_modes: int) -> npt.NDArray[np.complexfloating]: 

73 r"""Azimuthal coupling matrix ``M[f]`` from samples of ``f(theta)`` on a uniform grid. 

74 

75 ``M[f]_{a,b} = c_{m_a - m_b}`` where ``c_k`` is the ``k``-th Fourier coefficient of ``f`` (the 

76 multiplication-by-``f`` operator in the ``e^{i m theta}`` basis), ``m`` running ``-n_modes .. n_modes``. 

77 The grid must satisfy ``len(f_theta) >= 4 n_modes + 1`` so every needed coefficient ``|k| <= 2 n_modes`` 

78 is unaliased. 

79 

80 Returns 

81 ------- 

82 ndarray of complex, shape (2 n_modes + 1, 2 n_modes + 1) 

83 The Toeplitz coupling matrix. 

84 """ 

85 n = f_theta.shape[-1] 

86 coeffs = np.fft.fft(f_theta, axis=-1) / n # coeffs[k] = c_k, with c_{-k} at index n-k 

87 modes = np.arange(-n_modes, n_modes + 1) 

88 diff = modes[:, None] - modes[None, :] # m_a - m_b in [-2M, 2M] 

89 return coeffs[..., diff % n] 

90 

91 

92def _theta_grid(n_modes: int) -> npt.NDArray[np.floating]: 

93 r"""Uniform azimuthal FFT grid sized for the eps-banded harmonics. 

94 

95 The needed Fourier coefficients ``|k| <= 2 n_modes`` of the tensor components decay like 

96 ``eps^|k|`` and coefficient ``k`` aliases the ``k +- n_theta`` harmonic; ``8 n_modes + 48`` keeps 

97 that alias below ~1e-12 across the slow-drift envelope (``eps <= 0.6`` at the grid edge). 

98 

99 Returns 

100 ------- 

101 ndarray 

102 Grid angles, shape ``(8 n_modes + 48,)``. 

103 """ 

104 nth = 8 * n_modes + 48 

105 return np.arange(nth) * (2.0 * np.pi / nth) 

106 

107 

108def _tensor_components( 

109 rr: npt.NDArray[np.floating], theta: npt.NDArray[np.floating], *, alpha_l: float, a0: float, v_d: float, d_m: float 

110) -> tuple[npt.NDArray[np.floating], ...]: 

111 r"""Velocity and rank-1 Scheidegger tensor components on an ``(r, theta)`` product grid. 

112 

113 The single source of the drift field's tensor sampling, shared by the interior coupling matrices 

114 and the well-face operators so the formulas cannot drift apart. ``rr`` broadcasts as ``(n_r, 1)``. 

115 

116 Returns 

117 ------- 

118 v_r, v_th, speed, d_rr, d_rt, d_tt : ndarray 

119 Components on the product grid, each shape ``(n_r, n_theta)`` (``v_th`` is ``(1, n_theta)``). 

120 """ 

121 cos_t, sin_t = np.cos(theta), np.sin(theta) 

122 v_r = a0 / rr + v_d * cos_t[None, :] 

123 v_th = -v_d * sin_t[None, :] # theta-only (r-independent); broadcasts against the (n_r, .) factors 

124 speed = np.sqrt(v_r**2 + v_th**2) 

125 speed = np.where(speed == 0.0, 1.0, speed) # a0=v_d=0 never reached here (alpha_L>0 guard upstream) 

126 d_rr = alpha_l * v_r**2 / speed + d_m 

127 d_rt = alpha_l * v_r * v_th / speed 

128 d_tt = alpha_l * v_th**2 / speed + d_m 

129 return v_r, v_th, speed, d_rr, d_rt, d_tt 

130 

131 

132def block_coupling_matrices( 

133 r: npt.NDArray[np.floating], 

134 *, 

135 alpha_l: float, 

136 a0: float, 

137 v_d: float, 

138 d_m: float, 

139 n_modes: int, 

140) -> tuple[npt.NDArray[np.complexfloating], npt.NDArray[np.complexfloating], npt.NDArray[np.complexfloating]]: 

141 r"""Coefficient matrices ``A(r), B(r), S0(r)`` of the coupled-mode block ODE. 

142 

143 The block ODE is ``A c'' + B c' + (S0 + R s I) c = 0``; ``A, B, S0`` are ``s``-independent (the Laplace 

144 node enters only through ``R s I``). ``a0`` is the signed well strength ``A_0 = Q/(2 pi b n)`` (``> 0`` 

145 injection, ``< 0`` extraction); ``v_d = U/n`` the regional drift. Built by FFT-in-theta of the exact 

146 Scheidegger tensor components and their analytic radial derivatives. 

147 

148 Parameters 

149 ---------- 

150 r : ndarray 

151 Radii (m), shape ``(n_r,)``. 

152 alpha_l : float 

153 Longitudinal dispersivity (m). 

154 a0 : float 

155 Signed well strength ``A_0 = Q/(2 pi b n)`` (m^2/day). 

156 v_d : float 

157 Regional drift seepage velocity ``U/n`` (m/day). 

158 d_m : float 

159 Molecular diffusivity (m^2/day). 

160 n_modes : int 

161 Azimuthal truncation ``M`` (keeps modes ``-M .. M``). 

162 

163 Returns 

164 ------- 

165 a_mat, b_mat, s0_mat : ndarray of complex, each shape (n_r, 2 n_modes + 1, 2 n_modes + 1) 

166 The coefficient matrices at each radius. 

167 """ 

168 r = np.atleast_1d(np.asarray(r, dtype=float)) 

169 theta = _theta_grid(n_modes) 

170 nth = theta.size 

171 rr = r[:, None] # (n_r, 1) broadcast against theta 

172 v_r, v_th, speed, d_rr, d_rt, d_tt = _tensor_components(rr, theta, alpha_l=alpha_l, a0=a0, v_d=v_d, d_m=d_m) 

173 

174 # analytic radial derivatives (v_th is r-independent; d_r v_r = -a0/r^2) 

175 dv_r = -a0 / rr**2 

176 dspeed = v_r * dv_r / speed 

177 d_drr = alpha_l * (2.0 * v_r * dv_r / speed - v_r**2 * dspeed / speed**2) 

178 d_drt = alpha_l * (dv_r * v_th / speed - v_r * v_th * dspeed / speed**2) 

179 

180 # angular derivatives via spectral differentiation (exact on the grid) 

181 kth = np.fft.fftfreq(nth, d=1.0 / nth) # integer wavenumbers 

182 dth_drt = np.real(np.fft.ifft(1j * kth * np.fft.fft(d_rt, axis=-1), axis=-1)) 

183 dth_dtt = np.real(np.fft.ifft(1j * kth * np.fft.fft(d_tt, axis=-1), axis=-1)) 

184 

185 modes = np.arange(-n_modes, n_modes + 1) 

186 i_dm = 1j * modes # i * diag(m) as a row to scale columns 

187 m2 = modes**2 

188 

189 a_mat = -_toeplitz_from_theta(d_rr, n_modes) 

190 b_coeff = v_r - d_rr / rr - d_drr - dth_drt / rr 

191 b_mat = ( 

192 _toeplitz_from_theta(b_coeff, n_modes) + _toeplitz_from_theta(-2.0 * d_rt / rr, n_modes) * i_dm[None, None, :] 

193 ) 

194 s_coeff = v_th / rr - d_drt / rr - dth_dtt / rr**2 

195 s0_mat = ( 

196 _toeplitz_from_theta(s_coeff, n_modes) * i_dm[None, None, :] 

197 + _toeplitz_from_theta(-d_tt / rr**2, n_modes) * (-m2)[None, None, :] 

198 ) 

199 

200 # _toeplitz_from_theta always returns (n_r, nm, nm) and r is atleast_1d, so the shape is uniform. 

201 return a_mat, b_mat, s0_mat 

202 

203 

204def _face_matrices( 

205 r_w: float, *, alpha_l: float, a0_signed: float, v_d: float, d_m: float, n_modes: int 

206) -> tuple[npt.NDArray[np.complexfloating], npt.NDArray[np.complexfloating], npt.NDArray[np.complexfloating]]: 

207 r"""Well-face Toeplitz matrices ``M[v_r]``, ``M[D_rr]``, ``M[D_rtheta]`` at ``r = r_w``. 

208 

209 These carry the exact ``theta``-modulated well-face physics: the drift adds ``v_d cos(theta)`` to the 

210 face velocity and an ``O(eps_w)`` cross-dispersion ``D_rtheta``, which couple neighbouring modes in 

211 the face boundary conditions (Robin flux inlet / Danckwerts) and in the injected-flux source. Though 

212 individually ``O(eps_w)``, these couplings enter the drift recovery loss at ``O(eps^2)`` -- the order 

213 of the loss itself -- so the face conditions are built exactly. 

214 

215 Returns 

216 ------- 

217 m_vr, m_drr, m_drt : ndarray of complex, each shape (2 n_modes + 1, 2 n_modes + 1) 

218 Azimuthal coupling matrices of the face velocity and dispersion-tensor components. 

219 """ 

220 v_r, _, _, d_rr, d_rt, _ = _tensor_components( 

221 np.array([[r_w]]), _theta_grid(n_modes), alpha_l=alpha_l, a0=a0_signed, v_d=v_d, d_m=d_m 

222 ) 

223 return ( 

224 _toeplitz_from_theta(v_r, n_modes)[0], 

225 _toeplitz_from_theta(d_rr, n_modes)[0], 

226 _toeplitz_from_theta(d_rt, n_modes)[0], 

227 ) 

228 

229 

230def field_grid( 

231 flow: npt.NDArray[np.floating], 

232 dt_days: npt.NDArray[np.floating], 

233 c_geo: float, 

234 r_w: float, 

235 alpha_l: float, 

236 v_d: float, 

237 a0: float, 

238 n_quad: int, 

239 d_m: float = 0.0, 

240 drift_shift: float = 0.0, 

241) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating], float]: 

242 r"""Radial Gauss-Legendre grid and recessive-IC radius ``r_far`` for the drift block engine. 

243 

244 The grid spans the advective plume front ``r_front = sqrt(r_w^2 + V_peak/c_geo)`` plus a margin of 

245 breakthrough widths (radial std ``~ sqrt(alpha_L r_front)``), the molecular reach 

246 ``~ sqrt(D_m * total_time)`` (as the scalar engine's grid), and the total rest-phase drift 

247 displacement ``drift_shift`` (the rest kernel translates the plume down-gradient; the grid must 

248 contain it). For drift, the recessive initial condition at ``r_far`` must stay inside the stagnation 

249 radius ``r_s = |A_0|/v_d`` (the matrix Riccati has a coordinate finite-escape as 

250 ``eps(r) = v_d r/A_0 -> 1``), so ``r_far`` is hard-capped at ``_RS_FRAC * r_s``. The envelope is 

251 enforced **honestly**: the *significant* plume (front + ``_PLUME_WIDTHS`` breakthrough widths + the 

252 rest displacement, where the resident field has fallen to ~1% of peak) must fit below that cap -- 

253 otherwise a ``ValueError`` is raised rather than the plume being silently truncated. The quadrature 

254 spans the *significant* plume up to ``r_far`` (a generous ``_GRID_WIDTHS``-width field grid clipped at 

255 ``r_far``); only the negligible tail beyond ``r_far`` is dropped, which the recessive IC washes out 

256 anyway. At ``v_d = 0`` (``r_s = inf``) the cap is inactive and the grid is the scalar grid. ``a0`` 

257 should be the **smallest** per-phase ``|A_0|`` (the most restrictive ``r_s``). 

258 

259 Returns 

260 ------- 

261 r_nodes : ndarray 

262 Radial nodes (m), shape ``(n_quad,)``. 

263 dr_weights : ndarray 

264 Gauss-Legendre weights in ``r``. 

265 r_far : float 

266 Outer radius (m) where the recessive (decaying) block initial condition is imposed. 

267 

268 Raises 

269 ------ 

270 ValueError 

271 If the significant plume reaches the stagnation radius (``r_front + _PLUME_WIDTHS * (width + 

272 reach_dm) + drift_shift + r_w > _RS_FRAC * r_s``, with ``width`` carrying the rest kernel's 

273 mechanical spread in quadrature and ``reach_dm = sqrt(D_m * total_time)``) -- the drift is too 

274 strong for the radial engine's contained-plume assumption. 

275 """ 

276 net_volume = np.concatenate(([0.0], np.cumsum(flow * dt_days))) 

277 peak_volume = max(float(net_volume.max()), 0.0) 

278 total_time = float(np.sum(dt_days)) 

279 r_front = np.sqrt(r_w**2 + peak_volume / c_geo) 

280 # radial breakthrough variance ~ alpha_L r_front (+ alpha_L^2), plus -- in quadrature -- the rest 

281 # kernel's own mechanical spread sigma_x^2 = 2 alpha_L v_d t_rest / R = 2 alpha_L drift_shift 

282 width = np.sqrt(alpha_l * r_front + alpha_l**2 + 2.0 * alpha_l * drift_shift) 

283 reach_dm = np.sqrt(d_m * total_time) 

284 r_significant = r_front + _PLUME_WIDTHS * (width + reach_dm) + drift_shift + r_w 

285 r_grid = r_front + _GRID_WIDTHS * width + 2.0 * _PLUME_WIDTHS * reach_dm + drift_shift + r_w 

286 r_s = abs(a0) / abs(v_d) if v_d != 0.0 else np.inf 

287 r_far_cap = _RS_FRAC * r_s 

288 if r_significant > r_far_cap: 

289 eps_front = abs(v_d) * r_front / abs(a0) 

290 msg = ( 

291 f"drift too strong for the radial engine: the plume (front r_front={r_front:.2f} m, " 

292 f"eps_front={eps_front:.2f}) reaches the stagnation radius r_s={r_s:.2f} m; the slow-drift " 

293 "envelope requires the plume to stay well inside r_s = |A_0|/|v_d|" 

294 ) 

295 raise ValueError(msg) 

296 # r_far >= r_significant (the cap was not exceeded and the washout floor is >= r_grid >= r_significant), 

297 # so clipping the generous grid at r_far drops only the negligible far tail, not the significant plume. 

298 r_far = min(max(_RFAR_FIELD_MULT * r_grid, r_grid + _RFAR_DECAY * alpha_l), r_far_cap) 

299 r_max = min(r_grid, r_far) 

300 nodes, weights = np.polynomial.legendre.leggauss(n_quad) 

301 r_nodes = 0.5 * (r_max - r_w) * (nodes + 1.0) + r_w 

302 dr_weights = 0.5 * (r_max - r_w) * weights 

303 return r_nodes, dr_weights, r_far 

304 

305 

306def _interval_transitions( 

307 l_dense: Callable[[npt.NDArray[np.floating]], npt.NDArray[np.complexfloating]], 

308 r_from: npt.NDArray[np.floating], 

309 r_to: npt.NDArray[np.floating], 

310 n_s: int, 

311 nm: int, 

312) -> npt.NDArray[np.complexfloating]: 

313 r"""Fundamental transitions ``Psi(r_to[i], r_from[i])`` of ``Y' = L(r) Y`` for a batch of intervals. 

314 

315 All intervals are integrated in lockstep on a common ``tau in [0, 1]`` clock (``r = r_from + 

316 tau (r_to - r_from)``, signed widths), with the log-derivative field ``L`` evaluated through its dense 

317 ODE interpolant, chunked to bound memory. Every interval spans adjacent quadrature radii, so each 

318 transition stays within a few e-foldings of unity -- which is what keeps the prefix/suffix 

319 Green's-function recursions built from them unconditionally well-conditioned, where a globally 

320 pivoted fundamental matrix accumulates the full mode-split Sturm-Liouville exponent across the grid 

321 and overflows / colinearizes. 

322 

323 Returns 

324 ------- 

325 ndarray of complex, shape (n_intervals, n_s, nm, nm) 

326 The per-interval transition matrices. 

327 

328 Raises 

329 ------ 

330 RuntimeError 

331 If an interval-transition integration fails. 

332 """ 

333 n_int = r_from.size 

334 eye = np.eye(nm, dtype=complex) 

335 out = np.empty((n_int, n_s, nm, nm), dtype=complex) 

336 chunk = 64 # bounds the ODE state (and DOP853's stage copies) to ~chunk * n_s * nm^2 complex 

337 for lo in range(0, n_int, chunk): 

338 hi = min(lo + chunk, n_int) 

339 start, width = r_from[lo:hi], r_to[lo:hi] - r_from[lo:hi] 

340 nc = hi - lo 

341 

342 def rhs(tau: float, y: npt.NDArray[np.floating], start=start, width=width, nc=nc) -> npt.NDArray[np.floating]: 

343 psi = y.view(complex).reshape(nc, n_s, nm, nm) 

344 ld = l_dense(start + tau * width) # (nc, n_s, nm, nm) 

345 return (width[:, None, None, None] * (ld @ psi)).reshape(-1).view(float) 

346 

347 sol = solve_ivp( 

348 rhs, 

349 [0.0, 1.0], 

350 np.tile(eye.reshape(-1), nc * n_s).view(float), 

351 rtol=_RICCATI_RTOL, 

352 atol=_RICCATI_ATOL, 

353 method="DOP853", 

354 ) 

355 if not sol.success: 

356 msg = "block interval-transition integration failed" 

357 raise RuntimeError(msg) 

358 out[lo:hi] = np.ascontiguousarray(sol.y[:, -1]).view(complex).reshape(nc, n_s, nm, nm) 

359 return out 

360 

361 

362def _block_solutions( 

363 s: npt.NDArray[np.complexfloating], 

364 r_nodes: npt.NDArray[np.floating], 

365 r_w: float, 

366 *, 

367 alpha_l: float, 

368 a0: float, 

369 v_d: float, 

370 d_m: float, 

371 retardation_factor: float, 

372 n_modes: int, 

373 direction: str, 

374 r_far: float, 

375) -> dict[str, npt.NDArray[np.complexfloating]]: 

376 r"""Batched block log-derivative + per-interval transition solutions of one constant-Q phase. 

377 

378 Integrates, for every Laplace node ``s`` at once (the coefficient matrices ``A, B, S0`` are 

379 ``s``-independent -- only the ``R s A^{-1}`` term carries ``s`` -- so one vectorized ODE pass covers 

380 all nodes), the matrix Riccati ``L' = -L^2 - A^{-1} B L - A^{-1}(S0 + R s I)`` (``L = c' c^{-1}``) on 

381 two branches: 

382 

383 * **decaying** (recessive as ``r -> inf``): inward from ``r_far`` with a block-diagonal recessive IC 

384 (the scalar per-mode decaying log-derivative on every diagonal -- exact at ``v_d = 0`` and washed in 

385 by the inward attractor otherwise). Gives ``L_-`` at ``r_nodes`` and ``r_w``. 

386 * **regular** (well boundary): outward from ``r_w`` with the **exact block well IC** built from the 

387 face Toeplitz matrices (:func:`_face_matrices`): injection Robin 

388 ``L_+ = M[D_rr]^{-1}(M[v_r] - M[D_rt](i m)/r_w)``, extraction Danckwerts 

389 ``L_+ = -M[D_rr]^{-1} M[D_rt](i m)/r_w``. The ``O(eps_w)`` face couplings these carry look small 

390 but enter the drift loss at ``O(eps^2)``, the order of the loss itself. Gives ``L_+``. 

391 

392 The recessive / regular fundamental solutions enter the interior resolvent only through 

393 **per-interval transitions** ``Psi_-(r_i, r_{i-1})`` (outward hops) and ``Psi_+(r_{i-1}, r_i)`` 

394 (inward hops), with ``r_{-1} = r_w``, integrated through the dense ``L`` interpolants 

395 (:func:`_interval_transitions`). Each hop spans a few e-foldings at most, so the prefix/suffix 

396 resolvent recursions (:func:`_resolvent_field_laplace`, :func:`_readout_laplace`, 

397 :func:`_resident_laplace`) stay bounded and well-conditioned at any Peclet, Laplace frequency, and 

398 azimuthal truncation. A single fundamental matrix pivoted at ``r_w`` -- even de-scaled by a scalar 

399 log-amplitude -- accumulates the full mode-split Sturm-Liouville exponent across the grid and 

400 overflows / colinearizes for phases with small ``A_0`` or short durations, which is why the hops 

401 are the stored representation. 

402 

403 ``a0`` is the unsigned flow scale ``|A_0|``; the phase ``direction`` sets the operator orientation. 

404 Retardation enters as the explicit ``R`` in the ``R s`` term (the ODE keeps the physical ``A_0``/``D_m``). 

405 

406 Returns 

407 ------- 

408 dict of ndarray 

409 ``Lm_w`` (``L_-`` at ``r_w``, shape ``(n_s, nm, nm)``), ``H`` (the Wronskian blocks 

410 ``[A (L_- - L_+)]^{-1}`` at ``r_nodes``, ``(n_quad, n_s, nm, nm)``) and ``Tm``/``Tp`` (recessive 

411 outward / regular inward per-interval transitions, ``(n_quad, n_s, nm, nm)``). 

412 

413 Raises 

414 ------ 

415 RuntimeError 

416 If a matrix Riccati or interval-transition integration does not succeed (the log-derivative hit 

417 a pole, typically the coordinate finite-escape near the stagnation radius). 

418 """ 

419 s = np.asarray(s, dtype=complex).reshape(-1) 

420 n_s = s.size 

421 nm = 2 * n_modes + 1 

422 eye = np.eye(nm) 

423 sigma_a = 1.0 if direction == _INJECTION else -1.0 

424 a0_signed = sigma_a * abs(a0) 

425 r_max = float(np.max(r_nodes)) 

426 

427 def riccati_rhs(r: float, y: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]: 

428 ld = y.view(complex).reshape(n_s, nm, nm) 

429 a_m, b_m, s0_m = block_coupling_matrices( 

430 np.array([r]), alpha_l=alpha_l, a0=a0_signed, v_d=v_d, d_m=d_m, n_modes=n_modes 

431 ) 

432 a_inv = np.linalg.inv(a_m[0]) 

433 d_ld = ( 

434 -(ld @ ld) 

435 - ((a_inv @ b_m[0])[None] @ ld) 

436 - (a_inv @ s0_m[0])[None] 

437 - (retardation_factor * s)[:, None, None] * a_inv[None] 

438 ) 

439 return d_ld.reshape(-1).view(float) 

440 

441 # recessive IC at r_far (block-diagonal; scalar per-mode decaying log-derivative) 

442 if d_m == 0.0: 

443 beta = retardation_factor * s / (alpha_l * abs(a0)) 

444 b13 = beta ** (1.0 / 3.0) 

445 zeta = b13 * r_far + beta ** (-2.0 / 3.0) / (4.0 * alpha_l * alpha_l) 

446 aie, aipe, _, _ = airye(zeta) 

447 l0 = 1.0 / (2.0 * alpha_l) + b13 * (aipe / aie) 

448 else: 

449 astar = alpha_l * abs(a0) / d_m 

450 kappa = np.sqrt(retardation_factor * s / d_m) 

451 a_coef = (1.0 - sigma_a * abs(a0) / d_m) / 2.0 - kappa * astar / 2.0 

452 l0 = -kappa - a_coef / (r_far + astar) 

453 lm0 = (l0[:, None, None] * eye[None]).astype(complex) 

454 sol_m = solve_ivp( 

455 riccati_rhs, 

456 [r_far, r_w], 

457 lm0.reshape(-1).view(float), 

458 rtol=_RICCATI_RTOL, 

459 atol=_RICCATI_ATOL, 

460 dense_output=True, 

461 method="DOP853", 

462 ) 

463 # Exact block well-face IC for the regular branch: the theta-modulated face velocity and the 

464 # D_rtheta cross-dispersion couple the modes in the face condition (a mode-decoupled face IC would 

465 # mis-state the drift loss at O(eps^2), the order of the loss itself): 

466 # injection (Robin flux inlet, homogeneous part): M[v_r] c - M[D_rr] c' - M[D_rt] (i m / r_w) c = 0 

467 # extraction (Danckwerts, zero dispersive flux): M[D_rr] c' + M[D_rt] (i m / r_w) c = 0 

468 m_vr, m_drr, m_drt = _face_matrices(r_w, alpha_l=alpha_l, a0_signed=a0_signed, v_d=v_d, d_m=d_m, n_modes=n_modes) 

469 i_dm_face = 1j * np.arange(-n_modes, n_modes + 1) 

470 cross = (m_drt * i_dm_face[None, :]) / r_w 

471 lp_w = np.linalg.solve(m_drr, (m_vr - cross) if sigma_a > 0 else -cross) 

472 lp0 = np.broadcast_to(lp_w, (n_s, nm, nm)).astype(complex) 

473 sol_p = solve_ivp( 

474 riccati_rhs, 

475 [r_w, r_max], 

476 lp0.reshape(-1).view(float), 

477 rtol=_RICCATI_RTOL, 

478 atol=_RICCATI_ATOL, 

479 dense_output=True, 

480 method="DOP853", 

481 ) 

482 if not (sol_m.success and sol_p.success): 

483 msg = "block Riccati integration failed (the matrix log-derivative likely hit a pole near stagnation)" 

484 raise RuntimeError(msg) 

485 

486 def l_minus_at(r: npt.NDArray[np.floating]) -> npt.NDArray[np.complexfloating]: 

487 return np.ascontiguousarray(sol_m.sol(r).T).view(complex).reshape(np.size(r), n_s, nm, nm) 

488 

489 def l_plus_at(r: npt.NDArray[np.floating]) -> npt.NDArray[np.complexfloating]: 

490 return np.ascontiguousarray(sol_p.sol(r).T).view(complex).reshape(np.size(r), n_s, nm, nm) 

491 

492 prev = np.concatenate(([r_w], r_nodes[:-1])) 

493 a_n = block_coupling_matrices(r_nodes, alpha_l=alpha_l, a0=a0_signed, v_d=v_d, d_m=d_m, n_modes=n_modes)[0] 

494 # The log-derivative branches enter the resolvent only through the Wronskian block 

495 # H(r') = [A(r')(L_-(r') - L_+(r'))]^{-1} -- bounded wherever the recessive and regular subspaces are 

496 # transverse -- so H is materialized once here (and cached with the phase) instead of its factors. 

497 h = np.linalg.inv(a_n[:, None] @ (l_minus_at(r_nodes) - l_plus_at(r_nodes))) 

498 return { 

499 "Lm_w": l_minus_at(np.array([r_w]))[0], 

500 "H": h, 

501 "Tm": _interval_transitions(l_minus_at, prev, r_nodes, n_s, nm), # Psi_-(r_i, r_{i-1}), outward hops 

502 "Tp": _interval_transitions(l_plus_at, r_nodes, prev, n_s, nm), # Psi_+(r_{i-1}, r_i), inward hops 

503 } 

504 

505 

506def _resident_laplace( 

507 d: dict[str, npt.NDArray[np.complexfloating]], 

508 *, 

509 alpha_l: float, 

510 d_m: float, 

511 r_w: float, 

512 a0: float, 

513 v_d: float, 

514 n_modes: int, 

515) -> npt.NDArray[np.complexfloating]: 

516 r"""Laplace-domain resident mode-field per unit uniform-``cin`` flux injection at the well. 

517 

518 The injected water enters through the exact Kreft-Zuber flux boundary 

519 ``v_r c - D_rr d_r c - D_rt (1/r_w) d_th c = v_r cin``: the ``theta``-modulated face velocity gives 

520 the injected flux an ``O(eps_w)`` ``m = +-1`` modulation and the ``D_rtheta`` cross term couples the 

521 modes in the face operator, so with the recessive branch (``c' = L_-(r_w) c``) the face field is 

522 ``c(r_w) = F_w^{-1} M[v_r] e_0`` with ``F_w = M[v_r] - M[D_rr] L_-(r_w) - M[D_rt] (i m)/r_w``, carried 

523 outward by the stable hops ``c(r_i) = Psi_-(r_i, r_{i-1}) c(r_{i-1})``. Reduces to the scalar FR 

524 resident transfer ``E / f_w`` at ``v_d = 0``. 

525 

526 Returns 

527 ------- 

528 ndarray 

529 ``c(r_nodes; s)``, shape ``(n_s, n_quad, nm)``. 

530 """ 

531 nm = 2 * n_modes + 1 

532 n_s = d["Lm_w"].shape[0] 

533 e0 = np.zeros(nm, dtype=complex) 

534 e0[n_modes] = 1.0 

535 m_vr, m_drr, m_drt = _face_matrices(r_w, alpha_l=alpha_l, a0_signed=abs(a0), v_d=v_d, d_m=d_m, n_modes=n_modes) 

536 cross = (m_drt * (1j * np.arange(-n_modes, n_modes + 1))[None, :]) / r_w 

537 fw = (m_vr - cross)[None] - m_drr[None] @ d["Lm_w"] # (n_s, nm, nm) 

538 y = np.linalg.solve(fw, np.broadcast_to(m_vr @ e0, (n_s, nm))[..., None]) # (n_s, nm, 1) 

539 tm = d["Tm"] 

540 c = np.empty((tm.shape[0], n_s, nm), dtype=complex) 

541 for i in range(tm.shape[0]): 

542 y = tm[i] @ y 

543 c[i] = y[..., 0] 

544 return np.transpose(c, (1, 0, 2)) 

545 

546 

547def _source_blocks( 

548 d: dict[str, npt.NDArray[np.complexfloating]], 

549 field: npt.NDArray[np.floating], 

550 dr_weights: npt.NDArray[np.floating], 

551 retardation_factor: float, 

552) -> npt.NDArray[np.complexfloating]: 

553 r"""Wronskian-weighted source contributions ``H_j (R dr_j) field_j`` of the interior resolvent. 

554 

555 ``H(r') = [A(r')(L_-(r') - L_+(r'))]^{-1}`` is the matrix Wronskian block of the interior Green's 

556 function, materialized with the per-phase solutions (:func:`_block_solutions`). 

557 

558 Returns 

559 ------- 

560 ndarray 

561 ``H_j source_j``, shape ``(n_quad, n_s, nm, k)`` for a ``(n_quad, nm, k)`` mode-field batch. 

562 """ 

563 src = ((retardation_factor * dr_weights)[:, None, None] * field).astype(complex) # (n_quad, nm, k) 

564 return d["H"] @ src[:, None] # (n_quad, n_s, nm, k) 

565 

566 

567def _resolvent_field_laplace( 

568 d: dict[str, npt.NDArray[np.complexfloating]], 

569 field: npt.NDArray[np.floating], 

570 dr_weights: npt.NDArray[np.floating], 

571 retardation_factor: float, 

572) -> npt.NDArray[np.complexfloating]: 

573 r"""Laplace-domain resident field after one constant-Q phase propagates ``field`` (all nodes). 

574 

575 The interior Green's function of the constant-Q block operator, with recessive ``Y_-`` and well-regular 

576 ``Y_+`` and the matrix Wronskian ``H(r') = [A(r')(L_-(r') - L_+(r'))]^{-1}``, is 

577 ``Ghat(r, r') = Y_-(r) Y_-(r')^{-1} H(r')`` for ``r >= r'`` and ``Y_+(r) Y_+(r')^{-1} H(r')`` for 

578 ``r <= r'``. Applied to the source measure ``(R dr_j) field_j`` it separates into a recessive prefix 

579 (``j <= i``) and a regular suffix (``j > i``), both evaluated as first-order recursions over the 

580 per-interval transitions (``hs_j = H_j source_j``): 

581 

582 P_0 = hs_0, P_i = Psi_-(r_i, r_{i-1}) P_{i-1} + hs_i, 

583 S_{n-1} = 0, S_{i-1} = Psi_+(r_{i-1}, r_i) (S_i + hs_i), 

584 F_i = P_i + S_i. 

585 

586 Every factor is a short-interval transition (a few e-foldings) or a bounded Wronskian block, so the 

587 recursion is well-conditioned at any Peclet and Laplace frequency: distant contributions fade by 

588 repeated bounded multiplication exactly as the physical Green's function does, with no globally 

589 accumulated exponent to overflow or colinearize. 

590 

591 Returns 

592 ------- 

593 ndarray 

594 Propagated mode-field batch, shape ``(n_s, n_quad, nm, k)`` for a ``(n_quad, nm, k)`` ``field``. 

595 """ 

596 hs = _source_blocks(d, field, dr_weights, retardation_factor) # (n_quad, n_s, nm, k) 

597 tm, tp = d["Tm"], d["Tp"] 

598 n_quad = hs.shape[0] 

599 f = np.empty_like(hs) 

600 p = hs[0] 

601 f[0] = p 

602 for i in range(1, n_quad): 

603 p = tm[i] @ p + hs[i] 

604 f[i] = p 

605 s_acc = np.zeros_like(p) 

606 for i in range(n_quad - 1, 0, -1): 

607 s_acc = tp[i] @ (s_acc + hs[i]) 

608 f[i - 1] += s_acc 

609 return np.transpose(f, (1, 0, 2, 3)) # (n_s, n_quad, nm, k) 

610 

611 

612def _readout_laplace( 

613 d: dict[str, npt.NDArray[np.complexfloating]], 

614 field: npt.NDArray[np.floating], 

615 dr_weights: npt.NDArray[np.floating], 

616 retardation_factor: float, 

617 n_modes: int, 

618 eps_w: float, 

619) -> npt.NDArray[np.complexfloating]: 

620 r"""Laplace-domain extracted flux concentration from a resident ``field`` under extraction. 

621 

622 Under the exact Danckwerts boundary (zero dispersive flux across the face, cross term included in the 

623 regular-branch IC) the extracted flux concentration equals the resident concentration at the face 

624 pointwise in ``theta``, i.e. the extraction interior resolvent evaluated at ``r_w``. Since ``r_w`` 

625 lies below every node, only the regular (suffix) branch contributes: 

626 ``c(r_w) = sum_j Y_+(r_w) Y_+(r_j)^{-1} H_j (R dr_j) field_j``, evaluated by running the suffix 

627 recursion of :func:`_resolvent_field_laplace` one extra inward hop to ``r_w``. The extracted mixture 

628 is the **inflow-flux-weighted** azimuthal average -- ``|v_r(r_w, theta)| = (|A_0|/r_w)(1 - eps_w 

629 cos(theta))`` with ``eps_w = v_d r_w / |A_0|`` -- applied explicitly on the face modes: 

630 

631 cout_hat = c_0(r_w) - (eps_w / 2) (c_{+1}(r_w) + c_{-1}(r_w)). 

632 

633 (The face flux weighting and the ``D_rtheta`` face coupling are ``O(eps_w)`` individually but enter 

634 the drift loss at ``O(eps^2)``, the order of the loss itself; the engine-vs-FV loss tests pin them.) 

635 

636 Returns 

637 ------- 

638 ndarray 

639 ``cout_hat(s)``, shape ``(n_s, k)`` for a ``(n_quad, nm, k)`` ``field``. 

640 """ 

641 hs = _source_blocks(d, field, dr_weights, retardation_factor) # (n_quad, n_s, nm, k) 

642 tp = d["Tp"] 

643 s_acc = np.zeros_like(hs[0]) 

644 for i in range(hs.shape[0] - 1, 0, -1): 

645 s_acc = tp[i] @ (s_acc + hs[i]) 

646 face = tp[0] @ (s_acc + hs[0]) # (n_s, nm, k): the resident face modes 

647 return face[:, n_modes, :] - 0.5 * eps_w * (face[:, n_modes + 1, :] + face[:, n_modes - 1, :]) 

648 

649 

650def _rest_drift_field( 

651 field: npt.NDArray[np.floating], 

652 r_nodes: npt.NDArray[np.floating], 

653 dr_weights: npt.NDArray[np.floating], 

654 r_w: float, 

655 *, 

656 alpha_l: float, 

657 v_d: float, 

658 d_m: float, 

659 retardation_factor: float, 

660 t_rest: float, 

661 n_modes: int, 

662) -> npt.NDArray[np.floating]: 

663 r"""Advance the resident mode-field through a rest phase (``Q = 0``) under drift. 

664 

665 With the well shut the velocity field is the bare uniform drift ``v = v_d x_hat``, so the 

666 advection-dispersion kernel is exact in free space and constant-coefficient: the field translates 

667 down-gradient by ``delta = v_d t / R`` and spreads by the anisotropic Gaussian with 

668 ``sigma_x^2 = 2 (alpha_L |v_d| + D_m) t / R`` along the drift and ``sigma_y^2 = 2 D_m t / R`` across 

669 it (rank-1 Scheidegger tensor, ``alpha_T = 0``). The kernel is applied in real space: 

670 

671 1. the mode-field is evaluated at the Gauss-Hermite-shifted source points of every polar target node 

672 (barycentric interpolation on the radial Legendre nodes -- spectrally accurate for the smooth 

673 resident field -- times the azimuthal phase sum); 

674 2. Gauss-Hermite quadrature (``_REST_HERMITE`` nodes per Gaussian axis) averages the spread; 

675 3. the updated real-space samples are reprojected onto the modes by FFT over a uniform theta grid. 

676 

677 The shut well is closed by a **radial Neumann image**: source points falling inside the well disk are 

678 folded back across the face (``r' -> 2 r_w - r'``), the leading-order zero-dispersive-flux closure, 

679 which conserves the near-well mass (a zeroed disk would swallow ``O(sigma r_w / R_b^2)`` of the plume). 

680 The residual is the circle-vs-line curvature of the image and the neglected ``O(r_w^2/r^2)`` dipole 

681 distortion of the drift around the cylinder -- at ``v_d = 0``, ``D_m > 0`` the kernel differs from 

682 the scalar engine's exact well-respecting Bessel rest kernel only by this ``r_w``-scale closure 

683 residual for a stored plume (the public API dispatches ``v_d = 0`` to the scalar path anyway). 

684 Source points outside the radial grid carry no mass by the grid's containment guarantee 

685 (:func:`field_grid` provisions for the rest displacement). 

686 

687 An **honest truncation guard** protects the azimuthal representation: after reprojection, the energy 

688 in the harmonics just above the kept band (``M < |m| <= 2M``, available from the FFT grid) measures 

689 the translated field's spectral tail; if it exceeds ``_REST_TAIL_MAX`` of the field, the translated 

690 plume is too eccentric for the kept modes and a ``ValueError`` asks for a larger ``n_modes`` rather 

691 than silently folding the tail. 

692 

693 Returns 

694 ------- 

695 ndarray 

696 The advanced mode-field, shape ``(n_quad, nm, k)`` (matching ``field``). 

697 

698 Raises 

699 ------ 

700 ValueError 

701 If the translated field's azimuthal spectral tail exceeds the truncation guard (increase 

702 ``n_modes``). 

703 """ 

704 delta = v_d * t_rest / retardation_factor 

705 sig_x = np.sqrt(2.0 * (alpha_l * abs(v_d) + d_m) * t_rest / retardation_factor) 

706 sig_y = np.sqrt(2.0 * d_m * t_rest / retardation_factor) 

707 if delta == 0.0 and sig_x == 0.0: 

708 return field # v_d = 0 and D_m = 0: a rest phase is the identity 

709 n_quad, nm, k = field.shape 

710 modes = np.arange(-n_modes, n_modes + 1) 

711 theta = _theta_grid(n_modes) 

712 nth = theta.size 

713 x = r_nodes[:, None] * np.cos(theta)[None, :] # (n_quad, nth) 

714 y = r_nodes[:, None] * np.sin(theta)[None, :] 

715 interp = BarycentricInterpolator(r_nodes, field.reshape(n_quad, nm * k)) 

716 r_hi = float(r_nodes[-1]) 

717 

718 def field_at(xp: npt.NDArray[np.floating], yp: npt.NDArray[np.floating]) -> npt.NDArray[np.floating]: 

719 rp = np.hypot(xp, yp).ravel() 

720 thp = np.arctan2(yp, xp).ravel() 

721 # Source points inside the shut well are folded back across the face (radial Neumann image, 

722 # r' -> 2 r_w - r'): the leading-order zero-dispersive-flux closure at the well, which conserves 

723 # the near-well mass that a zeroed disk would silently swallow. Folded points land in 

724 # [r_w, 2 r_w], always inside the grid (field_grid's margins are many multiples of r_w). 

725 rp = np.where(rp < r_w, 2.0 * r_w - rp, rp) 

726 inside = rp <= r_hi # beyond the grid: no mass (containment guaranteed by field_grid) 

727 vals = np.zeros((rp.size, nm * k)) 

728 vals[inside] = interp(rp[inside]) 

729 phase = np.exp(1j * thp[:, None] * modes[None, :]) # (n_pts, nm) 

730 return np.einsum("pmk,pm->pk", vals.reshape(-1, nm, k), phase).real.reshape(*xp.shape, k) 

731 

732 zh, wh = np.polynomial.hermite.hermgauss(_REST_HERMITE) 

733 f_new = np.zeros((n_quad, nth, k)) 

734 if sig_y == 0.0: # D_m = 0: the Gaussian spread is 1-D along the drift 

735 for za, wa in zip(zh, wh, strict=True): 

736 f_new += wa * field_at(x - delta - np.sqrt(2.0) * sig_x * za, y) 

737 f_new /= np.sqrt(np.pi) 

738 else: 

739 for za, wa in zip(zh, wh, strict=True): 

740 x_shift = x - delta - np.sqrt(2.0) * sig_x * za 

741 for zb, wb in zip(zh, wh, strict=True): 

742 f_new += (wa * wb) * field_at(x_shift, y - np.sqrt(2.0) * sig_y * zb) 

743 f_new /= np.pi 

744 coeffs = np.fft.fft(f_new, axis=1) / nth # (n_quad, nth, k); c_m at index m mod nth 

745 measure = (r_nodes * dr_weights)[:, None, None] # radial area measure for the spectral-tail energy 

746 tail_idx = np.concatenate([np.arange(n_modes + 1, 2 * n_modes + 1), -np.arange(n_modes + 1, 2 * n_modes + 1)])