Coverage for src/gwtransport/diffusion.py: 98%

1r""" 

2Analytical solutions for 1D advection-dispersion transport. 

3 

4Water infiltrates and is transported in parallel along multiple aquifer pore volumes to 

5extraction. For each aquifer pore volume, transport is 1D advection with microdispersion, 

6molecular diffusion, and linear sorption; the spread across aquifer pore volumes provides 

7macrodispersion. Forward and backward modeling are supported. The flow is assumed orthogonal. 

8 

9The orthogonal-flow (Cartesian) geometry is what makes the Kreft-Zuber breakthrough the exact 

101D solution used below. 

11 

12Key functions: 

13 

14- :func:`infiltration_to_extraction` - Main transport function combining advection, 

15 microdispersion, and molecular diffusion with explicit pore volume distribution and 

16 streamline lengths. 

17 

18- :func:`extraction_to_infiltration` - Inverse operation (deconvolution with dispersion). 

19 

20- :func:`gamma_infiltration_to_extraction` - Gamma-distributed pore volumes with dispersion. 

21 Models aquifer heterogeneity with 2-parameter gamma distribution. Parameterizable via 

22 (alpha, beta) or (mean, std). Discretizes gamma distribution into equal-probability bins. 

23 

24- :func:`gamma_extraction_to_infiltration` - Gamma-distributed pore volumes, deconvolution 

25 with dispersion. Symmetric inverse of gamma_infiltration_to_extraction. 

26 

27When to choose this module vs :mod:`gwtransport.diffusion_fast` 

28--------------------------------------------------------------- 

29 

30This is the reference implementation: it evaluates the bin-averaged Kreft-Zuber flux 

31concentration by resolution-aware composite Gauss-Legendre quadrature (splitting at 

32flow-bin boundaries, with extra front-centred panels wherever a sharp breakthrough front 

33is otherwise under-resolved). 

34Prefer it only when the output grid is coarser than the flow detail -- it integrates the 

35full within-bin flow, which the closed-form :mod:`gwtransport.diffusion_fast` approximates as 

36constant per output bin. Otherwise that module computes the same physics to machine 

37precision for *every* parameter regime (including ``retardation_factor != 1`` with non-zero 

38molecular diffusivity, whose flux correction it also evaluates in closed form) and is 

39~80-90x faster (no quadrature, no residence-time inversion). Both modules accept 

40per-streamtube ``streamline_length`` / ``molecular_diffusivity`` / 

41``longitudinal_dispersivity`` arrays (heterogeneous flow paths -- partially-penetrating 

42wells, wedge-shaped capture zones). 

43 

44Reported outlet concentration: Kreft-Zuber (1978) flux concentration 

45--------------------------------------------------------------------- 

46 

47The outlet concentration reported by this module is the **flux concentration** 

48 

49 C_F(L, t) = C_R(L, t) - (D_s / v_s) * dC_R/dx \|_{x=L} 

50 

51with the solute-front (retarded-frame) velocity v_s = Q L / (R V_pore) and the 

52dispersion D_s = D_m + alpha_L * v_s, so the flux coefficient is 

53D_s / v_s = D_m / v_s + alpha_L = R D_m / v_fluid + alpha_L (with the fluid 

54velocity v_fluid = Q L / V_pore). The resident profile C_R solves the retarded 

55ADE with advection v_s and dispersion D_s, so its flux-vs-resident correction 

56must use v_s — not v_fluid; pairing v_s with the moving-frame variance below is 

57what conserves mass for R > 1 with D_m > 0. 

58 

59— the solute mass flux at the outlet divided by the volumetric fluid flux. This 

60is what is measured when sampling the extracted fluid. The resident 

61concentration ``C_R`` is Bear (1972) eq. 10.6.4, the variable-flow moving-frame 

62Ogata-Banks solution 

63 

64 C_R(L, V; t_j) = 0.5 * erfc((L - xi_j(V)) / (2 * sqrt(D_t(V)))) 

65 

66with the dispersion variance accumulated in the moving (Lagrangian) frame: 

67 

68 D_t(V) = sigma^2(V) / 2 = D_m * tau(V) + alpha_L * xi(V) 

69 

70where: 

71 

72- D_m is the effective molecular (or thermal) diffusivity [m²/day] 

73- alpha_L is the longitudinal dispersivity [m] 

74- tau(V) is the elapsed time since infiltration [day], with V the cumulative 

75 extracted volume 

76- xi(V) = L (V - V_j) / (R V_pore) is the distance the parcel has actually 

77 travelled [m] 

78 

79The K-Z flux-correction term is what makes the column-sum invariant 

80``integral Q c_out dt = integral Q c_in dt`` hold under arbitrary variable Q. 

81Without it, the leading-order C_R loses O(1/Pe) per column under variable Q + 

82pure D_m (issue #180). 

83 

84Implementation: the bin-averaged C_F is computed by resolution-aware composite 

85Gauss-Legendre quadrature in volume space, split at flow-bin boundaries so each 

86sub-interval sees a linear t(V). Within a sub-interval the erf-like front has 

87width ``sqrt(4*D_t)`` (in volume units); for near-zero dispersivity this can be 

88orders of magnitude below the flow-bin width, so a single fixed-order rule 

89cannot resolve it. Sub-intervals whose front is under-resolved are therefore 

90tiled with front-centred panels (fine near the front, flat tails outside), 

91which restores the column-mass invariant to ~1e-11 for every dispersion regime; 

92smooth/already-resolved sub-intervals keep the plain single 16-point rule. The 

93variance is evaluated at each quadrature node from the parcel's own tau and xi 

94histories — never capped at the residence time. The K-Z identity requires 

95Bear's formula to satisfy the variable-coefficient ADE exactly, which holds only 

96when D_t is allowed to keep growing past breakthrough. 

97 

98Macrodispersion vs microdispersion 

99---------------------------------- 

100 

101This module adds microdispersion (alpha_L) and molecular diffusion (D_m) on top of 

102macrodispersion captured by the pore volume distribution (APVD). Both represent velocity 

103heterogeneity at different scales. Microdispersion is an aquifer property; macrodispersion 

104depends additionally on hydrological boundary conditions. See :ref:`concept-dispersion-scales` 

105for guidance on when to use each approach and how to avoid double-counting spreading effects. 

106 

107Streamtube assumption (no cross-sectional area parameter) 

108--------------------------------------------------------- 

109 

110Each entry in ``aquifer_pore_volumes`` is treated as an independent 1D streamtube. There is 

111no cross-sectional area parameter: the variance budget uses ``2 D_m tau`` (molecular 

112diffusion in time) and ``2 alpha_L xi`` (microdispersion in travelled distance), with 

113the streamline length ``L`` and the pore volume ``V_pore`` together fixing the implicit 

114streamtube cross-section ``A = V_pore / L``. Callers who need distributed-area effects must 

115provide multiple streamtubes (via ``aquifer_pore_volumes`` or the gamma-parameterised 

116wrappers). 

117 

118References 

119---------- 

120Bear, J. (1972). Dynamics of Fluids in Porous Media. American Elsevier 

121Publishing Company. Equation 10.6.4 (variable-flow Ogata-Banks form). Provides 

122the resident concentration ``C_R``. 

123 

124Kreft, A., & Zuber, A. (1978). On the physical meaning of the dispersion 

125equation and its solutions for different initial and boundary conditions. 

126Chemical Engineering Science, 33(11), 1471-1480. Eq. 2 gives the resident-to- 

127flux concentration transformation; Eq. 1 is the mass-balance identity that 

128makes the column-sum invariant exact. 

129""" 

