Coverage for src/gwtransport/_diffusion_shared.py: 100%
89 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-07 21:17 +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 breakthrough antiderivative, input validation, the
6banded Tikhonov reverse solve, and the small per-streamtube / spin-up helpers. Both modules import
7from here so these primitives are defined once
8and evaluate bit-identically in either module (the modules' overall transport is *not* identical:
9diffusion_fast is exact, diffusion_fast_fast approximate).
11This file is part of gwtransport which is released under AGPL-3.0 license.
12See the ./LICENSE file or go to https://github.com/gwtransport/gwtransport/blob/main/LICENSE for full license details.
13"""
15import numpy as np
16import numpy.typing as npt
17import pandas as pd
18from scipy.special import erf
20from gwtransport._time import dt_to_days
21from gwtransport._validation import (
22 _validate_no_nan,
23 _validate_non_negative_array,
24 _validate_positive_array,
25 _validate_retardation_factor,
26)
27from gwtransport.utils import solve_inverse_transport_banded
29# Minimum coefficient sum to consider an output bin valid.
30_EPSILON_COEFF_SUM = 1e-10
32# sqrt(pi), used in the closed-form breakthrough antiderivative.
33_SQRT_PI = np.sqrt(np.pi)
35# Floor on the moving-frame dispersion product D_t [m^2] to keep the erf argument finite
36# for pre-breakthrough / zero-dispersion edges (where D_t -> 0).
37_DT_FLOOR = 1e-30
40def _breakthrough_antideriv(
41 step_widths: npt.NDArray[np.floating], dt_var: npt.NDArray[np.floating]
42) -> npt.NDArray[np.floating]:
43 r"""Closed-form antiderivative of the resident concentration, evaluated per edge.
45 Returns :math:`I(x) = \tfrac12 x + \tfrac12[x\,\operatorname{erf}(x/s) + (s/\sqrt\pi)e^{-(x/s)^2}]`
46 with :math:`s = 2\sqrt{D_t}`. Shared by the banded ``C_F`` build
47 (:func:`gwtransport.diffusion_fast._pv_band_values`) and the slow quadrature, so both compute
48 identical floating-point results for the same inputs. Because ``dD_t/dx = D_s/v_s`` is the
49 Kreft-Zuber flux coefficient at the solute-front velocity, differencing ``I`` across a cout bin
50 yields the flux concentration ``C_F`` directly.
52 Returns
53 -------
54 antideriv : ndarray
55 The antiderivative :math:`I(x)`.
56 """
57 s = 2.0 * np.sqrt(dt_var)
58 with np.errstate(over="ignore", invalid="ignore"):
59 u = step_widths / s
60 gaussian = np.exp(-(u * u))
61 return 0.5 * step_widths + 0.5 * (step_widths * erf(u) + (s / _SQRT_PI) * gaussian)
64def _cout_cumulative_volume(
65 *,
66 flow_out: npt.NDArray[np.floating] | None,
67 cout_tedges: pd.DatetimeIndex,
68 cout_tedges_days: npt.NDArray[np.floating],
69 tedges_days: npt.NDArray[np.floating],
70 cumulative_volume_at_cin: npt.NDArray[np.floating],
71) -> npt.NDArray[np.floating]:
72 """Cumulative through-flow volume at each cout edge, on the infiltration volume axis.
74 When ``flow_out`` is given (the user-specified extraction-side flow) the cout-edge volumes are
75 its cumulative integral, anchored at the first cout edge inside the flow record (so an output
76 window starting before the input data stays correctly aligned). Otherwise the cout edges are
77 interpolated from the infiltration cumulative-volume curve. Shared by
78 :func:`gwtransport.diffusion_fast._closed_form_coeff_matrix` and
79 :func:`gwtransport.diffusion_fast_fast._build_forward_operator` so every path --
80 diffusion_fast (exact) and diffusion_fast_fast's forward and reverse (both approximate, and built
81 from the same operator) -- places the cout grid on identical volume coordinates.
83 Parameters
84 ----------
85 flow_out : ndarray or None
86 Extraction flow rate [m3/day] on the output grid (length ``len(cout_tedges) - 1``), or None
87 to interpolate the cout edges from the infiltration curve.
88 cout_tedges : DatetimeIndex
89 Output time-bin edges (used only for the ``flow_out`` bin widths).
90 cout_tedges_days : ndarray
91 Output edges as days relative to the (work) infiltration reference.
92 tedges_days : ndarray
93 Infiltration edges as days relative to the same reference.
94 cumulative_volume_at_cin : ndarray
95 Cumulative infiltrated volume at each infiltration edge.
97 Returns
98 -------
99 ndarray
100 Cumulative volume at each cout edge (length ``len(cout_tedges)``).
101 """
102 if flow_out is None:
103 return np.interp(cout_tedges_days, tedges_days, cumulative_volume_at_cin)
104 cumsum_out = np.concatenate(([0.0], np.cumsum(flow_out * dt_to_days(cout_tedges))))
105 in_range = (cout_tedges_days >= tedges_days[0]) & (cout_tedges_days <= tedges_days[-1])
106 # np.argmax returns 0 for an all-False mask, the same fallback the guard provided.
107 i0 = int(np.argmax(in_range))
108 v_at_i0 = float(np.interp(cout_tedges_days[i0], tedges_days, cumulative_volume_at_cin))
109 return v_at_i0 + (cumsum_out - cumsum_out[i0])
112def _extend_tedges_flag(spinup: str | float | None) -> bool:
113 """Translate the public ``spinup`` parameter to the internal extend flag.
115 ``"constant"`` (default) extends ``tedges`` by 100 years on each side so a constant
116 warm-start fills the left-edge spin-up region; ``None`` disables the extension (spin-up
117 cout becomes NaN). Mirrors :func:`gwtransport.diffusion._diffusion_extend_tedges_flag`.
119 Returns
120 -------
121 bool
122 True if ``tedges`` should be extended (warm-start), False otherwise.
124 Raises
125 ------
126 ValueError
127 If ``spinup`` is a string other than ``"constant"``.
128 NotImplementedError
129 If ``spinup`` is a float (fraction-threshold mode is not implemented).
130 """
131 if spinup is None:
132 return False
133 if isinstance(spinup, str):
134 if spinup != "constant":
135 msg = f"spinup string must be 'constant'; got {spinup!r}"
136 raise ValueError(msg)
137 return True
138 msg = f"spinup only supports None or 'constant'; float thresholds are not implemented (got {spinup!r})"
139 raise NotImplementedError(msg)
142def _broadcast_to_pore_volumes(
143 values: npt.NDArray[np.floating] | float, n_pore_volumes: int
144) -> npt.NDArray[np.floating]:
145 """Return a per-pore-volume array: a scalar broadcasts to all streamtubes, an array passes through.
147 Length is validated upstream by :func:`_validate_inputs`, so a non-scalar array is
148 returned as-is (assumed length ``n_pore_volumes``).
150 Returns
151 -------
152 ndarray, shape (n_pore_volumes,)
153 Per-streamtube values.
154 """
155 arr = np.atleast_1d(np.asarray(values, dtype=float))
156 return np.broadcast_to(arr, (n_pore_volumes,)) if arr.size == 1 else arr
159def _validate_inputs(
160 *,
161 cin_or_cout: np.ndarray,
162 flow: np.ndarray,
163 tedges: pd.DatetimeIndex,
164 cout_tedges: pd.DatetimeIndex,
165 aquifer_pore_volumes: np.ndarray,
166 streamline_length: npt.NDArray[np.floating] | float,
167 molecular_diffusivity: npt.NDArray[np.floating] | float,
168 longitudinal_dispersivity: npt.NDArray[np.floating] | float,
169 retardation_factor: float,
170 is_forward: bool,
171 flow_out: np.ndarray | None = None,
172) -> None:
173 """Validate inputs for infiltration_to_extraction and extraction_to_infiltration.
175 ``streamline_length`` / ``molecular_diffusivity`` / ``longitudinal_dispersivity`` may be
176 a scalar or an array of length ``len(aquifer_pore_volumes)`` (one value per streamtube).
178 Raises
179 ------
180 ValueError
181 If array lengths are inconsistent, molecular_diffusivity or
182 longitudinal_dispersivity are negative or non-finite, cin or cout or flow contain NaN
183 values, aquifer_pore_volumes contains non-positive or non-finite values,
184 streamline_length is non-positive or non-finite, or retardation_factor is NaN or below 1
185 (anti-retardation is not physical for the supported sorption isotherms).
186 """
187 if is_forward:
188 if len(tedges) != len(cin_or_cout) + 1:
189 msg = "tedges must have one more element than cin"
190 raise ValueError(msg)
191 elif len(cout_tedges) != len(cin_or_cout) + 1:
192 msg = "cout_tedges must have one more element than cout"
193 raise ValueError(msg)
194 if len(tedges) != len(flow) + 1:
195 msg = "tedges must have one more element than flow"
196 raise ValueError(msg)
197 n_pore_volumes = len(aquifer_pore_volumes)
198 for name, arr in (
199 ("streamline_length", streamline_length),
200 ("molecular_diffusivity", molecular_diffusivity),
201 ("longitudinal_dispersivity", longitudinal_dispersivity),
202 ):
203 if np.size(arr) not in {1, n_pore_volumes}:
204 msg = f"{name} must be a scalar or have length len(aquifer_pore_volumes) = {n_pore_volumes}"
205 raise ValueError(msg)
206 # Delegate the finite+sign invariants to the shared _validation atoms so the NaN/+inf guards
207 # (which the bare ``< 0`` / ``<= 0`` / ``< 1.0`` comparisons here would let slip through) are
208 # enforced in exactly one place. Order and messages match the historical inline checks.
209 _validate_non_negative_array(molecular_diffusivity, name="molecular_diffusivity")
210 _validate_non_negative_array(longitudinal_dispersivity, name="longitudinal_dispersivity")
211 _validate_no_nan(cin_or_cout, name="cin" if is_forward else "cout")
212 _validate_no_nan(flow, name="flow")
213 _validate_non_negative_array(flow, name="flow", message="flow must be non-negative (negative flow not supported)")
214 _validate_positive_array(aquifer_pore_volumes, name="aquifer_pore_volumes")
215 _validate_positive_array(streamline_length, name="streamline_length")
216 _validate_retardation_factor(retardation_factor)
217 if flow_out is None:
218 # The output-grid extraction flow is only unambiguous when the cout grid matches
219 # the flow grid; otherwise it must be supplied (it defines the cout-bin volumes and
220 # the outlet velocity used by the retardation correction).
221 if not tedges.equals(cout_tedges):
222 msg = "flow_out is required when cout_tedges differs from tedges"
223 raise ValueError(msg)
224 else:
225 n_cout = len(cout_tedges) - 1
226 if len(flow_out) != n_cout:
227 msg = f"flow_out must have length len(cout_tedges) - 1 = {n_cout}, got {len(flow_out)}"
228 raise ValueError(msg)
229 if np.any(np.isnan(flow_out)):
230 msg = "flow_out contains NaN values, which are not allowed"
231 raise ValueError(msg)
232 if np.any(flow_out < 0):
233 msg = "flow_out must be non-negative (negative flow not supported)"
234 raise ValueError(msg)
237def _solve_reverse_banded(
238 *,
239 band_vals: npt.NDArray[np.floating],
240 col_start: npt.NDArray[np.intp],
241 valid_cout_bins: npt.NDArray[np.bool_],
242 cout: npt.NDArray[np.floating],
243 n_cin: int,
244 regularization_strength: float,
245) -> npt.NDArray[np.floating]:
246 """Normalize, decouple the warm-start tail, and solve the banded Tikhonov inverse.
248 Shared by :func:`gwtransport.diffusion_fast.extraction_to_infiltration` (which feeds the exact
249 closed-form banded operator) and :func:`gwtransport.diffusion_fast_fast.extraction_to_infiltration`
250 (which feeds the approximate banded ``G . M`` operator). The steps are:
252 1. Zero invalid rows (incomplete breakthrough) and normalize the remaining rows to sum to 1 --
253 the banded solver's ``W x ~= observed`` precondition.
254 2. Decouple the warm-start data-start tail. With a leading zero-flow plateau and ``D_m > 0`` the
255 forward operator carries a negative warm-start coefficient at the data-start columns (kept in
256 the forward band so the forward ``C_F`` is reproduced); their net column sum is ``<= 0``, so the
257 banded normal-equation solver would leave them unregularized -- a large, unregularized
258 ``WᵀW`` diagonal coupled to the spin-up nullspace, which is indefinite and breaks the Cholesky
259 factorisation. These columns are the unrecoverable spin-up region (NaN in the dense and the
260 slow-module inverses alike), so zeroing their band entries decouples them: the solver's
261 zero-diagonal path returns them as NaN and the remaining system is symmetric positive definite.
263 Parameters
264 ----------
265 band_vals : ndarray, shape (n_cout_bins, full_band)
266 Banded forward weights -- from :func:`gwtransport.diffusion_fast._closed_form_coeff_matrix`
267 (exact) or :func:`gwtransport.diffusion_fast_fast._banded_forward_matrix` (approximate ``G . M``).
268 col_start : ndarray of int, shape (n_cout_bins,)
269 First cin-bin column of each cout row's band.
270 valid_cout_bins : ndarray of bool, shape (n_cout_bins,)
271 Output bins with complete breakthrough information.
272 cout : ndarray, shape (n_cout_bins,)
273 Observed extraction concentration.
274 n_cin : int
275 Number of infiltration bins (output length).
276 regularization_strength : float
277 Tikhonov parameter.
279 Returns
280 -------
281 ndarray, shape (n_cin,)
282 Recovered infiltration concentration; NaN for unconstrained / spin-up bins.
283 """
284 row_sums = band_vals.sum(axis=1)
285 valid = valid_cout_bins & (row_sums > _EPSILON_COEFF_SUM)
286 bn = band_vals.copy()
287 bn[~valid] = 0.0
288 bn[valid] /= row_sums[valid, None]
290 full_band = bn.shape[1]
291 band_cols = col_start[:, None] + np.arange(full_band)
292 in_range = band_cols < n_cin
293 band_cols_clipped = np.clip(band_cols, 0, n_cin - 1)
294 col_sum = np.zeros(n_cin)
295 np.add.at(col_sum, band_cols_clipped[in_range], bn[in_range])
296 inactive_col = col_sum <= _EPSILON_COEFF_SUM
297 bn[inactive_col[band_cols_clipped] & in_range] = 0.0
299 return solve_inverse_transport_banded(
300 band_vals=bn,
301 col_start=col_start,
302 observed=cout,
303 n_output=n_cin,
304 regularization_strength=regularization_strength,
305 )