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