130 

131import numpy as np 

132import numpy.typing as npt 

133import pandas as pd 

134from scipy import special 

135 

136from gwtransport import gamma 

137from gwtransport._time import dt_to_days, tedges_to_days 

138from gwtransport._validation import ( 

139 _validate_no_nan, 

140 _validate_non_negative_array, 

141 _validate_positive_array, 

142 _validate_retardation_factor, 

143 _validate_scalar_or_matching_length, 

144 _validate_tedges_parity, 

145) 

146from gwtransport.residence_time import fraction_explained_full 

147from gwtransport.utils import cumulative_flow_volume, solve_inverse_transport 

148 

149# Numerical tolerance for coefficient sum to determine valid output bins 

150EPSILON_COEFF_SUM = 1e-10 

151 

152# Gauss-Legendre quadrature nodes and weights for volume-space integration 

153_GL_NODES, _GL_WEIGHTS = np.polynomial.legendre.leggauss(16) 

154 

155# Resolution-aware composite-quadrature template for the erf-like breakthrough front. 

156# When the front width sqrt(4*D_t) (in volume units) is much smaller than a 

157# (cell, flow-bin) sub-interval, plain 16-point GL cannot resolve it. Cells flagged 

158# as under-resolved get panels placed at ``V_front + front_width * _FRONT_OFFSETS`` 

159# (clipped to the sub-interval); the two outer panels then cover the flat erf tails, 

160# where 16-point GL is already exact. Panels near the front are ~1 front-width wide, 

161# spanning +-6 front-widths (erf and the flux-correction Gaussian are flat to ~1e-16 

162# beyond that). A cell is refined only when its front lies inside the sub-interval and 

163# the sub-interval is wider than _REFINE_RATIO front-widths, so adaptivity triggers 

164# only near sharp fronts (smooth regimes keep the plain single-panel cost and answer). 

165_FRONT_REACH = 6.0 

166_FRONT_OFFSETS = np.arange(-_FRONT_REACH, _FRONT_REACH + 0.5, 1.0) 

167_REFINE_RATIO = 4.0 

168 

169 

