Coverage for src/gwtransport/_radial_asr_kernels.py: 100%
114 statements
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-04 21:13 +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``. An exact flint/Arb Whittaker evaluation is the machine-precision test oracle
28 (``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# Detection boundary type (Kreft-Zuber mode). "flux" applies the flux operator
45# F[psi] = psi - (G/A_0) psi'; any other value ("resident") uses psi directly. Injection is always
46# through the Kreft-Zuber flux boundary at the well.
47_FLUX = "flux"
49# Phase orientation for the interior two-point resolvent. "injection" is the divergent operator
50# (flow pushes outward, Robin/flux well BC); "extraction" is the convergent operator (flow pulls
51# inward, Danckwerts/Neumann well BC).
52_INJECTION = "injection"
55def _airy_amplitudes(
56 s: npt.NDArray[np.complexfloating], r: float, alpha_l: float, a0_eff: float
57) -> tuple[npt.NDArray[np.complexfloating], npt.NDArray[np.complexfloating], npt.NDArray[np.complexfloating]]:
58 r"""Scaled Airy building blocks at radius ``r`` for the ``D_m = 0`` branch.
60 Returns ``(log_amp, psi_resident, psi_flux)`` such that, with the Airy scaling factored out,
62 * ``phi_s(r) = exp(log_amp) * psi_resident``
63 * ``F[phi_s](r) = exp(log_amp) * psi_flux``
65 where ``psi_resident = Aie(zeta)`` and ``psi_flux = 0.5 Aie(zeta) - alpha_L beta^{1/3} Aipe(zeta)``
66 are O(1) (``Aie``/``Aipe`` are the exponentially scaled Airy functions, ``scipy.special.airye``),
67 and ``log_amp = r/(2 alpha_L) - (2/3) zeta^{3/2}`` carries the (bounded, once differenced between
68 ``r`` and ``r_w``) exponent. Here ``beta = s/(alpha_L a0_eff)`` and
69 ``zeta = beta^{1/3} r + beta^{-2/3}/(4 alpha_L^2)``.
71 Keeping the amplitude as a log and the Airy parts scaled is what prevents the high-Peclet
72 overflow: the raw ``phi_s`` over/under-flows, but every transfer function is a ratio in which the
73 ``exp(log_amp)`` factors difference to a bounded exponent.
75 Returns
76 -------
77 log_amp : ndarray of complex
78 Bounded log-amplitude ``r/(2 alpha_L) - (2/3) zeta^{3/2}`` per node.
79 psi_resident : ndarray of complex
80 Scaled resident amplitude ``Aie(zeta)``.
81 psi_flux : ndarray of complex
82 Scaled flux amplitude ``0.5 Aie(zeta) - alpha_L beta^{1/3} Aipe(zeta)``.
83 """
84 beta = s / (alpha_l * a0_eff)
85 b13 = beta ** (1.0 / 3.0) # principal cube root; s on the Bromwich contour has Re(s) > 0
86 zeta = b13 * r + 1.0 / (4.0 * alpha_l * alpha_l) * beta ** (-2.0 / 3.0)
87 eai, eaip, _, _ = airye(zeta)
88 psi_resident = eai
89 psi_flux = 0.5 * eai - alpha_l * b13 * eaip
90 log_amp = r / (2.0 * alpha_l) - (2.0 / 3.0) * zeta**1.5
91 return log_amp, psi_resident, psi_flux
94def transfer_function(
95 *,
96 s: npt.NDArray[np.complexfloating],
97 r: float,
98 r_w: float,
99 alpha_l: float,
100 a0: float,
101 d_m: float = 0.0,
102 retardation_factor: float = 1.0,
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 Injection is through the Kreft-Zuber flux boundary at the well; the detection mode selects
108 ``F[phi_s]`` (flux) or ``phi_s`` (resident) at the detection radius ``r``:
110 ============ ===================================
111 detect ``g_hat``
112 ============ ===================================
113 flux (FF) ``F[phi_s](r) / F[phi_s](r_w)``
114 resident (FR) ``phi_s(r) / F[phi_s](r_w)``
115 ============ ===================================
117 ``g_hat(0) = 1`` (mass conservation). The flux-flux (FF) mode is the package observable.
119 Parameters
120 ----------
121 s : ndarray of complex
122 Laplace nodes (conjugate to time). The Bromwich contour has ``Re(s) > 0``; do not pass
123 ``s = 0`` (``g_hat(0) = 1`` is the limit).
124 r : float
125 Detection radius (m), ``r >= r_w``.
126 r_w : float
127 Well (screen) radius (m).
128 alpha_l : float
129 Longitudinal dispersivity (m), ``> 0`` for the Airy branch.
130 a0 : float
131 Physical flow scale ``A_0 = |Q| / (2 pi b n)`` (m^2/day).
132 d_m : float, optional
133 Molecular diffusivity (m^2/day). ``0`` selects the Airy branch; ``> 0`` the Riccati
134 log-derivative branch.
135 retardation_factor : float, optional
136 Linear retardation ``R >= 1``. Default 1.
137 detect : {'flux', 'resident'}, optional
138 Detection boundary type at ``r``. Default flux (FF).
140 Returns
141 -------
142 ndarray of complex
143 ``g_hat(s)``, same shape as ``s``.
145 Notes
146 -----
147 Validity floor: for ``|s|`` below ~1e-10 the scaled-Airy building blocks of the
148 ``D_m = 0`` branch degenerate and the result is NaN instead of the analytic limit
149 ``g_hat -> 1``. This is unreachable through the public API: the smallest de Hoog
150 node is ``gamma ~ -ln(tol) / (2 * 1.3 * max(mu, tau)) ~ 8 / max(mu, tau)``, which
151 stays orders of magnitude above 1e-10 for any physical phase volume.
152 """
153 # Linear retardation rescales the constant-Q operator: A_0 -> A_0/R, D_m -> D_m/R (alpha_L is
154 # geometric and unchanged); the residence time then scales by R.
155 a0_eff, d_m_eff = a0 / retardation_factor, d_m / retardation_factor
156 s = np.asarray(s, dtype=complex)
158 if d_m_eff == 0.0:
159 # Airy branch: vectorized, overflow-safe (scaled-Airy amplitudes, bounded log-difference).
160 log_r, res_r, flux_r = _airy_amplitudes(s, r, alpha_l, a0_eff)
161 log_w, _, flux_w = _airy_amplitudes(s, r_w, alpha_l, a0_eff)
162 numer = flux_r if detect == _FLUX else res_r
163 return np.exp(log_r - log_w) * (numer / flux_w)
165 # D_m > 0: Riccati log-derivative (numerical ODE on the log-derivative L = phi'/phi; exact to the
166 # de Hoog floor at any A_0/D_m, no special-function precision cap). Divergent orientation sigma_A = +1.
167 return _transfer_riccati(s.reshape(-1), r, r_w, alpha_l, a0_eff, d_m_eff, detect).reshape(s.shape)
170def _resolvent_airy_pieces(
171 s: npt.NDArray[np.complexfloating],
172 r: npt.NDArray[np.floating] | float,
173 alpha_l: float,
174 a0_eff: float,
175 gauge_sign: float,
176) -> dict[str, npt.NDArray[np.complexfloating]]:
177 r"""Scaled-Airy building blocks for the interior two-point resolvent at radius ``r`` (``D_m = 0``).
179 The two homogeneous solutions of the constant-Q ODE, in the gauge
180 ``e^{gauge_sign * r/(2 alpha_L)}`` (``+1`` divergent/injection, ``-1`` convergent/extraction):
182 * ``u_inf = s_inf * exp(gauge_sign r/2alpha_L - xi)`` -- the decaying branch (``Ai``),
183 * ``u_reg = s_reg * exp(gauge_sign r/2alpha_L + xiR)`` -- the growing branch (``Bi``),
185 with ``zeta = beta^{1/3} r + beta^{-2/3}/(4 alpha_L^2)``, ``beta = s/(alpha_L a0_eff)``,
186 ``xi = (2/3) zeta^{3/2}``. The scaled amplitudes ``s_inf = Aie``, ``s_reg = Bie`` and the
187 derivative amplitudes are O(1); the (possibly huge) gauge/Airy exponent is carried as the log
188 quantities ``xi`` and ``xiR``. **Crucially** ``scipy.special.airye`` scales ``Ai`` by
189 ``exp(+xi)`` but ``Bi`` by ``exp(-|Re xi|)``, so ``Ai = Aie exp(-xi)`` while ``Bi = Bie exp(+xiR)``
190 with ``xiR = |Re xi|``: the two differ for complex ``s`` (every de Hoog node), so they are tracked
191 separately. The caller forms only bounded exponent *differences* (no overflow to Pe ~ 600+).
193 Returns
194 -------
195 dict
196 ``s_inf, s_infp, s_reg, s_regp`` (scaled value and r-derivative amplitudes of ``u_inf``,
197 ``u_reg``) and the log-exponents ``xi`` (complex, for ``Ai``) and ``xiR`` (``|Re xi|``, ``Bi``).
198 """
199 beta = s / (alpha_l * a0_eff)
200 b13 = beta ** (1.0 / 3.0)
201 zeta = b13 * r + beta ** (-2.0 / 3.0) / (4.0 * alpha_l * alpha_l)
202 aie, aipe, bie, bipe = airye(zeta)
203 xi = (2.0 / 3.0) * zeta**1.5
204 g = gauge_sign / (2.0 * alpha_l)
205 return {
206 "s_inf": aie,
207 "s_infp": g * aie + b13 * aipe,
208 "s_reg": bie,
209 "s_regp": g * bie + b13 * bipe,
210 "xi": xi,
211 "xiR": np.abs(xi.real),
212 }
215def assemble_airy_resolvent(
216 piece_a: dict[str, npt.NDArray[np.complexfloating]],
217 piece_b: dict[str, npt.NDArray[np.complexfloating]],
218 piece_w: dict[str, npt.NDArray[np.complexfloating]],
219 r_sum: npt.NDArray[np.floating],
220 alpha_l: float,
221 gauge_sign: float,
222 source_log_weight: npt.NDArray[np.floating] | float = 0.0,
223) -> npt.NDArray[np.complexfloating]:
224 r"""Assemble ``Ghat(r, r'; s) = -(pref_a e^{ea} - pref_b e^{eb})`` from precomputed scaled-Airy pieces.
226 ``Ghat`` is the kernel of the spatial resolvent ``(s - L)^{-1}`` of the per-phase generator ``L``:
227 the field after propagating an initial resident profile ``f`` for flushed volume ``tau`` is
228 ``f_resid(r) = L^{-1}_s[ int Ghat(r, r'; s) f(r') w(r') dr' ](tau)``, with the Sturm-Liouville
229 weight ``w(r') = (2 c_geo r'/alpha_L) e^{-gauge_sign r'/alpha_L} dr'`` supplied by the caller. Built
230 from the convergent/divergent Airy solutions with the physical well boundary condition
231 (Danckwerts/Neumann for extraction, Robin/flux for injection) and outgoing decay,
233 ``Ghat(r, r'; s) = -u_0(r_<) u_inf(r_>) / N(s)``, ``N(s) = P(r)[u_0 u_inf' - u_0' u_inf]``,
235 ``u_inf`` the decaying solution, ``u_0`` the well-BC solution, ``P = e^{-gauge_sign r/alpha_L}``
236 (``N`` is constant in ``r`` -- the SL Abel identity). The leading minus sign and ``N`` are pinned by
237 the extraction duality: the well-face trace ``Ghat(r_w, r'; s) w(r')`` equals the extraction arrival
238 kernel.
240 ``piece_a``, ``piece_b`` are :func:`_resolvent_airy_pieces` at ``r_< = min(r, r')`` and
241 ``r_> = max(r, r')``; ``piece_w`` at ``r_w``; ``r_sum = r_< + r_> = r + r'`` (the radii enter the
242 bounded exponents only through their sum). The normalization ``N`` (with its huge exponent) is
243 factored into the bounded exponents ``ea, eb``, so the result is overflow-safe. Splitting piece
244 computation from assembly lets a caller evaluate the scaled Airy on a grid of radii once and
245 assemble every output node from prefix selection -- the ``O(n^2) -> O(n)`` saving the field
246 propagator relies on.
248 The Laplace variable enters only through ``beta = s/(alpha_L a0)`` (``D_m = 0``); for the
249 flushed-volume clock the caller passes the canonical ``s = 2 c_geo p``, ``a0 = 1`` so that
250 ``beta = 2 c_geo p/alpha_L`` is flow-magnitude independent. Retardation is a pure clock rescale
251 handled by the caller (propagate over ``tau/R``); ``Ghat`` itself is retardation-free.
253 The gauge term ``g * r_sum = gauge_sign (r + r')/(2 alpha_L)`` grows with the radii, so for the
254 field propagator its ``e^{g r_sum}`` factor is divergent (``+r/alpha_L`` injection) or its Airy
255 counterpart is (``-r/alpha_L`` extraction); on its own it overflows/underflows double precision at
256 Peclet ``r/alpha_L`` beyond ~700. It is tamed by the caller's Sturm-Liouville source weight
257 ``e^{-gauge_sign r'/alpha_L}``, whose LOG (``source_log_weight = -gauge_sign r'/alpha_L``, per
258 source node) must therefore be folded into the exponents *before* ``np.exp`` so the divergent parts
259 cancel to ``gauge_sign (r - r')/(2 alpha_L)`` (bounded by ``r_max/(2 alpha_L)``, and dominated by the
260 Airy decay) -- rather than overflowing to ``Inf`` and then meeting the taming factor as ``Inf * 0``.
261 The default ``0.0`` reproduces the bare interior resolvent (no source weight).
263 Parameters
264 ----------
265 piece_a, piece_b, piece_w : dict of ndarray
266 :func:`_resolvent_airy_pieces` at ``r_<``, ``r_>`` and ``r_w`` respectively.
267 r_sum : ndarray
268 ``r_< + r_> = r + r'`` (the radii enter the bounded exponents only through their sum).
269 alpha_l : float
270 Longitudinal dispersivity (m).
271 gauge_sign : float
272 ``+1`` divergent (injection, Robin well BC) / ``-1`` convergent (extraction, Neumann well BC).
273 source_log_weight : ndarray or float, optional
274 Log of the caller's Sturm-Liouville source weight per source node ``r'``
275 (``-gauge_sign r'/alpha_L``), broadcast over the source axis and folded into both exponents so
276 the divergent gauge cancels before ``np.exp``. Default ``0.0`` (bare resolvent, no weight).
278 Returns
279 -------
280 ndarray of complex
281 ``Ghat(r, r'; s)`` (times the source weight when ``source_log_weight`` is given), same broadcast
282 shape as the input pieces.
283 """
284 g = gauge_sign / (2.0 * alpha_l)
285 if gauge_sign < 0: # extraction: Danckwerts -> zero dispersive flux -> Neumann u_0'(r_w) = 0
286 bc_inf, bc_reg = piece_w["s_infp"], piece_w["s_regp"]
287 else: # injection: Robin/flux F[u_0](r_w) = 0, F[u] = u - alpha_L u'
288 bc_inf = piece_w["s_inf"] - alpha_l * piece_w["s_infp"]
289 bc_reg = piece_w["s_reg"] - alpha_l * piece_w["s_regp"]
290 # denom0 = scaled Wronskian piece at r_w (= b13/pi up to gauge).
291 denom0 = piece_w["s_regp"] * piece_w["s_inf"] - piece_w["s_infp"] * piece_w["s_reg"]
292 pref_a = bc_reg * piece_a["s_inf"] * piece_b["s_inf"] / (bc_inf * denom0)
293 pref_b = piece_a["s_reg"] * piece_b["s_inf"] / denom0
294 exp_a = g * r_sum + source_log_weight - (piece_a["xi"] + piece_b["xi"] - 2.0 * piece_w["xi"])
295 exp_b = g * r_sum + source_log_weight + (piece_a["xiR"] - piece_w["xiR"]) - (piece_b["xi"] - piece_w["xi"])
296 return -(pref_a * np.exp(exp_a) - pref_b * np.exp(exp_b))
299def rest_resolvent(
300 *,
301 s: npt.NDArray[np.complexfloating],
302 r: float,
303 r_prime: npt.ArrayLike,
304 r_w: float,
305 d_m: float,
306) -> npt.NDArray[np.complexfloating]:
307 r"""Interior two-point resolvent ``Ghat(r, r'; s)`` of a rest (``Q = 0``) phase -- pure diffusion.
309 With no flow the constant-Q ODE loses its advective and mechanical-dispersion terms and
310 collapses to the order-0 modified Bessel equation ``C'' + C'/r - (s/D_m) C = 0`` with
311 ``kappa = sqrt(s/D_m)``. The resident solution decaying as ``r -> inf`` is ``u_inf = K_0(kappa r)``;
312 the no-dispersive-flux (Danckwerts/Neumann) well solution is
313 ``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
314 Sturm-Liouville Wronskian normalization is ``N(s) = r [u_0 u_inf' - u_0' u_inf] = -K_1(kappa r_w)``
315 (constant in ``r``), giving
317 ``Ghat(r, r'; s) = -u_0(r_<) u_inf(r_>) / N(s) = u_0(r_<) K_0(kappa r_>) / K_1(kappa r_w)``,
319 ``r_< = min(r, r')``, ``r_> = max(r, r')``. It is evaluated overflow-safe with the exponentially
320 scaled modified Bessel functions (``scipy.special.ive``/``kve``): each term is a ratio whose scaling
321 exponents difference to a bounded value, so the growing ``I_0`` never overflows at high ``kappa r``.
322 The clock is wall-clock time (molecular diffusion is autonomous in ``t``); pair with the source
323 measure ``w(r') dr' = (r'/D_m) dr'`` (the Sturm-Liouville weight) when superposing a resident field.
325 Parameters
326 ----------
327 s : ndarray of complex
328 Laplace nodes (conjugate to wall-clock time). Shape ``(n_s,)``.
329 r : float
330 Output radius (m), ``>= r_w``.
331 r_prime : array-like
332 Source radius/radii (m), ``>= r_w``. Scalar or shape ``(n_r',)``.
333 r_w : float
334 Well radius (m).
335 d_m : float
336 Molecular diffusivity (m^2/day), ``> 0``.
338 Returns
339 -------
340 ndarray of complex
341 ``Ghat(r, r'; s)``, shape ``(n_s, n_r')`` (broadcast of ``s`` and ``r_prime``).
342 """
343 s = np.asarray(s, dtype=complex).reshape(-1, 1)
344 rp = np.atleast_1d(np.asarray(r_prime, dtype=float)).reshape(1, -1)
345 kappa = np.sqrt(s / d_m) # principal root; Re(s) > 0 on the Bromwich contour gives Re(kappa) > 0
346 r_lt = np.minimum(r, rp)
347 r_gt = np.maximum(r, rp)
348 z_lt, z_gt, z_w = kappa * r_lt, kappa * r_gt, kappa * r_w
349 # Ghat = [I_0(z_<) + (I_1(z_w)/K_1(z_w)) K_0(z_<)] K_0(z_>); split so the scaled-Bessel scaling
350 # exponents (ive scales by e^{-|Re z|}, kve by e^{+z}) difference to bounded values -- no overflow.
351 # Both terms carry the scaling exponents in a SINGLE np.exp of the combined (bounded) sum: the outer
352 # term's ``|Re z_w| + z_w`` alone overflows at Re(z_w) > ~354, but ``|Re z_w| + z_w - z_lt - z_gt``
353 # <= 0 (since r_w <= r_< <= r_>) so exponentiating the sum is overflow-safe.
354 term_inner = ive(0, z_lt) * kve(0, z_gt) * np.exp(np.abs(z_lt.real) - z_gt)
355 term_outer = (
356 (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)
357 )
358 return term_inner + term_outer
361# ---------------------------------------------------------------------------
362# D_m > 0 branch: Riccati log-derivative (numerical ODE, double precision)
363# ---------------------------------------------------------------------------
364# The decaying resident solution and the well-regular solution are tracked by their log-derivative
365# L = C'/C of the constant-Q ODE, integrated as a vector ODE over the Laplace nodes. L is O(kappa)
366# (bounded -- no 10^900 special-function magnitudes), so the transfer functions and the
367# Sturm-Liouville interior resolvent are assembled from O(1) quantities. This is exact to the de Hoog
368# inversion floor at ANY A_0/D_m -- no special-function precision blow-up, no tractability cap -- and
369# continuously becomes the Airy branch as D_m -> 0.
370_RICCATI_RTOL = 1e-12
371_RICCATI_ATOL = 1e-13
372# Outer boundary for the inward (decaying) integration. The truncated asymptotic IC at r_far washes out
373# because the recessive solution is the inward attractor (damping ~ e^{-2 Re(kappa)(r_far - r)}), but the
374# slowest (smallest Re(kappa)) Laplace node also needs r_far far enough that its z = 2 kappa r_far is in
375# the large-z asymptotic. So r_far is extended by ~_RICCATI_RFAR_DECAY decay lengths 1/Re(kappa_min)
376# beyond the field, floored at _RICCATI_RFAR_MULT * r_max and cost-capped at _RICCATI_RFAR_CAP * r_max
377# (the floor on Re(kappa_min) keeps z at the cap ~ 2 * _RICCATI_RFAR_DECAY, large enough for any node).
378_RICCATI_RFAR_MULT = 8.0
379_RICCATI_RFAR_DECAY = 22.0
380_RICCATI_RFAR_CAP = 500.0
383def _integrate_logderiv(
384 s: npt.NDArray[np.complexfloating],
385 radii: npt.ArrayLike,
386 r_w: float,
387 alpha_l: float,
388 a0_eff: float,
389 d_m_eff: float,
390 sigma_a: int,
391 branch: str,
392) -> tuple[npt.NDArray[np.complexfloating], npt.NDArray[np.complexfloating]]:
393 r"""Vector Riccati integration of the log-derivative ``L = C'/C`` for the ``D_m > 0`` branch.
395 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
396 Laplace nodes ``s`` at once, carrying the running integral ``J = int_{r_w}^{r} L dr``.
398 * ``branch='decaying'``: inward from ``r_far`` with the recessive asymptotic IC
399 ``L(r_far) = -kappa - a/x`` (``kappa = sqrt(s/D_m)``, ``x = r_far + a*``, ``a* = alpha_L A_0/D_m``,
400 ``a = b/2 - kappa a*/2``, ``b = 1 - sigma_A A_0/D_m``). The decaying solution is the inward
401 attractor, so the result is insensitive to ``r_far``.
402 * ``branch='regular'``: outward from ``r_w`` with the well-BC IC ``L(r_w) = A_0/G(r_w)`` (injection,
403 Robin ``F[u_0](r_w)=0``) or ``0`` (extraction, Neumann ``u_0'(r_w)=0``) -- both ``s``-independent.
405 ``sigma_a`` is ``+1`` divergent (injection) / ``-1`` convergent (extraction); ``radii`` (all ``>= r_w``)
406 are where ``L`` and ``J`` are returned.
408 Returns
409 -------
410 ld : ndarray of complex, shape (n_s, n_radii)
411 Log-derivative ``L`` at each requested radius.
412 jj : ndarray of complex, shape (n_s, n_radii)
413 ``int_{r_w}^{r} L dr`` at each requested radius.
414 """
415 s = np.asarray(s, dtype=complex).reshape(-1)
416 n = s.size
417 radii = np.atleast_1d(np.asarray(radii, dtype=float))
418 r_max = max(float(radii.max()), r_w)
420 def rhs(r: float, y: npt.NDArray[np.complexfloating]) -> npt.NDArray[np.complexfloating]:
421 ld = y[:n]
422 g = alpha_l * a0_eff + d_m_eff * r
423 d_ld = -(ld * ld) - ((d_m_eff - sigma_a * a0_eff) / g) * ld + s * r / g
424 return np.concatenate([d_ld, ld])
426 if branch == "decaying":
427 astar = alpha_l * a0_eff / d_m_eff
428 kappa = np.sqrt(s / d_m_eff)
429 # r_far must put the slowest node deep in the large-z asymptotic so the truncated IC washes out:
430 # extend by ~_RICCATI_RFAR_DECAY decay lengths 1/Re(kappa_min) (Re(kappa) floored so the extension
431 # is cost-capped at _RICCATI_RFAR_CAP * r_max while z = 2 kappa r_far stays large there).
432 re_kmin = max(float(kappa.real.min()), _RICCATI_RFAR_DECAY / (_RICCATI_RFAR_CAP * r_max))
433 r_far = max(_RICCATI_RFAR_MULT * r_max, r_max + _RICCATI_RFAR_DECAY / re_kmin)
434 a = (1.0 - sigma_a * a0_eff / d_m_eff) / 2.0 - kappa * astar / 2.0
435 y0 = np.concatenate([-kappa - a / (r_far + astar), np.zeros(n, dtype=complex)])
436 sol = solve_ivp(
437 rhs, [r_far, r_w], y0, rtol=_RICCATI_RTOL, atol=_RICCATI_ATOL, dense_output=True, method="DOP853"
438 )
439 y = sol.sol(radii) # dense output at all radii at once -> shape (2n, n_radii)
440 j_w = sol.sol(r_w)[n:, None] # re-anchor J to int_{r_w}^{r} (the IC put J(r_far) = 0)
441 # The result object sits in a reference cycle, so refcounting alone never frees the per-step
442 # interpolant arrays (the dominant allocation); drop them now that all evaluations are done.
443 sol.sol.interpolants.clear()
444 return y[:n], y[n:] - j_w
446 # regular branch: outward from r_w to r_max -- the growing solution is the stable outward attractor and
447 # the well-BC IC (s-independent) is exact, so no washout is needed; it need only reach the field.
448 ld0 = a0_eff / (alpha_l * a0_eff + d_m_eff * r_w) if sigma_a > 0 else 0.0
449 y0 = np.concatenate([np.full(n, ld0, dtype=complex), np.zeros(n, dtype=complex)])
450 sol = solve_ivp(rhs, [r_w, r_max], y0, rtol=_RICCATI_RTOL, atol=_RICCATI_ATOL, dense_output=True, method="DOP853")
451 y = sol.sol(radii)
452 sol.sol.interpolants.clear()
453 return y[:n], y[n:]
456def _transfer_riccati(
457 s: npt.NDArray[np.complexfloating],
458 r: float,
459 r_w: float,
460 alpha_l: float,
461 a0_eff: float,
462 d_m_eff: float,
463 detect: str,
464) -> npt.NDArray[np.complexfloating]:
465 r"""Kreft-Zuber flux-injection transfer modes for the ``D_m > 0`` branch via the decaying log-derivative.
467 With ``E = phi(r)/phi(r_w) = exp(int_{r_w}^{r} L)`` and the flux factor ``f(r) = 1 - (alpha_L +
468 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)`` and
469 ``FR = E/f(r_w)``. ``sigma_A = +1`` (the divergent operator); ``detect`` selects the Kreft-Zuber
470 detection boundary ('flux' or 'resident').
472 Returns
473 -------
474 ndarray of complex
475 ``g_hat(s)`` for the requested detection mode, shape ``(n_s,)``.
476 """
477 ld, jj = _integrate_logderiv(s, [r, r_w], r_w, alpha_l, a0_eff, d_m_eff, +1, "decaying")
478 l_r, l_w = ld[:, 0], ld[:, 1]
479 e = np.exp(jj[:, 0]) # phi(r)/phi(r_w)
480 f_r = 1.0 - (alpha_l + d_m_eff * r / a0_eff) * l_r
481 f_w = 1.0 - (alpha_l + d_m_eff * r_w / a0_eff) * l_w
482 return e * (f_r if detect == _FLUX else 1.0) / f_w
485def resolvent_riccati(
486 *,
487 s: npt.NDArray[np.complexfloating],
488 field: npt.NDArray[np.floating],
489 r_nodes: npt.NDArray[np.floating],
490 dr_weights: npt.NDArray[np.floating],
491 r_w: float,
492 alpha_l: float,
493 a0_eff: float,
494 d_m_eff: float,
495 direction: str,
496) -> npt.NDArray[np.complexfloating]:
497 r"""Interior Sturm-Liouville resolvent applied to a source ``field``, ``D_m > 0``, via log-derivatives.
499 Returns ``F(s)_{k,i} = sum_j Ghat(r_i, r_j; s_k) field_j w_j`` with the SL measure
500 ``w_j = r_j G(r_j)^{b-1} dr_j`` (``b = 1 - sigma_A A_0/D_m``). The Green's function
501 ``Ghat = phi_+(r_<) phi_-(r_>)/(-pW)`` is built from the decaying (``phi_-``, inward) and
502 well-regular (``phi_+``, outward) solutions normalized at ``r_w``; ``pW = G(r_w)^b (L_-(r_w) -
503 L_+(r_w))``. The divergent gauge ``G(r_w)^b`` (``b ~ -A_0/D_m``) is carried in LOG space
504 (``LG_j = b ln(G(r_j)/G(r_w)) - ln G(r_j)``, bounded as ``A_0/D_m -> inf``) so it never
505 underflows -- the assembly stays in double precision at any ``A_0/D_m``.
507 Parameters
508 ----------
509 s : ndarray of complex
510 Laplace nodes (conjugate to wall-clock time).
511 field : ndarray
512 Source resident-deviation profile on ``r_nodes``.
513 r_nodes : ndarray
514 Radial quadrature nodes (m), increasing, ``> r_w``.
515 dr_weights : ndarray
516 Quadrature weights for ``r_nodes`` (the ``dr`` measure).
517 r_w : float
518 Well radius (m).
519 alpha_l, a0_eff, d_m_eff : float
520 Dispersivity and the retardation-effective ``A_0`` / ``D_m``.
521 direction : {'injection', 'extraction'}
522 Phase orientation (divergent Robin / convergent Neumann well BC).
524 Returns
525 -------
526 ndarray of complex
527 ``F(s)_{k,i}`` -- the resolvent applied to ``field``, shape ``(n_s, n_nodes)``.
528 """
529 sigma_a = 1 if direction == _INJECTION else -1
530 s = np.asarray(s, dtype=complex).reshape(-1)
531 b = 1.0 - sigma_a * a0_eff / d_m_eff
532 rad = np.concatenate(([r_w], r_nodes))
533 ld_m, jj_m = _integrate_logderiv(s, rad, r_w, alpha_l, a0_eff, d_m_eff, sigma_a, "decaying")
534 lm_w = ld_m[:, 0] # L_-(r_w)
535 im = jj_m[:, 1:] # int_{r_w}^{r_i} L_- (n_s, n)
536 _, ip = _integrate_logderiv(s, r_nodes, r_w, alpha_l, a0_eff, d_m_eff, sigma_a, "regular") # int L_+
537 lp_w = a0_eff / (alpha_l * a0_eff + d_m_eff * r_w) if sigma_a > 0 else 0.0
538 g_nodes = alpha_l * a0_eff + d_m_eff * r_nodes
539 g_w = alpha_l * a0_eff + d_m_eff * r_w
540 lg = b * np.log(g_nodes / g_w) - np.log(g_nodes) # bounded gauge in log space
541 c = field * r_nodes * dr_weights
542 pmat = c[None, :] * np.exp(ip + lg[None, :]) # (n_s, n)
543 smat = c[None, :] * np.exp(im + lg[None, :])
544 prefix = np.cumsum(pmat, axis=1) # sum_{j<=i}
545 suffix = np.cumsum(smat[:, ::-1], axis=1)[:, ::-1] - smat # sum_{j>i}
546 denom = (lm_w - lp_w)[:, None]
547 return -(np.exp(im) * prefix + np.exp(ip) * suffix) / denom