Coverage for src/gwtransport/_radial_asr_kernels.py: 100%
125 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-07 21:17 +0000
1r"""Exact per-phase (constant-Q) Laplace kernels for radial advection-dispersion.
3This private module holds the closed-form Laplace-domain transfer functions for a single
4fully-penetrating well in an infinite aquifer (the theory of the radial ASR knowledge base). For a
5constant-Q phase the volume-coordinate PDE ``d_t C + Q d_V C = d_V(D_V d_V C)`` has, in the Laplace
6domain ``t -> s``, the ODE ``G(r) C'' + (D_m - sigma_A A_0) C' - s r C = 0`` with
7``G(r) = alpha_L A_0 + D_m r``. The decaying branch on ``[r_w, inf)`` gives the resident solution
8``phi_s`` (Airy when ``D_m = 0``; Tricomi-U / Whittaker when ``D_m > 0``), and the Kreft-Zuber flux
9operator ``F[psi] = psi - (G/A_0) psi'`` builds the four injection/detection transfer functions.
11Two evaluation regimes
12----------------------
13* ``D_m = 0`` (mechanical dispersion only): Airy functions of complex argument via
14 ``scipy.special.airye`` (exponentially scaled), vectorized over the Laplace nodes. The scaling is
15 essential -- the raw ``phi_s = e^{r/2 alpha_L} Ai(zeta)`` overflows/underflows to NaN for
16 Peclet ``r/alpha_L`` beyond ~200 (the prefactor overflows while ``Ai`` underflows). All transfer
17 functions are *ratios* of ``phi_s`` / ``F[phi_s]`` at ``r`` and ``r_w``; evaluating them with the
18 Airy scaling factored into a single bounded log-amplitude keeps the ratio finite at any Peclet.
19* ``D_m > 0`` (molecular diffusion present): the decaying solution is the confluent-hypergeometric
20 (Tricomi-U / Whittaker) function, but it is evaluated through its LOG-DERIVATIVE ``L = phi'/phi`` -- a
21 vector Riccati ODE ``L' = -L^2 - ((D_m - sigma_A A_0)/G) L + s r/G`` integrated over the Laplace nodes
22 (:func:`_integrate_logderiv`). ``L`` is ``O(kappa)`` (bounded -- no 10^900 special-function
23 magnitudes), so the transfer functions (:func:`_transfer_riccati`) and the interior resolvent
24 (:func:`resolvent_riccati`) are assembled from O(1) quantities, with the divergent Sturm-Liouville
25 gauge carried in log space. This is exact to the de Hoog inversion floor at ANY ``A_0/D_m`` -- no
26 special-function precision cap, no arbitrary-precision dependency -- and continuously becomes the Airy
27 branch as ``D_m -> 0``. The exact flint/Arb Whittaker evaluation it replaced is retained as a
28 machine-precision test oracle (``tests/src/_radial_asr_whittaker_oracle.py``).
30Retardation enters by the standard linear-sorption rescaling of the constant-Q operator: dividing the
31retarded equation ``R d_t C + ... `` by ``R`` is the unretarded equation with ``A_0 -> A_0/R`` and
32``D_m -> D_m/R`` (mechanical dispersivity ``alpha_L`` is geometric and unchanged). Callers pass the
33physical ``A_0`` / ``D_m`` and a retardation factor; the rescaling is applied here.
35This file is part of gwtransport which is released under AGPL-3.0 license.
36See the ./LICENSE file or go to https://github.com/gwtransport/gwtransport/blob/main/LICENSE for full license details.
37"""
39import numpy as np
40import numpy.typing as npt
41from scipy.integrate import solve_ivp
42from scipy.special import airye, ive, kve
44# Injection / detection boundary types (Kreft-Zuber modes). "flux" applies the flux operator
45# F[psi] = psi - (G/A_0) psi'; any other value ("resident") uses psi directly.
46_FLUX = "flux"
48# Phase orientation for the interior two-point resolvent. "injection" is the divergent operator
49# (flow pushes outward, Robin/flux well BC); "extraction" is the convergent operator (flow pulls
50# inward, Danckwerts/Neumann well BC).
51_INJECTION = "injection"
54def _airy_amplitudes(
55 s: npt.NDArray[np.complexfloating], r: float, alpha_l: float, a0_eff: float
56) -> tuple[npt.NDArray[np.complexfloating], npt.NDArray[np.complexfloating], npt.NDArray[np.complexfloating]]:
57 r"""Scaled Airy building blocks at radius ``r`` for the ``D_m = 0`` branch.
59 Returns ``(log_amp, psi_resident, psi_flux)`` such that, with the Airy scaling factored out,
61 * ``phi_s(r) = exp(log_amp) * psi_resident``
62 * ``F[phi_s](r) = exp(log_amp) * psi_flux``
64 where ``psi_resident = Aie(zeta)`` and ``psi_flux = 0.5 Aie(zeta) - alpha_L beta^{1/3} Aipe(zeta)``
65 are O(1) (``Aie``/``Aipe`` are the exponentially scaled Airy functions, ``scipy.special.airye``),
66 and ``log_amp = r/(2 alpha_L) - (2/3) zeta^{3/2}`` carries the (bounded, once differenced between
67 ``r`` and ``r_w``) exponent. Here ``beta = s/(alpha_L a0_eff)`` and
68 ``zeta = beta^{1/3} r + beta^{-2/3}/(4 alpha_L^2)``.
70 Keeping the amplitude as a log and the Airy parts scaled is what prevents the high-Peclet
71 overflow: the raw ``phi_s`` over/under-flows, but every transfer function is a ratio in which the
72 ``exp(log_amp)`` factors difference to a bounded exponent.
74 Returns
75 -------
76 log_amp : ndarray of complex
77 Bounded log-amplitude ``r/(2 alpha_L) - (2/3) zeta^{3/2}`` per node.
78 psi_resident : ndarray of complex
79 Scaled resident amplitude ``Aie(zeta)``.
80 psi_flux : ndarray of complex
81 Scaled flux amplitude ``0.5 Aie(zeta) - alpha_L beta^{1/3} Aipe(zeta)``.
82 """
83 beta = s / (alpha_l * a0_eff)
84 b13 = beta ** (1.0 / 3.0) # principal cube root; s on the Bromwich contour has Re(s) > 0
85 zeta = b13 * r + 1.0 / (4.0 * alpha_l * alpha_l) * beta ** (-2.0 / 3.0)
86 eai, eaip, _, _ = airye(zeta)
87 psi_resident = eai
88 psi_flux = 0.5 * eai - alpha_l * b13 * eaip
89 log_amp = r / (2.0 * alpha_l) - (2.0 / 3.0) * zeta**1.5
90 return log_amp, psi_resident, psi_flux
93def transfer_function(
94 *,
95 s: npt.NDArray[np.complexfloating],
96 r: float,
97 r_w: float,
98 alpha_l: float,
99 a0: float,
100 d_m: float = 0.0,
101 retardation_factor: float = 1.0,
102 inject: str = _FLUX,
103 detect: str = _FLUX,
104) -> npt.NDArray[np.complexfloating]:
105 r"""Laplace-domain transfer function ``g_hat(s)`` for a constant-Q divergent phase.
107 Implements the four Kreft-Zuber injection/detection modes (KB Sec. 5) as ratios of ``phi_s``
108 (resident) and ``F[phi_s]`` (flux) evaluated at the detection radius ``r`` and the well ``r_w``:
110 =========== ============ ===================================
111 inject detect ``g_hat``
112 =========== ============ ===================================
113 flux (FF) flux ``F[phi_s](r) / F[phi_s](r_w)``
114 flux (FR) resident ``phi_s(r) / F[phi_s](r_w)``
115 resident RF flux ``F[phi_s](r) / phi_s(r_w)``
116 resident RR resident ``phi_s(r) / phi_s(r_w)``
117 =========== ============ ===================================
119 ``g_hat(0) = 1`` (mass conservation). The flux-flux (FF) mode is the package observable.
121 Parameters
122 ----------
123 s : ndarray of complex
124 Laplace nodes (conjugate to time). The Bromwich contour has ``Re(s) > 0``; do not pass
125 ``s = 0`` (``g_hat(0) = 1`` is the limit).
126 r : float
127 Detection radius (m), ``r >= r_w``.
128 r_w : float
129 Well (screen) radius (m).
130 alpha_l : float
131 Longitudinal dispersivity (m), ``> 0`` for the Airy branch.
132 a0 : float
133 Physical flow scale ``A_0 = |Q| / (2 pi b n)`` (m^2/day).
134 d_m : float, optional
135 Molecular diffusivity (m^2/day). ``0`` selects the Airy branch; ``> 0`` the Riccati
136 log-derivative branch.
137 retardation_factor : float, optional
138 Linear retardation ``R >= 1``. Default 1.
139 inject, detect : {'flux', 'resident'}, optional
140 Boundary type at the well (injection) and at ``r`` (detection). Default flux/flux (FF).
142 Returns
143 -------
144 ndarray of complex
145 ``g_hat(s)``, same shape as ``s``.
146 """
147 # Linear retardation rescales the constant-Q operator: A_0 -> A_0/R, D_m -> D_m/R (alpha_L is
148 # geometric and unchanged); the residence time then scales by R.
149 a0_eff, d_m_eff = a0 / retardation_factor, d_m / retardation_factor
150 s = np.asarray(s, dtype=complex)
152 if d_m_eff == 0.0:
153 # Airy branch: vectorized, overflow-safe (scaled-Airy amplitudes, bounded log-difference).
154 log_r, res_r, flux_r = _airy_amplitudes(s, r, alpha_l, a0_eff)
155 log_w, res_w, flux_w = _airy_amplitudes(s, r_w, alpha_l, a0_eff)
156 numer = flux_r if detect == _FLUX else res_r
157 denom = flux_w if inject == _FLUX else res_w
158 return np.exp(log_r - log_w) * (numer / denom)
160 # D_m > 0: Riccati log-derivative (numerical ODE on the log-derivative L = phi'/phi; exact to the
161 # de Hoog floor at any A_0/D_m, no special-function precision cap). Divergent orientation sigma_A = +1.
162 return _transfer_riccati(s.reshape(-1), r, r_w, alpha_l, a0_eff, d_m_eff, inject, detect).reshape(s.shape)
165def _resolvent_airy_pieces(
166 s: npt.NDArray[np.complexfloating],
167 r: npt.NDArray[np.floating] | float,
168 alpha_l: float,
169 a0_eff: float,
170 gauge_sign: float,
171) -> dict[str, npt.NDArray[np.complexfloating]]:
172 r"""Scaled-Airy building blocks for the interior two-point resolvent at radius ``r`` (``D_m = 0``).
174 The two homogeneous solutions of the constant-Q ODE (KB Sec. 4), in the gauge
175 ``e^{gauge_sign * r/(2 alpha_L)}`` (``+1`` divergent/injection, ``-1`` convergent/extraction):
177 * ``u_inf = s_inf * exp(gauge_sign r/2alpha_L - xi)`` -- the decaying branch (``Ai``),
178 * ``u_reg = s_reg * exp(gauge_sign r/2alpha_L + xiR)`` -- the growing branch (``Bi``),
180 with ``zeta = beta^{1/3} r + beta^{-2/3}/(4 alpha_L^2)``, ``beta = s/(alpha_L a0_eff)``,
181 ``xi = (2/3) zeta^{3/2}``. The scaled amplitudes ``s_inf = Aie``, ``s_reg = Bie`` and the
182 derivative amplitudes are O(1); the (possibly huge) gauge/Airy exponent is carried as the log
183 quantities ``xi`` and ``xiR``. **Crucially** ``scipy.special.airye`` scales ``Ai`` by
184 ``exp(+xi)`` but ``Bi`` by ``exp(-|Re xi|)``, so ``Ai = Aie exp(-xi)`` while ``Bi = Bie exp(+xiR)``
185 with ``xiR = |Re xi|``: the two differ for complex ``s`` (every de Hoog node), so they are tracked
186 separately. The caller forms only bounded exponent *differences* (no overflow to Pe ~ 600+).
188 Returns
189 -------
190 dict
191 ``s_inf, s_infp, s_reg, s_regp`` (scaled value and r-derivative amplitudes of ``u_inf``,
192 ``u_reg``) and the log-exponents ``xi`` (complex, for ``Ai``) and ``xiR`` (``|Re xi|``, ``Bi``).
193 """
194 beta = s / (alpha_l * a0_eff)
195 b13 = beta ** (1.0 / 3.0)
196 zeta = b13 * r + beta ** (-2.0 / 3.0) / (4.0 * alpha_l * alpha_l)
197 aie, aipe, bie, bipe = airye(zeta)
198 xi = (2.0 / 3.0) * zeta**1.5
199 g = gauge_sign / (2.0 * alpha_l)
200 return {
201 "s_inf": aie,
202 "s_infp": g * aie + b13 * aipe,
203 "s_reg": bie,
204 "s_regp": g * bie + b13 * bipe,
205 "xi": xi,
206 "xiR": np.abs(xi.real),
207 }
210def interior_resolvent(
211 *,
212 s: npt.NDArray[np.complexfloating],
213 r: float,
214 r_prime: npt.ArrayLike,
215 r_w: float,
216 alpha_l: float,
217 a0: float,
218 direction: str,
219) -> npt.NDArray[np.complexfloating]:
220 r"""Interior two-point Laplace resolvent ``Ghat(r, r'; s)`` of a constant-Q phase (``D_m = 0``).
222 ``Ghat`` is the kernel of the spatial resolvent ``(s - L)^{-1}`` of the per-phase generator ``L``
223 (KB Sec. 7 / addendum Sec. A3): the field after propagating an initial resident profile ``f`` for
224 flushed volume ``tau`` is ``f_resid(r) = L^{-1}_s[ int Ghat(r, r'; s) f(r') w(r') dr' ](tau)``, with
225 the Sturm-Liouville weight ``w(r') = (2 c_geo r'/alpha_L) e^{-gauge_sign r'/alpha_L} dr'`` supplied
226 by the caller. Built from the convergent/divergent Airy solutions with the physical well boundary
227 condition (Danckwerts/Neumann for extraction, Robin/flux for injection) and outgoing decay:
229 ``Ghat(r, r'; s) = -u_0(r_<) u_inf(r_>) / N(s)``, ``N(s) = P(r)[u_0 u_inf' - u_0' u_inf]``,
231 ``u_inf`` the decaying solution, ``u_0`` the well-BC solution, ``r_< = min(r, r')``,
232 ``r_> = max(r, r')``, ``P = e^{-gauge_sign r/alpha_L}`` (``N`` is constant in ``r`` -- the SL Abel
233 identity). The leading minus sign and ``N`` are pinned by the KB Sec. 7 duality: the well-face
234 trace ``Ghat(r_w, r'; s) w(r')`` equals the extraction arrival kernel. The normalization is
235 factored out before exponentiating, so the scaled-Airy form is overflow-safe.
237 The Laplace variable enters only through ``beta = s/(alpha_L a0)`` (``D_m = 0``); for the
238 flushed-volume clock pass ``s = flow_scale * p``, ``a0 = flow_scale/(2 c_geo)`` so that
239 ``beta = 2 c_geo p/alpha_L`` is flow-magnitude independent. Retardation is a pure clock rescale
240 handled by the caller (propagate over ``tau/R``); ``Ghat`` itself is retardation-free.
242 Parameters
243 ----------
244 s : ndarray of complex
245 Laplace nodes (conjugate to flushed volume). Shape ``(n_s,)``.
246 r : float
247 Output radius (m), ``>= r_w``.
248 r_prime : array-like
249 Source radius/radii (m), ``>= r_w``. Scalar or shape ``(n_r',)``.
250 r_w : float
251 Well radius (m).
252 alpha_l : float
253 Longitudinal dispersivity (m).
254 a0 : float
255 Flow scale ``A_0`` setting ``beta = s/(alpha_L a0)``.
256 direction : {'injection', 'extraction'}
257 Phase orientation: divergent (Robin well BC) or convergent (Neumann well BC).
259 Returns
260 -------
261 ndarray of complex
262 ``Ghat(r, r'; s)``, shape ``(n_s, n_r')`` (broadcast of ``s`` and ``r_prime``).
263 """
264 gauge_sign = 1.0 if direction == _INJECTION else -1.0
265 s = np.asarray(s, dtype=complex).reshape(-1, 1)
266 rp = np.atleast_1d(np.asarray(r_prime, dtype=float)).reshape(1, -1)
267 r_a = np.minimum(r, rp)
268 r_b = np.maximum(r, rp)
269 piece_a = _resolvent_airy_pieces(s, r_a, alpha_l, a0, gauge_sign)
270 piece_b = _resolvent_airy_pieces(s, r_b, alpha_l, a0, gauge_sign)
271 piece_w = _resolvent_airy_pieces(s, r_w, alpha_l, a0, gauge_sign)
272 return assemble_airy_resolvent(piece_a, piece_b, piece_w, r_a + r_b, alpha_l, gauge_sign)
275def assemble_airy_resolvent(
276 piece_a: dict[str, npt.NDArray[np.complexfloating]],
277 piece_b: dict[str, npt.NDArray[np.complexfloating]],
278 piece_w: dict[str, npt.NDArray[np.complexfloating]],
279 r_sum: npt.NDArray[np.floating],
280 alpha_l: float,
281 gauge_sign: float,
282 source_log_weight: npt.NDArray[np.floating] | float = 0.0,
283) -> npt.NDArray[np.complexfloating]:
284 r"""Assemble ``Ghat(r, r'; s) = -(pref_a e^{ea} - pref_b e^{eb})`` from precomputed scaled-Airy pieces.
286 ``piece_a``, ``piece_b`` are :func:`_resolvent_airy_pieces` at ``r_< = min(r, r')`` and
287 ``r_> = max(r, r')``; ``piece_w`` at ``r_w``; ``r_sum = r_< + r_> = r + r'`` (the radii enter the
288 bounded exponents only through their sum). The normalization ``N`` (with its huge exponent) is
289 factored into the bounded exponents ``ea, eb``, so the result is overflow-safe (Sec. 1b of the
290 plan). Splitting piece computation from assembly lets a caller evaluate the scaled Airy on a grid
291 of radii once and assemble every output node from prefix selection -- the ``O(n^2) -> O(n)``
292 saving the field propagator relies on.
294 The gauge term ``g * r_sum = gauge_sign (r + r')/(2 alpha_L)`` grows with the radii, so for the
295 field propagator its ``e^{g r_sum}`` factor is divergent (``+r/alpha_L`` injection) or its Airy
296 counterpart is (``-r/alpha_L`` extraction); on its own it overflows/underflows double precision at
297 Peclet ``r/alpha_L`` beyond ~700. It is tamed by the caller's Sturm-Liouville source weight
298 ``e^{-gauge_sign r'/alpha_L}``, whose LOG (``source_log_weight = -gauge_sign r'/alpha_L``, per
299 source node) must therefore be folded into the exponents *before* ``np.exp`` so the divergent parts
300 cancel to ``gauge_sign (r - r')/(2 alpha_L)`` (bounded by ``r_max/(2 alpha_L)``, and dominated by the
301 Airy decay) -- rather than overflowing to ``Inf`` and then meeting the taming factor as ``Inf * 0``.
302 The default ``0.0`` reproduces the bare interior resolvent (no source weight).
304 Parameters
305 ----------
306 piece_a, piece_b, piece_w : dict of ndarray
307 :func:`_resolvent_airy_pieces` at ``r_<``, ``r_>`` and ``r_w`` respectively.
308 r_sum : ndarray
309 ``r_< + r_> = r + r'`` (the radii enter the bounded exponents only through their sum).
310 alpha_l : float
311 Longitudinal dispersivity (m).
312 gauge_sign : float
313 ``+1`` divergent (injection, Robin well BC) / ``-1`` convergent (extraction, Neumann well BC).
314 source_log_weight : ndarray or float, optional
315 Log of the caller's Sturm-Liouville source weight per source node ``r'``
316 (``-gauge_sign r'/alpha_L``), broadcast over the source axis and folded into both exponents so
317 the divergent gauge cancels before ``np.exp``. Default ``0.0`` (bare resolvent, no weight).
319 Returns
320 -------
321 ndarray of complex
322 ``Ghat(r, r'; s)`` (times the source weight when ``source_log_weight`` is given), same broadcast
323 shape as the input pieces.
324 """
325 g = gauge_sign / (2.0 * alpha_l)
326 if gauge_sign < 0: # extraction: Danckwerts -> zero dispersive flux -> Neumann u_0'(r_w) = 0
327 bc_inf, bc_reg = piece_w["s_infp"], piece_w["s_regp"]
328 else: # injection: Robin/flux F[u_0](r_w) = 0, F[u] = u - alpha_L u'
329 bc_inf = piece_w["s_inf"] - alpha_l * piece_w["s_infp"]
330 bc_reg = piece_w["s_reg"] - alpha_l * piece_w["s_regp"]
331 # denom0 = scaled Wronskian piece at r_w (= b13/pi up to gauge).
332 denom0 = piece_w["s_regp"] * piece_w["s_inf"] - piece_w["s_infp"] * piece_w["s_reg"]
333 pref_a = bc_reg * piece_a["s_inf"] * piece_b["s_inf"] / (bc_inf * denom0)
334 pref_b = piece_a["s_reg"] * piece_b["s_inf"] / denom0
335 exp_a = g * r_sum + source_log_weight - (piece_a["xi"] + piece_b["xi"] - 2.0 * piece_w["xi"])
336 exp_b = g * r_sum + source_log_weight + (piece_a["xiR"] - piece_w["xiR"]) - (piece_b["xi"] - piece_w["xi"])
337 return -(pref_a * np.exp(exp_a) - pref_b * np.exp(exp_b))
340def rest_resolvent(
341 *,
342 s: npt.NDArray[np.complexfloating],
343 r: float,
344 r_prime: npt.ArrayLike,
345 r_w: float,
346 d_m: float,
347) -> npt.NDArray[np.complexfloating]:
348 r"""Interior two-point resolvent ``Ghat(r, r'; s)`` of a rest (``Q = 0``) phase -- pure diffusion.
350 With no flow the constant-Q ODE (KB Sec. 4) loses its advective and mechanical-dispersion terms and
351 collapses to the order-0 modified Bessel equation ``C'' + C'/r - (s/D_m) C = 0`` with
352 ``kappa = sqrt(s/D_m)``. The resident solution decaying as ``r -> inf`` is ``u_inf = K_0(kappa r)``;
353 the no-dispersive-flux (Danckwerts/Neumann) well solution is
354 ``u_0(r) = K_1(kappa r_w) I_0(kappa r) + I_1(kappa r_w) K_0(kappa r)`` (so ``u_0'(r_w) = 0``). The
355 Sturm-Liouville Wronskian normalization is ``N(s) = r [u_0 u_inf' - u_0' u_inf] = -K_1(kappa r_w)``
356 (constant in ``r``), giving
358 ``Ghat(r, r'; s) = -u_0(r_<) u_inf(r_>) / N(s) = u_0(r_<) K_0(kappa r_>) / K_1(kappa r_w)``,
360 ``r_< = min(r, r')``, ``r_> = max(r, r')``. It is evaluated overflow-safe with the exponentially
361 scaled modified Bessel functions (``scipy.special.ive``/``kve``): each term is a ratio whose scaling
362 exponents difference to a bounded value, so the growing ``I_0`` never overflows at high ``kappa r``.
363 The clock is wall-clock time (molecular diffusion is autonomous in ``t``); pair with the source
364 measure ``w(r') dr' = (r'/D_m) dr'`` (the Sturm-Liouville weight) when superposing a resident field.
366 Parameters
367 ----------
368 s : ndarray of complex
369 Laplace nodes (conjugate to wall-clock time). Shape ``(n_s,)``.
370 r : float
371 Output radius (m), ``>= r_w``.
372 r_prime : array-like
373 Source radius/radii (m), ``>= r_w``. Scalar or shape ``(n_r',)``.
374 r_w : float
375 Well radius (m).
376 d_m : float
377 Molecular diffusivity (m^2/day), ``> 0``.
379 Returns
380 -------
381 ndarray of complex
382 ``Ghat(r, r'; s)``, shape ``(n_s, n_r')`` (broadcast of ``s`` and ``r_prime``).
383 """
384 s = np.asarray(s, dtype=complex).reshape(-1, 1)
385 rp = np.atleast_1d(np.asarray(r_prime, dtype=float)).reshape(1, -1)
386 kappa = np.sqrt(s / d_m) # principal root; Re(s) > 0 on the Bromwich contour gives Re(kappa) > 0
387 r_lt = np.minimum(r, rp)
388 r_gt = np.maximum(r, rp)
389 z_lt, z_gt, z_w = kappa * r_lt, kappa * r_gt, kappa * r_w
390 # Ghat = [I_0(z_<) + (I_1(z_w)/K_1(z_w)) K_0(z_<)] K_0(z_>); split so the scaled-Bessel scaling
391 # exponents (ive scales by e^{-|Re z|}, kve by e^{+z}) difference to bounded values -- no overflow.
392 # Both terms carry the scaling exponents in a SINGLE np.exp of the combined (bounded) sum: the outer
393 # term's ``|Re z_w| + z_w`` alone overflows at Re(z_w) > ~354, but ``|Re z_w| + z_w - z_lt - z_gt``
394 # <= 0 (since r_w <= r_< <= r_>) so exponentiating the sum is overflow-safe.
395 term_inner = ive(0, z_lt) * kve(0, z_gt) * np.exp(np.abs(z_lt.real) - z_gt)
396 term_outer = (
397 (ive(1, z_w) / kve(1, z_w)) * kve(0, z_lt) * kve(0, z_gt) * np.exp(np.abs(z_w.real) + z_w - z_lt - z_gt)
398 )
399 return term_inner + term_outer
402# ---------------------------------------------------------------------------
403# D_m > 0 branch: Riccati log-derivative (numerical ODE, double precision)
404# ---------------------------------------------------------------------------
405# The decaying resident solution and the well-regular solution are tracked by their log-derivative
406# L = C'/C of the constant-Q ODE, integrated as a vector ODE over the Laplace nodes. L is O(kappa)
407# (bounded -- no 10^900 special-function magnitudes), so the transfer functions and the
408# Sturm-Liouville interior resolvent are assembled from O(1) quantities. This is exact to the de Hoog
409# inversion floor at ANY A_0/D_m -- no special-function precision blow-up, no tractability cap -- and
410# continuously becomes the Airy branch as D_m -> 0.
411_RICCATI_RTOL = 1e-12
412_RICCATI_ATOL = 1e-13
413# Outer boundary for the inward (decaying) integration. The truncated asymptotic IC at r_far washes out
414# because the recessive solution is the inward attractor (damping ~ e^{-2 Re(kappa)(r_far - r)}), but the
415# slowest (smallest Re(kappa)) Laplace node also needs r_far far enough that its z = 2 kappa r_far is in
416# the large-z asymptotic. So r_far is extended by ~_RICCATI_RFAR_DECAY decay lengths 1/Re(kappa_min)
417# beyond the field, floored at _RICCATI_RFAR_MULT * r_max and cost-capped at _RICCATI_RFAR_CAP * r_max
418# (the floor on Re(kappa_min) keeps z at the cap ~ 2 * _RICCATI_RFAR_DECAY, large enough for any node).
419_RICCATI_RFAR_MULT = 8.0
420_RICCATI_RFAR_DECAY = 22.0
421_RICCATI_RFAR_CAP = 500.0
424def _integrate_logderiv(
425 s: npt.NDArray[np.complexfloating],
426 radii: npt.ArrayLike,
427 r_w: float,
428 alpha_l: float,
429 a0_eff: float,
430 d_m_eff: float,
431 sigma_a: int,
432 branch: str,
433) -> tuple[npt.NDArray[np.complexfloating], npt.NDArray[np.complexfloating]]:
434 r"""Vector Riccati integration of the log-derivative ``L = C'/C`` for the ``D_m > 0`` branch.
436 Solves ``L' = -L^2 - ((D_m - sigma_A A_0)/G) L + s r/G`` (``G = alpha_L A_0 + D_m r``) over all
437 Laplace nodes ``s`` at once, carrying the running integral ``J = int_{r_w}^{r} L dr``.
439 * ``branch='decaying'``: inward from ``r_far`` with the recessive asymptotic IC
440 ``L(r_far) = -kappa - a/x`` (``kappa = sqrt(s/D_m)``, ``x = r_far + a*``, ``a* = alpha_L A_0/D_m``,
441 ``a = b/2 - kappa a*/2``, ``b = 1 - sigma_A A_0/D_m``). The decaying solution is the inward
442 attractor, so the result is insensitive to ``r_far``.
443 * ``branch='regular'``: outward from ``r_w`` with the well-BC IC ``L(r_w) = A_0/G(r_w)`` (injection,
444 Robin ``F[u_0](r_w)=0``) or ``0`` (extraction, Neumann ``u_0'(r_w)=0``) -- both ``s``-independent.
446 ``sigma_a`` is ``+1`` divergent (injection) / ``-1`` convergent (extraction); ``radii`` (all ``>= r_w``)
447 are where ``L`` and ``J`` are returned.
449 Returns
450 -------
451 ld : ndarray of complex, shape (n_s, n_radii)
452 Log-derivative ``L`` at each requested radius.
453 jj : ndarray of complex, shape (n_s, n_radii)
454 ``int_{r_w}^{r} L dr`` at each requested radius.
455 """
456 s = np.asarray(s, dtype=complex).reshape(-1)
457 n = s.size
458 radii = np.atleast_1d(np.asarray(radii, dtype=float))
459 r_max = max(float(radii.max()), r_w)
461 def rhs(r: float, y: npt.NDArray[np.complexfloating]) -> npt.NDArray[np.complexfloating]:
462 ld = y[:n]
463 g = alpha_l * a0_eff + d_m_eff * r
464 d_ld = -(ld * ld) - ((d_m_eff - sigma_a * a0_eff) / g) * ld + s * r / g
465 return np.concatenate([d_ld, ld])
467 if branch == "decaying":
468 astar = alpha_l * a0_eff / d_m_eff
469 kappa = np.sqrt(s / d_m_eff)
470 # r_far must put the slowest node deep in the large-z asymptotic so the truncated IC washes out:
471 # extend by ~_RICCATI_RFAR_DECAY decay lengths 1/Re(kappa_min) (Re(kappa) floored so the extension
472 # is cost-capped at _RICCATI_RFAR_CAP * r_max while z = 2 kappa r_far stays large there).
473 re_kmin = max(float(kappa.real.min()), _RICCATI_RFAR_DECAY / (_RICCATI_RFAR_CAP * r_max))
474 r_far = max(_RICCATI_RFAR_MULT * r_max, r_max + _RICCATI_RFAR_DECAY / re_kmin)
475 a = (1.0 - sigma_a * a0_eff / d_m_eff) / 2.0 - kappa * astar / 2.0
476 y0 = np.concatenate([-kappa - a / (r_far + astar), np.zeros(n, dtype=complex)])
477 sol = solve_ivp(
478 rhs, [r_far, r_w], y0, rtol=_RICCATI_RTOL, atol=_RICCATI_ATOL, dense_output=True, method="DOP853"
479 )
480 y = sol.sol(radii) # dense output at all radii at once -> shape (2n, n_radii)
481 j_w = sol.sol(r_w)[n:, None] # re-anchor J to int_{r_w}^{r} (the IC put J(r_far) = 0)
482 return y[:n], y[n:] - j_w
484 # regular branch: outward from r_w to r_max -- the growing solution is the stable outward attractor and
485 # the well-BC IC (s-independent) is exact, so no washout is needed; it need only reach the field.
486 ld0 = a0_eff / (alpha_l * a0_eff + d_m_eff * r_w) if sigma_a > 0 else 0.0
487 y0 = np.concatenate([np.full(n, ld0, dtype=complex), np.zeros(n, dtype=complex)])
488 sol = solve_ivp(rhs, [r_w, r_max], y0, rtol=_RICCATI_RTOL, atol=_RICCATI_ATOL, dense_output=True, method="DOP853")
489 y = sol.sol(radii)
490 return y[:n], y[n:]
493def _transfer_riccati(
494 s: npt.NDArray[np.complexfloating],
495 r: float,
496 r_w: float,
497 alpha_l: float,
498 a0_eff: float,
499 d_m_eff: float,
500 inject: str,
501 detect: str,
502) -> npt.NDArray[np.complexfloating]:
503 r"""Four Kreft-Zuber transfer modes for the ``D_m > 0`` branch via the decaying log-derivative.
505 With ``E = phi(r)/phi(r_w) = exp(int_{r_w}^{r} L)`` and the flux factor ``f(r) = 1 - (alpha_L +
506 D_m r/A_0) L(r)`` (so ``F[phi](r) = phi(r) f(r)``), the modes are ``FF = E f(r)/f(r_w)``,
507 ``FR = E/f(r_w)``, ``RF = E f(r)``, ``RR = E``. ``sigma_A = +1`` (the divergent operator). ``inject``
508 / ``detect`` select the Kreft-Zuber well / detection boundary ('flux' or 'resident').
510 Returns
511 -------
512 ndarray of complex
513 ``g_hat(s)`` for the requested ``(inject, detect)`` mode, shape ``(n_s,)``.
514 """
515 ld, jj = _integrate_logderiv(s, [r, r_w], r_w, alpha_l, a0_eff, d_m_eff, +1, "decaying")
516 l_r, l_w = ld[:, 0], ld[:, 1]
517 e = np.exp(jj[:, 0]) # phi(r)/phi(r_w)
518 f_r = 1.0 - (alpha_l + d_m_eff * r / a0_eff) * l_r
519 f_w = 1.0 - (alpha_l + d_m_eff * r_w / a0_eff) * l_w
520 numer = e * (f_r if detect == _FLUX else 1.0)
521 denom = f_w if inject == _FLUX else 1.0
522 return numer / denom
525def resolvent_riccati(
526 *,
527 s: npt.NDArray[np.complexfloating],
528 field: npt.NDArray[np.floating],
529 r_nodes: npt.NDArray[np.floating],
530 dr_weights: npt.NDArray[np.floating],
531 r_w: float,
532 alpha_l: float,
533 a0_eff: float,
534 d_m_eff: float,
535 direction: str,
536) -> npt.NDArray[np.complexfloating]:
537 r"""Interior Sturm-Liouville resolvent applied to a source ``field``, ``D_m > 0``, via log-derivatives.
539 Returns ``F(s)_{k,i} = sum_j Ghat(r_i, r_j; s_k) field_j w_j`` with the SL measure
540 ``w_j = r_j G(r_j)^{b-1} dr_j`` (``b = 1 - sigma_A A_0/D_m``). The Green's function
541 ``Ghat = phi_+(r_<) phi_-(r_>)/(-pW)`` is built from the decaying (``phi_-``, inward) and
542 well-regular (``phi_+``, outward) solutions normalized at ``r_w``; ``pW = G(r_w)^b (L_-(r_w) -
543 L_+(r_w))``. The divergent gauge ``G(r_w)^b`` (``b ~ -A_0/D_m``) is carried in LOG space
544 (``LG_j = b ln(G(r_j)/G(r_w)) - ln G(r_j)``, bounded as ``A_0/D_m -> inf``) so it never
545 underflows -- the assembly stays in double precision at any ``A_0/D_m``.
547 Parameters
548 ----------
549 s : ndarray of complex
550 Laplace nodes (conjugate to wall-clock time).
551 field : ndarray
552 Source resident-deviation profile on ``r_nodes``.
553 r_nodes : ndarray
554 Radial quadrature nodes (m), increasing, ``> r_w``.
555 dr_weights : ndarray
556 Quadrature weights for ``r_nodes`` (the ``dr`` measure).
557 r_w : float
558 Well radius (m).
559 alpha_l, a0_eff, d_m_eff : float
560 Dispersivity and the retardation-effective ``A_0`` / ``D_m``.
561 direction : {'injection', 'extraction'}
562 Phase orientation (divergent Robin / convergent Neumann well BC).
564 Returns
565 -------
566 ndarray of complex
567 ``F(s)_{k,i}`` -- the resolvent applied to ``field``, shape ``(n_s, n_nodes)``.
568 """
569 sigma_a = 1 if direction == _INJECTION else -1
570 s = np.asarray(s, dtype=complex).reshape(-1)
571 b = 1.0 - sigma_a * a0_eff / d_m_eff
572 rad = np.concatenate(([r_w], r_nodes))
573 ld_m, jj_m = _integrate_logderiv(s, rad, r_w, alpha_l, a0_eff, d_m_eff, sigma_a, "decaying")
574 lm_w = ld_m[:, 0] # L_-(r_w)
575 im = jj_m[:, 1:] # int_{r_w}^{r_i} L_- (n_s, n)
576 _, ip = _integrate_logderiv(s, r_nodes, r_w, alpha_l, a0_eff, d_m_eff, sigma_a, "regular") # int L_+
577 lp_w = a0_eff / (alpha_l * a0_eff + d_m_eff * r_w) if sigma_a > 0 else 0.0
578 g_nodes = alpha_l * a0_eff + d_m_eff * r_nodes
579 g_w = alpha_l * a0_eff + d_m_eff * r_w
580 lg = b * np.log(g_nodes / g_w) - np.log(g_nodes) # bounded gauge in log space
581 c = field * r_nodes * dr_weights
582 pmat = c[None, :] * np.exp(ip + lg[None, :]) # (n_s, n)
583 smat = c[None, :] * np.exp(im + lg[None, :])
584 prefix = np.cumsum(pmat, axis=1) # sum_{j<=i}
585 suffix = np.cumsum(smat[:, ::-1], axis=1)[:, ::-1] - smat # sum_{j>i}
586 denom = (lm_w - lp_w)[:, None]
587 return -(np.exp(im) * prefix + np.exp(ip) * suffix) / denom