170def _cfrac_mean_volume( 

171 *, 

172 step_widths: npt.NDArray[np.floating], 

173 cumulative_volume_at_cout_tedges: npt.NDArray[np.floating], 

174 cumulative_volume_at_cin_tedges: npt.NDArray[np.floating], 

175 tedges_days: npt.NDArray[np.floating], 

176 molecular_diffusivity: float, 

177 longitudinal_dispersivity: float, 

178 r_vpv: float, 

179 streamline_len: float, 

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

181 r"""Compute bin-averaged flux concentration at the outlet for each cell. 

182 

183 For each cell (cout-bin *i*, cin-edge *j*), computes the flow-weighted 

184 average of the Kreft-Zuber (1978) **flux concentration** at the outlet: 

185 

186 .. math:: 

187 

188 \text{frac}_{i,j} = \frac{1}{\Delta V_i} 

189 \int_{V_i}^{V_{i+1}} C_F\!\left(L,\,V;\,t_j\right) dV 

190 

191 where :math:`C_F = C_R - (D_s / v_s) \, \partial_x C_R\big|_{x=L}` and 

192 :math:`C_R` is Bear's (1972) moving-frame solution: 

193 

194 .. math:: 

195 

196 C_R(L, V; t_j) &= \tfrac{1}{2}\, 

197 \mathrm{erfc}\!\left( \frac{L - \xi_j(V)}{2\sqrt{D_t(V)}} \right) \\ 

198 \xi_j(V) &= L \cdot (V - V_j) \,/\, (R\,V_\text{pore}) \\ 

199 D_t(V) &= D_m\,\tau_j(V) + \alpha_L\,\xi_j(V), 

200 \quad \tau_j(V) = t(V) - t_j \\ 

201 D_s &= D_m + \alpha_L\,v_s(t), \quad v_s(t) = Q(t)\,L\,/\,(R\,V_\text{pore}). 

202 

203 The solute-front velocity :math:`v_s` (advection speed of the retarded ADE 

204 that :math:`C_R` solves), not the fluid velocity :math:`Q L / V_\text{pore}`, 

205 sets the flux coefficient :math:`D_s/v_s = D_m/v_s + \alpha_L`. The added 

206 flux-correction term 

207 

208 .. math:: 

209 

210 \frac{D_s}{v_s(t(V))} \cdot 

211 \frac{1}{\sqrt{4\pi\,D_t(V)}}\, 

212 \exp\!\left( -\frac{(L - \xi_j(V))^2}{4\,D_t(V)} \right) 

213 

214 converts Bear's *resident* concentration to a *flux* concentration. This 

215 makes the coefficient matrix conserve mass under the criterion 

216 ``integral Q c_out dt = integral Q c_in dt`` — the relevant invariant for 

217 tracer measurements taken in the extracted fluid (Kreft & Zuber, 1978, 

218 Eq. 5 and Eq. 1). Without this correction, Bear's leading-order kernel 

219 misses the dispersive boundary flux at the outlet and column-sum mass 

220 conservation fails by O(1/Pe) under variable Q. 

221 

222 Implementation: resolution-aware composite Gauss-Legendre quadrature in 

223 volume space, split at flow-bin boundaries so that within each sub-interval 

224 :math:`t(V)` is linear. The erf-like front has width :math:`\sqrt{4 D_t}` 

225 (in volume units); a sub-interval whose front is under-resolved by a single 

226 16-point rule (front width far below the sub-interval width) is tiled with 

227 front-centred panels (see ``_FRONT_OFFSETS``), while smooth/already-resolved 

228 sub-intervals keep the plain single 16-point rule (bit-identical to it). No 

229 "fully capped" branch: the moving-frame variance keeps growing past 

230 breakthrough, and the K-Z identity requires Bear's formula to satisfy the 

231 variable-coefficient ADE exactly (which it does only without capping). 

232 

233 Parameters 

234 ---------- 

235 step_widths : ndarray, shape (n_cout_edges, n_cin_edges) 

236 x-position ``x(V_cout, V_cin) = (V_cout - V_cin - r_vpv) * L / r_vpv`` 

237 at each (cout-edge, cin-edge). NaN for inactive cells. Equals 

238 :math:`\xi - L`. 

239 cumulative_volume_at_cout_tedges : ndarray, shape (n_cout_edges,) 

240 Cumulative extracted volume at each cout time edge [m³]. 

241 cumulative_volume_at_cin_tedges : ndarray, shape (n_cin_edges,) 

242 Cumulative volume at each cin (flow) time edge [m³]. 

243 tedges_days : ndarray, shape (n_cin_edges,) 

244 Flow time edges in days. 

245 molecular_diffusivity : float 

246 Effective (retarded-frame) molecular diffusivity D_m [m²/day]. 

247 Contributes ``D_m * tau`` to the dispersion product ``D_t``. 

248 longitudinal_dispersivity : float 

249 Longitudinal dispersivity alpha_L [m]. Contributes ``alpha_L * xi`` 

250 to the dispersion product ``D_t``. 

251 r_vpv : float 

252 Retardation factor times pore volume = R * V_pore [m³]. 

253 streamline_len : float 

254 Streamline length L [m]. 

255 

256 Returns 

257 ------- 

258 ndarray, shape (n_cout_bins, n_cin_edges) 

259 Bin-averaged flux concentration for each cell. NaN for inactive cells. 

260 

261 References 

262 ---------- 

263 Kreft, A., & Zuber, A. (1978). On the physical meaning of the dispersion 

264 equation and its solutions for different initial and boundary conditions. 

265 Chemical Engineering Science, 33(11), 1471-1480. 

266 """ 

267 n_cout_edges, n_cin_edges = step_widths.shape 

268 n_cout_bins = n_cout_edges - 1 

269 

270 x_lo = step_widths[:-1] 

271 x_hi = step_widths[1:] 

272 dx = x_hi - x_lo 

273 

274 v_lo_arr = cumulative_volume_at_cout_tedges[:-1] 

275 v_hi_arr = cumulative_volume_at_cout_tedges[1:] 

276 

277 is_valid = ~np.isnan(x_lo) & ~np.isnan(x_hi) 

278 

279 frac = np.full((n_cout_bins, n_cin_edges), np.nan) 

280 

281 # --- No dispersion: C_F = C_R = step function (no dispersive flux) --- 

282 if molecular_diffusivity == 0.0 and longitudinal_dispersivity == 0.0: 

283 with np.errstate(divide="ignore", invalid="ignore"): 

284 cr_no_disp = 0.5 + 0.5 * (np.abs(x_hi) - np.abs(x_lo)) / dx 

285 cr_no_disp = np.where(dx == 0.0, 0.5 + 0.5 * np.sign(x_lo), cr_no_disp) 

286 return np.where(is_valid, cr_no_disp, frac) 

287 

288 # --- Pre-compute solute-front velocity and K-Z coefficient (D_s/v_s) per flow bin --- 

289 dv_per_bin = np.diff(cumulative_volume_at_cin_tedges) 

290 dt_per_bin = np.diff(tedges_days) 

291 with np.errstate(divide="ignore", invalid="ignore"): 

292 q_per_bin = np.where(dt_per_bin > 0, dv_per_bin / dt_per_bin, 0.0) 

293 # Solute-front velocity v_s = Q L / (R V_pore) -- the advection speed of the retarded ADE 

294 # that C_R actually solves. Kreft-Zuber requires the flux coefficient D_s/v_s = D_m/v_s + 

295 # alpha_L to use THAT velocity (not the fluid velocity Q L / V_pore); pairing it with the 

296 # moving-frame variance D_t = D_m tau + alpha_L xi is what conserves mass for R>1, D_m>0. 

297 v_per_bin = q_per_bin * streamline_len / r_vpv 

298 # (D_s/v_s) = D_m/v_s + alpha_L. At v_s=0 the bin has dV=0 and is skipped below; the 

299 # surrounding errstate suppresses the divide warning for those lanes. 

300 dl_over_v_per_bin = np.where( 

301 v_per_bin > 0, 

302 molecular_diffusivity / v_per_bin + longitudinal_dispersivity, 

303 0.0, 

304 ) 

305 

306 # --- Resolution-aware composite Gauss-Legendre quadrature, split by flow bins --- 

307 # The integration window per (cell, flow-bin) is the intersection of 

308 # [V_lo, V_hi] (cell), [ve_lo, ve_hi] (flow bin), and [V_j, infty) (parcel 

309 # entered). Within each sub-interval, t(V) is linear. Where the sub-interval is 

310 # much wider than the front width sqrt(4*D_t), the erf-like front is under-resolved 

311 # by a single 16-point GL rule; such sub-intervals are tiled with front-centred 

312 # panels (see _FRONT_OFFSETS). Smooth sub-intervals keep the single panel and are 

313 # bit-identical to the plain 16-point rule. 

314 idx_i, idx_j = np.nonzero(is_valid) 

315 if len(idx_i) == 0: 

316 return frac 

317 

318 v_lo_cells = v_lo_arr[idx_i] 

319 v_hi_cells = v_hi_arr[idx_i] 

320 v_cin_cells = cumulative_volume_at_cin_tedges[idx_j] 

321 t_j_cells = tedges_days[idx_j] 

322 total_dv = v_hi_cells - v_lo_cells 

323 valid_cells = total_dv > 0 

324 

325 # Post-injection lower bound is loop-invariant: max(V_lo, V_cin) (D2 hoist). 

326 v_lo_or_cin = np.maximum(v_lo_cells, v_cin_cells) 

327 

328 integral_cf = np.zeros(len(idx_i)) 

329 

330 vol_edges = cumulative_volume_at_cin_tedges 

331 for k in range(len(vol_edges) - 1): 

332 ve_lo, ve_hi = vol_edges[k], vol_edges[k + 1] 

333 if v_per_bin[k] <= 0.0: 

334 continue 

335 dl_over_v_k = dl_over_v_per_bin[k] 

336 dt_sub_bin = tedges_days[k + 1] - tedges_days[k] 

337 dv_sub_edge = ve_hi - ve_lo 

338 

339 # Intersection of cell, flow-bin, and post-injection range 

340 sub_lo = np.maximum(v_lo_or_cin, ve_lo) 

341 sub_hi = np.minimum(v_hi_cells, ve_hi) 

342 overlap = (sub_hi > sub_lo) & valid_cells 

343 if not np.any(overlap): 

344 continue 

345 

346 lo = sub_lo[overlap] 

347 hi = sub_hi[overlap] 

348 vcin = v_cin_cells[overlap] 

349 tj = t_j_cells[overlap] 

350 

351 # Front centre (x = 0 => xi = L => V = V_cin + r_vpv) and its width in volume. 

352 # front_width = sqrt(4*D_t_front) * r_vpv / L, with D_t_front = D_m*tau_front + 

353 # alpha_L*L (xi = L at the front). t(V) is linear within flow bin k. 

354 v_front = vcin + r_vpv 

355 t_front = tedges_days[k] + (v_front - ve_lo) * (dt_sub_bin / dv_sub_edge) 

356 tau_front = np.maximum(t_front - tj, 0.0) 

357 dt_front = molecular_diffusivity * tau_front + longitudinal_dispersivity * streamline_len 

358 front_width = np.sqrt(4.0 * dt_front) * r_vpv / streamline_len 

359 

360 # Refine only where a sharp front region intersects this sub-interval and is 

361 # under-resolved by a single 16-point rule; elsewhere the integrand is flat 

362 # or already resolved, so one panel is exact. The front-region test (not just 

363 # "centre in this bin") also catches a front whose tail spills across a 

364 # flow-bin boundary. A spuriously large extrapolated front_width far from the 

365 # front fails ``underresolved`` and so cannot trigger refinement. 

366 front_hits = (v_front + _FRONT_REACH * front_width > lo) & (v_front - _FRONT_REACH * front_width < hi) 

367 underresolved = (hi - lo) > _REFINE_RATIO * front_width 

368 if np.any(front_hits & underresolved): 

369 inner = np.clip( 

370 v_front[:, np.newaxis] + front_width[:, np.newaxis] * _FRONT_OFFSETS[np.newaxis, :], 

371 lo[:, np.newaxis], 

372 hi[:, np.newaxis], 

373 ) 

374 edges = np.concatenate([lo[:, np.newaxis], inner, hi[:, np.newaxis]], axis=1) 

375 else: 

376 edges = np.stack([lo, hi], axis=1) 

377 

378 p_lo = edges[:, :-1] 

379 p_hi = edges[:, 1:] 

380 p_mid = 0.5 * (p_lo + p_hi) 

381 p_half = 0.5 * (p_hi - p_lo) 

382 

383 # GL nodes over every panel: shape (n_cell, n_panel, n_gl) 

384 v_nodes = p_mid[:, :, np.newaxis] + p_half[:, :, np.newaxis] * _GL_NODES[np.newaxis, np.newaxis, :] 

385 

386 # Geometry: x = xi - L = (V - V_j - r_vpv) * L / r_vpv (parcel position 

387 # relative to outlet); xi = parcel travel distance. 

388 x_nodes = (v_nodes - vcin[:, np.newaxis, np.newaxis] - r_vpv) * streamline_len / r_vpv 

389 xi_nodes = x_nodes + streamline_len 

390 

391 t_nodes = tedges_days[k] + (v_nodes - ve_lo) * (dt_sub_bin / dv_sub_edge) 

392 # tau >= 0 by construction (lo >= v_cin); clip for safety. 

393 tau_nodes = np.maximum(t_nodes - tj[:, np.newaxis, np.newaxis], 0.0) 

394 

395 # Bear's variance accumulator (sigma^2/2) — NO capping at RT/L 

396 dt_nodes = molecular_diffusivity * tau_nodes + longitudinal_dispersivity * xi_nodes 

397 

398 with np.errstate(divide="ignore", invalid="ignore"): 

399 arg = x_nodes / (2.0 * np.sqrt(dt_nodes)) 

400 # C_R = 0.5 * (1 + erf(arg)) = 0.5 * erfc((L-xi)/(2*sqrt(Dt))) 

401 erf_vals = np.where(np.isfinite(arg), special.erf(arg), np.sign(x_nodes)) 

402 cr_vals = 0.5 * (1.0 + erf_vals) 

403 

404 # K-Z flux correction: FC = (D_s/v_s) * (1/sqrt(4 pi D_t)) * exp(-arg^2) 

405 with np.errstate(divide="ignore", invalid="ignore"): 

406 gauss_vals = np.where( 

407 dt_nodes > 0.0, 

408 np.exp(-(arg**2)) / np.sqrt(4.0 * np.pi * dt_nodes), 

409 0.0, 

410 ) 

411 cf_vals = cr_vals + dl_over_v_k * gauss_vals 

412 

413 # Integrate: GL-weight over nodes, then sum panel contributions per cell. 

414 # The weight contraction is done in 2D so a single-panel (non-refined) 

415 # sub-interval is bit-identical to a plain 16-point rule. 

416 cf_weighted = (cf_vals.reshape(-1, cf_vals.shape[-1]) @ _GL_WEIGHTS).reshape(cf_vals.shape[:-1]) 

417 integral_cf[overlap] += (p_half * cf_weighted).sum(axis=1) 

418 

419 with np.errstate(divide="ignore", invalid="ignore"): 

420 frac_cells = np.where(valid_cells, integral_cf / total_dv, np.nan) 

421 frac[idx_i, idx_j] = frac_cells 

422 

423 return frac 

424 

425 

426def _diffusion_extend_tedges_flag(spinup: object) -> bool: 

427 """Translate the public ``spinup`` parameter to the internal extend flag. 

428 

429 The diffusion module's existing warm-start behavior is to extend 

430 ``tedges`` by 100 years on each side. The public ``spinup`` parameter 

431 maps onto this binary toggle: ``"constant"`` enables the extension 

432 (default; preserves legacy behavior), ``None`` disables it (cout in 

433 spin-up region becomes NaN). The float fraction-threshold mode of 

434 other modules is not implemented here. 

435 

436 Returns 

437 ------- 

438 bool 

439 True if tedges should be extended (warm-start), False if not. 

440 

441 Raises 

442 ------ 

443 ValueError 

444 If ``spinup`` is a string other than ``"constant"``. 

445 NotImplementedError 

446 If ``spinup`` is a float (fraction-threshold mode is not 

447 implemented for the diffusion module). 

448 """ 

449 if spinup is None: 

450 return False 

451 if isinstance(spinup, str): 

452 if spinup != "constant": 

453 msg = f"spinup string must be 'constant'; got {spinup!r}" 

454 raise ValueError(msg) 

455 return True 

456 msg = ( 

457 "diffusion's spinup parameter only supports None or 'constant'; " 

458 f"float thresholds are not yet implemented (got {spinup!r})" 

459 ) 

460 raise NotImplementedError(msg) 

461 

462 

463def _validate_diffusion_inputs( 

464 *, 

465 tedges: pd.DatetimeIndex, 

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

467 aquifer_pore_volumes: npt.NDArray[np.floating], 

468 streamline_length: npt.NDArray[np.floating], 

469 molecular_diffusivity: npt.NDArray[np.floating], 

470 longitudinal_dispersivity: npt.NDArray[np.floating], 

471 retardation_factor: float, 

472 cin_values: npt.NDArray[np.floating] | None = None, 

473 cout_values: npt.NDArray[np.floating] | None = None, 

474 cout_tedges: pd.DatetimeIndex | None = None, 

475) -> None: 

476 """Validate inputs common to diffusion forward / reverse entry points. 

477 

478 Path selection via mutually-exclusive kwargs: 

479 

480 - ``cin_values`` provided => forward path. ``tedges`` parities cin and flow. 

481 - ``cout_values`` + ``cout_tedges`` provided => reverse path. ``tedges`` parities 

482 flow; ``cout_tedges`` parities cout. 

483 

484 Raises 

485 ------ 

486 ValueError 

487 If any check fails. The message identifies which invariant was violated. 

488 """ 

489 n_pore_volumes = len(aquifer_pore_volumes) 

490 

491 if cin_values is not None: 

492 _validate_tedges_parity(tedges, cin_values, tedges_name="tedges", values_name="cin") 

493 _validate_tedges_parity(tedges, flow, tedges_name="tedges", values_name="flow") 

494 elif cout_values is not None and cout_tedges is not None: 

495 _validate_tedges_parity(tedges, flow, tedges_name="tedges", values_name="flow") 

496 _validate_tedges_parity(cout_tedges, cout_values, tedges_name="cout_tedges", values_name="cout") 

497 else: 

498 msg = "must provide cin_values (forward) or both cout_values and cout_tedges (reverse)" 

499 raise ValueError(msg) 

500 if len(aquifer_pore_volumes) != len(streamline_length): 

501 msg = "aquifer_pore_volumes and streamline_length must have the same length" 

502 raise ValueError(msg) 

503 _validate_scalar_or_matching_length( 

504 molecular_diffusivity, 

505 name="molecular_diffusivity", 

506 expected_len=n_pore_volumes, 

507 ref_name="aquifer_pore_volumes", 

508 ) 

509 _validate_scalar_or_matching_length( 

510 longitudinal_dispersivity, 

511 name="longitudinal_dispersivity", 

512 expected_len=n_pore_volumes, 

513 ref_name="aquifer_pore_volumes", 

514 ) 

515 _validate_non_negative_array(molecular_diffusivity, name="molecular_diffusivity") 

516 _validate_non_negative_array(longitudinal_dispersivity, name="longitudinal_dispersivity") 

517 if cin_values is not None: 

518 _validate_no_nan(cin_values, name="cin") 

519 elif cout_values is not None: 

520 _validate_no_nan(cout_values, name="cout") 

521 _validate_no_nan(flow, name="flow") 

522 _validate_non_negative_array(flow, name="flow", message="flow must be non-negative (negative flow not supported)") 

523 _validate_positive_array(aquifer_pore_volumes, name="aquifer_pore_volumes") 

524 _validate_positive_array(streamline_length, name="streamline_length") 

525 _validate_retardation_factor(retardation_factor) 

526 

527 

528def _prepare_diffusion_arrays( 

529 *, 

530 flow: npt.ArrayLike, 

531 aquifer_pore_volumes: npt.ArrayLike, 

532 streamline_length: npt.ArrayLike, 

533 molecular_diffusivity: npt.ArrayLike, 

534 longitudinal_dispersivity: npt.ArrayLike, 

535) -> tuple[ 

536 npt.NDArray[np.floating], 

537 npt.NDArray[np.floating], 

538 npt.NDArray[np.floating], 

539 npt.NDArray[np.floating], 

540 npt.NDArray[np.floating], 

541]: 

542 """Coerce flow / geometry / dispersion inputs to broadcasted float arrays. 

543 

544 Each per-streamtube parameter (``streamline_length``, ``molecular_diffusivity``, 

545 ``longitudinal_dispersivity``) may be passed as a scalar; it is broadcast to one 

546 value per pore volume. The returned arrays are read-only views when broadcast (none 

547 is mutated downstream). 

548 

549 Returns 

550 ------- 

551 tuple of ndarray 

552 ``(flow, aquifer_pore_volumes, streamline_length, molecular_diffusivity, 

553 longitudinal_dispersivity)`` as float arrays. 

554 """ 

555 flow = np.asarray(flow, dtype=float) 

556 aquifer_pore_volumes = np.asarray(aquifer_pore_volumes, dtype=float) 

557 streamline_length = np.atleast_1d(np.asarray(streamline_length, dtype=float)) 

558 molecular_diffusivity = np.atleast_1d(np.asarray(molecular_diffusivity, dtype=float)) 

559 longitudinal_dispersivity = np.atleast_1d(np.asarray(longitudinal_dispersivity, dtype=float)) 

560 

561 n_pore_volumes = len(aquifer_pore_volumes) 

562 if streamline_length.size == 1: 

563 streamline_length = np.broadcast_to(streamline_length, (n_pore_volumes,)) 

564 if molecular_diffusivity.size == 1: 

565 molecular_diffusivity = np.broadcast_to(molecular_diffusivity, (n_pore_volumes,)) 

566 if longitudinal_dispersivity.size == 1: 

567 longitudinal_dispersivity = np.broadcast_to(longitudinal_dispersivity, (n_pore_volumes,)) 

568 

569 return flow, aquifer_pore_volumes, streamline_length, molecular_diffusivity, longitudinal_dispersivity 

570 

571 

572def _infiltration_to_extraction_coeff_matrix( 

573 *, 

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

575 tedges: pd.DatetimeIndex, 

576 cout_tedges: pd.DatetimeIndex, 

577 aquifer_pore_volumes: npt.NDArray[np.floating], 

578 streamline_length: npt.NDArray[np.floating], 

579 molecular_diffusivity: npt.NDArray[np.floating], 

580 longitudinal_dispersivity: npt.NDArray[np.floating], 

581 retardation_factor: float, 

582 extend_tedges: bool = True, 

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

584 """Build the forward coefficient matrix for diffusion transport. 

585 

586 Constructs the matrix W such that ``cout = W @ cin``, accounting for 

587 advection, microdispersion, and molecular diffusion. NaN entries in the raw coefficient 

588 matrix are replaced with zero. 

589 

590 Parameters 

591 ---------- 

592 flow : ndarray 

593 Flow rate of water [m³/day]. Already validated. 

594 tedges : DatetimeIndex 

595 Cin/flow time edges (not yet extended for spin-up). 

596 cout_tedges : DatetimeIndex 

597 Cout time edges. 

598 aquifer_pore_volumes : ndarray 

599 Pore volumes [m³]. Already validated. 

600 streamline_length : ndarray 

601 Travel distances [m]. Already validated. 

602 molecular_diffusivity : ndarray 

603 Effective molecular diffusivities [m²/day]. Already broadcasted. 

604 See :func:`infiltration_to_extraction` for physical interpretation. 

605 longitudinal_dispersivity : ndarray 

606 Longitudinal dispersivities [m]. Already broadcasted. 

607 retardation_factor : float 

608 Retardation factor. 

609 

610 Returns 

611 ------- 

612 coeff_matrix : ndarray 

613 Filled coefficient matrix of shape (n_cout, n_cin). NaN replaced 

614 with zero. 

615 valid_cout_bins : ndarray 

616 Boolean mask of shape (n_cout,) indicating valid output bins. 

617 """ 

618 if extend_tedges: 

619 # Extend tedges by 100 years on each side to provide warm-start cin 

620 # and post-data flow for cout bins near the boundaries. Equivalent to 

621 # the public ``spinup="constant"`` policy in other modules. 

622 tedges = pd.DatetimeIndex([ 

623 tedges[0] - pd.Timedelta("36500D"), 

624 *list(tedges[1:-1]), 

625 tedges[-1] + pd.Timedelta("36500D"), 

626 ]) 

627 

628 # Compute the cumulative flow at tedges 

629 cumulative_volume_at_cin_tedges = cumulative_flow_volume(flow, dt_to_days(tedges)) # m³ 

630 

631 # Compute the cumulative flow at cout_tedges. Both edge arrays are first reduced to a shared 

632 # day axis: np.interp coerces each datetime64 operand to int64 in its own resolution, so a 

633 # cout_tedges / tedges unit mismatch (e.g. ns vs us) would send every query out of range and 

634 # silently return all-NaN. 

635 tedges_days_arr = tedges_to_days(tedges) 

636 cout_tedges_days = tedges_to_days(cout_tedges, ref=tedges[0]) 

637 cumulative_volume_at_cout_tedges = np.interp( 

638 cout_tedges_days, tedges_days_arr, cumulative_volume_at_cin_tedges 

639 ).astype(float) 

640 

641 # Output bin valid where every streamtube's advective look-back is in-record across the whole 

642 # bin -- i.e. advective coverage == 1 for all pore volumes (NaN outside the record -> invalid). 

643 # This is the advective validity gate only; the dispersive informedness is the captured kernel 

644 # mass (total_coeff) applied downstream. 

645 valid_cout_bins = np.all( 

646 fraction_explained_full( 

647 flow=flow, 

648 tedges=tedges, 

649 cout_tedges=cout_tedges, 

650 aquifer_pore_volumes=aquifer_pore_volumes, 

651 retardation_factor=retardation_factor, 

652 direction="extraction_to_infiltration", 

653 ) 

654 >= 1.0, 

655 axis=0, 

656 ) 

657 

658 # Initialize coefficient matrix accumulator 

659 n_cout_bins = len(cout_tedges) - 1 

660 n_cin_bins = len(flow) 

661 accumulated_coeff = np.zeros((n_cout_bins, n_cin_bins)) 

662 

663 # Determine when infiltration has occurred: cout_tedge must be >= tedge (infiltration time) 

664 isactive = cout_tedges.to_numpy()[:, None] >= tedges.to_numpy()[None, :] 

665 

666 # Loop over each pore volume 

667 for i_pv in range(len(aquifer_pore_volumes)): 

668 r_vpv = retardation_factor * aquifer_pore_volumes[i_pv] 

669 

670 delta_volume = cumulative_volume_at_cout_tedges[:, None] - cumulative_volume_at_cin_tedges[None, :] - r_vpv 

671 delta_volume[~isactive] = np.nan 

672 

673 step_widths = delta_volume / r_vpv * streamline_length[i_pv] 

674 

675 frac = _cfrac_mean_volume( 

676 step_widths=step_widths, 

677 cumulative_volume_at_cout_tedges=cumulative_volume_at_cout_tedges, 

678 cumulative_volume_at_cin_tedges=cumulative_volume_at_cin_tedges, 

679 tedges_days=tedges_days_arr, 

680 molecular_diffusivity=float(molecular_diffusivity[i_pv]), 

681 longitudinal_dispersivity=float(longitudinal_dispersivity[i_pv]), 

682 r_vpv=r_vpv, 

683 streamline_len=streamline_length[i_pv], 

684 ) 

685 

686 frac_start = frac[:, :-1] 

687 frac_end = frac[:, 1:] 

688 frac_end_filled = np.where(np.isnan(frac_end) & ~np.isnan(frac_start), 0.0, frac_end) 

689 coeff = frac_start - frac_end_filled 

690 

691 accumulated_coeff += coeff 

692 

693 coeff_matrix_filled = np.nan_to_num(accumulated_coeff / len(aquifer_pore_volumes), nan=0.0) 

694 

695 return coeff_matrix_filled, valid_cout_bins 

696 

697 

698def infiltration_to_extraction( 

699 *, 

700 cin: npt.ArrayLike, 

701 flow: npt.ArrayLike, 

702 tedges: pd.DatetimeIndex, 

703 cout_tedges: pd.DatetimeIndex, 

704 aquifer_pore_volumes: npt.ArrayLike, 

705 streamline_length: npt.ArrayLike, 

706 molecular_diffusivity: npt.ArrayLike, 

707 longitudinal_dispersivity: npt.ArrayLike, 

708 retardation_factor: float = 1.0, 

709 spinup: str | None = "constant", 

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

711 """ 

712 Compute extracted concentration with advection, microdispersion, and molecular diffusion. 

713 

714 This function models 1D solute transport through an aquifer system along orthogonal 

715 (Cartesian) flow paths. Each aquifer pore volume is an independent streamline carrying 

716 advection with microdispersion (alpha_L) and molecular diffusion (D_m); the spread across 

717 the pore volume distribution provides macrodispersion. Linear sorption enters via the 

718 retardation factor. 

719 

720 The physical model assumes: 

721 

722 1. Water infiltrates with concentration cin at time t_in 

723 2. Water travels distance L through aquifer with residence time tau = V_pore / Q 

724 3. During transport, microdispersion and molecular diffusion spread each streamline, 

725 while the spread across pore volumes provides macrodispersion 

726 4. At extraction, the concentration is a dispersed breakthrough curve 

727 

728 The reported extracted concentration is the Kreft-Zuber (1978) **flux 

729 concentration** at the outlet, defined as the solute mass flux divided 

730 by the volumetric fluid flux. This is what is measured when sampling the 

731 outflowing fluid. Compared to Bear's leading-order resident concentration, 

732 it includes the dispersive boundary flux ``-D_s * dC_R/dx`` at 

733 ``x = L`` (with the solute-front dispersion ``D_s = D_m + alpha_L * v_s`` 

734 and velocity ``v_s = Q L / (R V_pore)``), which is what makes the column-sum 

735 invariant ``integral Q c_out dt = integral Q c_in dt`` hold exactly under 

736 variable flow. 

737 

738 Microdispersion and molecular diffusion enter as the moving-frame variance 

739 

740 sigma^2(V) = 2 * D_m * tau(V) + 2 * alpha_L * xi(V), 

741 

742 where ``tau(V)`` is the elapsed time since infiltration and ``xi(V)`` is 

743 the distance the parcel has actually travelled. Evaluating sigma^2 at each 

744 quadrature node — and avoiding any artificial capping past breakthrough — 

745 keeps Bear's formula an exact solution of the variable-coefficient ADE, 

746 which the Kreft-Zuber identity relies on. 

747 

748 Parameters 

749 ---------- 

750 cin : array-like 

751 Concentration of the compound in infiltrating water [concentration units]. 

752 Length must match the number of time bins defined by tedges. The model assumes 

753 this value is constant over each interval ``[tedges[i], tedges[i+1])``. 

754 flow : array-like 

755 Flow rate of water in the aquifer [m³/day]. 

756 Length must match cin and the number of time bins defined by tedges. The model 

757 assumes this value is constant over each interval ``[tedges[i], tedges[i+1])``. 

758 tedges : pandas.DatetimeIndex 

759 Time edges defining bins for both cin and flow data. Has length of 

760 len(cin) + 1. 

761 cout_tedges : pandas.DatetimeIndex 

762 Time edges for output data bins. Has length of desired output + 1. 

763 The output concentration is averaged over each bin. 

764 aquifer_pore_volumes : array-like 

765 Array of aquifer pore volumes [m³] representing the distribution 

766 of flow paths. Each pore volume determines the residence time for 

767 that flow path: tau = V_pore / Q. 

768 streamline_length : array-like 

769 Array of travel distances [m] corresponding to each pore volume. 

770 Must have the same length as aquifer_pore_volumes. 

771 molecular_diffusivity : float or array-like 

772 Effective (retarded-frame) molecular diffusivity [m²/day]. Can be a 

773 scalar (same for all pore volumes) or an array with the same length as 

774 aquifer_pore_volumes. Must be non-negative. For solute transport, this is 

775 the molecular diffusion coefficient D_m [m²/day] — typically ~1e-5 m²/day, 

776 negligible compared to microdispersion. For heat transport, pass the 

777 thermal diffusivity D_th = lambda / (rho*c)_eff [m²/day], typically 

778 0.01-0.1 m²/day. 

779 

780 Internally, this contributes ``2 * molecular_diffusivity * tau`` to the 

781 variance, where ``tau`` is the elapsed time in days (no extra factor of 

782 R). The retardation factor instead enters the flux coefficient 

783 ``D_s/v_s = R D_m / v_fluid + alpha_L`` through the solute-front velocity 

784 ``v_s = Q L / (R V_pore)``. For heat transport, the thermal diffusivity 

785 already represents the effective diffusivity D_eff in the porous matrix; 

786 for solutes the contribution is typically negligible. 

787 longitudinal_dispersivity : float or array-like 

788 Longitudinal dispersivity [m]. Can be a scalar (same for all pore 

789 volumes) or an array with the same length as aquifer_pore_volumes. 

790 Must be non-negative. Represents microdispersion from pore-scale velocity variations. 

791 Set to 0 for pure molecular diffusion. 

792 retardation_factor : float, optional 

793 Retardation factor of the compound in the aquifer (default 1.0). 

794 Values > 1.0 indicate slower transport due to sorption. 

795 spinup : {'constant'} or None, optional 

796 Spin-up policy (default ``'constant'``). ``'constant'`` extends tedges by 

797 100 years on each side so that output bins near the boundary are always 

798 informed. ``None`` disables the extension; output bins without sufficient 

799 upstream data become NaN. Float fraction-threshold mode is not implemented 

800 and raises ``NotImplementedError``. 

801 

802 Returns 

803 ------- 

804 numpy.ndarray 

805 Bin-averaged concentration in the extracted water. Same units as cin. 

806 Length equals len(cout_tedges) - 1. NaN values indicate time periods 

807 with no valid contributions from the infiltration data. 

808 

809 Raises 

810 ------ 

811 ValueError 

812 If input dimensions are inconsistent, if diffusivity is negative, 

813 or if aquifer_pore_volumes and streamline_length have different lengths. 

814 

815 See Also 

816 -------- 

817 extraction_to_infiltration : Inverse operation (deconvolution) 

818 gwtransport.advection.infiltration_to_extraction : Pure advection (no dispersion) 

819 gwtransport.diffusion_fast.infiltration_to_extraction : Fast closed-form equivalent 

820 :ref:`concept-dispersion-scales` : Macrodispersion vs microdispersion 

821 

822 Notes 

823 ----- 

824 The algorithm constructs a coefficient matrix W where cout = W @ cin: 

825 

826 1. For each pore volume, build a cell grid in cumulative volume space: 

827 

828 - cells span ``(V_cout[i], V_cout[i+1]) x V_cin[j]`` for each 

829 (cout-bin i, cin-edge j) 

830 - delta_volume = V_cout - V_cin - r_vpv encodes the parcel's offset 

831 from the outlet at each (cout-edge, cin-edge) 

832 

833 2. For each cell, compute the bin-averaged Kreft-Zuber flux concentration 

834 ``frac[i, j] = (1/dV_i) * integral C_F(L, V; t_j) dV`` by resolution-aware 

835 composite Gauss-Legendre quadrature in volume space, split at flow-bin 

836 boundaries so that ``t(V)`` is linear within each sub-interval. Where the 

837 erf-like front (width ``sqrt(4*D_t)`` in volume units) is under-resolved by 

838 a single 16-point rule -- as for near-zero dispersivity -- the sub-interval 

839 is tiled with front-centred panels; smooth sub-intervals keep the single 

840 rule. The moving-frame variance ``D_t = D_m*tau + alpha_L*xi`` is evaluated 

841 at each quadrature node (never capped at the residence time). 

842 

843 3. Coefficient for bin: ``coeff[i,j] = frac[i, j] - frac[i, j+1]``. This 

844 is the contribution of cin[j] to cout[i] in the W matrix. 

845 

846 4. Average coefficients across all pore volumes. 

847 

848 The K-Z flux-correction term in C_F = C_R - (D_s/v_s) * dC_R/dx (solute-front 

849 velocity v_s = Q L / (R V_pore), dispersion D_s = D_m + alpha_L * v_s) is what 

850 makes the column-sum invariant exact under variable Q; see the module 

851 docstring for the derivation. 

852 

853 Examples 

854 -------- 

855 Basic usage with constant flow: 

856 

857 >>> import pandas as pd 

858 >>> import numpy as np 

859 >>> from gwtransport.diffusion import infiltration_to_extraction 

860 >>> 

861 >>> # Create time edges 

862 >>> tedges = pd.date_range(start="2020-01-01", end="2020-01-20", freq="D") 

863 >>> cout_tedges = pd.date_range(start="2020-01-05", end="2020-01-25", freq="D") 

864 >>> 

865 >>> # Input concentration (step function) and constant flow 

866 >>> cin = np.zeros(len(tedges) - 1) 

867 >>> cin[5:10] = 1.0 # Pulse of concentration 

868 >>> flow = np.ones(len(tedges) - 1) * 100.0 # 100 m³/day 

869 >>> 

870 >>> # Single pore volume of 500 m³, travel distance 100 m 

871 >>> aquifer_pore_volumes = np.array([500.0]) 

872 >>> streamline_length = np.array([100.0]) 

873 >>> 

874 >>> # Compute with dispersion (molecular diffusion + dispersivity) 

875 >>> # Scalar values broadcast to all pore volumes 

876 >>> cout = infiltration_to_extraction( 

877 ... cin=cin, 

878 ... flow=flow, 

879 ... tedges=tedges, 

880 ... cout_tedges=cout_tedges, 

881 ... aquifer_pore_volumes=aquifer_pore_volumes, 

882 ... streamline_length=streamline_length, 

883 ... molecular_diffusivity=1e-4, # m²/day, same for all pore volumes 

884 ... longitudinal_dispersivity=1.0, # m, same for all pore volumes 

885 ... ) 

886 

887 With multiple pore volumes (heterogeneous aquifer): 

888 

889 >>> # Distribution of pore volumes and corresponding travel distances 

890 >>> aquifer_pore_volumes = np.array([400.0, 500.0, 600.0]) 

891 >>> streamline_length = np.array([80.0, 100.0, 120.0]) 

892 >>> 

893 >>> # Scalar diffusion parameters broadcast to all pore volumes 

894 >>> cout = infiltration_to_extraction( 

895 ... cin=cin, 

896 ... flow=flow, 

897 ... tedges=tedges, 

898 ... cout_tedges=cout_tedges, 

899 ... aquifer_pore_volumes=aquifer_pore_volumes, 

900 ... streamline_length=streamline_length, 

901 ... molecular_diffusivity=1e-4, # m²/day 

902 ... longitudinal_dispersivity=1.0, # m 

903 ... ) 

904 """ 

905 cout_tedges = pd.DatetimeIndex(cout_tedges) 

906 tedges = pd.DatetimeIndex(tedges) 

907 

908 cin = np.asarray(cin, dtype=float) 

909 flow, aquifer_pore_volumes, streamline_length, molecular_diffusivity, longitudinal_dispersivity = ( 

910 _prepare_diffusion_arrays( 

911 flow=flow, 

912 aquifer_pore_volumes=aquifer_pore_volumes, 

913 streamline_length=streamline_length, 

914 molecular_diffusivity=molecular_diffusivity, 

915 longitudinal_dispersivity=longitudinal_dispersivity, 

916 ) 

917 ) 

918 

919 _validate_diffusion_inputs( 

920 tedges=tedges, 

921 flow=flow, 

922 aquifer_pore_volumes=aquifer_pore_volumes, 

923 streamline_length=streamline_length, 

924 molecular_diffusivity=molecular_diffusivity, 

925 longitudinal_dispersivity=longitudinal_dispersivity, 

926 retardation_factor=retardation_factor, 

927 cin_values=cin, 

928 ) 

929 

930 extend_tedges = _diffusion_extend_tedges_flag(spinup) 

931 coeff_matrix, valid_cout_bins = _infiltration_to_extraction_coeff_matrix( 

932 flow=flow, 

933 tedges=tedges, 

934 cout_tedges=cout_tedges, 

935 aquifer_pore_volumes=aquifer_pore_volumes, 

936 streamline_length=streamline_length, 

937 molecular_diffusivity=molecular_diffusivity, 

938 longitudinal_dispersivity=longitudinal_dispersivity, 

939 retardation_factor=retardation_factor, 

940 extend_tedges=extend_tedges, 

941 ) 

942 

943 cout = coeff_matrix @ cin 

944 

945 # Output bins are invalid where the coefficient sum is near zero (no cin has broken 

946 # through yet) or the bin extends beyond the input data range (valid_cout_bins). 

947 total_coeff = np.sum(coeff_matrix, axis=1) 

948 no_valid_contribution = (total_coeff < EPSILON_COEFF_SUM) | ~valid_cout_bins 

949 cout[no_valid_contribution] = np.nan 

950 

951 return cout 

952 

953 

954def extraction_to_infiltration( 

955 *, 

956 cout: npt.ArrayLike, 

957 flow: npt.ArrayLike, 

958 tedges: pd.DatetimeIndex, 

959 cout_tedges: pd.DatetimeIndex, 

960 aquifer_pore_volumes: npt.ArrayLike, 

961 streamline_length: npt.ArrayLike, 

962 molecular_diffusivity: npt.ArrayLike, 

963 longitudinal_dispersivity: npt.ArrayLike, 

964 retardation_factor: float = 1.0, 

965 regularization_strength: float = 1e-10, 

966 spinup: str | None = "constant", 

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

968 """ 

969 Compute infiltration concentration from extracted water (deconvolution with dispersion). 

970 

971 Inverts the forward transport model by building the forward coefficient 

972 matrix ``W_forward`` from :func:`infiltration_to_extraction` and solving 

973 ``W_forward @ cin = cout`` via Tikhonov regularization. Well-determined 

974 modes are dominated by the data; poorly-determined modes are pulled 

975 toward the physically motivated target (transpose-and-normalize of the