Coverage for src/gwtransport/diffusion_fast_fast.py: 0%
125 statements
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-04 20:54 +0000
1"""
2Fast *approximate* 1D advection-dispersion transport (Kreft-Zuber flux concentration).
4This module shares the conceptual model of :mod:`gwtransport.diffusion` and
5:mod:`gwtransport.diffusion_fast` -- advection with microdispersion (``alpha_L``) and molecular
6diffusion (``D_m``) along orthogonal (Cartesian) flow paths, one independent streamtube per aquifer
7pore volume, the spread across the pore volume distribution providing macrodispersion, and linear
8sorption via the retardation factor. It targets the bin-averaged Kreft-Zuber (1978) flux
9concentration ``C_F`` on the streamtube bundle, but trades exactness for a single fast (~1.5 ms)
10native-grid evaluation that does not depend on the flow being constant. It is **approximate** --
11see Accuracy below for the error budget and when to reach for :mod:`gwtransport.diffusion_fast`
12instead.
14The three modules form an accuracy/speed ladder: :mod:`gwtransport.diffusion` is the dense reference
15implementation and the ground-truth oracle (composite Gauss-Legendre quadrature, slowest),
16:mod:`gwtransport.diffusion_fast` is its banded closed-form equivalent (agrees to machine precision,
17~80-90x faster), and this module is the approximate rung -- fastest, and exact only up to the error
18budget below.
20How it works -- one skewed breakthrough on the native volume grid
21-----------------------------------------------------------------
23The moving-frame dispersion product ``D_t = D_m*tau + alpha_L*xi`` mixes a *time* term (molecular
24diffusion ``D_m*tau``) and a *volume* term (microdispersion ``alpha_L*xi``). Under constant flow the
25two coincide: elapsed time and breakthrough coordinate are locked, ``tau = r_vpv*xi/(L*Q)``
26(``r_vpv = R*V_pore``), so ``D_m*tau = (D_m*r_vpv/(L*Q))*xi`` -- the same ``xi``-proportional form as
27``alpha_L*xi``. Molecular diffusion is therefore an **effective dispersivity**
28``alpha_eff = alpha_L + D_m*r_vpv/(L*q_mean)`` per streamtube (``q_mean`` the record-mean flow), and
29the whole method is a single skewed ``D_t = alpha_eff*xi`` Kreft-Zuber breakthrough applied banded on
30the **native cumulative-volume grid**: the whole aquifer pore volume distribution (APVD) is pre-summed
31into one 1D antiderivative ``Ibar(dV)`` -- exact for any APVD shape -- finely sampled once and read
32back by interpolation.
34``tedges`` need **not** be regularly spaced and ``cout_tedges`` need not equal ``tedges`` (supply
35``flow_out`` when they differ): the build runs on the native cumulative-volume grid for any spacing.
36The **only** approximations are the ``Ibar`` interpolation (~1e-4, below) and -- under *variable* flow
37-- freezing ``q_mean`` at the record mean: the ``tau = r_vpv*xi/(L*Q)`` map holds exactly only at
38constant flow, so the microdispersion part stays exact but the molecular part picks up a commutator
39residual.
41Accuracy (vs :mod:`gwtransport.diffusion_fast`)
42-----------------------------------------------
44- **Constant flow:** exact to the ``Ibar`` interpolation floor -- ~1e-6 for smooth inputs, degrading
45 to ~1e-4 for sharp inputs at ``alpha_L = 0`` (the kink limit; ~1e-3 for a single large pore volume).
46 This holds across ``D_m`` and ``R`` at realistic Peclet numbers ``Pe = L/alpha_eff >> 1``,
47 **including heat transport** (``R > 1``, large ``D_m``): the effective-dispersivity fold reproduces
48 the skewed molecular breakthrough exactly, not merely its second moment. Only in the extreme
49 low-Peclet corner (``alpha_eff`` approaching ``L`` -- a short streamline with very strong molecular
50 diffusion, ``Pe <~ 2``) does the fold's wide dispersive band exceed the fine-grid sample cap
51 (``_MAX_KERNEL_SAMPLES``), coarsening the breakthrough shape; use :mod:`gwtransport.diffusion_fast`
52 there.
53- **Variable flow:** the microdispersion part stays exact; the molecular part carries the frozen-
54 ``q_mean`` commutator residual. It is small when molecular diffusion is sub-dominant -- the typical
55 ``alpha_L > 0`` regime, <~1e-3 for realistic solute ``D_m`` even through strong flow -- and grows in
56 the molecular-diffusion-dominated corner (``alpha_L`` ~ 0) with sharp inputs under strongly variable
57 or seasonal flow, reaching the multi-percent range where no fast approximation is reliable.
58 **Use** :mod:`gwtransport.diffusion_fast` **there** (in particular for heat transport under strongly
59 variable flow).
61The inverse (:func:`extraction_to_infiltration`) deconvolves the *same* approximate operator ``W`` the
62forward applies. It assembles ``W`` directly in banded form (one ``Ibar`` gather -- no per-pore-volume
63closed-form loop, no dense ``(n_cout, n_cin)`` matrix) and solves it with banded Tikhonov
64regularisation (banded Cholesky, ``O(n * band**2)``), so it is much faster than
65:mod:`gwtransport.diffusion_fast`'s reverse, especially for many streamtubes. Inverting exactly the
66forward operator makes a round trip self-consistent (recovering the input up to the deconvolution
67conditioning); on *real* extraction data the reverse carries the forward's error budget above,
68amplified by the deconvolution conditioning.
70Available functions:
72- :func:`infiltration_to_extraction` - Forward transport from an explicit ``aquifer_pore_volumes``
73 distribution: returns the approximate bin-averaged Kreft-Zuber flux concentration on ``cout_tedges``,
74 from one banded breakthrough on the native cumulative-volume grid with molecular diffusion folded
75 into the effective dispersivity ``alpha_eff = alpha_L + D_m * R * V_pore / (L * q_mean)`` per
76 streamtube. ``streamline_length``, ``molecular_diffusivity`` and ``longitudinal_dispersivity`` are
77 either scalars shared by all streamtubes or one value per pore volume. ``flow_out`` (extraction flow
78 on the ``cout_tedges`` grid) is required whenever ``cout_tedges`` differs from ``tedges``. Output
79 bins without complete breakthrough information are NaN.
81- :func:`extraction_to_infiltration` - Reverse direction: assembles the same approximate banded
82 operator and deconvolves it with banded Tikhonov regularization (banded Cholesky on the normal
83 equations, ``O(n * band**2)``), returning the bin-averaged infiltration concentration on ``tedges``.
84 NaN entries in ``cout`` mark measurement gaps and are excluded from the solve; spin-up and
85 unconstrained cin bins come back NaN.
87- :func:`gamma_infiltration_to_extraction` - :func:`infiltration_to_extraction` with the pore volume
88 distribution given as a (shifted) gamma -- either (mean, std) or (alpha, beta), plus ``loc`` --
89 discretized into ``n_bins`` equal-probability streamtubes that share one ``streamline_length`` and
90 one pair of dispersion parameters.
92- :func:`gamma_extraction_to_infiltration` - :func:`extraction_to_infiltration` with the same gamma
93 parameterization of the pore volume distribution: reconstructs ``cin`` on ``tedges`` from ``cout``.
95References
96----------
97Kreft, A., & Zuber, A. (1978). On the physical meaning of the dispersion equation and its
98solutions for different initial and boundary conditions. Chemical Engineering Science,
9933(11), 1471-1480.
101This file is part of gwtransport which is released under AGPL-3.0 license.
102See the ./LICENSE file or go to https://github.com/gwtransport/gwtransport/blob/main/LICENSE for full license details.
103"""
105import numpy as np
106import numpy.typing as npt
107import pandas as pd
108from scipy.sparse import coo_array
110from gwtransport import gamma
111from gwtransport._diffusion_shared import (
112 _DT_FLOOR,
113 _EPSILON_COEFF_SUM,
114 _advective_valid_cout_bins,
115 _breakthrough_antideriv,
116 _coerce_and_validate,
117 _cout_cumulative_volume,
118 _extend_tedges,
119 _extend_tedges_flag,
120 _solve_reverse_banded,
121)
122from gwtransport._time import dt_to_days, tedges_to_days
123from gwtransport.diffusion_fast import _DEFAULT_SATURATION_THRESHOLD
124from gwtransport.utils import cumulative_flow_volume
126# Samples per native bin used to discretise the 1D breakthrough antiderivative ``Ibar``. Higher =
127# more accurate ``Ibar`` interpolation (breakthrough error ~ O(1/_KERNEL_FINE^2), or O(1/_KERNEL_FINE)
128# at the alpha_eff=0 kink) at the cost of a larger one-time precompute; 16 gives ~1e-4 -- the
129# constant-flow accuracy floor of the method -- so it is not exposed as a user knob.
130_KERNEL_FINE = 16
132# Upper bound on the number of fine ``Ibar`` samples. Caps the one-time precompute when the
133# breakthrough band is very wide -- either tiny flow (enormous front offset; those bins are masked by
134# residence time anyway, so coarsening is benign) or the extreme low-Peclet corner where molecular
135# diffusion folds into a large ``alpha_eff`` (a valid, unmasked band). In the latter the cap coarsens
136# ``dv_fine`` below one sample per output bin and the breakthrough shape loses accuracy -- use
137# :mod:`gwtransport.diffusion_fast` for a short streamline with very strong molecular diffusion
138# (Pe = L/alpha_eff <~ 2). No realistic groundwater/heat regime reaches that corner.
139_MAX_KERNEL_SAMPLES = 20000
142def _summed_antideriv(
143 *,
144 aquifer_pore_volumes: npt.NDArray[np.floating],
145 streamline_length: npt.NDArray[np.floating],
146 molecular_diffusivity: npt.NDArray[np.floating],
147 longitudinal_dispersivity: npt.NDArray[np.floating],
148 retardation_factor: float,
149 q_mean: float,
150 mean_bin_volume: float,
151) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating], float, float, float]:
152 r"""Precompute the APVD-summed breakthrough antiderivative ``Ibar(dV)`` on a fine 1D grid.
154 For one streamtube the bin-averaged flux fraction over a cout bin equals the second difference of
155 the antiderivative ``I(x)`` (:func:`gwtransport._diffusion_shared._breakthrough_antideriv`) of the
156 resident concentration in the breakthrough coordinate ``x = (dV - r_vpv)*L/r_vpv``
157 (``r_vpv = R*V_pore``). The antiderivative *with respect to cumulative volume* ``dV`` is
158 ``(r_vpv/L)*I(x(dV))``; averaging it over the APVD gives a single 1D function
160 .. math::
162 \bar I(\Delta V) = \operatorname{mean}_{pv}\Bigl[\tfrac{r_{vpv}}{L}\,
163 I\bigl((\Delta V - r_{vpv})L/r_{vpv}\bigr)\Bigr],\quad D_t = \alpha_{eff}\,\xi,
165 whose edge-differences reproduce the per-streamtube-averaged ``C_F`` **exactly** for any APVD
166 (the ``r_vpv/L`` Jacobian and the cout-bin volume normalisation cancel the per-streamtube ``dx``).
167 Only the interpolation of ``Ibar`` is approximate.
169 Molecular diffusion is folded into ``D_t`` as an **effective dispersivity**. The moving-frame
170 variance is ``D_t = D_m*tau + alpha_L*xi``; under constant flow ``tau = r_vpv*xi/(L*Q)``, so
171 ``D_m*tau = (D_m*r_vpv/(L*Q))*xi`` -- the same ``xi``-proportional form as ``alpha_L*xi``. Hence
172 ``D_t = alpha_eff*xi`` with ``alpha_eff = alpha_L + D_m*r_vpv/(L*q_mean)`` per streamtube, which
173 reproduces the skewed Kreft-Zuber molecular breakthrough exactly at constant flow (``q_mean`` is
174 the record-mean flow; freezing it is the only approximation, and only under variable flow).
176 The grid is **uniform** over the breakthrough band ``[off_lo, off_hi]`` (front ``r_vpv`` plus the
177 conservative ``alpha_eff`` dispersion smear, unioned over streamtubes) plus a margin, sampled at
178 ``mean_bin_volume / _KERNEL_FINE`` (uniformity lets :func:`_eval_antideriv` interpolate by
179 fractional indexing instead of a per-point search). For ``alpha_eff > 0`` ``Ibar`` is smooth and
180 the interpolation error is ``O(1/_KERNEL_FINE^2)``; only at ``alpha_L = 0`` *and* ``D_m = 0`` does
181 it have a kink at each ``r_vpv`` (error ``O(1/_KERNEL_FINE)``, sub-1e-3).
183 Returns
184 -------
185 grid : ndarray
186 Cumulative-volume offsets ``dV`` (uniformly spaced, strictly increasing).
187 ibar : ndarray
188 ``Ibar`` sampled at ``grid``.
189 mean_r_vpv : float
190 ``R * mean(V_pore)`` -- the saturated offset (``Ibar -> dV - mean_r_vpv`` above the band).
191 off_lo, off_hi : float
192 Lower / upper cumulative-volume offset of the breakthrough band.
193 """
194 r_vpv = retardation_factor * aquifer_pore_volumes
195 mean_r_vpv = float(r_vpv.mean())
196 u = _DEFAULT_SATURATION_THRESHOLD
197 # Molecular diffusion as an effective dispersivity (see docstring): D_t = alpha_eff*xi. q_mean > 0
198 # is guaranteed -- _build_forward_operator returns early when total_volume <= 0, and
199 # tedges_days[-1] - tedges_days[0] > 0 for any valid strictly-increasing tedges.
200 alpha_eff = longitudinal_dispersivity + molecular_diffusivity * r_vpv / (streamline_length * q_mean)
201 # Dispersion half-widths in the breakthrough coordinate x (front D_t = alpha_eff*L): the pre side
202 # shrinks (|x| = U*2*sqrt(alpha_eff*L)); the post side grows with slope alpha_eff, giving the
203 # quadratic root x = 2U^2*alpha_eff + 2U*sqrt(U^2*alpha_eff^2 + alpha_eff*L). Mapped to volume
204 # offsets via r_vpv/L and unioned over streamtubes (conservative, never under-covers the band).
205 pre_x = u * 2.0 * np.sqrt(alpha_eff * streamline_length)
206 post_x = 2.0 * u * u * alpha_eff + 2.0 * u * np.sqrt(u * u * alpha_eff**2 + alpha_eff * streamline_length)
207 off_lo = float(np.min(r_vpv - (r_vpv / streamline_length) * pre_x))
208 off_hi = float(np.max(r_vpv + (r_vpv / streamline_length) * post_x))
210 margin = 4.0 * mean_bin_volume
211 span = (off_hi - off_lo) + 2.0 * margin
212 dv_fine = mean_bin_volume / _KERNEL_FINE
213 if span / dv_fine > _MAX_KERNEL_SAMPLES:
214 dv_fine = span / _MAX_KERNEL_SAMPLES
215 grid = np.arange(off_lo - margin, off_hi + margin + dv_fine, dv_fine)
217 x = (grid[None, :] - r_vpv[:, None]) * streamline_length[:, None] / r_vpv[:, None]
218 dt_var = np.maximum(alpha_eff[:, None] * np.maximum(x + streamline_length[:, None], 0.0), _DT_FLOOR)
219 antideriv = _breakthrough_antideriv(x, dt_var)
220 ibar = ((r_vpv[:, None] / streamline_length[:, None]) * antideriv).mean(axis=0)
221 return grid, ibar, mean_r_vpv, off_lo, off_hi
224def _eval_antideriv(
225 dv: npt.NDArray[np.floating],
226 grid: npt.NDArray[np.floating],
227 ibar: npt.NDArray[np.floating],
228 mean_r_vpv: float,
229) -> npt.NDArray[np.floating]:
230 """Evaluate ``Ibar`` at arbitrary cumulative-volume offsets.
232 Linear interpolation on the *uniform* precomputed grid by fractional indexing (no per-point
233 search -- the gather is evaluated at ``O(N*band)`` points, so this dominates the runtime).
234 Below the grid ``Ibar = 0`` (not broken through); above it ``Ibar = dV - mean_r_vpv`` (saturated,
235 the exact linear asymptote).
237 Returns
238 -------
239 ndarray
240 ``Ibar(dv)`` with the same shape as ``dv``.
241 """
242 g0 = grid[0]
243 dstep = grid[1] - grid[0]
244 f = (dv - g0) / dstep
245 # astype truncates toward zero, which equals floor on f >= 0 -- the only regime that survives
246 # the dv < g0 override below and the lower clip. Precompute the segment slopes once so the
247 # O(N*band) gather reads a single array instead of differencing two.
248 slopes = np.diff(ibar)
249 i0 = np.clip(f.astype(np.intp), 0, grid.size - 2)
250 out = ibar[i0] + (f - i0) * slopes[i0]
251 out = np.where(dv < g0, 0.0, out)
252 return np.where(dv > grid[-1], dv - mean_r_vpv, out)
255def _breakthrough_band(
256 *,
257 cumulative_volume_at_cin: npt.NDArray[np.floating],
258 cumulative_volume_at_cout: npt.NDArray[np.floating],
259 grid: npt.NDArray[np.floating],
260 ibar: npt.NDArray[np.floating],
261 mean_r_vpv: float,
262 off_lo: float,
263 off_hi: float,
264 extend: bool,
265) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.intp]]:
266 """Banded breakthrough coefficients (advection + macro + micro + molecular) on the volume grid.
268 For each cout bin, gathers the band of cin edges whose breakthrough offset falls in
269 ``[off_lo, off_hi]`` (a conservative fixed window from ``searchsorted``, mirroring
270 :func:`gwtransport.diffusion_fast._closed_form_coeff_matrix`), evaluates ``Ibar`` at the native
271 edge offsets, and telescopes the edge-differences into per-cin-bin coefficients. With
272 ``extend`` the cin axis is extended by one wide virtual bin each side carrying the constant
273 boundary value (``cin[0]`` / ``cin[-1]``), reproducing the 100-year warm-start.
275 The coefficients depend only on the volume grid (not on ``cin``), so this band is built once and
276 shared: :func:`infiltration_to_extraction` applies it to the (extended) ``cin``, and
277 :func:`extraction_to_infiltration` folds it onto the real cin axis to assemble the banded
278 operator it deconvolves -- guaranteeing forward and reverse use the same operator.
280 Returns
281 -------
282 coeff : ndarray, shape (n_cout, band)
283 Per-(cout bin, band slot) coefficient.
284 cin_bin : ndarray of int, shape (n_cout, band)
285 Column index of each coefficient on the (warm-start-extended when ``extend``) cin axis:
286 ``cin_ext[cin_bin]`` for the forward, ``clip(cin_bin - int(extend), 0, n_cin - 1)`` for the
287 real cin axis in the reverse.
288 """
289 vi = cumulative_volume_at_cin
290 vc = cumulative_volume_at_cout
291 if extend:
292 big = abs(off_hi) + abs(off_lo) + (vc[-1] - vc[0]) + (vi[-1] - vi[0]) + 1.0
293 vi_ext = np.concatenate([[vi[0] - big], vi, [vi[-1] + big]])
294 n_cin_ext = vi.size + 1 # (vi.size - 1) real bins + 2 virtual boundary bins
295 else:
296 vi_ext = vi
297 n_cin_ext = vi.size - 1
299 vc_lo, vc_hi = vc[:-1], vc[1:]
300 w = vc_hi - vc_lo
301 vc_mid = 0.5 * (vc_lo + vc_hi)
302 # Conservative two-sided fixed window: center on the front, widen to the farthest cin edge whose
303 # offset is still inside the band on either side. The +1 absorbs the edge consumed by the
304 # telescoping difference and searchsorted rounding (an under-sized window silently drops mass).
305 center = np.searchsorted(vi_ext, vc_mid - mean_r_vpv)
306 j_lo = np.searchsorted(vi_ext, vc_lo - off_hi, side="left")
307 j_hi = np.searchsorted(vi_ext, vc_hi - off_lo, side="right")
308 hw_lo = int(np.max(center - j_lo)) + 1
309 hw_hi = int(np.max(j_hi - center)) + 1
311 cols = center[:, None] + np.arange(-hw_lo, hw_hi + 1)[None, :]
312 vi_band = vi_ext[np.clip(cols, 0, len(vi_ext) - 1)]
313 ibar_hi = _eval_antideriv(vc_hi[:, None] - vi_band, grid, ibar, mean_r_vpv)
314 ibar_lo = _eval_antideriv(vc_lo[:, None] - vi_band, grid, ibar, mean_r_vpv)
315 with np.errstate(divide="ignore", invalid="ignore"):
316 frac_edge = np.where(w[:, None] > 0.0, (ibar_hi - ibar_lo) / w[:, None], 0.0)
317 coeff = frac_edge[:, :-1] - frac_edge[:, 1:]
318 cin_bin = np.clip(cols[:, :-1], 0, n_cin_ext - 1)
319 return coeff, cin_bin
322def _build_forward_operator(
323 *,
324 flow: npt.NDArray[np.floating],
325 tedges: pd.DatetimeIndex,
326 cout_tedges: pd.DatetimeIndex,
327 flow_out: npt.NDArray[np.floating] | None,
328 aquifer_pore_volumes: npt.NDArray[np.floating],
329 streamline_length: npt.NDArray[np.floating],
330 molecular_diffusivity: npt.NDArray[np.floating],
331 longitudinal_dispersivity: npt.NDArray[np.floating],
332 retardation_factor: float,
333 extend: bool,
334) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.intp], npt.NDArray[np.bool_]] | None:
335 r"""Build the cin-independent pieces of the approximate banded forward operator ``W``.
337 Both directions share this build: the banded breakthrough operator ``W`` (``coeff`` / ``cin_bin``,
338 :func:`_breakthrough_band` -- advection + macro + micro + molecular, with molecular diffusion
339 folded in as an effective dispersivity) and the residence-time ``valid`` mask. Because the band
340 depends only on the volume grid (not on ``cin`` / ``cout``), forward transport and reverse
341 deconvolution operate on exactly the same operator.
343 Returns
344 -------
345 coeff : ndarray, shape (n_cout, band)
346 Banded ``W`` coefficients.
347 cin_bin : ndarray of int, shape (n_cout, band)
348 Column indices of ``coeff`` on the warm-start-extended cin axis.
349 valid : ndarray of bool, shape (n_cout,)
350 Output bins with residence time finite at both edges (complete breakthrough).
352 None
353 Returned instead of the tuple when there is no through-flow (nothing breaks through).
354 """
355 tedges_days = tedges_to_days(tedges)
356 cout_tedges_days = tedges_to_days(cout_tedges, ref=tedges[0])
357 cumulative_volume_at_cin = cumulative_flow_volume(flow, dt_to_days(tedges))
358 total_volume = float(cumulative_volume_at_cin[-1])
359 if total_volume <= 0.0:
360 return None
362 cumulative_volume_at_cout = _cout_cumulative_volume(
363 flow_out=flow_out,
364 cout_tedges=cout_tedges,
365 cout_tedges_days=cout_tedges_days,
366 tedges_days=tedges_days,
367 cumulative_volume_at_cin=cumulative_volume_at_cin,
368 )
370 mean_bin_volume = total_volume / len(flow)
371 q_mean = total_volume / float(tedges_days[-1] - tedges_days[0])
372 grid, ibar, mean_r_vpv, off_lo, off_hi = _summed_antideriv(
373 aquifer_pore_volumes=aquifer_pore_volumes,
374 streamline_length=streamline_length,
375 molecular_diffusivity=molecular_diffusivity,
376 longitudinal_dispersivity=longitudinal_dispersivity,
377 retardation_factor=retardation_factor,
378 q_mean=q_mean,
379 mean_bin_volume=mean_bin_volume,
380 )
381 coeff, cin_bin = _breakthrough_band(
382 cumulative_volume_at_cin=cumulative_volume_at_cin,
383 cumulative_volume_at_cout=cumulative_volume_at_cout,
384 grid=grid,
385 ibar=ibar,
386 mean_r_vpv=mean_r_vpv,
387 off_lo=off_lo,
388 off_hi=off_hi,
389 extend=extend,
390 )
392 # Mask bins beyond the data range (and, without warm-start, incompletely-broken-through spin-up
393 # bins). The look-back runs on the extended grid when warm-starting so spin-up bins stay valid.
394 valid = _advective_valid_cout_bins(
395 flow=flow,
396 tedges=_extend_tedges(tedges) if extend else tedges,
397 cout_tedges=cout_tedges,
398 aquifer_pore_volumes=aquifer_pore_volumes,
399 retardation_factor=retardation_factor,
400 )
401 return coeff, cin_bin, valid
404def _banded_forward_matrix(
405 *,
406 coeff: npt.NDArray[np.floating],
407 cin_bin: npt.NDArray[np.intp],
408 extend: bool,
409 n_cin: int,
410) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.intp]]:
411 """Assemble the banded forward operator ``W`` as a per-row contiguous band for ``_solve_reverse_banded``.
413 The breakthrough coefficients are scattered from the native band onto the real cin axis, folding
414 the warm-start virtual columns into the boundary bins (``clip(cin_bin - int(extend), 0, n_cin - 1)``,
415 so ``W @ cin`` equals the forward's ``coeff @ cin_ext`` exactly). The returned band carries the
416 forward operator verbatim; ``_solve_reverse_banded`` masks the spin-up rows/columns and normalizes,
417 so a forward-then-inverse round trip is self-consistent.
419 Returns
420 -------
421 band_vals : ndarray, shape (n_cout, full_band)
422 Forward weights in banded layout (explicit zeros outside each row's support).
423 col_start : ndarray of int, shape (n_cout,)
424 First real-cin column of each row's band.
425 """
426 n_cout = coeff.shape[0]
427 rows = np.broadcast_to(np.arange(n_cout)[:, None], coeff.shape)
428 real_col = np.clip(cin_bin - int(extend), 0, n_cin - 1)
429 # COO -> CSR sums the warm-start virtual columns folded onto the boundary real columns.
430 w_mat = coo_array((coeff.ravel(), (rows.ravel(), real_col.ravel())), shape=(n_cout, n_cin)).tocsr()
432 # CSR -> contiguous per-row band. Each row spans one contiguous cin band (the fold only saturates
433 # the ends onto the boundary columns), so the banded layout carries no spurious interior gaps.
434 w_mat.sort_indices()
435 indptr, indices, data = w_mat.indptr, w_mat.indices, w_mat.data
436 row_counts = np.diff(indptr)
437 nonempty = row_counts > 0
438 col_start = np.zeros(n_cout, dtype=np.intp)
439 last_col = np.zeros(n_cout, dtype=np.intp)
440 col_start[nonempty] = indices[indptr[:-1][nonempty]]
441 last_col[nonempty] = indices[indptr[1:][nonempty] - 1]
442 full_band = int((last_col[nonempty] - col_start[nonempty] + 1).max()) if nonempty.any() else 1
444 band_vals = np.zeros((n_cout, full_band))
445 rows_of_nz = np.repeat(np.arange(n_cout), row_counts)
446 band_vals[rows_of_nz, indices - col_start[rows_of_nz]] = data
447 return band_vals, col_start
450def infiltration_to_extraction(
451 *,
452 cin: npt.ArrayLike,
453 flow: npt.ArrayLike,
454 tedges: pd.DatetimeIndex,
455 cout_tedges: pd.DatetimeIndex,
456 aquifer_pore_volumes: npt.ArrayLike,
457 streamline_length: npt.NDArray[np.floating] | float,
458 molecular_diffusivity: npt.NDArray[np.floating] | float,
459 longitudinal_dispersivity: npt.NDArray[np.floating] | float,
460 retardation_factor: float = 1.0,
461 flow_out: npt.ArrayLike | None = None,
462 spinup: str | None = "constant",
463) -> npt.NDArray[np.floating]:
464 """Compute extracted concentration with advection, microdispersion, and molecular diffusion (approximate).
466 Fast *approximate* counterpart of :func:`gwtransport.diffusion_fast.infiltration_to_extraction`.
467 Advection + macrodispersion + microdispersion (``alpha_L``) and molecular diffusion (``D_m``,
468 folded in as an effective dispersivity ``alpha_L + D_m*r_vpv/(L*q_mean)`` per streamtube) form a
469 single exact skewed Kreft-Zuber breakthrough on the native cumulative-volume grid. At **constant
470 flow** the result reproduces :func:`gwtransport.diffusion_fast.infiltration_to_extraction` to the
471 ``Ibar`` interpolation floor (~1e-6 smooth, ~1e-4 sharp) at realistic Peclet numbers
472 (``Pe = L/alpha_eff >> 1``), including heat (``R > 1``, large ``D_m``); only the extreme low-Peclet
473 corner (``Pe <~ 2``) loses breakthrough-shape accuracy to the fine-grid sample cap. Under
474 **variable flow** the molecular part carries a commutator residual
475 from the frozen ``q_mean``: small when molecular diffusion is sub-dominant (``alpha_L > 0``, <~1e-3
476 for realistic solute ``D_m``), growing to the multi-percent range for sharp inputs in the
477 molecular-diffusion-dominated corner (``alpha_L`` ~ 0) under strongly variable flow. For machine
478 precision -- or that corner -- use :mod:`gwtransport.diffusion_fast`.
480 Parameters
481 ----------
482 cin : array-like
483 Concentration of the compound in the infiltrating water. Length ``len(tedges) - 1``.
484 flow : array-like
485 Flow rate of water in the aquifer [m³/day]. Length ``len(tedges) - 1``.
486 tedges : pandas.DatetimeIndex
487 Time edges for cin and flow data. Length ``len(cin) + 1``.
488 cout_tedges : pandas.DatetimeIndex
489 Time edges for output data bins. Length ``len(output) + 1``.
490 aquifer_pore_volumes : array-like
491 Aquifer pore volumes [m³] -- one independent streamtube per entry. Any distribution shape
492 (the APVD is pre-summed exactly).
493 streamline_length : float or ndarray
494 Travel distance L [m]: a scalar (shared by all streamtubes) or an array with one
495 value per aquifer pore volume. Must be positive.
496 molecular_diffusivity : float or ndarray
497 Effective molecular diffusivity D_m [m²/day]: scalar or one value per pore volume.
498 Must be non-negative.
499 longitudinal_dispersivity : float or ndarray
500 Longitudinal dispersivity alpha_L [m] (microdispersion): scalar or one value per pore volume.
501 Must be non-negative.
502 retardation_factor : float, optional
503 Retardation factor (default 1.0). Values > 1.0 indicate slower transport.
504 flow_out : array-like or None, optional
505 Extraction flow rate [m³/day] on the output grid (aligned to ``cout_tedges``,
506 length ``len(cout_tedges) - 1``). Required when ``cout_tedges`` differs from ``tedges``;
507 may be omitted only when ``cout_tedges`` equals ``tedges``. Default None.
508 spinup : {"constant"} | None, optional
509 ``"constant"`` (default) extends ``tedges`` by 100 years on each side so a constant
510 warm-start fills the left-edge spin-up region; ``None`` leaves spin-up cout as NaN.
512 Returns
513 -------
514 numpy.ndarray
515 Bin-averaged Kreft-Zuber flux concentration ``C_F`` in the extracted water. Length
516 ``len(cout_tedges) - 1``. NaN where no infiltration data has broken through.
518 See Also
519 --------
520 gwtransport.diffusion_fast.infiltration_to_extraction : Exact (machine-precision) counterpart;
521 use it when approximation is unacceptable, especially in the molecular-dominant regime.
522 gwtransport.diffusion.infiltration_to_extraction : Quadrature reference.
523 extraction_to_infiltration : Inverse operation (deconvolves this same operator).
524 :ref:`concept-dispersion-scales` : Macrodispersion vs microdispersion.
525 """
526 cin, transport = _coerce_and_validate(
527 cin_or_cout=cin,
528 flow=flow,
529 tedges=tedges,
530 cout_tedges=cout_tedges,
531 aquifer_pore_volumes=aquifer_pore_volumes,
532 streamline_length=streamline_length,
533 molecular_diffusivity=molecular_diffusivity,
534 longitudinal_dispersivity=longitudinal_dispersivity,
535 retardation_factor=retardation_factor,
536 is_forward=True,
537 flow_out=flow_out,
538 )
539 extend = _extend_tedges_flag(spinup)
540 operator = _build_forward_operator(**transport, retardation_factor=retardation_factor, extend=extend)
541 if operator is None:
542 # No through-flow: nothing breaks through (matches diffusion_fast's all-NaN result).
543 return np.full(len(transport["cout_tedges"]) - 1, np.nan)
544 coeff, cin_bin, valid = operator
546 # Apply the banded breakthrough operator to the (warm-start-extended) cin. Output bins with no
547 # through-flow carry coeff.sum() ~= 0, so the support mask NaNs them out; the band adds zero volume
548 # across a zero-flow gap, so a constant cin stays constant across it (no smear to leak).
549 cin_ext = np.concatenate([[cin[0]], cin, [cin[-1]]]) if extend else cin
550 cout = np.einsum("kb,kb->k", coeff, cin_ext[cin_bin])
551 support = coeff.sum(axis=1) >= _EPSILON_COEFF_SUM
552 return np.where(support & valid, cout, np.nan)
555def extraction_to_infiltration(
556 *,
557 cout: npt.ArrayLike,
558 flow: npt.ArrayLike,
559 tedges: pd.DatetimeIndex,
560 cout_tedges: pd.DatetimeIndex,
561 aquifer_pore_volumes: npt.ArrayLike,
562 streamline_length: npt.NDArray[np.floating] | float,
563 molecular_diffusivity: npt.NDArray[np.floating] | float,
564 longitudinal_dispersivity: npt.NDArray[np.floating] | float,
565 retardation_factor: float = 1.0,
566 regularization_strength: float = 1e-10,
567 flow_out: npt.ArrayLike | None = None,
568 spinup: str | None = "constant",
569) -> npt.NDArray[np.floating]:
570 """Reconstruct infiltration concentration from extracted water (fast approximate deconvolution).
572 Inverts the **same** approximate operator ``W`` the forward applies: it assembles the banded
573 breakthrough operator ``W`` (advection + macro + micro + molecular, molecular folded in as an
574 effective dispersivity) directly in banded form and deconvolves it with banded Tikhonov
575 regularization (``_solve_reverse_banded`` -- banded Cholesky on the normal equations,
576 ``O(n * band**2)``). It builds ``W`` from one ``Ibar`` gather -- no per-pore-volume closed-form
577 loop and no dense ``(n_cout, n_cin)`` matrix -- so it is much faster than
578 :func:`gwtransport.diffusion_fast.extraction_to_infiltration` (which evaluates the exact
579 breakthrough per streamtube), especially for many streamtubes. Because the deconvolved operator is
580 exactly the forward operator, a forward-then-inverse round trip recovers ``cin`` up to the
581 deconvolution conditioning and regularization. On real extraction data the reverse is as accurate
582 as the forward: at constant flow it matches
583 :func:`gwtransport.diffusion_fast.extraction_to_infiltration` to ~1e-6, while the
584 molecular-diffusion-dominated corner under strongly variable flow amplifies the forward's
585 commutator residual (use :mod:`gwtransport.diffusion_fast` there).
587 Parameters
588 ----------
589 cout : array-like
590 Concentration of the compound in extracted water. Length ``len(cout_tedges) - 1``.
591 flow : array-like
592 Flow rate of water in the aquifer [m³/day]. Length ``len(tedges) - 1``.
593 tedges : pandas.DatetimeIndex
594 Time edges for cin (output) and flow data. Length ``len(flow) + 1``.
595 cout_tedges : pandas.DatetimeIndex
596 Time edges for cout data bins. Length ``len(cout) + 1``.
597 aquifer_pore_volumes : array-like
598 Aquifer pore volumes [m³] -- one independent streamtube per entry.
599 streamline_length : float or ndarray
600 Travel distance L [m]: scalar or one value per pore volume. Must be positive.
601 molecular_diffusivity : float or ndarray
602 Effective molecular diffusivity D_m [m²/day]: scalar or one value per pore volume.
603 Must be non-negative.
604 longitudinal_dispersivity : float or ndarray
605 Longitudinal dispersivity alpha_L [m] (microdispersion): scalar or one value per pore volume.
606 Must be non-negative.
607 retardation_factor : float, optional
608 Retardation factor (default 1.0).
609 regularization_strength : float, optional
610 Tikhonov regularization parameter (default 1e-10). Must be strictly positive: the banded
611 solver relies on it to make the normal equations positive-definite (it cannot return the
612 dense ``lambda = 0`` minimum-norm solution).
613 flow_out : array-like or None, optional
614 Extraction flow rate [m³/day] on the output grid (aligned to ``cout_tedges``). See
615 :func:`infiltration_to_extraction`. Default None.
616 spinup : {"constant"} | None, optional
617 See :func:`infiltration_to_extraction`. Default ``"constant"``.
619 Returns
620 -------
621 numpy.ndarray
622 Bin-averaged concentration in the infiltrating water. Length ``len(tedges) - 1``.
623 NaN where no extraction data constrains the bin.
625 See Also
626 --------
627 infiltration_to_extraction : Forward operation (the operator inverted here).
628 gwtransport.diffusion_fast.extraction_to_infiltration : Exact (machine-precision) counterpart;
629 use it when the approximation is unacceptable.
630 :ref:`concept-dispersion-scales` : Macrodispersion vs microdispersion.
631 """
632 cout, transport = _coerce_and_validate(
633 cin_or_cout=cout,
634 flow=flow,
635 tedges=tedges,
636 cout_tedges=cout_tedges,
637 aquifer_pore_volumes=aquifer_pore_volumes,
638 streamline_length=streamline_length,
639 molecular_diffusivity=molecular_diffusivity,
640 longitudinal_dispersivity=longitudinal_dispersivity,
641 retardation_factor=retardation_factor,
642 is_forward=False,
643 flow_out=flow_out,
644 )
645 n_cin = len(transport["flow"])
646 extend = _extend_tedges_flag(spinup)
647 operator = _build_forward_operator(**transport, retardation_factor=retardation_factor, extend=extend)
648 if operator is None:
649 # No through-flow: nothing constrains the infiltration signal.
650 return np.full(n_cin, np.nan)
651 coeff, cin_bin, valid = operator
653 band_vals, col_start = _banded_forward_matrix(coeff=coeff, cin_bin=cin_bin, extend=extend, n_cin=n_cin)
654 return _solve_reverse_banded(
655 band_vals=band_vals,
656 col_start=col_start,
657 valid_cout_bins=valid,
658 cout=cout,
659 n_cin=n_cin,
660 regularization_strength=regularization_strength,
661 )
664def gamma_infiltration_to_extraction(
665 *,
666 cin: npt.ArrayLike,
667 flow: npt.ArrayLike,
668 tedges: pd.DatetimeIndex,
669 cout_tedges: pd.DatetimeIndex,
670 mean: float | None = None,
671 std: float | None = None,
672 loc: float = 0.0,
673 alpha: float | None = None,
674 beta: float | None = None,
675 n_bins: int = 100,
676 streamline_length: float,
677 molecular_diffusivity: float,
678 longitudinal_dispersivity: float,
679 retardation_factor: float = 1.0,
680 flow_out: npt.ArrayLike | None = None,
681 spinup: str | None = "constant",
682) -> npt.NDArray[np.floating]:
683 """Compute extracted concentration for a gamma-distributed pore volume distribution (approximate).
685 Convenience wrapper around :func:`infiltration_to_extraction` that discretizes a (shifted)
686 gamma aquifer pore-volume distribution into ``n_bins`` equal-probability streamtubes. Provide
687 either (mean, std) or (alpha, beta); ``loc`` defaults to 0. Approximate -- see
688 :func:`infiltration_to_extraction`.
690 Parameters
691 ----------
692 cin : array-like
693 Concentration of the compound in infiltrating water.
694 flow : array-like
695 Flow rate of water in the aquifer [m³/day].
696 tedges : pandas.DatetimeIndex
697 Time edges for cin and flow data. Length ``len(cin) + 1``.
698 cout_tedges : pandas.DatetimeIndex
699 Time edges for output data bins. Length ``len(result) + 1``.
700 mean, std : float, optional
701 Mean and standard deviation of the gamma pore-volume distribution.
702 loc : float, optional
703 Location (minimum pore volume), ``0 <= loc < mean``. Default 0.0.
704 alpha, beta : float, optional
705 Shape and scale parameters of the gamma distribution (alternative to mean/std).
706 n_bins : int, optional
707 Number of equal-probability streamtubes. Default 100.
708 streamline_length : float
709 Travel distance L [m], applied to all gamma streamtubes. Must be positive.
710 molecular_diffusivity : float
711 Effective molecular diffusivity D_m [m²/day], applied to all streamtubes. Must be
712 non-negative.
713 longitudinal_dispersivity : float
714 Longitudinal dispersivity alpha_L [m] (microdispersion), applied to all streamtubes. Must be non-negative.
715 retardation_factor : float, optional
716 Retardation factor (default 1.0).
717 flow_out : array-like or None, optional
718 Extraction flow rate [m³/day] on the output grid. See
719 :func:`infiltration_to_extraction`. Default None.
720 spinup : {"constant"} | None, optional
721 See :func:`infiltration_to_extraction`. Default ``"constant"``.
723 Returns
724 -------
725 numpy.ndarray
726 Bin-averaged Kreft-Zuber flux concentration ``C_F`` in the extracted water.
727 Length ``len(cout_tedges) - 1``. NaN where no infiltration data has broken through.
729 See Also
730 --------
731 infiltration_to_extraction : Transport with an explicit pore volume distribution.
732 gamma_extraction_to_infiltration : Reverse operation.
733 gwtransport.gamma.bins : Create gamma distribution bins.
734 :ref:`concept-gamma-distribution` : Two-parameter pore volume model.
735 """
736 bins = gamma.bins(mean=mean, std=std, loc=loc, alpha=alpha, beta=beta, n_bins=n_bins)
737 return infiltration_to_extraction(
738 cin=cin,
739 flow=flow,
740 tedges=tedges,
741 cout_tedges=cout_tedges,
742 aquifer_pore_volumes=bins["expected_values"],
743 streamline_length=streamline_length,
744 molecular_diffusivity=molecular_diffusivity,
745 longitudinal_dispersivity=longitudinal_dispersivity,
746 retardation_factor=retardation_factor,
747 flow_out=flow_out,
748 spinup=spinup,
749 )
752def gamma_extraction_to_infiltration(
753 *,
754 cout: npt.ArrayLike,
755 flow: npt.ArrayLike,
756 tedges: pd.DatetimeIndex,
757 cout_tedges: pd.DatetimeIndex,
758 mean: float | None = None,
759 std: float | None = None,
760 loc: float = 0.0,
761 alpha: float | None = None,
762 beta: float | None = None,
763 n_bins: int = 100,
764 streamline_length: float,
765 molecular_diffusivity: float,
766 longitudinal_dispersivity: float,
767 retardation_factor: float = 1.0,
768 regularization_strength: float = 1e-10,
769 flow_out: npt.ArrayLike | None = None,
770 spinup: str | None = "constant",
771) -> npt.NDArray[np.floating]:
772 """Reconstruct infiltration concentration for a gamma-distributed pore volume distribution.
774 Convenience wrapper around :func:`extraction_to_infiltration` that discretizes a (shifted)
775 gamma aquifer pore-volume distribution into ``n_bins`` equal-probability streamtubes. Provide
776 either (mean, std) or (alpha, beta); ``loc`` defaults to 0. Fast approximate banded deconvolution
777 (see :func:`extraction_to_infiltration`).
779 Parameters
780 ----------
781 cout : array-like
782 Concentration of the compound in extracted water.
783 flow : array-like
784 Flow rate of water in the aquifer [m³/day].
785 tedges : pandas.DatetimeIndex
786 Time edges for cin (output) and flow data. Length ``len(flow) + 1``.
787 cout_tedges : pandas.DatetimeIndex
788 Time edges for cout data bins. Length ``len(cout) + 1``.
789 mean, std : float, optional
790 Mean and standard deviation of the gamma pore-volume distribution.
791 loc : float, optional
792 Location (minimum pore volume), ``0 <= loc < mean``. Default 0.0.
793 alpha, beta : float, optional
794 Shape and scale parameters of the gamma distribution (alternative to mean/std).
795 n_bins : int, optional
796 Number of equal-probability streamtubes. Default 100.
797 streamline_length : float
798 Travel distance L [m], applied to all gamma streamtubes. Must be positive.
799 molecular_diffusivity : float
800 Effective molecular diffusivity D_m [m²/day], applied to all streamtubes. Must be
801 non-negative.
802 longitudinal_dispersivity : float
803 Longitudinal dispersivity alpha_L [m] (microdispersion), applied to all streamtubes. Must be non-negative.
804 retardation_factor : float, optional
805 Retardation factor (default 1.0).
806 regularization_strength : float, optional
807 Tikhonov regularization parameter (default 1e-10).
808 flow_out : array-like or None, optional
809 Extraction flow rate [m³/day] on the output grid. See
810 :func:`infiltration_to_extraction`. Default None.
811 spinup : {"constant"} | None, optional
812 See :func:`infiltration_to_extraction`. Default ``"constant"``.
814 Returns
815 -------
816 numpy.ndarray
817 Bin-averaged concentration in the infiltrating water. Length ``len(tedges) - 1``.
819 See Also
820 --------
821 extraction_to_infiltration : Deconvolution with an explicit pore volume distribution.
822 gamma_infiltration_to_extraction : Forward operation.
823 gwtransport.gamma.bins : Create gamma distribution bins.
824 :ref:`concept-gamma-distribution` : Two-parameter pore volume model.
825 """
826 bins = gamma.bins(mean=mean, std=std, loc=loc, alpha=alpha, beta=beta, n_bins=n_bins)
827 return extraction_to_infiltration(
828 cout=cout,
829 flow=flow,
830 tedges=tedges,
831 cout_tedges=cout_tedges,
832 aquifer_pore_volumes=bins["expected_values"],
833 streamline_length=streamline_length,
834 molecular_diffusivity=molecular_diffusivity,
835 longitudinal_dispersivity=longitudinal_dispersivity,
836 retardation_factor=retardation_factor,
837 regularization_strength=regularization_strength,
838 flow_out=flow_out,
839 spinup=spinup,
840 )