Coverage for src/gwtransport/advection_utils.py: 98%
115 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-07 21:17 +0000
1"""
2Private helper functions for advective transport modeling.
4This module contains internal helper functions used by the advection module.
5These functions implement various algorithms for computing transport weights
6and handling nonlinear sorption.
8This file is part of gwtransport which is released under AGPL-3.0 license.
9See the ./LICENSE file or go to https://github.com/gwtransport/gwtransport/blob/main/LICENSE for full license details.
10"""
12import numpy as np
13import numpy.typing as npt
14import pandas as pd
16from gwtransport._time import tedges_to_days
17from gwtransport.utils import cumulative_flow_volume
19# Target number of (streamtube x cout-bin) pairs per tile in the banded weight build. The
20# per-tile working set is O(_WEIGHT_BUILD_BLOCK x band), so peak memory is bounded
21# independent of record length; chosen to keep the build under ~30 MB while spanning most
22# records in one or a few tiles. See _infiltration_to_extraction_weights.
23_WEIGHT_BUILD_BLOCK = 100_000
26def _infiltration_to_extraction_weights(
27 *,
28 tedges: pd.DatetimeIndex,
29 cout_tedges: pd.DatetimeIndex,
30 aquifer_pore_volumes: npt.NDArray[np.floating],
31 flow: npt.NDArray[np.floating],
32 retardation_factor: float,
33) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.intp], npt.NDArray[np.intp], npt.NDArray[np.bool_]]:
34 """
35 Compute raw streamtube-bundle weights for infiltration to extraction transformation.
37 Builds the per-cout-bin sum of streamtube-normalized overlap rows in a
38 **compact banded layout**, plus the count of contributing streamtubes per
39 cout bin and the zero-flow cout-bin mask. The caller decides how to convert
40 these into final weights (strict-validity, count-mean, or padded "constant"
41 spin-up); see :func:`_resolve_spinup_mask` and :func:`_resolve_spinup_inputs`.
42 Reconstruct the dense ``(n_cout, n_cin)`` matrix with :func:`_densify_weights`.
44 The per-streamtube weight is the literal mass-flux / water-flux ratio for
45 that streamtube and cout bin: each contributing row sums to 1 to ULP.
46 Equal-mass pore-volume bins from a gamma APVD discretization carry equal
47 flow at the outlet (steady-streamline assumption), so the bundle output is
48 the arithmetic mean over contributing streamtubes.
50 Pure advection (``D_m = 0``, ``alpha_L = 0``) is volume-stationary. Let
51 ``Vi`` be the cumulative throughflow volume at the cin edges and ``Vc`` the
52 same cumulative volume sampled at the cout edges. A streamtube of retarded
53 pore volume ``r = R * V_pore`` carries each cout edge back to the
54 infiltration time at cumulative volume ``Vc_edge - r``. The cout bin's
55 source window in infiltration time spans one cout bin's worth of volume, so
56 it overlaps only a few cin bins. The nonzeros of cout row ``k`` therefore
57 span only the residence-time spread of the APVD, ``[col_start[k],
58 col_start[k] + full_band)``, and are accumulated into one banded buffer over
59 a **pore-volume loop** -- ``O(n_cout * full_band)`` memory regardless of
60 record length or ``n_pv``. The overlap is the flow-weighted time overlap
61 normalized by the window's total in-range flow, so each contributing row
62 sums to 1 exactly: the exact flow-weighted overlap, matching the dense
63 ``residence_time`` + ``partial_isin`` reference to machine precision, not an
64 approximation.
66 A streamtube whose source volume leaves ``[Vi[0], Vi[-1]]`` -- or a cout
67 edge outside the cin time range -- maps to NaN and is **dropped, not
68 clipped**: its whole window for that cout bin is discarded (mirroring the
69 dense build, where a single NaN residence-time edge poisons the row).
71 Parameters
72 ----------
73 tedges : pandas.DatetimeIndex
74 Time edges for infiltration bins.
75 cout_tedges : pandas.DatetimeIndex
76 Time edges for extraction bins.
77 aquifer_pore_volumes : array-like
78 Distribution of pore volumes [m³].
79 flow : array-like
80 Flow rate values [m³/day].
81 retardation_factor : float
82 Constant retardation factor.
84 Returns
85 -------
86 band_vals : numpy.ndarray
87 Sum over streamtubes of per-streamtube normalized overlap rows in
88 banded layout. Shape: (len(cout_tedges) - 1, full_band). Slot
89 ``band_vals[k, b]`` is the weight on cin bin ``col_start[k] + b``;
90 for a cout bin ``k`` with ``c`` contributing streamtubes the row
91 sums to ``c`` (not ``n_pv`` and not 1).
92 col_start : numpy.ndarray of int
93 Shape: (len(cout_tedges) - 1,). First cin bin index of each cout
94 row's band. Defaults to 0 for rows with no contributing streamtube.
95 contributing_bins : numpy.ndarray of int
96 Shape: (len(cout_tedges) - 1,). Number of streamtubes that
97 actually contributed to each cout bin (had a source window fully
98 inside the cin volume range).
99 zero_flow_cout : numpy.ndarray of bool
100 Shape: (len(cout_tedges) - 1,). True for cout bins with zero
101 time-averaged extraction flow over their interval.
103 See Also
104 --------
105 _densify_weights : Reconstruct the dense (n_cout, n_cin) matrix.
106 """
107 cin_tedges_days = tedges_to_days(tedges)
108 cout_tedges_days = tedges_to_days(cout_tedges, ref=tedges[0])
109 flow = np.asarray(flow, dtype=float)
111 n_cout = len(cout_tedges) - 1
112 n_cin = len(tedges) - 1
113 n_pv = len(aquifer_pore_volumes)
115 # Cumulative throughflow volume at cin edges (Vi) and sampled at cout edges (Vc).
116 vi = cumulative_flow_volume(flow, np.diff(cin_tedges_days))
117 vc = np.interp(cout_tedges_days, cin_tedges_days, vi)
119 # A cout bin's time-averaged extraction flow is zero exactly when no throughflow
120 # volume passes during the bin, i.e. its cumulative-volume width is zero. This is
121 # bit-identical to the dense flow-overlap-matrix product but is O(n_cout), not O(N^2).
122 zero_flow_cout = np.diff(vc) == 0
124 if n_pv == 0:
125 return np.zeros((n_cout, 1)), np.zeros(n_cout, dtype=np.intp), np.zeros(n_cout, dtype=np.intp), zero_flow_cout
127 r = np.sort(np.asarray(aquifer_pore_volumes, dtype=float) * retardation_factor)
129 # Cumulative source volume at each cout edge. Out-of-range cout edges have no source
130 # volume; NaN there propagates so the adjacent windows are dropped (not clipped),
131 # matching the dense build.
132 edge_in_range = (cout_tedges_days >= cin_tedges_days[0]) & (cout_tedges_days <= cin_tedges_days[-1])
133 src_edge = np.where(edge_in_range, vc, np.nan)
135 # Tile the cout axis so the (n_pv, tile, band) overlap tensor never materialises for the
136 # whole record: peak memory stays O(block) regardless of record length, while each tile
137 # stays fully vectorised over its streamtubes. Most records span one or a few tiles. Each
138 # tile sizes its own band independently and is right-padded to the global width at the end.
139 block = max(1, _WEIGHT_BUILD_BLOCK // n_pv)
140 col_start = np.zeros(n_cout, dtype=np.intp)
141 contributing_bins = np.zeros(n_cout, dtype=np.intp)
142 tiles: list[tuple[int, npt.NDArray[np.floating]]] = []
143 full_band = 1
144 for start in range(0, n_cout, block):
145 stop = min(start + block, n_cout)
146 m = stop - start
148 # Back-project this tile's cout edges by every streamtube to infiltration time.
149 infil = np.interp(
150 (src_edge[start : stop + 1] - r[:, None]).ravel(), vi, cin_tedges_days, left=np.nan, right=np.nan
151 ).reshape(n_pv, m + 1)
152 win_lo, win_hi = infil[:, :-1], infil[:, 1:] # (n_pv, m) infiltration-time window per cout bin
153 contained = np.isfinite(win_lo) & np.isfinite(win_hi)
154 j_lo = np.clip(np.searchsorted(cin_tedges_days, win_lo, side="right") - 1, 0, n_cin - 1)
155 j_hi = np.clip(np.searchsorted(cin_tedges_days, win_hi, side="left"), 0, n_cin)
157 # Tile band = union of contained streamtube windows; per_band = widest single window.
158 any_contained = contained.any(axis=0)
159 tile_start = np.where(any_contained, np.where(contained, j_lo, n_cin).min(axis=0), 0)
160 row_hi = np.where(contained, j_hi, 0).max(axis=0)
161 tile_band = min(int(np.max(np.where(any_contained, row_hi - tile_start, 0))) + 1, n_cin)
162 per_band = min(int(np.max(np.where(contained, j_hi - j_lo, 0))) + 1, n_cin)
164 # Per-streamtube flow-weighted overlap on its own narrow window, normalised so each
165 # contributing row sums to 1, then scatter-summed over streamtubes into the tile's
166 # banded buffer (offset by each row's tile_start).
167 cols = j_lo[:, :, None] + np.arange(per_band)[None, None, :] # (n_pv, m, per_band)
168 in_cin = contained[:, :, None] & (cols < n_cin)
169 cols_clipped = np.clip(cols, 0, n_cin - 1)
170 overlap = np.maximum(
171 0.0,
172 np.minimum(win_hi[:, :, None], cin_tedges_days[cols_clipped + 1])
173 - np.maximum(win_lo[:, :, None], cin_tedges_days[cols_clipped]),
174 )
175 flux = np.where(in_cin, flow[cols_clipped] * overlap, 0.0)
176 window_flux = flux.sum(axis=2) # (n_pv, m) total in-range flow per window
177 contributes = contained & (window_flux > 0)
178 np.divide(flux, window_flux[:, :, None], out=flux, where=contributes[:, :, None])
180 # Only nonzero in-window slots scatter; their banded offset is < tile_band by
181 # construction (a contributing slot lies in [tile_start, row_hi)).
182 keep = flux > 0
183 offset = cols_clipped - tile_start[None, :, None]
184 row_idx = np.broadcast_to(np.arange(m)[None, :, None], cols.shape)
185 tile_vals = (
186 np
187 .bincount((row_idx * tile_band + offset)[keep], weights=flux[keep], minlength=m * tile_band)
188 .astype(float, copy=False)
189 .reshape(m, tile_band)
190 )
192 col_start[start:stop] = tile_start
193 contributing_bins[start:stop] = contributes.sum(axis=0)
194 tiles.append((start, tile_vals))
195 full_band = max(full_band, tile_band)
197 band_vals = np.zeros((n_cout, full_band))
198 for start, tile_vals in tiles:
199 band_vals[start : start + tile_vals.shape[0], : tile_vals.shape[1]] = tile_vals
201 return band_vals, col_start, contributing_bins, zero_flow_cout
204def _resolve_spinup_mask(
205 *,
206 band_vals: npt.NDArray[np.floating],
207 col_start: npt.NDArray[np.intp],
208 contributing_bins: npt.NDArray[np.intp],
209 zero_flow_cout: npt.NDArray[np.bool_],
210 n_pv: int,
211 spinup: float | None,
212) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.intp], npt.NDArray[np.bool_]]:
213 """Convert raw banded bundle outputs into final banded weights + invalid mask.
215 Parameters
216 ----------
217 band_vals : numpy.ndarray
218 Per-cout-bin sum of streamtube-normalized rows in banded layout from
219 :func:`_infiltration_to_extraction_weights`. Shape (n_cout, full_band).
220 col_start : numpy.ndarray of int
221 First cin bin index of each cout row's band. Passed through unchanged.
222 contributing_bins : numpy.ndarray of int
223 Number of streamtubes that contributed to each cout bin.
224 zero_flow_cout : numpy.ndarray of bool
225 Mask of zero-extraction-flow cout bins.
226 n_pv : int
227 Total number of streamtubes (length of aquifer_pore_volumes).
228 spinup : float in [0, 1] or None
229 Spin-up policy. ``None`` requires every streamtube to have
230 contributed (strict-validity, mass-conserving across cin → cout
231 per row). A float in [0, 1] is the minimum fraction of
232 contributing streamtubes for the row to be emitted; the bundle
233 is then a count-mean over the contributing subset (NOT
234 mass-conserving across the bundle when contributing < n_pv —
235 this is the pre-fix behavior of issue #161 made explicit).
237 Returns
238 -------
239 weights : numpy.ndarray
240 Final banded weight matrix of the same shape as ``band_vals``.
241 Rows where the policy is not satisfied are zero.
242 col_start : numpy.ndarray of int
243 The input ``col_start``, returned unchanged for caller convenience.
244 invalid_mask : numpy.ndarray of bool
245 True for cout bins where the policy is not satisfied.
247 Notes
248 -----
249 ``spinup`` is assumed already validated: :func:`_resolve_spinup_inputs`
250 resolves the string ``"constant"`` and rejects out-of-range values before
251 this function is reached, so ``spinup`` is always ``None`` or a float in
252 ``[0, 1]`` here.
253 """
254 if n_pv == 0:
255 # No streamtubes means no transport: every cout bin is invalid.
256 return np.zeros_like(band_vals), col_start, np.ones(band_vals.shape[0], dtype=bool)
258 if spinup is None:
259 # Strict validity: every streamtube must have contributed, so
260 # contributing_bins == n_pv on valid rows and the divisor below
261 # (contributing_bins) coincides with n_pv there.
262 valid = (contributing_bins == n_pv) & ~zero_flow_cout
263 else:
264 # ``_resolve_spinup_inputs`` has already narrowed a non-None threshold to a
265 # float in [0, 1], so no further type/range check is reachable here.
266 valid = (contributing_bins >= spinup * n_pv) & ~zero_flow_cout & (contributing_bins > 0)
268 # Every valid row has contributing_bins > 0 (== n_pv > 0 for strict validity,
269 # > 0 for the float threshold), so the divisor is guaranteed positive.
270 weights = np.zeros_like(band_vals)
271 weights[valid, :] = band_vals[valid, :] / contributing_bins[valid, None]
272 return weights, col_start, ~valid
275def _densify_weights(
276 band_vals: npt.NDArray[np.floating], col_start: npt.NDArray[np.intp], n_cin: int
277) -> npt.NDArray[np.floating]:
278 """Reconstruct the dense (n_cout, n_cin) weight matrix from the banded layout.
280 Inverse of the banded packing produced by
281 :func:`_infiltration_to_extraction_weights` and
282 :func:`_resolve_spinup_mask`: row ``k`` places ``band_vals[k]`` at cin
283 columns ``[col_start[k], col_start[k] + full_band)``, dropping band slots
284 that fall past ``n_cin`` (the right-edge padding).
286 Parameters
287 ----------
288 band_vals : numpy.ndarray
289 Banded weights of shape (n_cout, full_band).
290 col_start : numpy.ndarray of int
291 First cin bin index of each row's band, shape (n_cout,).
292 n_cin : int
293 Number of cin bins (dense column count).
295 Returns
296 -------
297 numpy.ndarray
298 Dense weight matrix of shape (n_cout, n_cin).
299 """
300 n_cout, full_band = band_vals.shape
301 dense = np.zeros((n_cout, n_cin))
302 cols = col_start[:, None] + np.arange(full_band)[None, :]
303 in_range = cols < n_cin
304 rows = np.broadcast_to(np.arange(n_cout)[:, None], cols.shape)
305 dense[rows[in_range], cols[in_range]] = band_vals[in_range]
306 return dense
309def _resolve_spinup_inputs(
310 spinup: object,
311 *,
312 tedges: pd.DatetimeIndex,
313 flow: npt.NDArray[np.floating],
314 aquifer_pore_volumes: npt.ArrayLike,
315 retardation_factor: float,
316 cin: npt.NDArray[np.floating] | None = None,
317) -> tuple[pd.DatetimeIndex, npt.NDArray[np.floating], npt.NDArray[np.floating] | None, float | None, int]:
318 """Validate ``spinup`` and apply its input-side effects.
320 Returns the (possibly padded) tedges, flow, and cin to use for
321 weight computation, plus the threshold for the output-side policy
322 and the number of bins prepended.
324 Modes:
326 - ``spinup is None`` — strict-validity. Returns inputs unchanged
327 and ``threshold = None``. Mass-conserving but NaN where any
328 streamtube has not broken through.
329 - ``spinup == "constant"`` — warm-start. Prepends ``n_pad`` bins
330 (each of width ``tedges[1] - tedges[0]``) so the total prepended
331 duration covers ``retardation_factor * max(aquifer_pore_volumes)
332 / flow[0]``. cin and flow are extended with their first observed
333 value (constant warm start). Strict-validity on the padded system
334 yields no spin-up NaN for cout bins at or after the original
335 ``tedges[0]``. Right-edge spin-up (cout extending past
336 ``tedges[-1]``) is not addressed.
337 - ``spinup`` is a float in ``[0, 1]`` — fraction threshold. Returns
338 inputs unchanged and ``threshold = float(spinup)``. *Warning:*
339 with ``spinup < 1.0`` the bundle is a count-mean over contributing
340 streamtubes; this conserves mass per row but NOT cin → cout (it
341 is exactly the issue #161 over-attribution made explicit).
343 The first bin width ``tedges[1] - tedges[0]`` is used as the unit
344 for prepended bins; this preserves uniformity if the input
345 ``tedges`` are uniform (required by the smooth-then-advect path of
346 ``diffusion_fast``).
348 Parameters
349 ----------
350 spinup : None, "constant", or float in [0, 1]
351 Public spin-up policy.
352 tedges : pandas.DatetimeIndex
353 Original cin/flow time edges (length n_cin + 1).
354 flow : numpy.ndarray
355 Flow values; ``flow[0]`` sets the warm-start flow.
356 aquifer_pore_volumes : array-like
357 Pore volumes; the maximum sets the warm-start duration.
358 retardation_factor : float
359 Retardation factor.
360 cin : numpy.ndarray, optional
361 Concentration values to prepend with ``cin[0]``. Pass ``None``
362 for the inverse direction (cin is the unknown to recover); the
363 returned ``new_cin`` is then ``None``.
365 Returns
366 -------
367 new_tedges : pandas.DatetimeIndex
368 Tedges to pass to :func:`_infiltration_to_extraction_weights`.
369 Length is ``len(tedges) + n_pad``.
370 new_flow : numpy.ndarray
371 Flow values aligned with ``new_tedges``; length ``n_cin + n_pad``.
372 new_cin : numpy.ndarray or None
373 cin values aligned with ``new_tedges`` if ``cin`` was provided;
374 otherwise ``None``.
375 threshold : float in [0, 1] or None
376 Threshold for :func:`_resolve_spinup_mask`. ``None`` (strict)
377 for ``spinup is None`` and ``spinup="constant"``.
378 n_pad : int
379 Number of bins prepended. ``0`` for non-padding modes. Callers
380 of the inverse direction must drop the first ``n_pad`` entries
381 from the recovered cin to align with the original ``tedges``.
383 Raises
384 ------
385 TypeError
386 If ``spinup`` has an unsupported type.
387 ValueError
388 If ``spinup`` is a string other than ``"constant"`` or a float
389 outside ``[0, 1]``.
390 """
391 if spinup is None:
392 return tedges, np.asarray(flow, dtype=float), cin, None, 0
393 flow_arr = np.asarray(flow, dtype=float)
394 if isinstance(spinup, str):
395 if spinup != "constant":
396 msg = f"spinup string must be 'constant'; got {spinup!r}"
397 raise ValueError(msg)
398 # Determine whether padding is feasible. We fall back to strict-validity
399 # (no padding) silently when the warm-start is undefined (zero or NaN
400 # initial flow) or when the implied padding would be unreasonably large
401 # (extreme pore volumes). This keeps the default usable for edge cases
402 # while still triggering the strict-validity NaN when the warm-start
403 # assumption cannot meaningfully be applied.
404 if len(tedges) < 2: # noqa: PLR2004
405 return tedges, flow_arr, cin, None, 0
406 q0 = float(flow_arr[0])
407 apvs = np.asarray(aquifer_pore_volumes, dtype=float)
408 if apvs.size == 0:
409 return tedges, flow_arr, cin, None, 0
410 v_max = float(np.max(apvs))
411 if not (q0 > 0 and v_max > 0):
412 return tedges, flow_arr, cin, None, 0
413 bin_width = tedges[1] - tedges[0]
414 bin_width_days = bin_width / pd.Timedelta(days=1)
415 if not bin_width_days > 0:
416 return tedges, flow_arr, cin, None, 0
417 pad_days = retardation_factor * v_max / q0
418 # Add 1 extra bin so the longest streamtube's source window for the
419 # earliest original cout bin lies strictly inside the padded range
420 # (avoids strict-validity NaN due to floating-point edge alignment).
421 n_pad_float = np.ceil(pad_days / bin_width_days) + 1
422 # Cap to keep memory bounded; beyond this, "constant" is no longer a
423 # meaningful warm-start (the user probably has unphysical pore volumes
424 # or extreme retardation), so fall through to strict-validity.
425 max_n_pad = max(10_000, 10 * len(flow_arr))
426 if not np.isfinite(n_pad_float) or n_pad_float > max_n_pad:
427 return tedges, flow_arr, cin, None, 0
428 n_pad = int(n_pad_float)
429 offsets = pd.TimedeltaIndex(bin_width * np.arange(n_pad, 0, -1))
430 new_tedges = (tedges[0] - offsets).append(tedges)
431 new_flow = np.concatenate([np.full(n_pad, flow_arr[0]), flow_arr])
432 new_cin = np.concatenate([np.full(n_pad, cin[0]), cin]) if cin is not None else None
433 return new_tedges, new_flow, new_cin, None, n_pad
434 if isinstance(spinup, bool) or not isinstance(spinup, (int, float)):
435 msg = f"spinup must be None, 'constant', or float in [0, 1]; got {spinup!r}"
436 raise TypeError(msg)
437 if not (0.0 <= spinup <= 1.0):
438 msg = f"spinup float must be in [0, 1]; got {spinup!r}"
439 raise ValueError(msg)
440 return tedges, flow_arr, cin, float(spinup), 0