Coverage for src/gwtransport/_diffusion_shared.py: 100%
107 statements
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-04 21:13 +0000
1"""
2Shared closed-form helpers for the Kreft-Zuber flux-concentration transport modules.
4This private module holds the pieces common to :mod:`gwtransport.diffusion_fast` and
5:mod:`gwtransport.diffusion_fast_fast`: the closed-form breakthrough antiderivative, the input
6coerce/validate/broadcast preamble, the cout-edge cumulative-volume axis, the warm-start (spin-up)
7grid and its policy flag, the advective-validity gate, and the banded Tikhonov reverse solve. Both
8modules import from here so these primitives are defined once and evaluate bit-identically in either
9module (the modules' overall transport is *not* identical: diffusion_fast is exact,
10diffusion_fast_fast approximate). :mod:`gwtransport.diffusion`, the quadrature reference, shares only
11the parameter broadcasting and the spin-up policy flag; its kernel is its own.
13This file is part of gwtransport which is released under AGPL-3.0 license.
14See the ./LICENSE file or go to https://github.com/gwtransport/gwtransport/blob/main/LICENSE for full license details.
15"""
17import numpy as np
18import numpy.typing as npt
19import pandas as pd
20from scipy.special import erf
22from gwtransport._time import dt_to_days
23from gwtransport._validation import (
24 _validate_no_nan,
25 _validate_non_negative_array,
26 _validate_positive_array,
27 _validate_retardation_factor,
28)
29from gwtransport.residence_time import fraction_explained_full
30from gwtransport.utils import solve_inverse_transport_banded
32# Minimum coefficient sum to consider an output bin valid.
33_EPSILON_COEFF_SUM = 1e-10
35# sqrt(pi), used in the closed-form breakthrough antiderivative.
36_SQRT_PI = np.sqrt(np.pi)
38# Floor on the moving-frame dispersion product D_t [m^2] to keep the erf argument finite
39# for pre-breakthrough / zero-dispersion edges (where D_t -> 0).
40_DT_FLOOR = 1e-30
43def _breakthrough_antideriv(
44 step_widths: npt.NDArray[np.floating], dt_var: npt.NDArray[np.floating]
45) -> npt.NDArray[np.floating]:
46 r"""Closed-form antiderivative of the resident concentration, evaluated per edge.
48 Returns :math:`I(x) = \tfrac12 x + \tfrac12[x\,\operatorname{erf}(x/s) + (s/\sqrt\pi)e^{-(x/s)^2}]`
49 with :math:`s = 2\sqrt{D_t}`. Shared by the banded ``C_F`` build
50 (:func:`gwtransport.diffusion_fast._pv_band_values`) and the slow quadrature, so both compute
51 identical floating-point results for the same inputs. Because ``dD_t/dx = D_s/v_s`` is the
52 Kreft-Zuber flux coefficient at the solute-front velocity, differencing ``I`` across a cout bin
53 yields the flux concentration ``C_F`` directly.
55 Returns
56 -------
57 antideriv : ndarray
58 The antiderivative :math:`I(x)`.
59 """
60 s = 2.0 * np.sqrt(dt_var)
61 with np.errstate(over="ignore", invalid="ignore"):
62 u = step_widths / s
63 gaussian = np.exp(-(u * u))
64 return 0.5 * step_widths + 0.5 * (step_widths * erf(u) + (s / _SQRT_PI) * gaussian)
67def _cout_cumulative_volume(
68 *,
69 flow_out: npt.NDArray[np.floating] | None,
70 cout_tedges: pd.DatetimeIndex,
71 cout_tedges_days: npt.NDArray[np.floating],
72 tedges_days: npt.NDArray[np.floating],
73 cumulative_volume_at_cin: npt.NDArray[np.floating],
74) -> npt.NDArray[np.floating]:
75 """Cumulative through-flow volume at each cout edge, on the infiltration volume axis.
77 When ``flow_out`` is given (the user-specified extraction-side flow) the cout-edge volumes are
78 its cumulative integral, anchored at the first cout edge inside the flow record (so an output
79 window starting before the input data stays correctly aligned). Otherwise the cout edges are
80 interpolated from the infiltration cumulative-volume curve. Shared by
81 :func:`gwtransport.diffusion_fast._closed_form_coeff_matrix` and
82 :func:`gwtransport.diffusion_fast_fast._build_forward_operator` so every path --
83 diffusion_fast (exact) and diffusion_fast_fast's forward and reverse (both approximate, and built
84 from the same operator) -- places the cout grid on identical volume coordinates.
86 Parameters
87 ----------
88 flow_out : ndarray or None
89 Extraction flow rate [m3/day] on the output grid (length ``len(cout_tedges) - 1``), or None
90 to interpolate the cout edges from the infiltration curve.
91 cout_tedges : DatetimeIndex
92 Output time-bin edges (used only for the ``flow_out`` bin widths).
93 cout_tedges_days : ndarray
94 Output edges as days relative to the (work) infiltration reference.
95 tedges_days : ndarray
96 Infiltration edges as days relative to the same reference.
97 cumulative_volume_at_cin : ndarray
98 Cumulative infiltrated volume at each infiltration edge.
100 Returns
101 -------
102 ndarray
103 Cumulative volume at each cout edge (length ``len(cout_tedges)``).
104 """
105 if flow_out is None:
106 return np.interp(cout_tedges_days, tedges_days, cumulative_volume_at_cin)
107 cumsum_out = np.concatenate(([0.0], np.cumsum(flow_out * dt_to_days(cout_tedges))))
108 in_range = (cout_tedges_days >= tedges_days[0]) & (cout_tedges_days <= tedges_days[-1])
109 # np.argmax returns 0 for an all-False mask (no cout edge inside the flow record).
110 i0 = int(np.argmax(in_range))
111 v_at_i0 = float(np.interp(cout_tedges_days[i0], tedges_days, cumulative_volume_at_cin))
112 return v_at_i0 + (cumsum_out - cumsum_out[i0])
115def _extend_tedges_flag(spinup: str | float | None) -> bool:
116 """Translate the public ``spinup`` parameter to the internal extend flag.
118 ``"constant"`` (default) extends ``tedges`` by 100 years on each side so a constant
119 warm-start fills the left-edge spin-up region; ``None`` disables the extension (spin-up
120 cout becomes NaN).
122 Returns
123 -------
124 bool
125 True if ``tedges`` should be extended (warm-start), False otherwise.
127 Raises
128 ------
129 ValueError
130 If ``spinup`` is a string other than ``"constant"``.
131 NotImplementedError
132 If ``spinup`` is a float (fraction-threshold mode is not implemented).
133 """
134 if spinup is None:
135 return False
136 if isinstance(spinup, str):
137 if spinup != "constant":
138 msg = f"spinup string must be 'constant'; got {spinup!r}"
139 raise ValueError(msg)
140 return True
141 msg = f"spinup only supports None or 'constant'; float thresholds are not implemented (got {spinup!r})"
142 raise NotImplementedError(msg)
145def _extend_tedges(tedges: pd.DatetimeIndex) -> pd.DatetimeIndex:
146 """Pad ``tedges`` by 100 years on each side, the warm-start grid of ``spinup="constant"``.
148 Timestamp arithmetic keeps the input timezone (tz-naive stays naive, tz-aware stays
149 tz-aware); going through ``.to_numpy()`` would strip it.
151 Returns
152 -------
153 DatetimeIndex
154 Padded edges, same length as ``tedges``.
155 """
156 pad = pd.Timedelta(days=36500)
157 return (tedges[:1] - pad).append(tedges[1:-1]).append(tedges[-1:] + pad)
160def _advective_valid_cout_bins(
161 *,
162 flow: npt.NDArray[np.floating],
163 tedges: pd.DatetimeIndex,
164 cout_tedges: pd.DatetimeIndex,
165 aquifer_pore_volumes: npt.NDArray[np.floating],
166 retardation_factor: float,
167) -> npt.NDArray[np.bool_]:
168 """Output bins whose advective look-back stays in-record for every streamtube.
170 A bin is valid where the advective coverage is 1 for all pore volumes (NaN outside the
171 record -> invalid). This is the advective gate only; the dispersive informedness is the
172 captured kernel mass, applied downstream.
174 Returns
175 -------
176 ndarray of bool, shape (len(cout_tedges) - 1,)
177 True where the bin is fully informed by the infiltration record.
178 """
179 return np.all(
180 fraction_explained_full(
181 flow=flow,
182 tedges=tedges,
183 cout_tedges=cout_tedges,
184 aquifer_pore_volumes=aquifer_pore_volumes,
185 retardation_factor=retardation_factor,
186 direction="extraction_to_infiltration",
187 )
188 >= 1.0,
189 axis=0,
190 )
193def _broadcast_to_pore_volumes(values: npt.ArrayLike, n_pore_volumes: int) -> npt.NDArray[np.floating]:
194 """Return a per-pore-volume array: a scalar broadcasts to all streamtubes, an array passes through.
196 Callers validate the length separately, so a non-scalar array is returned as-is (assumed
197 length ``n_pore_volumes``). Broadcast results are read-only views.
199 Returns
200 -------
201 ndarray, shape (n_pore_volumes,)
202 Per-streamtube values.
203 """
204 arr = np.atleast_1d(np.asarray(values, dtype=float))
205 return np.broadcast_to(arr, (n_pore_volumes,)) if arr.size == 1 else arr
208def _validate_inputs(
209 *,
210 cin_or_cout: np.ndarray,
211 flow: np.ndarray,
212 tedges: pd.DatetimeIndex,
213 cout_tedges: pd.DatetimeIndex,
214 aquifer_pore_volumes: np.ndarray,
215 streamline_length: npt.NDArray[np.floating] | float,
216 molecular_diffusivity: npt.NDArray[np.floating] | float,
217 longitudinal_dispersivity: npt.NDArray[np.floating] | float,
218 retardation_factor: float,
219 is_forward: bool,
220 flow_out: np.ndarray | None = None,
221) -> None:
222 """Validate inputs for infiltration_to_extraction and extraction_to_infiltration.
224 ``streamline_length`` / ``molecular_diffusivity`` / ``longitudinal_dispersivity`` may be
225 a scalar or an array of length ``len(aquifer_pore_volumes)`` (one value per streamtube).
227 Raises
228 ------
229 ValueError
230 If array lengths are inconsistent, molecular_diffusivity or
231 longitudinal_dispersivity are negative or non-finite, cin (forward) or flow contain NaN
232 values, aquifer_pore_volumes contains non-positive or non-finite values,
233 streamline_length is non-positive or non-finite, or retardation_factor is NaN or below 1
234 (anti-retardation is not physical for the supported sorption isotherms).
235 """
236 if is_forward:
237 if len(tedges) != len(cin_or_cout) + 1:
238 msg = "tedges must have one more element than cin"
239 raise ValueError(msg)
240 elif len(cout_tedges) != len(cin_or_cout) + 1:
241 msg = "cout_tedges must have one more element than cout"
242 raise ValueError(msg)
243 if len(tedges) != len(flow) + 1:
244 msg = "tedges must have one more element than flow"
245 raise ValueError(msg)
246 n_pore_volumes = len(aquifer_pore_volumes)
247 for name, arr in (
248 ("streamline_length", streamline_length),
249 ("molecular_diffusivity", molecular_diffusivity),
250 ("longitudinal_dispersivity", longitudinal_dispersivity),
251 ):
252 if np.size(arr) not in {1, n_pore_volumes}:
253 msg = f"{name} must be a scalar or have length len(aquifer_pore_volumes) = {n_pore_volumes}"
254 raise ValueError(msg)
255 # The finite+sign invariants go through the shared _validation atoms so the NaN/+inf guards
256 # (which bare ``< 0`` / ``<= 0`` / ``< 1.0`` comparisons would let slip through) are enforced
257 # in exactly one place.
258 _validate_non_negative_array(molecular_diffusivity, name="molecular_diffusivity")
259 _validate_non_negative_array(longitudinal_dispersivity, name="longitudinal_dispersivity")
260 # Reverse cout may contain NaN (measurement gaps); gapped rows are excluded from the solve.
261 if is_forward:
262 _validate_no_nan(cin_or_cout, name="cin")
263 _validate_no_nan(flow, name="flow")
264 _validate_non_negative_array(flow, name="flow", message="flow must be non-negative (negative flow not supported)")
265 _validate_positive_array(aquifer_pore_volumes, name="aquifer_pore_volumes")
266 _validate_positive_array(streamline_length, name="streamline_length")
267 _validate_retardation_factor(retardation_factor)
268 if flow_out is None:
269 # The output-grid extraction flow is only unambiguous when the cout grid matches
270 # the flow grid; otherwise it must be supplied (it defines the cout-bin volumes and
271 # the outlet velocity used by the retardation correction).
272 if not tedges.equals(cout_tedges):
273 msg = "flow_out is required when cout_tedges differs from tedges"
274 raise ValueError(msg)
275 else:
276 n_cout = len(cout_tedges) - 1
277 if len(flow_out) != n_cout:
278 msg = f"flow_out must have length len(cout_tedges) - 1 = {n_cout}, got {len(flow_out)}"
279 raise ValueError(msg)
280 if np.any(np.isnan(flow_out)):
281 msg = "flow_out contains NaN values, which are not allowed"
282 raise ValueError(msg)
283 if np.any(flow_out < 0):
284 msg = "flow_out must be non-negative (negative flow not supported)"
285 raise ValueError(msg)
288def _coerce_and_validate(
289 *,
290 cin_or_cout: npt.ArrayLike,
291 flow: npt.ArrayLike,
292 tedges: pd.DatetimeIndex,
293 cout_tedges: pd.DatetimeIndex,
294 aquifer_pore_volumes: npt.ArrayLike,
295 streamline_length: npt.NDArray[np.floating] | float,
296 molecular_diffusivity: npt.NDArray[np.floating] | float,
297 longitudinal_dispersivity: npt.NDArray[np.floating] | float,
298 retardation_factor: float,
299 is_forward: bool,
300 flow_out: npt.ArrayLike | None,
301) -> tuple[npt.NDArray[np.floating], dict]:
302 """Coerce, validate, and broadcast the public transport inputs of the fast modules.
304 Validation runs on the coerced-but-not-yet-broadcast per-streamtube parameters, so a
305 length mismatch against ``aquifer_pore_volumes`` is still caught.
307 Returns
308 -------
309 values : ndarray
310 ``cin`` (forward) or ``cout`` (reverse) as a float array.
311 transport : dict
312 Coefficient-matrix builder arguments, keyed by parameter name: the coerced ``flow`` /
313 ``tedges`` / ``cout_tedges`` / ``flow_out`` / ``aquifer_pore_volumes`` plus the three
314 per-streamtube arrays broadcast to one value per pore volume.
315 """
316 tedges = pd.DatetimeIndex(tedges)
317 cout_tedges = pd.DatetimeIndex(cout_tedges)
318 values = np.asarray(cin_or_cout, dtype=float)
319 flow = np.asarray(flow, dtype=float)
320 aquifer_pore_volumes = np.asarray(aquifer_pore_volumes, dtype=float)
321 if flow_out is not None:
322 flow_out = np.asarray(flow_out, dtype=float)
324 _validate_inputs(
325 cin_or_cout=values,
326 flow=flow,
327 tedges=tedges,
328 cout_tedges=cout_tedges,
329 aquifer_pore_volumes=aquifer_pore_volumes,
330 streamline_length=streamline_length,
331 molecular_diffusivity=molecular_diffusivity,
332 longitudinal_dispersivity=longitudinal_dispersivity,
333 retardation_factor=retardation_factor,
334 is_forward=is_forward,
335 flow_out=flow_out,
336 )
338 n_pore_volumes = len(aquifer_pore_volumes)
339 return values, {
340 "flow": flow,
341 "tedges": tedges,
342 "cout_tedges": cout_tedges,
343 "flow_out": flow_out,
344 "aquifer_pore_volumes": aquifer_pore_volumes,
345 "streamline_length": _broadcast_to_pore_volumes(streamline_length, n_pore_volumes),
346 "molecular_diffusivity": _broadcast_to_pore_volumes(molecular_diffusivity, n_pore_volumes),
347 "longitudinal_dispersivity": _broadcast_to_pore_volumes(longitudinal_dispersivity, n_pore_volumes),
348 }
351def _solve_reverse_banded(
352 *,
353 band_vals: npt.NDArray[np.floating],
354 col_start: npt.NDArray[np.intp],
355 valid_cout_bins: npt.NDArray[np.bool_],
356 cout: npt.NDArray[np.floating],
357 n_cin: int,
358 regularization_strength: float,
359) -> npt.NDArray[np.floating]:
360 """Normalize, decouple the warm-start tail, and solve the banded Tikhonov inverse.
362 Shared by :func:`gwtransport.diffusion_fast.extraction_to_infiltration` (which feeds the exact
363 closed-form banded operator) and :func:`gwtransport.diffusion_fast_fast.extraction_to_infiltration`
364 (which feeds the approximate banded breakthrough operator). The steps are:
366 1. Zero invalid rows (incomplete breakthrough) and normalize the remaining rows to sum to 1 --
367 the banded solver's ``W x ~= observed`` precondition.
368 2. Decouple the warm-start data-start tail. With a leading zero-flow plateau and ``D_m > 0`` the
369 forward operator carries a negative warm-start coefficient at the data-start columns (kept in
370 the forward band so the forward ``C_F`` is reproduced); their net column sum is ``<= 0``, so the
371 banded normal-equation solver would leave them unregularized -- a large, unregularized
372 ``WᵀW`` diagonal coupled to the spin-up nullspace, which is indefinite and breaks the Cholesky
373 factorisation. These columns are the unrecoverable spin-up region (NaN in the dense and the
374 slow-module inverses alike), so zeroing their band entries decouples them: the solver's
375 zero-diagonal path returns them as NaN and the remaining system is symmetric positive definite.
377 Parameters
378 ----------
379 band_vals : ndarray, shape (n_cout_bins, full_band)
380 Banded forward weights -- from :func:`gwtransport.diffusion_fast._closed_form_coeff_matrix`
381 (exact) or :func:`gwtransport.diffusion_fast_fast._banded_forward_matrix` (approximate breakthrough).
382 col_start : ndarray of int, shape (n_cout_bins,)
383 First cin-bin column of each cout row's band.
384 valid_cout_bins : ndarray of bool, shape (n_cout_bins,)
385 Output bins with complete breakthrough information.
386 cout : ndarray, shape (n_cout_bins,)
387 Observed extraction concentration.
388 n_cin : int
389 Number of infiltration bins (output length).
390 regularization_strength : float
391 Tikhonov parameter.
393 Returns
394 -------
395 ndarray, shape (n_cin,)
396 Recovered infiltration concentration; NaN for unconstrained / spin-up bins.
397 """
398 row_sums = band_vals.sum(axis=1)
399 valid = valid_cout_bins & (row_sums > _EPSILON_COEFF_SUM)
400 bn = band_vals.copy()
401 bn[~valid] = 0.0
402 bn[valid] /= row_sums[valid, None]
404 full_band = bn.shape[1]
405 band_cols = col_start[:, None] + np.arange(full_band)
406 in_range = band_cols < n_cin
407 band_cols_clipped = np.clip(band_cols, 0, n_cin - 1)
408 col_sum = np.zeros(n_cin)
409 np.add.at(col_sum, band_cols_clipped[in_range], bn[in_range])
410 inactive_col = col_sum <= _EPSILON_COEFF_SUM
411 bn[inactive_col[band_cols_clipped] & in_range] = 0.0
413 return solve_inverse_transport_banded(
414 band_vals=bn,
415 col_start=col_start,
416 observed=cout,
417 n_output=n_cin,
418 regularization_strength=regularization_strength,
419 )