Coverage for src/gwtransport/advection_utils.py: 0%
101 statements
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-04 20:54 +0000
1"""
2Private helper functions for advective transport modeling.
4This module contains internal helper functions used by :mod:`gwtransport.advection`. It has no
5public API; every name is private and the module is not part of the documented interface.
7The helpers build and post-process the linear infiltration-to-extraction operator that both the
8forward and the reverse advection entry points share. One helper computes the raw
9streamtube-bundle weights on the cumulative-throughflow-volume axis, in a compact banded layout
10together with the per-cout-bin count of contributing streamtubes and the zero-flow cout mask; a
11second turns those raw outputs into the final banded weights plus the mask of cout bins the
12spin-up policy rejects; a third reconstructs the dense ``(n_cout, n_cin)`` matrix from the banded
13layout; and a fourth validates the ``spinup`` argument and applies its input-side effect of
14prepending warm-start bins to tedges, flow and cin.
16Nonlinear sorption is not handled here -- the front-tracking solver lives in
17:mod:`gwtransport.fronttracking`.
19This file is part of gwtransport which is released under AGPL-3.0 license.
20See the ./LICENSE file or go to https://github.com/gwtransport/gwtransport/blob/main/LICENSE for full license details.
21"""
23import numpy as np
24import numpy.typing as npt
25import pandas as pd
27from gwtransport._time import tedges_to_days
28from gwtransport.utils import cumulative_flow_volume
30# Target number of (streamtube x cout-bin) pairs per tile in the banded weight build. The
31# per-tile working set is O(_WEIGHT_BUILD_BLOCK x band), so peak memory is bounded
32# independent of record length; chosen to keep the build under ~30 MB while spanning most
33# records in one or a few tiles. See _infiltration_to_extraction_weights.
34_WEIGHT_BUILD_BLOCK = 100_000
37def _infiltration_to_extraction_weights(
38 *,
39 tedges: pd.DatetimeIndex,
40 cout_tedges: pd.DatetimeIndex,
41 aquifer_pore_volumes: npt.NDArray[np.floating],
42 flow: npt.NDArray[np.floating],
43 retardation_factor: float,
44) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.intp], npt.NDArray[np.intp], npt.NDArray[np.bool_]]:
45 """
46 Compute raw streamtube-bundle weights for infiltration to extraction transformation.
48 Builds the per-cout-bin sum of streamtube-normalized overlap rows in a
49 **compact banded layout**, plus the count of contributing streamtubes per
50 cout bin and the zero-flow cout-bin mask. The caller converts these into
51 final weights; see :func:`_resolve_spinup_mask` and
52 :func:`_resolve_spinup_inputs`. Reconstruct the dense ``(n_cout, n_cin)``
53 matrix with :func:`_densify_weights`.
55 The per-streamtube weight is the literal mass-flux / water-flux ratio for
56 that streamtube and cout bin: each contributing row sums to 1 to ULP.
57 Equal-mass pore-volume bins from a gamma APVD discretization carry equal
58 flow at the outlet (steady-streamline assumption), so the bundle output is
59 the arithmetic mean over contributing streamtubes.
61 Pure advection (``D_m = 0``, ``alpha_L = 0``) is volume-stationary. Let
62 ``Vi`` be the cumulative throughflow volume at the cin edges and ``Vc`` the
63 same cumulative volume sampled at the cout edges. A streamtube of retarded
64 pore volume ``r = R * V_pore`` carries each cout edge back to the
65 infiltration time at cumulative volume ``Vc_edge - r``. The cout bin's
66 source window in infiltration time spans one cout bin's worth of volume, so
67 it overlaps only a few cin bins. The nonzeros of cout row ``k`` therefore
68 span only the residence-time spread of the APVD, ``[col_start[k],
69 col_start[k] + full_band)``, and are accumulated into one banded buffer over
70 a **pore-volume loop** -- ``O(n_cout * full_band)`` memory regardless of
71 record length or ``n_pv``. The overlap is the flow-weighted time overlap
72 normalized by the window's total in-range flow, so each contributing row
73 sums to 1 exactly: the exact flow-weighted overlap, not an approximation.
75 A streamtube whose source volume leaves ``[Vi[0], Vi[-1]]`` -- or a cout
76 edge outside the cin time range -- maps to NaN and is **dropped, not
77 clipped**: its whole window for that cout bin is discarded (mirroring the
78 dense build, where a single NaN residence-time edge poisons the row).
80 Parameters
81 ----------
82 tedges : pandas.DatetimeIndex
83 Time edges for infiltration bins.
84 cout_tedges : pandas.DatetimeIndex
85 Time edges for extraction bins.
86 aquifer_pore_volumes : array-like
87 Distribution of pore volumes [m³].
88 flow : array-like
89 Flow rate values [m³/day].
90 retardation_factor : float
91 Constant retardation factor.
93 Returns
94 -------
95 band_vals : numpy.ndarray
96 Sum over streamtubes of per-streamtube normalized overlap rows in
97 banded layout. Shape: (len(cout_tedges) - 1, full_band). Slot
98 ``band_vals[k, b]`` is the weight on cin bin ``col_start[k] + b``;
99 for a cout bin ``k`` with ``c`` contributing streamtubes the row
100 sums to ``c`` (not ``n_pv`` and not 1).
101 col_start : numpy.ndarray of int
102 Shape: (len(cout_tedges) - 1,). First cin bin index of each cout
103 row's band. Defaults to 0 for rows with no contributing streamtube.
104 contributing_bins : numpy.ndarray of int
105 Shape: (len(cout_tedges) - 1,). Number of streamtubes that
106 actually contributed to each cout bin (had a source window fully
107 inside the cin volume range).
108 zero_flow_cout : numpy.ndarray of bool
109 Shape: (len(cout_tedges) - 1,). True for cout bins with zero
110 time-averaged extraction flow over their interval.
112 See Also
113 --------
114 _densify_weights : Reconstruct the dense (n_cout, n_cin) matrix.
115 """
116 cin_tedges_days = tedges_to_days(tedges)
117 cout_tedges_days = tedges_to_days(cout_tedges, ref=tedges[0])
118 flow = np.asarray(flow, dtype=float)
120 n_cout = len(cout_tedges) - 1
121 n_cin = len(tedges) - 1
122 n_pv = len(aquifer_pore_volumes)
124 # Cumulative throughflow volume at cin edges (Vi) and sampled at cout edges (Vc).
125 vi = cumulative_flow_volume(flow, np.diff(cin_tedges_days))
126 vc = np.interp(cout_tedges_days, cin_tedges_days, vi)
128 # A cout bin's time-averaged extraction flow is zero exactly when no throughflow
129 # volume passes during the bin, i.e. its cumulative-volume width is zero. This is
130 # bit-identical to the dense flow-overlap-matrix product but is O(n_cout), not O(N^2).
131 zero_flow_cout = np.diff(vc) == 0
133 r = np.sort(np.asarray(aquifer_pore_volumes, dtype=float) * retardation_factor)
135 # Cumulative source volume at each cout edge. Out-of-range cout edges have no source
136 # volume; NaN there propagates so the adjacent windows are dropped (not clipped),
137 # matching the dense build.
138 edge_in_range = (cout_tedges_days >= cin_tedges_days[0]) & (cout_tedges_days <= cin_tedges_days[-1])
139 src_edge = np.where(edge_in_range, vc, np.nan)
141 # Tile the cout axis so the (n_pv, tile, band) overlap tensor never materialises for the
142 # whole record: peak memory stays O(block) regardless of record length, while each tile
143 # stays fully vectorised over its streamtubes. Most records span one or a few tiles. Each
144 # tile sizes its own band independently and is right-padded to the global width at the end.
145 block = max(1, _WEIGHT_BUILD_BLOCK // n_pv)
146 col_start = np.zeros(n_cout, dtype=np.intp)
147 contributing_bins = np.zeros(n_cout, dtype=np.intp)
148 tiles: list[tuple[int, npt.NDArray[np.floating]]] = []
149 full_band = 1
150 for start in range(0, n_cout, block):
151 stop = min(start + block, n_cout)
152 m = stop - start
154 # Back-project this tile's cout edges by every streamtube to infiltration time.
155 infil = np.interp(
156 (src_edge[start : stop + 1] - r[:, None]).ravel(), vi, cin_tedges_days, left=np.nan, right=np.nan
157 ).reshape(n_pv, m + 1)
158 win_lo, win_hi = infil[:, :-1], infil[:, 1:] # (n_pv, m) infiltration-time window per cout bin
159 contained = np.isfinite(win_lo) & np.isfinite(win_hi)
160 j_lo = np.clip(np.searchsorted(cin_tedges_days, win_lo, side="right") - 1, 0, n_cin - 1)
161 j_hi = np.clip(np.searchsorted(cin_tedges_days, win_hi, side="left"), 0, n_cin)
163 # Tile band = union of contained streamtube windows; per_band = widest single window.
164 any_contained = contained.any(axis=0)
165 tile_start = np.where(any_contained, np.where(contained, j_lo, n_cin).min(axis=0), 0)
166 row_hi = np.where(contained, j_hi, 0).max(axis=0)
167 tile_band = min(int(np.max(np.where(any_contained, row_hi - tile_start, 0))) + 1, n_cin)
168 per_band = min(int(np.max(np.where(contained, j_hi - j_lo, 0))) + 1, n_cin)
170 # Per-streamtube flow-weighted overlap on its own narrow window, normalised so each
171 # contributing row sums to 1, then scatter-summed over streamtubes into the tile's
172 # banded buffer (offset by each row's tile_start).
173 cols = j_lo[:, :, None] + np.arange(per_band)[None, None, :] # (n_pv, m, per_band)
174 in_cin = contained[:, :, None] & (cols < n_cin)
175 cols_clipped = np.clip(cols, 0, n_cin - 1)
176 overlap = np.maximum(
177 0.0,
178 np.minimum(win_hi[:, :, None], cin_tedges_days[cols_clipped + 1])
179 - np.maximum(win_lo[:, :, None], cin_tedges_days[cols_clipped]),
180 )
181 flux = np.where(in_cin, flow[cols_clipped] * overlap, 0.0)
182 window_flux = flux.sum(axis=2) # (n_pv, m) total in-range flow per window
183 contributes = contained & (window_flux > 0)
184 np.divide(flux, window_flux[:, :, None], out=flux, where=contributes[:, :, None])
186 # Only nonzero in-window slots scatter; their banded offset is < tile_band by
187 # construction (a contributing slot lies in [tile_start, row_hi)).
188 keep = flux > 0
189 offset = cols_clipped - tile_start[None, :, None]
190 row_idx = np.broadcast_to(np.arange(m)[None, :, None], cols.shape)
191 tile_vals = (
192 np
193 .bincount((row_idx * tile_band + offset)[keep], weights=flux[keep], minlength=m * tile_band)
194 .astype(float, copy=False)
195 .reshape(m, tile_band)
196 )
198 col_start[start:stop] = tile_start
199 contributing_bins[start:stop] = contributes.sum(axis=0)
200 tiles.append((start, tile_vals))
201 full_band = max(full_band, tile_band)
203 band_vals = np.zeros((n_cout, full_band))
204 for start, tile_vals in tiles:
205 band_vals[start : start + tile_vals.shape[0], : tile_vals.shape[1]] = tile_vals
207 return band_vals, col_start, contributing_bins, zero_flow_cout
210def _resolve_spinup_mask(
211 *,
212 band_vals: npt.NDArray[np.floating],
213 contributing_bins: npt.NDArray[np.intp],
214 zero_flow_cout: npt.NDArray[np.bool_],
215 n_pv: int,
216) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.bool_]]:
217 """Convert raw banded bundle outputs into final banded weights + invalid mask.
219 A cout bin is valid only when every streamtube has contributed and the bin
220 carries throughflow; its bundle row is then the arithmetic mean over the
221 ``n_pv`` streamtube rows and sums to 1, so the operator conserves cin → cout
222 mass. Rows failing either condition are zeroed and flagged invalid.
224 Parameters
225 ----------
226 band_vals : numpy.ndarray
227 Per-cout-bin sum of streamtube-normalized rows in banded layout from
228 :func:`_infiltration_to_extraction_weights`. Shape (n_cout, full_band).
229 contributing_bins : numpy.ndarray of int
230 Number of streamtubes that contributed to each cout bin.
231 zero_flow_cout : numpy.ndarray of bool
232 Mask of zero-extraction-flow cout bins.
233 n_pv : int
234 Total number of streamtubes (length of aquifer_pore_volumes), at least 1.
236 Returns
237 -------
238 weights : numpy.ndarray
239 Final banded weight matrix of the same shape as ``band_vals``.
240 Rows where the policy is not satisfied are zero.
241 invalid_mask : numpy.ndarray of bool
242 True for cout bins where the policy is not satisfied.
243 """
244 valid = (contributing_bins == n_pv) & ~zero_flow_cout
245 # Every valid row has contributing_bins == n_pv > 0, so the divisor is positive.
246 weights = np.zeros_like(band_vals)
247 weights[valid, :] = band_vals[valid, :] / contributing_bins[valid, None]
248 return weights, ~valid
251def _densify_weights(
252 band_vals: npt.NDArray[np.floating], col_start: npt.NDArray[np.intp], n_cin: int
253) -> npt.NDArray[np.floating]:
254 """Reconstruct the dense (n_cout, n_cin) weight matrix from the banded layout.
256 Inverse of the banded packing produced by
257 :func:`_infiltration_to_extraction_weights` and
258 :func:`_resolve_spinup_mask`: row ``k`` places ``band_vals[k]`` at cin
259 columns ``[col_start[k], col_start[k] + full_band)``, dropping band slots
260 that fall past ``n_cin`` (the right-edge padding).
262 Parameters
263 ----------
264 band_vals : numpy.ndarray
265 Banded weights of shape (n_cout, full_band).
266 col_start : numpy.ndarray of int
267 First cin bin index of each row's band, shape (n_cout,).
268 n_cin : int
269 Number of cin bins (dense column count).
271 Returns
272 -------
273 numpy.ndarray
274 Dense weight matrix of shape (n_cout, n_cin).
275 """
276 n_cout, full_band = band_vals.shape
277 dense = np.zeros((n_cout, n_cin))
278 cols = col_start[:, None] + np.arange(full_band)[None, :]
279 in_range = cols < n_cin
280 rows = np.broadcast_to(np.arange(n_cout)[:, None], cols.shape)
281 dense[rows[in_range], cols[in_range]] = band_vals[in_range]
282 return dense
285def _resolve_spinup_inputs(
286 spinup: object,
287 *,
288 tedges: pd.DatetimeIndex,
289 flow: npt.NDArray[np.floating],
290 aquifer_pore_volumes: npt.ArrayLike,
291 retardation_factor: float,
292 cin: npt.NDArray[np.floating] | None = None,
293) -> tuple[pd.DatetimeIndex, npt.NDArray[np.floating], npt.NDArray[np.floating] | None, None, int]:
294 """Validate ``spinup`` and apply its input-side effects.
296 Returns the (possibly padded) tedges, flow, and cin to use for
297 weight computation, and the number of bins prepended.
299 Modes:
301 - ``spinup is None`` — strict-validity. Returns inputs unchanged.
302 Mass-conserving but NaN where any streamtube has not broken through.
303 - ``spinup == "constant"`` — warm-start. Prepends ``n_pad`` bins
304 (each of width ``tedges[1] - tedges[0]``) so the total prepended
305 duration covers ``retardation_factor * max(aquifer_pore_volumes)
306 / flow[0]``. cin and flow are extended with their first observed
307 value (constant warm start). Strict-validity on the padded system
308 yields no spin-up NaN for cout bins at or after the original
309 ``tedges[0]``. Right-edge spin-up (cout extending past
310 ``tedges[-1]``) is not addressed.
312 The first bin width ``tedges[1] - tedges[0]`` is used as the unit
313 for prepended bins; this preserves uniformity if the input
314 ``tedges`` are uniform (required by the smooth-then-advect path of
315 ``diffusion_fast``).
317 Parameters
318 ----------
319 spinup : None or "constant"
320 Public spin-up policy.
321 tedges : pandas.DatetimeIndex
322 Original cin/flow time edges (length n_cin + 1).
323 flow : numpy.ndarray
324 Flow values; ``flow[0]`` sets the warm-start flow.
325 aquifer_pore_volumes : array-like
326 Pore volumes; the maximum sets the warm-start duration.
327 retardation_factor : float
328 Retardation factor.
329 cin : numpy.ndarray, optional
330 Concentration values to prepend with ``cin[0]``. Pass ``None``
331 for the inverse direction (cin is the unknown to recover); the
332 returned ``new_cin`` is then ``None``.
334 Returns
335 -------
336 new_tedges : pandas.DatetimeIndex
337 Tedges to pass to :func:`_infiltration_to_extraction_weights`.
338 Length is ``len(tedges) + n_pad``.
339 new_flow : numpy.ndarray
340 Flow values aligned with ``new_tedges``; length ``n_cin + n_pad``.
341 new_cin : numpy.ndarray or None
342 cin values aligned with ``new_tedges`` if ``cin`` was provided;
343 otherwise ``None``.
344 threshold : None
345 Always ``None``: both policies are resolved on the input side and
346 :func:`_resolve_spinup_mask` needs no threshold.
347 n_pad : int
348 Number of bins prepended. ``0`` for the non-padding mode. Callers
349 of the inverse direction must drop the first ``n_pad`` entries
350 from the recovered cin to align with the original ``tedges``.
352 Raises
353 ------
354 TypeError
355 If ``spinup`` is neither ``None`` nor a string.
356 ValueError
357 If ``spinup`` is a string other than ``"constant"``.
358 """
359 if spinup is None:
360 return tedges, np.asarray(flow, dtype=float), cin, None, 0
361 if not isinstance(spinup, str):
362 msg = f"spinup must be None or 'constant'; got {spinup!r}"
363 raise TypeError(msg)
364 if spinup != "constant":
365 msg = f"spinup string must be 'constant'; got {spinup!r}"
366 raise ValueError(msg)
367 flow_arr = np.asarray(flow, dtype=float)
368 # Determine whether padding is feasible. We fall back to strict-validity
369 # (no padding) silently when the warm-start is undefined (zero or NaN
370 # initial flow) or when the implied padding would be unreasonably large
371 # (extreme pore volumes). This keeps the default usable for edge cases
372 # while still triggering the strict-validity NaN when the warm-start
373 # assumption cannot meaningfully be applied.
374 if len(tedges) < 2: # noqa: PLR2004
375 return tedges, flow_arr, cin, None, 0
376 q0 = float(flow_arr[0])
377 v_max = float(np.max(np.asarray(aquifer_pore_volumes, dtype=float)))
378 if not (q0 > 0 and v_max > 0):
379 return tedges, flow_arr, cin, None, 0
380 bin_width = tedges[1] - tedges[0]
381 bin_width_days = bin_width / pd.Timedelta(days=1)
382 if not bin_width_days > 0:
383 return tedges, flow_arr, cin, None, 0
384 pad_days = retardation_factor * v_max / q0
385 # Add 1 extra bin so the longest streamtube's source window for the
386 # earliest original cout bin lies strictly inside the padded range
387 # (avoids strict-validity NaN due to floating-point edge alignment).
388 n_pad_float = np.ceil(pad_days / bin_width_days) + 1
389 # Cap to keep memory bounded; beyond this, "constant" is not a meaningful
390 # warm-start (the user probably has unphysical pore volumes or extreme
391 # retardation), so fall through to strict-validity.
392 max_n_pad = max(10_000, 10 * len(flow_arr))
393 if not np.isfinite(n_pad_float) or n_pad_float > max_n_pad:
394 return tedges, flow_arr, cin, None, 0
395 n_pad = int(n_pad_float)
396 offsets = pd.TimedeltaIndex(bin_width * np.arange(n_pad, 0, -1))
397 new_tedges = (tedges[0] - offsets).append(tedges)
398 new_flow = np.concatenate([np.full(n_pad, flow_arr[0]), flow_arr])
399 new_cin = np.concatenate([np.full(n_pad, cin[0]), cin]) if cin is not None else None
400 return new_tedges, new_flow, new_cin, None, n_pad