Coverage for src/gwtransport/diffusion_fast.py: 100%
157 statements
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-04 21:13 +0000
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-04 21:13 +0000
1"""
2Fast closed-form 1D advection-dispersion transport (Kreft-Zuber flux concentration).
4This module shares the conceptual model of :mod:`gwtransport.diffusion` -- advection with
5microdispersion (``alpha_L``) and molecular diffusion (``D_m``) along orthogonal (Cartesian)
6flow paths, one independent streamtube per aquifer pore volume, with the spread across the
7pore volume distribution providing macrodispersion and linear sorption entering through the
8retardation factor. It reports the same Kreft-Zuber (1978) flux concentration ``C_F`` at the
9outlet of the streamtube bundle, but evaluates the bin-averaged breakthrough in closed form
10instead of by Gauss-Legendre quadrature.
12For each streamtube (one aquifer pore volume) the resident concentration in moving-frame
13cumulative-volume (V) coordinates is the Gaussian CDF
14``C_R = 0.5 * erfc((L - xi) / (2 * sqrt(D_t)))``, with ``D_t = D_m * tau + alpha_L * xi``
15the moving-frame dispersion product. Its bin-average over a cout bin has the closed-form
16antiderivative ``I(x) = 0.5*x + 0.5*[x*erf(x/s) + (s/sqrt(pi))*exp(-(x/s)^2)]``,
17``s = 2*sqrt(D_t)``. Evaluating ``I`` once per cout edge with ``D_t`` carried *per edge*
18and differencing yields the flux concentration ``C_F`` directly -- not merely ``C_R`` --
19because ``dD_t/dx = D_m/v_s + alpha_L = D_s/v_s`` is exactly the Kreft-Zuber flux coefficient
20at the solute-front velocity ``v_s = Q*L/(R*V_pore)`` (using ``d(tau)/dx = 1/v_s =
21R*V_pore/(L*Q)``). The dispersive boundary-flux correction therefore emerges from the
22``D_t`` variation across the bin; no explicit correction term is added.
24The elapsed time ``tau`` and travel distance ``xi`` are read directly from the time and
25cumulative-volume edges (``tau_ij = t_cout_i - t_cin_j``, ``xi`` geometric), so no per-cell
26quadrature and no residence-time inversion is needed. The coefficient matrix is built only on
27the breakthrough band -- the cumulative-volume band where the bin-averaged ``C_F`` is
28unsaturated, the only region with non-zero coefficients -- so the build cost scales with the
29band width (a few percent of the matrix at realistic dispersion) rather than with the full grid.
31Streamtube assumption (no cross-sectional area parameter)
32---------------------------------------------------------
34Each entry in ``aquifer_pore_volumes`` is an independent 1D streamtube; molecular diffusion
35enters the V-space variance through ``D_m * tau`` and microdispersion through
36``alpha_L * xi``. ``streamline_length`` / ``molecular_diffusivity`` /
37``longitudinal_dispersivity`` may be a scalar (shared by all streamtubes) or an array with
38one value per pore volume, exactly as in :mod:`gwtransport.diffusion`.
40Choosing among the three diffusion modules
41------------------------------------------
43Whenever the cout grid is at or finer than the flow grid, this module reproduces
44:mod:`gwtransport.diffusion` to machine precision for *every* parameter regime -- including
45``retardation_factor != 1`` with ``molecular_diffusivity > 0``, where the antiderivative's slope
46``dD_t/dx = D_s/v_s`` already carries the solute-front Kreft-Zuber flux coefficient natively --
47while being ~80-90x faster even before banding. So it is the right default. The only case that
48favours :mod:`gwtransport.diffusion` is a cout grid *coarser* than the flow detail: this module
49treats ``flow_out`` as constant within each cout bin, whereas :mod:`gwtransport.diffusion`
50integrates the full ``tedges``-resolution flow within each cout bin -- a ~0.1%-of-peak
51difference for a rapidly-varying ``cin`` over wide cout bins under variable flow.
53The third rung is :mod:`gwtransport.diffusion_fast_fast`, which is approximate by design: it folds
54molecular diffusion into an effective dispersivity per streamtube and evaluates one banded
55breakthrough on the native cumulative-volume grid, so it is faster still. It agrees with this module
56to ~1e-6 (smooth inputs) to ~1e-4 (sharp inputs) at constant flow and degrades in the
57molecular-diffusion-dominated corner (``alpha_L`` ~ 0) under strongly variable flow; this module is
58the one to use whenever machine precision against :mod:`gwtransport.diffusion` is required.
60Available functions:
62- :func:`infiltration_to_extraction` - Forward transport from an explicit ``aquifer_pore_volumes``
63 distribution: returns the bin-averaged Kreft-Zuber flux concentration on ``cout_tedges``, evaluated
64 from the closed-form antiderivative on the breakthrough band only and averaged with equal weight over
65 the streamtubes. ``streamline_length``, ``molecular_diffusivity`` and ``longitudinal_dispersivity``
66 are either scalars shared by all streamtubes or one value per pore volume. ``flow_out`` (extraction
67 flow on the ``cout_tedges`` grid) is required whenever ``cout_tedges`` differs from ``tedges``: it
68 sets the cout-bin volumes and the outlet velocity. Output bins without complete breakthrough
69 information are NaN.
71- :func:`extraction_to_infiltration` - Reverse direction: assembles the same banded forward operator
72 and deconvolves it with banded Tikhonov regularization (banded Cholesky on the normal equations),
73 returning the bin-averaged infiltration concentration on ``tedges``. NaN entries in ``cout`` mark
74 measurement gaps and are excluded from the solve; spin-up and unconstrained cin bins come back NaN.
76- :func:`gamma_infiltration_to_extraction` - :func:`infiltration_to_extraction` with the pore volume
77 distribution given as a (shifted) gamma -- either (mean, std) or (alpha, beta), plus ``loc`` --
78 discretized into ``n_bins`` equal-probability streamtubes that share one ``streamline_length`` and
79 one pair of dispersion parameters.
81- :func:`gamma_extraction_to_infiltration` - :func:`extraction_to_infiltration` with the same gamma
82 parameterization of the pore volume distribution: reconstructs ``cin`` on ``tedges`` from ``cout``.
84References
85----------
86Kreft, A., & Zuber, A. (1978). On the physical meaning of the dispersion equation and its
87solutions for different initial and boundary conditions. Chemical Engineering Science,
8833(11), 1471-1480.
90This file is part of gwtransport which is released under AGPL-3.0 license.
91See the ./LICENSE file or go to https://github.com/gwtransport/gwtransport/blob/main/LICENSE for full license details.
92"""
94import numpy as np
95import numpy.typing as npt
96import pandas as pd
98from gwtransport import gamma
99from gwtransport._diffusion_shared import (
100 _DT_FLOOR,
101 _EPSILON_COEFF_SUM,
102 _advective_valid_cout_bins,
103 _breakthrough_antideriv,
104 _coerce_and_validate,
105 _cout_cumulative_volume,
106 _extend_tedges,
107 _extend_tedges_flag,
108 _solve_reverse_banded,
109)
110from gwtransport._time import dt_to_days, tedges_to_days
111from gwtransport.utils import cumulative_flow_volume
113# Saturation threshold U of the banded build: a cout/cin pair is only evaluated while the
114# breakthrough |x|/(2*sqrt(D_t)) <= U; beyond it the bin-averaged C_F is saturated to 0 or 1.
115# At U >= ~6 the dense kernel itself rounds the dropped tail to exactly 0 or 1 (the Gaussian
116# term underflows below the ulp of x), so the banded matrix is bit-identical to the dense one.
117# A smaller U narrows the band -> faster, at the cost of dropping breakthrough tails of order
118# exp(-U^2); the builders take it as a parameter so that trade-off stays testable.
119_DEFAULT_SATURATION_THRESHOLD = 7.0
122def _pv_band_geometry(
123 *,
124 cumulative_volume_at_cout: npt.NDArray[np.floating],
125 cumulative_volume_at_cin: npt.NDArray[np.floating],
126 cout_tedges_days: npt.NDArray[np.floating],
127 tedges_days: npt.NDArray[np.floating],
128 r_vpv: float,
129 length: float,
130 molecular_diffusivity: float,
131 longitudinal_dispersivity: float,
132 min_cin_flow: float,
133 saturation_threshold: float,
134 n_cin_bins: int,
135 stagnant_time_at_cout: npt.NDArray[np.floating],
136) -> tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]]:
137 r"""Per-cout-row band bounds (lo, hi) in cin-bin columns for one streamtube (geometry only).
139 Locates the narrow cumulative-volume band where the breakthrough transitions between 0 and 1
140 -- the only cin bins with a non-zero coefficient. Within a streamtube the moving-frame
141 dispersion product is exactly linear in the breakthrough coordinate, ``D_t(x) = A + B*x``,
142 with slope ``B = dD_t/dx = R*D_m/v_fluid + alpha_L`` (``v_fluid = Q*L/V_pore``) and
143 intercept ``A`` the front value, so the
144 saturation edge ``|x| = saturation_threshold * 2*sqrt(D_t(x))`` is the root of a quadratic --
145 the band half-width is closed form, no iteration. The band is centred per cout bin on the
146 front ``V_cin = V_cout - R*V_pore`` and mapped to cin-edge columns with ``searchsorted`` (so
147 non-uniform / variable-flow grids need no special handling).
149 With a warm-start spin-up and ``D_m > 0`` the breakthrough is *also* unsaturated at the
150 data-start edge (cin edge 0): a leading zero-flow plateau holds the cumulative volume flat
151 there, so the front search lands the local band one or more columns inside the record while a
152 genuine, non-negligible coefficient remains at column 0 (the warm-start tail of a wide bin with
153 large ``tau`` -> large ``D_t``). Each row therefore additionally tests ``|x0| < U*2*sqrt(D_t0)``
154 at edge 0 and, when non-saturated, drops its band lower bound to 0 so that tail is kept.
156 ``stagnant_time_at_cout`` is the interior zero-flow (stagnation) time accumulated before each
157 cout bin's upper edge [days]. During a pumped-off gap the moving-frame variance keeps growing
158 (``D_m * tau``) at frozen cumulative volume, so post-side columns before the gap carry a ``D_t``
159 larger than the front intercept plus the flowing slope bound; the gap-accumulated product
160 ``D_m * stagnant_time`` is added to the post-side intercept so those columns stay in the band.
162 Returns
163 -------
164 lo : ndarray of int, shape (n_cout_bins,)
165 Per-row band lower bound (inclusive), clipped to ``[0, n_cin_bins - 1]``.
166 hi : ndarray of int, shape (n_cout_bins,)
167 Per-row band upper bound (inclusive), clipped to ``[0, n_cin_bins - 1]``.
168 """
169 v_cin = cumulative_volume_at_cin
170 n_cin_edges = v_cin.size
171 v_cout_lo, v_cout_hi = cumulative_volume_at_cout[:-1], cumulative_volume_at_cout[1:]
172 t_cout_lo, t_cout_hi = cout_tedges_days[:-1], cout_tedges_days[1:]
174 # Front locus per cout bin, then the band half-widths in closed form. In the breakthrough
175 # coordinate x (x > 0 broken through, x < 0 not), the moving-frame dispersion product is
176 # D_t(x) = D_m*tau(x) + alpha_L*xi(x), with intercept A := front D_t. Both bounds below are
177 # conservative -- they never under-cover the unsaturated band, where |x| < U*2*sqrt(D_t):
178 # PRE side (x < 0): tau and xi both shrink away from the front, so D_t <= a_pre := front D_t
179 # and the band reaches |x| = 2U*sqrt(a_pre). Flow-independent, no cancellation.
180 # POST side (x > 0): D_t grows; bounded by the steepest slope D_t <= A + B*x with
181 # B = alpha_L + D_m*r_vpv/(L*min_cin_flow) (dtau/dx = r_vpv/(L*flow) <= r_vpv/(L*min_cin_flow)),
182 # giving the root x_post = 2U^2*B + 2U*sqrt(U^2*B^2 + A). The flow term only matters for strong
183 # molecular diffusion (already the wide-band regime); the mechanical term is exact.
184 # The post extent anchors at the lower cout edge (smallest V_cin), whose front can sit in slower
185 # flow with larger tau, so it carries its own intercept a_post; max(a_pre, a_post) bounds the front
186 # D_t of both cout edges there. The pre extent anchors at the upper cout edge. a_pre uses the upper
187 # edge time / mid front (conservative for the pre side). +1 absorbs searchsorted rounding;
188 # min_cin_flow == 0 (no flow) -> the post slope is unbounded; b_max is set to span the whole axis
189 # only when D_m > 0, else 0 (no diffusion, no flow term), avoiding a 0/0 NaN.
190 u = saturation_threshold
191 v_front = 0.5 * (v_cout_lo + v_cout_hi) - r_vpv
192 center = np.searchsorted(v_cin, v_front)
193 center_post = np.searchsorted(v_cin, v_cout_lo - r_vpv) # searchsorted >= 0, so only the high clip binds
194 disp = longitudinal_dispersivity * length
195 a_pre = molecular_diffusivity * np.maximum(t_cout_hi - tedges_days[np.minimum(center, n_cin_edges - 1)], 0.0) + disp
196 a_post = (
197 molecular_diffusivity * np.maximum(t_cout_lo - tedges_days[np.minimum(center_post, n_cin_edges - 1)], 0.0)
198 + disp
199 )
200 # Interior stagnation: tau across a zero-flow gap grows without volume, so the flowing slope
201 # bound below misses D_m * (gap time) at post-side columns before the gap; add it to the
202 # intercept (conservative: all interior stagnant time before the row's upper edge).
203 a_post = np.maximum(a_post, a_pre) + molecular_diffusivity * stagnant_time_at_cout
204 pre_x = 2.0 * u * np.sqrt(a_pre)
205 if min_cin_flow > 0.0:
206 b_max = longitudinal_dispersivity + molecular_diffusivity * r_vpv / (length * min_cin_flow)
207 else:
208 b_max = np.inf if molecular_diffusivity > 0.0 else longitudinal_dispersivity
209 post_x = 2.0 * u * u * b_max + 2.0 * u * np.sqrt(u * u * b_max * b_max + a_post)
210 col_post = np.searchsorted(v_cin, (v_cout_lo - r_vpv) - r_vpv * post_x / length) # smallest V_cin
211 col_pre = np.searchsorted(v_cin, (v_cout_hi - r_vpv) + r_vpv * pre_x / length) # largest V_cin
212 # Scalar (global-max) half-widths, applied to every row via the band centre. Taking the max over
213 # all cout rows is conservative: a row whose front sits in a slow / zero-flow plateau (large tau,
214 # wide local transition) sets the half-width for all rows, so an interior zero-flow gap straddled
215 # by a misaligned cout bin is never under-covered. +1 absorbs the searchsorted rounding.
216 hw_post = max(int(np.max(center - col_post)), 1) + 1
217 hw_pre = max(int(np.max(col_pre - center)), 1) + 1
218 lo = np.clip(center - hw_post, 0, n_cin_bins - 1)
219 hi = np.clip(center + hw_pre - 1, 0, n_cin_bins - 1)
221 # Warm-start data-start tail: drop the band to column 0 where the breakthrough is still
222 # unsaturated at cin edge 0. x0 is the breakthrough coordinate of the lower cout edge measured
223 # from V_cin[0] (the smallest x0 over the bin, hence the most-broken-through / largest |D_t0|);
224 # D_t0 floors at _DT_FLOOR so the zero-dispersion limit gives x0 != 0 -> saturated -> no spurious
225 # widening. The leading zero-flow plateau that triggers this keeps full_band bounded by the
226 # plateau length, independent of the record length.
227 x0 = (v_cout_lo - v_cin[0] - r_vpv) * length / r_vpv
228 tau0 = np.maximum(t_cout_lo - tedges_days[0], 0.0)
229 dt0 = np.maximum(molecular_diffusivity * tau0 + longitudinal_dispersivity * np.maximum(x0 + length, 0.0), _DT_FLOOR)
230 non_saturated0 = np.abs(x0) < u * 2.0 * np.sqrt(dt0)
231 lo[non_saturated0] = 0
232 return lo, hi
235def _pv_band_values(
236 *,
237 col_start: npt.NDArray[np.intp],
238 full_band: int,
239 cumulative_volume_at_cout: npt.NDArray[np.floating],
240 cumulative_volume_at_cin: npt.NDArray[np.floating],
241 cout_tedges_days: npt.NDArray[np.floating],
242 tedges_days: npt.NDArray[np.floating],
243 r_vpv: float,
244 length: float,
245 molecular_diffusivity: float,
246 longitudinal_dispersivity: float,
247) -> npt.NDArray[np.floating]:
248 r"""Bin-averaged ``C_F`` stripe for one streamtube on the shared band (values pass).
250 ``C_F`` over a cout bin is ``(I(x_hi) - I(x_lo)) / dx`` with the closed-form antiderivative
251 ``I`` evaluated at the two cout edges bounding the bin. Because ``D_t = D_m*tau + alpha_L*xi``
252 with ``d(tau)/dx = 1/v_s = R*V_pore/(L*Q)``, the antiderivative's slope ``dD_t/dx = R*D_m/v_fluid + alpha_L =
253 D_s/v_s`` is exactly the Kreft-Zuber flux coefficient at the solute-front velocity
254 ``v_s = Q*L/(R*V_pore)``, so the flux concentration emerges natively -- no correction term is
255 added. The stripe is the band itself: each row spans cin edges
256 ``col_start[k] .. col_start[k] + full_band`` (``full_band + 1`` edges), so the coefficient for
257 band offset ``b`` (cin bin ``col_start[k] + b``) is ``frac[b] - frac[b + 1]``.
259 ``I`` is evaluated once per cout EDGE, not once per (row, edge) pair: an interior edge bounds two
260 adjacent rows (it is a row's upper edge and the next row's lower edge), so the two evaluations of
261 ``I`` there are the identical function of the identical breakthrough coordinate. The build
262 therefore evaluates ``I`` on the ``n_cout_bins + 1`` cout edges over a single per-edge cin-edge
263 window (anchored so both adjacent rows can read it), then gathers each row's lower/upper edge
264 values from it -- roughly halving the ``erf`` work relative to evaluating both edges per row. The
265 zero-dispersion limit is exact here too: ``D_t`` floors to ``_DT_FLOOR``, so ``C_F`` is a step
266 smoothed by ~1e-15.
268 Returns
269 -------
270 coeff : ndarray, shape (n_cout_bins, full_band)
271 Per-row, per-band-offset coefficient ``frac[:, :-1] - frac[:, 1:]`` already aligned to the
272 banded buffer (offset 0 -> cin bin ``col_start[k]``).
273 """
274 v_cin = cumulative_volume_at_cin
275 n_cin_edges = v_cin.size
276 n_cout_bins = cumulative_volume_at_cout.size - 1
277 width = full_band + 1
279 # Per-edge cin-edge window. Cout edge e bounds row e (as its lower edge, needing cin band anchor
280 # col_start[e]) and row e-1 (as its upper edge, needing col_start[e-1]); anchoring at the smaller
281 # of the two lets one evaluation serve both. col_start is non-decreasing here, but the min keeps
282 # both read offsets non-negative regardless. The sentinel (>= any column) drops the absent
283 # consumer at the two end edges. The window is widened past full_band + 1 only by the col_start
284 # jump between adjacent rows (typically 0-1), so it collapses to the row width when aligned.
285 big = n_cin_edges
286 estart = np.minimum(np.append(col_start, big), np.append(big, col_start))
287 lo_off = col_start - estart[:-1]
288 hi_off = col_start - estart[1:]
289 edge_width = width + int(max(lo_off.max(), hi_off.max()))
290 edge_cols = np.clip(estart[:, None] + np.arange(edge_width)[None, :], 0, n_cin_edges - 1)
291 v_c, t_c = v_cin[edge_cols], tedges_days[edge_cols]
292 sw_edge = (cumulative_volume_at_cout[:, None] - v_c - r_vpv) * length / r_vpv
294 # Gather each row's lower (cout edge k) and upper (cout edge k + 1) breakthrough coordinates from
295 # the per-edge stripe. Clipped cin edges fall in the warm-start-saturated tail, carrying frac 0/1
296 # that telescopes away in the coefficient difference.
297 rows = np.arange(n_cout_bins)[:, None]
298 band = np.arange(width)[None, :]
299 lo_idx = lo_off[:, None] + band
300 hi_idx = hi_off[:, None] + band
301 sw_lo = sw_edge[rows, lo_idx]
302 sw_hi = sw_edge[rows + 1, hi_idx]
304 # No dispersion: C_R is the step H(x); its exact bin-average is the fraction of the cout bin with
305 # x > 0, and at a zero-width (dv_cout = 0) cout bin it is the point value 0.5*(1 + sign(x_lo)).
306 # This matches the dense kernel's zero-dispersion branch bit-for-bit (the floored erf form would
307 # instead give 0 at dx = 0, dropping the step at a gap-straddling misaligned cout bin).
308 if molecular_diffusivity == 0.0 and longitudinal_dispersivity == 0.0:
309 dx = sw_hi - sw_lo
310 with np.errstate(divide="ignore", invalid="ignore"):
311 frac = (np.maximum(sw_hi, 0.0) - np.maximum(sw_lo, 0.0)) / dx
312 frac = np.where(dx > 0.0, frac, 0.5 + 0.5 * np.sign(sw_lo))
313 return frac[:, :-1] - frac[:, 1:]
315 dt_edge = np.maximum(
316 molecular_diffusivity * np.maximum(cout_tedges_days[:, None] - t_c, 0.0)
317 + longitudinal_dispersivity * np.maximum(sw_edge + length, 0.0),
318 _DT_FLOOR,
319 )
320 i_edge = _breakthrough_antideriv(sw_edge, dt_edge)
321 i_lo = i_edge[rows, lo_idx]
322 i_hi = i_edge[rows + 1, hi_idx]
323 dx = sw_hi - sw_lo
324 delta = i_hi - i_lo
326 # Stagnation (zero-flow gap) correction. The endpoint difference above integrates the exact
327 # 1-form dI = C_R dx + (dI/dD_t) dD_t along the bin; on flowing stretches dD_t/dx = D_s/v_s
328 # makes dI = C_F dx, but across a zero-flow gap x is frozen while tau (hence D_t) keeps
329 # growing, so a cout bin straddling a gap picks up a vertical int (dI/dD_t) dD_t that carries
330 # no extracted volume and does not belong in the flux-weighted bin average. Subtract, per
331 # (cout bin, zero-flow cin bin) overlap clipped to the bin's time span, the antiderivative
332 # jump at the gap's frozen breakthrough coordinate. Consecutive zero-flow bins telescope to
333 # the full gap jump; zero-length overlaps (grid-aligned cout edges) are skipped, keeping the
334 # aligned path bit-identical.
335 if molecular_diffusivity > 0.0:
336 gap_bins = np.nonzero((np.diff(v_cin) <= 0.0) & (np.diff(tedges_days) > 0.0))[0]
337 if gap_bins.size:
338 t_gap_lo, t_gap_hi = tedges_days[gap_bins], tedges_days[gap_bins + 1]
339 k_first = np.maximum(np.searchsorted(cout_tedges_days, t_gap_lo, side="right") - 1, 0)
340 k_last = np.minimum(np.searchsorted(cout_tedges_days, t_gap_hi, side="left") - 1, n_cout_bins - 1)
341 counts = np.maximum(k_last - k_first + 1, 0)
342 k_pair = np.repeat(k_first, counts) + (
343 np.arange(counts.sum()) - np.repeat(np.cumsum(counts) - counts, counts)
344 )
345 g_pair = np.repeat(np.arange(gap_bins.size), counts)
346 t_lo_pair = np.maximum(cout_tedges_days[k_pair], t_gap_lo[g_pair])
347 t_hi_pair = np.minimum(cout_tedges_days[k_pair + 1], t_gap_hi[g_pair])
348 keep = t_hi_pair > t_lo_pair
349 if np.any(keep):
350 k_pair, g_pair = k_pair[keep], g_pair[keep]
351 t_lo_pair, t_hi_pair = t_lo_pair[keep], t_hi_pair[keep]
352 cols = np.clip(col_start[k_pair][:, None] + band, 0, n_cin_edges - 1)
353 x_gap = (v_cin[gap_bins[g_pair], None] - v_cin[cols] - r_vpv) * length / r_vpv
354 disp_term = longitudinal_dispersivity * np.maximum(x_gap + length, 0.0)
355 t_c = tedges_days[cols]
356 dt_gap_lo = np.maximum(
357 molecular_diffusivity * np.maximum(t_lo_pair[:, None] - t_c, 0.0) + disp_term, _DT_FLOOR
358 )
359 dt_gap_hi = np.maximum(
360 molecular_diffusivity * np.maximum(t_hi_pair[:, None] - t_c, 0.0) + disp_term, _DT_FLOOR
361 )
362 jump = _breakthrough_antideriv(x_gap, dt_gap_hi) - _breakthrough_antideriv(x_gap, dt_gap_lo)
363 np.add.at(delta, k_pair, -jump)
365 with np.errstate(divide="ignore", invalid="ignore"):
366 frac = np.where(dx > 0.0, delta / dx, 0.0)
368 return frac[:, :-1] - frac[:, 1:]
371def _closed_form_coeff_matrix(
372 *,
373 flow: npt.NDArray[np.floating],
374 tedges: pd.DatetimeIndex,
375 cout_tedges: pd.DatetimeIndex,
376 flow_out: npt.NDArray[np.floating] | None,
377 aquifer_pore_volumes: npt.NDArray[np.floating],
378 streamline_length: npt.NDArray[np.floating],
379 molecular_diffusivity: npt.NDArray[np.floating],
380 longitudinal_dispersivity: npt.NDArray[np.floating],
381 retardation_factor: float,
382 extend_tedges: bool,
383 saturation_threshold: float = _DEFAULT_SATURATION_THRESHOLD,
384) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.intp], npt.NDArray[np.bool_]]:
385 """Build the banded forward operator (``cout = W @ cin``) via the closed-form C_F.
387 Mirrors :func:`gwtransport.diffusion._infiltration_to_extraction_coeff_matrix`
388 (per-streamtube loop over pore volumes, 100-year warm-start extension, residence-time
389 validity) but computes the bin-averaged flux concentration in closed form instead of
390 16-point Gauss-Legendre quadrature, and stores it in BANDED layout: row ``k`` of the dense
391 operator ``W`` is ``band_vals[k]`` placed at columns ``[col_start[k], col_start[k] + full_band)``.
392 The build runs in two passes over the pore-volume loop -- a cheap geometry pass that sizes the
393 per-row band union, then a values pass that scatters each streamtube's ``C_F`` stripe into the
394 banded buffer. ``streamline_length``, ``molecular_diffusivity``, and
395 ``longitudinal_dispersivity`` are per-pore-volume arrays (length ``len(aquifer_pore_volumes)``).
397 Returns
398 -------
399 band_vals : ndarray, shape (n_cout_bins, full_band)
400 Banded forward weights, NaN replaced with zero.
401 col_start : ndarray of int, shape (n_cout_bins,)
402 First cin-bin column of each cout row's band.
403 valid_cout_bins : ndarray of bool, shape (n_cout_bins,)
404 Output bins with complete breakthrough information for every streamtube.
405 """
406 work_tedges = _extend_tedges(tedges) if extend_tedges else tedges
408 tedges_days = tedges_to_days(work_tedges)
409 cout_tedges_days = tedges_to_days(cout_tedges, ref=work_tedges[0])
411 # Cumulative through-flow volume on a common axis. cout-edge volumes come from flow_out
412 # when provided (the user-specified extraction-side flow), placed on the infiltration
413 # volume axis by anchoring at the first cout edge inside the flow record (so an output
414 # window that starts before the input data stays correctly aligned); otherwise
415 # interpolated from the infiltration curve.
416 cumulative_volume_at_cin = cumulative_flow_volume(flow, dt_to_days(work_tedges))
417 cumulative_volume_at_cout = _cout_cumulative_volume(
418 flow_out=flow_out,
419 cout_tedges=cout_tedges,
420 cout_tedges_days=cout_tedges_days,
421 tedges_days=tedges_days,
422 cumulative_volume_at_cin=cumulative_volume_at_cin,
423 )
425 valid_cout_bins = _advective_valid_cout_bins(
426 flow=flow,
427 tedges=work_tedges,
428 cout_tedges=cout_tedges,
429 aquifer_pore_volumes=aquifer_pore_volumes,
430 retardation_factor=retardation_factor,
431 )
433 # Slowest cin-side flow rate, used to bound the broken-through band width (the slowest flow
434 # gives the steepest dD_t/dx). Zero when flow is everywhere zero -> the band widens (capped
435 # at n_cin_bins), and the resulting no-flow rows are masked invalid anyway.
436 positive_flow = flow[flow > 0.0]
437 min_cin_flow = float(positive_flow.min()) if positive_flow.size else 0.0
439 n_cout_bins = len(cout_tedges) - 1
440 n_cin_bins = len(flow)
442 # Interior stagnation (zero-flow gap) time accumulated before each cout bin's upper edge, for
443 # the post-side band intercept (see _pv_band_geometry). The data-start plateau (including the
444 # 100-year warm-start pad when flow[0] == 0) is excluded: its columns share the data-start
445 # cumulative volume, so the per-row x0 tail check already covers them without forcing a
446 # record-wide band.
447 dv_bins = np.diff(cumulative_volume_at_cin)
448 stagnant_dt = np.where(dv_bins > 0.0, 0.0, np.diff(tedges_days))
449 stagnant_dt[: int(np.argmax(dv_bins > 0.0))] = 0.0
450 stagnant_cum = np.concatenate(([0.0], np.cumsum(stagnant_dt)))
451 stagnant_time_at_cout = np.interp(cout_tedges_days[1:], tedges_days, stagnant_cum)
453 # PASS 1 (geometry): per-streamtube band bounds (lo, hi) in cin-bin columns; accumulate the
454 # per-row union over streamtubes. union_lo / union_hi are the min / max bounds per cout row.
455 union_lo = np.full(n_cout_bins, n_cin_bins - 1, dtype=np.intp)
456 union_hi = np.zeros(n_cout_bins, dtype=np.intp)
457 geometry = []
458 for i_pv, v_pore in enumerate(aquifer_pore_volumes):
459 r_vpv = retardation_factor * v_pore
460 length = float(streamline_length[i_pv])
461 d_m = float(molecular_diffusivity[i_pv])
462 alpha_l = float(longitudinal_dispersivity[i_pv])
463 lo, hi = _pv_band_geometry(
464 cumulative_volume_at_cout=cumulative_volume_at_cout,
465 cumulative_volume_at_cin=cumulative_volume_at_cin,
466 cout_tedges_days=cout_tedges_days,
467 tedges_days=tedges_days,
468 r_vpv=r_vpv,
469 length=length,
470 molecular_diffusivity=d_m,
471 longitudinal_dispersivity=alpha_l,
472 min_cin_flow=min_cin_flow,
473 saturation_threshold=saturation_threshold,
474 n_cin_bins=n_cin_bins,
475 stagnant_time_at_cout=stagnant_time_at_cout,
476 )
477 np.minimum(union_lo, lo, out=union_lo)
478 np.maximum(union_hi, hi, out=union_hi)
479 geometry.append((r_vpv, length, d_m, alpha_l, lo, hi))
481 col_start = union_lo
482 full_band = min(int(np.max(union_hi - union_lo)) + 1, n_cin_bins)
483 band_vals = np.zeros((n_cout_bins, full_band))
485 # PASS 2 (values): evaluate each streamtube's C_F stripe on its OWN band [lo, hi] -- typically
486 # much narrower than the shared union band when the APVD spread is wide -- then scatter it
487 # offset-shifted into the union buffer at per-row offset (lo - union_lo). The erf-heavy
488 # antiderivative in _pv_band_values then runs over own_band columns instead of full_band,
489 # removing the redundant erf evaluation across the union width beyond each streamtube's own
490 # front. This is bit-identical to a union-band build: the coefficient at (row, absolute cin bin)
491 # depends only on that cin bin's breakthrough coordinate, not on the band anchoring, and every
492 # union-band column outside [lo, hi] carries a saturated (exactly-0 at the default threshold)
493 # coefficient that contributes nothing to the sum.
494 row_idx = np.arange(n_cout_bins)[:, None]
495 for r_vpv, length, d_m, alpha_l, lo, hi in geometry:
496 own_band = min(int(np.max(hi - lo)) + 1, n_cin_bins)
497 stripe = _pv_band_values(
498 col_start=lo,
499 full_band=own_band,
500 cumulative_volume_at_cout=cumulative_volume_at_cout,
501 cumulative_volume_at_cin=cumulative_volume_at_cin,
502 cout_tedges_days=cout_tedges_days,
503 tedges_days=tedges_days,
504 r_vpv=r_vpv,
505 length=length,
506 molecular_diffusivity=d_m,
507 longitudinal_dispersivity=alpha_l,
508 )
509 # Per-row column offset of the own band inside the union band (>= 0 since union_lo <= lo).
510 cols = (lo - union_lo)[:, None] + np.arange(own_band)[None, :]
511 # Own-band stripes may extend one column past the union width into the saturated tail
512 # (all-zero at the default threshold); drop those so the scatter stays in bounds.
513 in_band = cols < full_band
514 rows = np.broadcast_to(row_idx, cols.shape)
515 np.add.at(band_vals, (rows[in_band], cols[in_band]), stripe[in_band])
517 band_vals /= len(aquifer_pore_volumes)
518 return np.nan_to_num(band_vals, nan=0.0), col_start, valid_cout_bins
521def infiltration_to_extraction(
522 *,
523 cin: npt.ArrayLike,
524 flow: npt.ArrayLike,
525 tedges: pd.DatetimeIndex,
526 cout_tedges: pd.DatetimeIndex,
527 aquifer_pore_volumes: npt.ArrayLike,
528 streamline_length: npt.NDArray[np.floating] | float,
529 molecular_diffusivity: npt.NDArray[np.floating] | float,
530 longitudinal_dispersivity: npt.NDArray[np.floating] | float,
531 retardation_factor: float = 1.0,
532 flow_out: npt.ArrayLike | None = None,
533 spinup: str | None = "constant",
534) -> npt.NDArray[np.floating]:
535 """Compute extracted concentration with advection, microdispersion, and molecular diffusion.
537 Fast closed-form counterpart of :func:`gwtransport.diffusion.infiltration_to_extraction`.
538 Reports the Kreft-Zuber (1978) flux concentration ``C_F`` and reproduces the slow module
539 to machine precision when the cout grid aligns with the flow grid (supply ``flow_out``).
541 Parameters
542 ----------
543 cin : array-like
544 Concentration of the compound in the infiltrating water. Length ``len(tedges) - 1``.
545 flow : array-like
546 Flow rate of water in the aquifer [m³/day]. Length ``len(tedges) - 1``.
547 tedges : pandas.DatetimeIndex
548 Time edges for cin and flow data. Length ``len(cin) + 1``.
549 cout_tedges : pandas.DatetimeIndex
550 Time edges for output data bins. Length ``len(output) + 1``.
551 aquifer_pore_volumes : array-like
552 Aquifer pore volumes [m³] -- one independent streamtube per entry.
553 streamline_length : float or ndarray
554 Travel distance L [m]: a scalar (shared by all streamtubes) or an array with one
555 value per aquifer pore volume. Must be positive.
556 molecular_diffusivity : float or ndarray
557 Effective (retarded-frame) molecular diffusivity D_m [m²/day]: scalar or one value per pore volume.
558 Must be non-negative.
559 longitudinal_dispersivity : float or ndarray
560 Longitudinal dispersivity alpha_L [m] (microdispersion): scalar or one value per pore volume.
561 Must be non-negative.
562 retardation_factor : float, optional
563 Retardation factor (default 1.0). Values > 1.0 indicate slower transport.
564 flow_out : array-like or None, optional
565 Extraction flow rate [m³/day] on the output grid (aligned to ``cout_tedges``,
566 length ``len(cout_tedges) - 1``); constant within each cout bin, like ``flow`` is
567 within each ``tedges`` bin. It defines the cout-bin volumes and the outlet velocity.
568 **Required when ``cout_tedges`` differs from ``tedges``**; may be omitted only when
569 ``cout_tedges`` equals ``tedges`` (then it equals ``flow``). Default None.
570 spinup : {"constant"} | None, optional
571 ``"constant"`` (default) extends ``tedges`` by 100 years on each side so a constant
572 warm-start fills the left-edge spin-up region; ``None`` leaves spin-up cout as NaN.
574 Returns
575 -------
576 numpy.ndarray
577 Bin-averaged Kreft-Zuber flux concentration ``C_F`` in the extracted water. Length
578 ``len(cout_tedges) - 1``. NaN where no infiltration data has broken through.
580 See Also
581 --------
582 gwtransport.diffusion.infiltration_to_extraction : Quadrature reference; prefer for cout
583 grids coarser than the flow detail.
584 extraction_to_infiltration : Inverse operation.
585 :ref:`concept-dispersion-scales` : Macrodispersion vs microdispersion.
586 """
587 cin, transport = _coerce_and_validate(
588 cin_or_cout=cin,
589 flow=flow,
590 tedges=tedges,
591 cout_tedges=cout_tedges,
592 aquifer_pore_volumes=aquifer_pore_volumes,
593 streamline_length=streamline_length,
594 molecular_diffusivity=molecular_diffusivity,
595 longitudinal_dispersivity=longitudinal_dispersivity,
596 retardation_factor=retardation_factor,
597 is_forward=True,
598 flow_out=flow_out,
599 )
600 band_vals, col_start, valid_cout_bins = _closed_form_coeff_matrix(
601 **transport,
602 retardation_factor=retardation_factor,
603 extend_tedges=_extend_tedges_flag(spinup),
604 )
606 n_cin = len(cin)
607 full_band = band_vals.shape[1]
608 cols = np.clip(col_start[:, None] + np.arange(full_band), 0, n_cin - 1)
609 cout = np.einsum("kb,kb->k", band_vals, cin[cols])
611 # Mark output bins invalid where no input has broken through (spin-up) or the output
612 # bin extends beyond the input data range.
613 total_coeff = band_vals.sum(axis=1)
614 cout[(total_coeff < _EPSILON_COEFF_SUM) | ~valid_cout_bins] = np.nan
615 return cout
618def extraction_to_infiltration(
619 *,
620 cout: npt.ArrayLike,
621 flow: npt.ArrayLike,
622 tedges: pd.DatetimeIndex,
623 cout_tedges: pd.DatetimeIndex,
624 aquifer_pore_volumes: npt.ArrayLike,
625 streamline_length: npt.NDArray[np.floating] | float,
626 molecular_diffusivity: npt.NDArray[np.floating] | float,
627 longitudinal_dispersivity: npt.NDArray[np.floating] | float,
628 retardation_factor: float = 1.0,
629 regularization_strength: float = 1e-10,
630 flow_out: npt.ArrayLike | None = None,
631 spinup: str | None = "constant",
632) -> npt.NDArray[np.floating]:
633 """Reconstruct infiltration concentration from extracted water (deconvolution).
635 Inverts the forward model by building the same closed-form flux-concentration matrix as
636 :func:`infiltration_to_extraction` and solving ``W @ cin = cout`` via Tikhonov
637 regularization. Fast closed-form counterpart of
638 :func:`gwtransport.diffusion.extraction_to_infiltration`.
640 Parameters
641 ----------
642 cout : array-like
643 Concentration of the compound in extracted water. Length ``len(cout_tedges) - 1``.
644 flow : array-like
645 Flow rate of water in the aquifer [m³/day]. Length ``len(tedges) - 1``.
646 tedges : pandas.DatetimeIndex
647 Time edges for cin (output) and flow data. Length ``len(flow) + 1``.
648 cout_tedges : pandas.DatetimeIndex
649 Time edges for cout data bins. Length ``len(cout) + 1``.
650 aquifer_pore_volumes : array-like
651 Aquifer pore volumes [m³] -- one independent streamtube per entry.
652 streamline_length : float or ndarray
653 Travel distance L [m]: a scalar (shared by all streamtubes) or an array with one
654 value per aquifer pore volume. Must be positive.
655 molecular_diffusivity : float or ndarray
656 Effective (retarded-frame) molecular diffusivity D_m [m²/day]: scalar or one value per pore volume.
657 Must be non-negative.
658 longitudinal_dispersivity : float or ndarray
659 Longitudinal dispersivity alpha_L [m] (microdispersion): scalar or one value per pore volume.
660 Must be non-negative.
661 retardation_factor : float, optional
662 Retardation factor (default 1.0).
663 regularization_strength : float, optional
664 Tikhonov regularization parameter (default 1e-10).
665 flow_out : array-like or None, optional
666 Extraction flow rate [m³/day] on the output grid (aligned to ``cout_tedges``).
667 See :func:`infiltration_to_extraction`. Default None.
668 spinup : {"constant"} | None, optional
669 See :func:`infiltration_to_extraction`. Default ``"constant"``.
671 Returns
672 -------
673 numpy.ndarray
674 Bin-averaged concentration in the infiltrating water. Length ``len(tedges) - 1``.
675 NaN where no extraction data constrains the bin.
677 See Also
678 --------
679 infiltration_to_extraction : Forward operation.
680 gwtransport.diffusion.extraction_to_infiltration : Quadrature reference.
681 :ref:`concept-dispersion-scales` : Macrodispersion vs microdispersion.
682 """
683 cout, transport = _coerce_and_validate(
684 cin_or_cout=cout,
685 flow=flow,
686 tedges=tedges,
687 cout_tedges=cout_tedges,
688 aquifer_pore_volumes=aquifer_pore_volumes,
689 streamline_length=streamline_length,
690 molecular_diffusivity=molecular_diffusivity,
691 longitudinal_dispersivity=longitudinal_dispersivity,
692 retardation_factor=retardation_factor,
693 is_forward=False,
694 flow_out=flow_out,
695 )
696 band_vals, col_start, valid_cout_bins = _closed_form_coeff_matrix(
697 **transport,
698 retardation_factor=retardation_factor,
699 extend_tedges=_extend_tedges_flag(spinup),
700 )
702 n_cin = len(transport["flow"])
703 return _solve_reverse_banded(
704 band_vals=band_vals,
705 col_start=col_start,
706 valid_cout_bins=valid_cout_bins,
707 cout=cout,
708 n_cin=n_cin,
709 regularization_strength=regularization_strength,
710 )
713def gamma_infiltration_to_extraction(
714 *,
715 cin: npt.ArrayLike,
716 flow: npt.ArrayLike,
717 tedges: pd.DatetimeIndex,
718 cout_tedges: pd.DatetimeIndex,
719 mean: float | None = None,
720 std: float | None = None,
721 loc: float = 0.0,
722 alpha: float | None = None,
723 beta: float | None = None,
724 n_bins: int = 100,
725 streamline_length: float,
726 molecular_diffusivity: float,
727 longitudinal_dispersivity: float,
728 retardation_factor: float = 1.0,
729 flow_out: npt.ArrayLike | None = None,
730 spinup: str | None = "constant",
731) -> npt.NDArray[np.floating]:
732 """Compute extracted concentration for a gamma-distributed pore volume distribution.
734 Convenience wrapper around :func:`infiltration_to_extraction` that discretizes a
735 (shifted) gamma aquifer pore-volume distribution into ``n_bins`` equal-probability
736 streamtubes. Provide either (mean, std) or (alpha, beta); ``loc`` defaults to 0.
738 Parameters
739 ----------
740 cin : array-like
741 Concentration of the compound in infiltrating water.
742 flow : array-like
743 Flow rate of water in the aquifer [m³/day].
744 tedges : pandas.DatetimeIndex
745 Time edges for cin and flow data. Length ``len(cin) + 1``.
746 cout_tedges : pandas.DatetimeIndex
747 Time edges for output data bins.
748 mean, std : float, optional
749 Mean and standard deviation of the gamma pore-volume distribution [m³].
750 loc : float, optional
751 Location (minimum pore volume) [m³], ``0 <= loc < mean``. Default 0.0.
752 alpha, beta : float, optional
753 Shape and scale parameters of the gamma distribution (alternative to mean/std).
754 n_bins : int, optional
755 Number of equal-probability streamtubes. Default 100.
756 streamline_length : float
757 Travel distance L [m], applied to all gamma streamtubes. Must be positive.
758 molecular_diffusivity : float
759 Effective (retarded-frame) molecular diffusivity D_m [m²/day], applied to all streamtubes. Must be
760 non-negative.
761 longitudinal_dispersivity : float
762 Longitudinal dispersivity alpha_L [m] (microdispersion), applied to all streamtubes. Must be
763 non-negative.
764 retardation_factor : float, optional
765 Retardation factor (default 1.0).
766 flow_out : array-like or None, optional
767 Extraction flow rate [m³/day] on the output grid. See
768 :func:`infiltration_to_extraction`. Default None.
769 spinup : {"constant"} | None, optional
770 See :func:`infiltration_to_extraction`. Default ``"constant"``.
772 Returns
773 -------
774 numpy.ndarray
775 Bin-averaged Kreft-Zuber flux concentration ``C_F`` in the extracted water.
776 Length ``len(cout_tedges) - 1``.
778 See Also
779 --------
780 infiltration_to_extraction : Transport with an explicit pore volume distribution.
781 gamma_extraction_to_infiltration : Reverse operation.
782 gwtransport.gamma.bins : Create gamma distribution bins.
783 :ref:`concept-gamma-distribution` : Two-parameter pore volume model.
784 """
785 bins = gamma.bins(mean=mean, std=std, loc=loc, alpha=alpha, beta=beta, n_bins=n_bins)
786 return infiltration_to_extraction(
787 cin=cin,
788 flow=flow,
789 tedges=tedges,
790 cout_tedges=cout_tedges,
791 aquifer_pore_volumes=bins["expected_values"],
792 streamline_length=streamline_length,
793 molecular_diffusivity=molecular_diffusivity,
794 longitudinal_dispersivity=longitudinal_dispersivity,
795 retardation_factor=retardation_factor,
796 flow_out=flow_out,
797 spinup=spinup,
798 )
801def gamma_extraction_to_infiltration(
802 *,
803 cout: npt.ArrayLike,
804 flow: npt.ArrayLike,
805 tedges: pd.DatetimeIndex,
806 cout_tedges: pd.DatetimeIndex,
807 mean: float | None = None,
808 std: float | None = None,
809 loc: float = 0.0,
810 alpha: float | None = None,
811 beta: float | None = None,
812 n_bins: int = 100,
813 streamline_length: float,
814 molecular_diffusivity: float,
815 longitudinal_dispersivity: float,
816 retardation_factor: float = 1.0,
817 regularization_strength: float = 1e-10,
818 flow_out: npt.ArrayLike | None = None,
819 spinup: str | None = "constant",
820) -> npt.NDArray[np.floating]:
821 """Reconstruct infiltration concentration for a gamma-distributed pore volume distribution.
823 Convenience wrapper around :func:`extraction_to_infiltration` that discretizes a
824 (shifted) gamma aquifer pore-volume distribution into ``n_bins`` equal-probability
825 streamtubes. Provide either (mean, std) or (alpha, beta); ``loc`` defaults to 0.
827 Parameters
828 ----------
829 cout : array-like
830 Concentration of the compound in extracted water.
831 flow : array-like
832 Flow rate of water in the aquifer [m³/day].
833 tedges : pandas.DatetimeIndex
834 Time edges for cin (output) and flow data. Length ``len(flow) + 1``.
835 cout_tedges : pandas.DatetimeIndex
836 Time edges for cout data bins. Length ``len(cout) + 1``.
837 mean, std : float, optional
838 Mean and standard deviation of the gamma pore-volume distribution [m³].
839 loc : float, optional
840 Location (minimum pore volume) [m³], ``0 <= loc < mean``. Default 0.0.
841 alpha, beta : float, optional
842 Shape and scale parameters of the gamma distribution (alternative to mean/std).
843 n_bins : int, optional
844 Number of equal-probability streamtubes. Default 100.
845 streamline_length : float
846 Travel distance L [m], applied to all gamma streamtubes. Must be positive.
847 molecular_diffusivity : float
848 Effective (retarded-frame) molecular diffusivity D_m [m²/day], applied to all streamtubes. Must be
849 non-negative.
850 longitudinal_dispersivity : float
851 Longitudinal dispersivity alpha_L [m] (microdispersion), applied to all streamtubes. Must be
852 non-negative.
853 retardation_factor : float, optional
854 Retardation factor (default 1.0).
855 regularization_strength : float, optional
856 Tikhonov regularization parameter (default 1e-10).
857 flow_out : array-like or None, optional
858 Extraction flow rate [m³/day] on the output grid. See
859 :func:`infiltration_to_extraction`. Default None.
860 spinup : {"constant"} | None, optional
861 See :func:`infiltration_to_extraction`. Default ``"constant"``.
863 Returns
864 -------
865 numpy.ndarray
866 Bin-averaged concentration in the infiltrating water. Length ``len(tedges) - 1``.
868 See Also
869 --------
870 extraction_to_infiltration : Deconvolution with an explicit pore volume distribution.
871 gamma_infiltration_to_extraction : Forward operation.
872 gwtransport.gamma.bins : Create gamma distribution bins.
873 :ref:`concept-gamma-distribution` : Two-parameter pore volume model.
874 """
875 bins = gamma.bins(mean=mean, std=std, loc=loc, alpha=alpha, beta=beta, n_bins=n_bins)
876 return extraction_to_infiltration(
877 cout=cout,
878 flow=flow,
879 tedges=tedges,
880 cout_tedges=cout_tedges,
881 aquifer_pore_volumes=bins["expected_values"],
882 streamline_length=streamline_length,
883 molecular_diffusivity=molecular_diffusivity,
884 longitudinal_dispersivity=longitudinal_dispersivity,
885 retardation_factor=retardation_factor,
886 regularization_strength=regularization_strength,
887 flow_out=flow_out,
888 spinup=spinup,
889 )