Coverage for src/gwtransport/utils.py: 90%
436 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-07 21:17 +0000
1"""
2General Utilities for 1D Groundwater Transport Modeling.
4This module provides general-purpose utility functions for time series manipulation,
5interpolation, numerical operations, and data processing used throughout the gwtransport
6package. Functions include linear interpolation/averaging, bin overlap calculations,
7underdetermined system solvers, and external data retrieval.
9Available functions:
11- :func:`step_plot_coords` - Compute step-plot coordinates from bin edges and
12 bin-averaged values. Returns paired x/y arrays for plotting piecewise-constant
13 functions with ``ax.plot(x, y)``.
15- ``_make_strictly_monotone`` (private) - Bump consecutive duplicates in a non-decreasing
16 array by ``k * ulp(max)`` so it becomes strictly monotone. Used before V → t inversions to
17 prevent ``np.interp`` from silently picking one limit at plateau levels.
19- :func:`cumulative_flow_volume` - Cumulative infiltrated/extracted volume from per-bin flow
20 rates and bin widths, prepended with a leading zero. Optionally bumped to strict
21 monotonicity for V → t inversions.
23- :func:`linear_interpolate` - Linear interpolation using numpy's optimized interp function.
24 Automatically handles unsorted data with configurable extrapolation (None for clamping,
25 float for constant values). Handles multi-dimensional query arrays.
27- :func:`linear_average` - Compute average values of piecewise linear time series between
28 specified x-edges. Supports 1D or 2D edge arrays for batch processing. Handles NaN values
29 and offers multiple extrapolation methods ('nan', 'outer', 'raise').
31- :func:`time_bin_overlap` - Calculate fraction of time bins overlapping with specified time
32 ranges. Similar to partial_isin but for time-based bin overlaps with list of (start, end)
33 tuples.
35- :func:`simplify_bins` - Simplify a piecewise-constant time series by merging adjacent bins
36 whose values are within a tolerance. Uses volume-weighted (flow x width) averaging when
37 flow is provided, otherwise width-weighted. Direction-independent via largest-jump splitting.
39- :func:`compute_time_edges` - Compute DatetimeIndex of time bin edges from explicit edges,
40 start times, or end times. Validates consistency with expected number of bins and handles
41 uniform spacing extrapolation.
43The inverse solvers below are two intentionally coexisting families: a Tikhonov family (the dense
44:func:`solve_inverse_transport` and its banded equivalent :func:`solve_inverse_transport_banded`,
45both fed by :func:`compute_reverse_target` and built on :func:`solve_tikhonov`) for the
46overdetermined deconvolution in advection/diffusion, and a separate nullspace solver
47(:func:`solve_underdetermined_system`) for the underdetermined deposition inverse.
49- :func:`solve_tikhonov` - Solve linear system with Tikhonov regularization toward a target.
50 Well-determined modes follow the data; poorly-determined modes are pulled toward the target.
52- :func:`compute_reverse_target` - Build the regularization target for the inverse problem by
53 transposing and row-normalizing the forward coefficient matrix. Consumed by
54 :func:`solve_tikhonov` and :func:`solve_inverse_transport`.
56- :func:`solve_inverse_transport` - Solve the inverse transport problem (deconvolution) via
57 Tikhonov regularization. Shared by advection, diffusion, and diffusion_fast
58 ``extraction_to_infiltration`` functions.
60- :func:`solve_inverse_transport_banded` - Memory-light banded equivalent of
61 :func:`solve_inverse_transport` for a forward operator stored in banded layout. Assembles the
62 Tikhonov normal equations directly in banded form and solves them via banded Cholesky.
64- :func:`solve_underdetermined_system` - Solve underdetermined linear system (Ax = b, m < n)
65 with nullspace regularization. Handles NaN values by row exclusion. Supports built-in
66 objectives ('squared_differences', 'summed_differences') or custom callable objectives.
67 Used by :mod:`gwtransport.deposition`.
69- :func:`get_soil_temperature` - Download soil temperature data from KNMI weather stations with
70 automatic caching. Supports stations 260 (De Bilt), 273 (Marknesse), 286 (Nieuw Beerta),
71 323 (Wilhelminadorp). Returns DataFrame with columns TB1-TB5, TNB1-TNB2, TXB1-TXB2 at various
72 depths. Daily cache prevents redundant downloads.
74- ``_generate_failed_coverage_badge`` (private) - Generate SVG badge indicating failed coverage
75 using genbadge library. Used in CI/CD workflows.
77This file is part of gwtransport which is released under AGPL-3.0 license.
78See the ./LICENSE file or go to https://github.com/gwtransport/gwtransport/blob/main/LICENSE for full license details.
79"""
81from __future__ import annotations
83import io
84import warnings
85from collections.abc import Callable
86from datetime import date
87from pathlib import Path
89import numpy as np
90import numpy.typing as npt
91import pandas as pd
92from scipy.linalg import cho_solve_banded, cholesky_banded, null_space
93from scipy.optimize import minimize
95cache_dir = Path(__file__).parent.parent.parent / "cache"
98def step_plot_coords(edges: npt.ArrayLike, values: npt.ArrayLike) -> tuple[npt.NDArray, npt.NDArray]:
99 """Compute step-plot coordinates from bin edges and bin-averaged values.
101 Converts bin edges (n+1) and bin values (n) into paired x/y arrays
102 suitable for plotting piecewise-constant (step) functions with
103 ``ax.plot(x, y)``.
105 Parameters
106 ----------
107 edges : array-like
108 Bin edges (n+1 elements for n bins). Can be numeric, datetime, or
109 any type accepted by :func:`numpy.repeat`.
110 values : array-like
111 Bin-averaged values (n elements), one per bin.
113 Returns
114 -------
115 x : ndarray
116 Step x-coordinates (2n elements). Same dtype as *edges*.
117 y : ndarray
118 Step y-coordinates (2n elements). Same dtype as *values*.
120 Examples
121 --------
122 >>> import numpy as np
123 >>> edges = np.array([0.0, 1.0, 3.0, 6.0])
124 >>> values = np.array([2.0, 5.0, 1.0])
125 >>> x, y = step_plot_coords(edges, values)
126 >>> x
127 array([0., 1., 1., 3., 3., 6.])
128 >>> y
129 array([2., 2., 5., 5., 1., 1.])
130 """
131 x = np.repeat(edges, 2)[1:-1]
132 y = np.repeat(values, 2)
133 return x, y
136_DUP_BUMP_ULPS = 16 # safety factor in ulps; see _make_strictly_monotone docstring
139def _make_strictly_monotone(arr: npt.ArrayLike) -> npt.NDArray[np.floating]:
140 """Bump consecutive duplicates so a non-decreasing array becomes strictly monotone.
142 Returns the input unchanged if no consecutive duplicates are present. Otherwise returns a
143 new array with each duplicate bumped up by ``k * step``, where ``k`` is its 1-based
144 position within the consecutive duplicate run and ``step`` is ``16 * ulp(max(arr))``
145 capped per run so the largest bump stays strictly below the next genuine value above the
146 plateau (``step = min(16 * ulp(max(arr)), gap / (run_len + 1))``). The cap prevents a long
147 run from overshooting a closely-spaced successor; a gap narrower than the run length in
148 ulps is unrepresentable and cannot be separated.
150 The factor of 16 is a safety margin against IEEE 754 rounding noise in ``np.interp``'s
151 linear-interpolation arithmetic, which differs subtly between Linux x86_64 (with FMA)
152 and ARM macOS. A 1-ulp gap, while strictly monotone, can place a downstream query value
153 on the wrong side of a bracket boundary if the intermediate arithmetic rounds 1 ulp away
154 from the exact value. 16 ulps ensures the bracket selection is unambiguous on every
155 platform we support. The perturbation is relative to the array scale:
156 ``bump ≈ 16 * ulp(max(arr)) ≈ 3.5e-15 * max(arr)``, i.e. about 15 significant digits
157 below the data scale and well below physical relevance. The absolute size therefore grows
158 with the cumulative-volume magnitude (e.g. ``~1e-13`` only for ``max(arr) ~ 30``).
160 Parameters
161 ----------
162 arr : array-like
163 1D non-decreasing array (e.g., a cumulative volume sequence ``flow_cum`` that contains
164 plateaus from ``Q = 0`` bins).
166 Returns
167 -------
168 ndarray
169 Strictly monotone array of the same length.
171 Notes
172 -----
173 Use this before passing ``arr`` as ``x_ref`` to a ``V → t`` inversion via
174 :func:`linear_interpolate` or :func:`numpy.interp`. Plateaus in ``arr`` make ``arr⁻¹``
175 multi-valued, and ``np.interp`` would silently pick one of the two limits, biasing
176 integrals over output bins that span the kink.
177 """
178 arr = np.asarray(arr, dtype=float)
179 diffs = np.diff(arr)
180 if not np.any(diffs == 0):
181 return arr
182 ulp_max = np.nextafter(arr.max(), np.inf) - arr.max()
183 n = len(arr)
184 idx = np.arange(n)
185 is_dup = np.concatenate(([False], diffs == 0))
186 # 1-based position of each duplicate within its consecutive run.
187 last_nondup = np.maximum.accumulate(np.where(is_dup, -1, idx))
188 cumcount = np.where(is_dup, idx - last_nondup, 0)
190 # Per-run headroom: each bumped value must stay strictly below the next genuine
191 # (non-duplicate) value above the plateau, otherwise a long run can overshoot a
192 # closely-spaced next value and break monotonicity. ``next_nondup_idx`` is the first
193 # non-duplicate index after each position (``n`` when the run reaches the array end, where
194 # there is no successor and hence no overshoot risk). The gap to that successor caps the
195 # bump step so the last (largest) bump in a run of length L is at most ``L/(L+1)`` of the
196 # gap. A gap narrower than the run length in ulps is unrepresentable and cannot be split.
197 next_nondup_idx = np.minimum.accumulate(np.where(is_dup, n, idx)[::-1])[::-1]
198 has_successor = next_nondup_idx < n
199 gap_to_next = arr[np.clip(next_nondup_idx, 0, n - 1)] - arr[idx]
200 run_len = next_nondup_idx - last_nondup - 1
201 full_step = _DUP_BUMP_ULPS * ulp_max
202 with np.errstate(invalid="ignore", divide="ignore"):
203 capped_step = np.where(has_successor, np.minimum(full_step, gap_to_next / (run_len + 1.0)), full_step)
204 bump = np.where(is_dup, cumcount * capped_step, 0.0)
205 return arr + bump
208def cumulative_flow_volume(
209 flow: npt.ArrayLike, dt_days: npt.ArrayLike, *, strictly_monotone: bool = False
210) -> npt.NDArray[np.floating]:
211 """Cumulative infiltrated/extracted volume from per-bin flow rates.
213 Multiplies each per-bin flow rate by its bin width and accumulates, with a
214 leading zero prepended so the result has one entry per bin edge (n+1 values
215 for n bins). The result is the cumulative volume ``V`` at each time edge.
217 Parameters
218 ----------
219 flow : array-like
220 Flow rate per bin (m³/day), length n.
221 dt_days : array-like
222 Bin widths in days, length n (e.g. ``numpy.diff`` of edge days).
223 strictly_monotone : bool, optional
224 When ``True``, bump consecutive duplicates (plateaus from ``Q = 0``
225 bins) via ``_make_strictly_monotone`` so the cumulative volume is
226 strictly increasing. Required before a V → t inversion; leave ``False``
227 when the plateaus must be preserved. Default is ``False``.
229 Returns
230 -------
231 ndarray
232 Cumulative volume at each edge (length ``len(flow) + 1``), starting at
233 zero.
235 See Also
236 --------
237 ``_make_strictly_monotone`` : Bump duplicates before V → t inversion.
238 """
239 flow_cum = np.concatenate(([0.0], np.cumsum(np.asarray(flow) * np.asarray(dt_days))))
240 return _make_strictly_monotone(flow_cum) if strictly_monotone else flow_cum
243def linear_interpolate(
244 *,
245 x_ref: npt.ArrayLike,
246 y_ref: npt.ArrayLike,
247 x_query: npt.ArrayLike,
248 left: float | None = None,
249 right: float | None = None,
250) -> npt.NDArray[np.floating]:
251 """
252 Linear interpolation using numpy's optimized interp function.
254 Automatically handles unsorted reference data by sorting it first.
256 Parameters
257 ----------
258 x_ref : array-like
259 Reference x-values. If unsorted, will be automatically sorted.
260 y_ref : array-like
261 Reference y-values corresponding to x_ref.
262 x_query : array-like
263 Query x-values where interpolation is needed. Array may have any shape.
264 left : float, optional
265 Value to return for x_query < x_ref[0].
267 - If ``left=None``: clamp to y_ref[0] (default)
268 - If ``left=float``: use specified value (e.g., ``np.nan``)
270 right : float, optional
271 Value to return for x_query > x_ref[-1].
273 - If ``right=None``: clamp to y_ref[-1] (default)
274 - If ``right=float``: use specified value (e.g., ``np.nan``)
276 Returns
277 -------
278 ndarray
279 Interpolated y-values with the same shape as x_query.
281 Examples
282 --------
283 Basic interpolation with clamping (default):
285 >>> import numpy as np
286 >>> from gwtransport.utils import linear_interpolate
287 >>> x_ref = np.array([1.0, 2.0, 3.0, 4.0])
288 >>> y_ref = np.array([10.0, 20.0, 30.0, 40.0])
289 >>> x_query = np.array([0.5, 1.5, 2.5, 3.5, 4.5])
290 >>> linear_interpolate(x_ref=x_ref, y_ref=y_ref, x_query=x_query)
291 array([10., 15., 25., 35., 40.])
293 Using NaN for extrapolation:
295 >>> linear_interpolate(
296 ... x_ref=x_ref, y_ref=y_ref, x_query=x_query, left=np.nan, right=np.nan
297 ... )
298 array([nan, 15., 25., 35., nan])
300 Handles unsorted reference data automatically:
302 >>> x_unsorted = np.array([3.0, 1.0, 4.0, 2.0])
303 >>> y_unsorted = np.array([30.0, 10.0, 40.0, 20.0])
304 >>> linear_interpolate(x_ref=x_unsorted, y_ref=y_unsorted, x_query=x_query)
305 array([10., 15., 25., 35., 40.])
306 """
307 x_ref = np.asarray(x_ref)
308 y_ref = np.asarray(y_ref)
309 x_query = np.asarray(x_query)
311 sort_idx = np.argsort(x_ref)
312 x_ref_sorted = x_ref[sort_idx]
313 y_ref_sorted = y_ref[sort_idx]
315 return np.interp(x_query, x_ref_sorted, y_ref_sorted, left=left, right=right)
318def linear_average(
319 *,
320 x_data: npt.ArrayLike,
321 y_data: npt.ArrayLike,
322 x_edges: npt.ArrayLike,
323 extrapolate_method: str = "nan",
324) -> npt.NDArray[np.floating]:
325 """
326 Compute the average value of a piecewise linear time series between specified x-edges.
328 Parameters
329 ----------
330 x_data : array-like
331 x-coordinates of the time series data points, must be in ascending order.
332 y_data : array-like
333 y-coordinates of the time series data points. Can be 1D or 2D.
335 - If 1D: shape ``(n_data,)`` -- a single series.
336 - If 2D: shape ``(n_series_y, n_data)`` -- multiple series sharing the same
337 ``x_data``. The leading axis is averaged independently per row. Cannot be
338 combined with 2D ``x_edges`` (each row of ``x_edges`` and each row of
339 ``y_data`` would otherwise have to broadcast against each other, which is
340 not supported).
341 x_edges : array-like
342 x-coordinates of the integration edges.
344 - If 1D: shape ``(n_edges,)``, must be in ascending order.
345 - If 2D: shape ``(n_series_x, n_edges)``, each row must be in ascending order.
346 extrapolate_method : str, optional
347 Method for handling bin edges that fall outside ``x_data``. Default
348 is ``'nan'``.
350 - ``'outer'``: average over the **in-range** portion of each bin
351 (clip-then-average). The bin width used for normalisation is the
352 clipped width, not the original width. For example,
353 ``x_data = y_data = [1, 2, 3]`` and ``x_edges = [0, 5]`` returns
354 ``2.0`` (integral over ``[1, 3]`` divided by clipped width 2),
355 **not** ``2.2`` (which a constant-extension scheme would give).
356 - ``'nan'``: bins that extend outside ``x_data`` are returned as ``nan``.
357 - ``'raise'``: raise an error if any bin edge falls outside ``x_data``.
359 Returns
360 -------
361 ndarray
362 2D array of average values between consecutive pairs of x_edges.
363 Shape is ``(n_series, n_bins)`` where ``n_bins = n_edges - 1`` and
364 ``n_series = max(n_series_x, n_series_y)``. Both ``x_edges`` and ``y_data``
365 being 1D yields ``n_series = 1``.
367 Raises
368 ------
369 ValueError
370 If ``x_edges`` is not 1D or 2D. If ``y_data`` is not 1D or 2D. If both
371 ``x_edges`` and ``y_data`` are 2D. If ``x_data`` and ``y_data`` have
372 incompatible shapes or are empty. If ``x_edges`` has fewer than 2 values per
373 row. If ``x_data`` is not in ascending order. If ``x_edges`` rows are not in
374 ascending order. If ``extrapolate_method`` is ``'raise'`` and any edge falls
375 outside the data range.
377 Notes
378 -----
379 **NaN handling is asymmetric between 1D and 2D ``y_data``.**
381 - 1D ``y_data`` is treated as a single series; internal NaN gaps are
382 silently bridged by linear interpolation across the gap (via
383 ``np.interp`` with ``left=nan, right=nan``).
384 - 2D ``y_data`` is treated row-wise; any output bin whose
385 ``[edge_left, edge_right]`` touches a NaN segment **in that row** is
386 set to NaN, while other rows are unaffected.
388 Callers that need NaN-bridging behaviour across multiple series must
389 pre-fill (e.g., ``pd.DataFrame.interpolate``) before calling.
391 Examples
392 --------
393 >>> import numpy as np
394 >>> from gwtransport.utils import linear_average
395 >>> x_data = [0, 1, 2, 3]
396 >>> y_data = [0, 1, 1, 0]
397 >>> x_edges = [0, 1.5, 3]
398 >>> linear_average(
399 ... x_data=x_data, y_data=y_data, x_edges=x_edges
400 ... ) # doctest: +ELLIPSIS
401 array([[0.666..., 0.666...]])
403 >>> x_edges_2d = [[0, 1.5, 3], [0.5, 2, 3]]
404 >>> linear_average(x_data=x_data, y_data=y_data, x_edges=x_edges_2d)
405 array([[0.66666667, 0.66666667],
406 [0.91666667, 0.5 ]])
408 Multiple y-series with shared x_data and x_edges:
410 >>> y_data_2d = [[0, 1, 1, 0], [0, 2, 2, 0]]
411 >>> linear_average(x_data=x_data, y_data=y_data_2d, x_edges=x_edges)
412 array([[0.66666667, 0.66666667],
413 [1.33333333, 1.33333333]])
414 """
415 # Convert inputs to numpy arrays
416 x_data = np.asarray(x_data, dtype=float)
417 y_data = np.asarray(y_data, dtype=float)
418 x_edges = np.asarray(x_edges, dtype=float)
420 # Ensure x_edges is always 2D
421 if x_edges.ndim == 1:
422 x_edges = x_edges[np.newaxis, :]
423 elif x_edges.ndim != 2: # noqa: PLR2004
424 msg = "x_edges must be 1D or 2D array"
425 raise ValueError(msg)
427 # Ensure y_data is always 2D internally with shape (n_series_y, n_data)
428 if y_data.ndim == 1:
429 y_data = y_data[np.newaxis, :]
430 elif y_data.ndim != 2: # noqa: PLR2004
431 msg = "y_data must be 1D or 2D array"
432 raise ValueError(msg)
434 # 2D y_data requires 1D x_edges (no per-row x_edges allowed). The combination would
435 # require an outer product over (n_series_x, n_series_y), which is intentionally
436 # not supported -- callers can loop or stack instead.
437 n_series_x = x_edges.shape[0]
438 n_series_y = y_data.shape[0]
439 if n_series_x > 1 and n_series_y > 1:
440 msg = "Cannot combine 2D x_edges with 2D y_data"
441 raise ValueError(msg)
442 n_series = max(n_series_x, n_series_y)
444 # Input validation
445 if y_data.shape[1] != x_data.shape[0] or x_data.shape[0] == 0:
446 msg = "x_data and y_data must have the same length and be non-empty"
447 raise ValueError(msg)
448 if x_edges.shape[1] < 2: # noqa: PLR2004
449 msg = "x_edges must contain at least 2 values in each row"
450 raise ValueError(msg)
451 if not np.all(np.diff(x_data) >= 0):
452 msg = "x_data must be in ascending order"
453 raise ValueError(msg)
454 if not np.all(np.diff(x_edges, axis=1) >= 0):
455 msg = "x_edges must be in ascending order along each row"
456 raise ValueError(msg)
458 # Filter out NaN values. With 2D y_data, a column is dropped only when all rows
459 # have NaN there; per-row NaNs are handled via segment masking below so that one
460 # series' NaNs do not contaminate the others.
461 x_nan = np.isnan(x_data)
462 y_any_finite = np.any(~np.isnan(y_data), axis=0)
463 show = ~x_nan & y_any_finite
464 if show.sum() < 2: # noqa: PLR2004
465 if show.sum() == 1 and extrapolate_method == "outer":
466 # For a single retained data point with outer extrapolation, use the
467 # row-wise value broadcast across all output bins.
468 constant_value = y_data[:, show][:, 0] # shape (n_series_y,)
469 return np.broadcast_to(constant_value[:, None], (n_series, x_edges.shape[1] - 1)).astype(
470 np.float64, copy=True
471 )
472 return np.full(shape=(n_series, x_edges.shape[1] - 1), fill_value=np.nan)
474 x_data_clean = x_data[show]
475 y_data_clean = y_data[:, show] # shape (n_series_y, n_clean)
477 # Handle extrapolation for all series at once (vectorized). The 'raise' and 'nan'
478 # branches never mutate edges_processed, so they alias x_edges directly; 'outer'
479 # produces a fresh clipped array.
480 if extrapolate_method == "outer":
481 edges_processed = np.clip(x_edges, x_data_clean[0], x_data_clean[-1])
482 elif extrapolate_method == "raise":
483 if np.any(x_edges < x_data_clean[0]) or np.any(x_edges > x_data_clean[-1]):
484 msg = "x_edges must be within the range of x_data"
485 raise ValueError(msg)
486 edges_processed = x_edges
487 else: # nan method
488 edges_processed = x_edges
490 # Create a combined grid of all unique x points (data + all edges)
491 all_unique_x = np.unique(np.concatenate([x_data_clean, edges_processed.ravel()]))
493 # Interpolate y values at all unique x points once. For 2D y_data we vectorize
494 # the linear interpolation manually since np.interp does not accept 2D y.
495 if n_series_y == 1:
496 all_unique_y_result = np.interp(all_unique_x, x_data_clean, y_data_clean[0], left=np.nan, right=np.nan)
497 all_unique_y: npt.NDArray[np.floating] = np.asarray(all_unique_y_result, dtype=np.float64)[np.newaxis, :]
498 else:
499 # Locate each query x in x_data_clean. For x within the data range, idx is in
500 # [1, len(x_data_clean) - 1] so left_idx = idx - 1 is the bracketing left index.
501 idx = np.searchsorted(x_data_clean, all_unique_x).clip(1, len(x_data_clean) - 1)
502 left_idx = idx - 1
503 right_idx = idx
504 x_left = x_data_clean[left_idx]
505 x_right = x_data_clean[right_idx]
506 denom = x_right - x_left
507 # Detect query points coincident with an x_data point. Handling them via a
508 # direct lookup avoids the IEEE 754 trap where NaN * 0 = NaN, which would
509 # otherwise contaminate exact-endpoint queries adjacent to a NaN sample.
510 on_left_node = denom == 0 # only happens if x_left == x_right (duplicate)
511 weights = np.where(on_left_node, 0.0, (all_unique_x - x_left) / np.where(on_left_node, 1.0, denom))
512 all_unique_y = y_data_clean[:, left_idx] * (1.0 - weights) + y_data_clean[:, right_idx] * weights
513 # Override at exact x_data positions to avoid NaN * 0 contamination.
514 is_left_match = all_unique_x == x_left
515 is_right_match = all_unique_x == x_right
516 all_unique_y[:, is_left_match] = y_data_clean[:, left_idx[is_left_match]]
517 all_unique_y[:, is_right_match] = y_data_clean[:, right_idx[is_right_match]]
518 # Mark out-of-range query points as NaN (matches np.interp(left=nan, right=nan)).
519 out_of_range = (all_unique_x < x_data_clean[0]) | (all_unique_x > x_data_clean[-1])
520 all_unique_y[:, out_of_range] = np.nan
522 # Compute cumulative integrals once using trapezoidal rule.
523 # Segments outside the data range carry NaN (from the interp step with left/right=NaN);
524 # those NaNs will be masked out later via the bin-range check, so we suppress
525 # them here only to keep the cumulative sum finite for in-range bins.
526 dx = np.diff(all_unique_x)
527 y_avg = (all_unique_y[:, :-1] + all_unique_y[:, 1:]) / 2
528 segment_integrals = np.where(np.isnan(y_avg), 0.0, dx[np.newaxis, :] * y_avg)
529 # Cumulative integral with leading 0 along the x axis.
530 cumulative_integral = np.concatenate([np.zeros((y_avg.shape[0], 1)), np.cumsum(segment_integrals, axis=1)], axis=1)
532 # Vectorized computation for all series
533 # Find indices of all edges in the combined grid
534 edge_indices_result = np.searchsorted(all_unique_x, edges_processed)
535 # Ensure it's a 2D array for type checker
536 edge_indices: npt.NDArray[np.intp] = np.asarray(edge_indices_result, dtype=np.intp).reshape(edges_processed.shape)
538 # Compute integral between consecutive edges. Broadcast over n_series via the leading axis
539 # of cumulative_integral. edge_indices is (n_series_x, n_bins+1); cumulative_integral is
540 # (n_series_y, n_unique_x). We rely on n_series_x == 1 or n_series_y == 1 (enforced above).
541 integral_values = cumulative_integral[:, edge_indices[:, 1:]] - cumulative_integral[:, edge_indices[:, :-1]]
542 # integral_values has shape (n_series_y, n_series_x, n_bins). Squeeze the singleton.
543 integral_values_2d = integral_values[0] if n_series_y == 1 else integral_values[:, 0, :]
545 # Compute widths between consecutive edges for all series (vectorized)
546 edge_widths = np.diff(edges_processed, axis=1) # shape (n_series_x, n_bins)
547 # Broadcast widths to match (n_series, n_bins)
548 edge_widths_b = np.broadcast_to(edge_widths, (n_series, edge_widths.shape[1])) if n_series_y > 1 else edge_widths
550 # Handle zero-width intervals (vectorized)
551 zero_width_mask = edge_widths_b == 0
552 result = np.zeros_like(edge_widths_b, dtype=np.float64)
554 # For non-zero width intervals, compute average = integral / width (vectorized)
555 non_zero_mask = ~zero_width_mask
556 result[non_zero_mask] = integral_values_2d[non_zero_mask] / edge_widths_b[non_zero_mask]
558 # For zero-width intervals, interpolate y-value directly (vectorized)
559 if np.any(zero_width_mask):
560 # Positions where zero width occurs; use the left edge's x position.
561 if n_series_y == 1:
562 zero_positions = edges_processed[:, :-1][zero_width_mask] # 1D
563 result[zero_width_mask] = np.interp(zero_positions, x_data_clean, y_data_clean[0])
564 else:
565 # zero_width_mask has shape (n_series_y, n_bins); positions vary per row.
566 # edges_processed is (1, n_bins+1) here since n_series_x == 1.
567 edges_left = np.broadcast_to(edges_processed[:, :-1], (n_series, edge_widths.shape[1]))
568 zero_positions = edges_left[zero_width_mask]
569 # Interpolate per series using the same x_data_clean. Find bracketing indices
570 # for each zero-width position, then index into the appropriate y row.
571 # Get the row index for each zero-width entry.
572 row_idx_grid = np.broadcast_to(np.arange(n_series)[:, None], (n_series, edge_widths.shape[1]))
573 zero_rows = row_idx_grid[zero_width_mask]
574 idx_z = np.searchsorted(x_data_clean, zero_positions).clip(1, len(x_data_clean) - 1)
575 xl = x_data_clean[idx_z - 1]
576 xr = x_data_clean[idx_z]
577 denom_z = np.where(xr == xl, 1.0, xr - xl)
578 w_z = (zero_positions - xl) / denom_z
579 yl = y_data_clean[zero_rows, idx_z - 1]
580 yr = y_data_clean[zero_rows, idx_z]
581 result[zero_width_mask] = yl * (1.0 - w_z) + yr * w_z
583 # Handle extrapolation when 'nan' method is used (vectorized).
584 # Bins must lie entirely within the data range; bins partially outside
585 # (straddling) are also set to NaN, since the integral over the missing
586 # portion is undefined and dividing by the full bin width would bias the
587 # average low. Bins fully outside are likewise NaN.
588 if extrapolate_method == "nan":
589 bins_within_range = (x_edges[:, :-1] >= x_data_clean[0]) & (x_edges[:, 1:] <= x_data_clean[-1])
590 if n_series_y > 1:
591 bins_within_range = np.broadcast_to(bins_within_range, (n_series, bins_within_range.shape[1]))
592 result[~bins_within_range] = np.nan
594 # With 2D y_data, propagate per-row NaNs from the y series itself: any output bin that
595 # touches an x_data segment with NaN y in this row must be NaN. This 2-D NaN contract is
596 # method-independent -- it also holds for 'outer'/'raise', which would otherwise return a
597 # silently wrong finite average (the NaN trapezoids were zeroed above). Per-row NaN info is
598 # preserved in y_avg; mark bins whose spanned segments contain a NaN segment for this row.
599 if n_series_y > 1:
600 seg_nan = np.isnan(y_avg) # shape (n_series_y, n_unique_x - 1)
601 seg_nan_cum = np.concatenate([np.zeros((n_series_y, 1)), np.cumsum(seg_nan, axis=1)], axis=1)
602 nan_count_per_bin = seg_nan_cum[:, edge_indices[0, 1:]] - seg_nan_cum[:, edge_indices[0, :-1]]
603 result[nan_count_per_bin > 0] = np.nan
605 return result
608def time_bin_overlap(*, tedges: npt.ArrayLike, bin_tedges: list[tuple]) -> npt.NDArray[np.floating]:
609 """
610 Calculate the fraction of each time bin that overlaps with each time range.
612 This function computes an array where element (i, j) represents the fraction
613 of time bin j that overlaps with time range i. The computation uses
614 vectorized operations to avoid loops.
616 Parameters
617 ----------
618 tedges : array-like
619 1D array of time bin edges in ascending order. For n bins, there
620 should be n+1 edges.
621 bin_tedges : list of tuple
622 List of tuples where each tuple contains ``(start_time, end_time)``
623 defining a time range.
625 Returns
626 -------
627 overlap_array : ndarray
628 Array of shape (len(bin_tedges), n_bins) where n_bins is the number of
629 time bins. Each element (i, j) represents the fraction of time bin j
630 that overlaps with time range i. Values range from 0 (no overlap) to
631 1 (complete overlap).
633 Raises
634 ------
635 ValueError
636 If ``tedges`` is not a 1D array, has fewer than 2 elements, or if
637 ``bin_tedges`` is empty.
639 Notes
640 -----
641 - tedges must be sorted in ascending order
642 - Uses vectorized operations to handle large arrays efficiently
643 - Time ranges in bin_tedges can be in any order and can overlap
645 Examples
646 --------
647 >>> import numpy as np
648 >>> from gwtransport.utils import time_bin_overlap
649 >>> tedges = np.array([0, 10, 20, 30])
650 >>> bin_tedges = [(5, 15), (25, 35)]
651 >>> time_bin_overlap(
652 ... tedges=tedges, bin_tedges=bin_tedges
653 ... ) # doctest: +NORMALIZE_WHITESPACE
654 array([[0.5, 0.5, 0. ],
655 [0. , 0. , 0.5]])
656 """
657 # Convert inputs to numpy arrays
658 tedges = np.asarray(tedges)
659 bin_tedges_array = np.asarray(bin_tedges)
661 # Validate inputs
662 if tedges.ndim != 1:
663 msg = "tedges must be a 1D array"
664 raise ValueError(msg)
665 if len(tedges) < 2: # noqa: PLR2004
666 msg = "tedges must have at least 2 elements"
667 raise ValueError(msg)
668 if bin_tedges_array.size == 0:
669 msg = "bin_tedges must be non-empty"
670 raise ValueError(msg)
672 # Normalize datetime-like inputs (datetime64 or object arrays of Timestamps/datetimes) to a
673 # common int64-nanosecond float scale so numeric, datetime64, and Timestamp inputs share one
674 # arithmetic path; ``np.maximum(0, Timedelta)`` on an object array would otherwise raise. Only
675 # differences enter the result, so the shared epoch origin cancels and the fractions are exact.
676 if not np.issubdtype(tedges.dtype, np.number):
677 tedges = pd.DatetimeIndex(tedges).asi8.astype(float)
678 if not np.issubdtype(bin_tedges_array.dtype, np.number):
679 flat = pd.DatetimeIndex(bin_tedges_array.ravel()).asi8.astype(float)
680 bin_tedges_array = flat.reshape(bin_tedges_array.shape)
682 # Calculate overlaps for all combinations using broadcasting
683 overlap_left = np.maximum(bin_tedges_array[:, [0]], tedges[None, :-1])
684 overlap_right = np.minimum(bin_tedges_array[:, [1]], tedges[None, 1:])
685 overlap_widths = np.maximum(0, overlap_right - overlap_left)
687 # Calculate fractions (handle division by zero for zero-width bins)
688 bin_width_bc = np.diff(tedges)[None, :] # Shape: (1, n_bins)
690 return np.divide(
691 overlap_widths, bin_width_bc, out=np.zeros_like(overlap_widths, dtype=float), where=bin_width_bc != 0.0
692 )
695def simplify_bins(
696 *,
697 edges: npt.ArrayLike,
698 values: npt.ArrayLike,
699 flow: npt.ArrayLike | None = None,
700 tol: float = 0.0,
701) -> tuple[
702 npt.NDArray[np.floating] | pd.DatetimeIndex,
703 npt.NDArray[np.floating],
704 npt.NDArray[np.floating] | None,
705]:
706 """Simplify a piecewise-constant time series by merging adjacent bins.
708 Splits at the largest value jump until the peak-to-peak range within
709 every group does not exceed `tol`. The result is independent of scan
710 direction.
712 Parameters
713 ----------
714 edges : array-like
715 Bin edges with shape ``(n+1,)``. May be numeric or pandas Timestamps.
716 values : array-like
717 Bin-averaged values with shape ``(n,)`` (e.g., concentrations).
718 flow : array-like, optional
719 Flow rate per bin with shape ``(n,)`` (e.g., m³/day). When provided,
720 merged-bin values are weighted by volume (flow x bin width) instead of
721 bin width alone.
722 tol : float, optional
723 Maximum peak-to-peak range within a merged group.
724 Default is 0.0, which merges only runs of identical values.
726 Returns
727 -------
728 new_edges : ndarray or DatetimeIndex
729 Simplified bin edges with shape ``(m+1,)``, preserving the type of
730 `edges`.
731 new_values : ndarray of float
732 Volume-weighted (or width-weighted) average values per simplified
733 bin, with shape ``(m,)``.
734 new_flow : ndarray of float or None
735 Time-weighted (width-weighted) average flow per simplified bin, with
736 shape ``(m,)``. None when `flow` is not provided.
737 """
738 edges = np.asarray(edges) if not isinstance(edges, pd.DatetimeIndex) else edges
739 values = np.asarray(values, dtype=float)
740 if len(values) == 0:
741 flow_out = np.asarray(flow, dtype=float) if flow is not None else None
742 return edges, values, flow_out
744 widths = np.asarray(np.diff(edges), dtype=float)
745 if flow is not None:
746 flow = np.asarray(flow, dtype=float)
747 weights = widths * flow
748 else:
749 weights = widths
751 # Iteratively split each segment at its largest value jump until every group's peak-to-peak
752 # range is within tol. An explicit LIFO stack replaces the natural recursion, which peels one
753 # element per level on smooth monotone data (argmax|diff| sits at a segment edge) and overflows
754 # the interpreter stack for a few thousand points. Every split index is interior to its
755 # (disjoint) segment, so sorting the collected splits reproduces the recursion's in-order
756 # output exactly -- the merged bins are identical.
757 splits: list[int] = []
758 stack: list[tuple[int, int]] = [(0, len(values))]
759 while stack:
760 lo, hi = stack.pop()
761 if np.ptp(values[lo:hi]) <= tol:
762 continue
763 i = lo + int(np.argmax(np.abs(np.diff(values[lo:hi])))) + 1
764 splits.append(i)
765 stack.extend(((lo, i), (i, hi)))
766 splits.sort()
767 s = np.array([0, *splits])
768 idx = np.append(s, len(values))
769 new_edges = edges[idx]
770 new_widths = np.add.reduceat(widths, s)
771 weight_sums = np.add.reduceat(weights, s)
772 new_values = np.add.reduceat(weights * values, s) / weight_sums
773 # When flow is given, weights == flow * widths, so weight_sums == reduceat(flow * widths, s) exactly.
774 new_flow = weight_sums / new_widths if flow is not None else None
776 return new_edges, new_values, new_flow
779def _generate_failed_coverage_badge() -> None:
780 """Generate a badge indicating failed coverage."""
781 from genbadge import Badge # type: ignore # noqa: PLC0415
783 b = Badge(left_txt="coverage", right_txt="failed", color="red")
784 b.write_to("coverage_failed.svg", use_shields=False)
787def compute_time_edges(
788 *,
789 tedges: pd.DatetimeIndex | None,
790 tstart: pd.DatetimeIndex | None,
791 tend: pd.DatetimeIndex | None,
792 number_of_bins: int,
793) -> pd.DatetimeIndex:
794 """
795 Compute time edges for binning data based on provided time parameters.
797 This function creates a DatetimeIndex of time bin edges from one of three possible
798 input formats: explicit edges, start times, or end times. The resulting edges
799 define the boundaries of time intervals for data binning.
801 Define either explicit time edges, or start and end times for each bin and leave the others at None.
803 Parameters
804 ----------
805 tedges : pandas.DatetimeIndex or None
806 Explicit time edges for the bins. If provided, must have one more element
807 than the number of bins (n_bins + 1). Takes precedence over tstart and tend.
808 tstart : pandas.DatetimeIndex or None
809 Start times for each bin. Must have the same number of elements as the
810 number of bins. Used when tedges is None.
811 tend : pandas.DatetimeIndex or None
812 End times for each bin. Must have the same number of elements as the
813 number of bins. Used when both tedges and tstart are None.
814 number_of_bins : int
815 The expected number of time bins. Used for validation against the provided
816 time parameters.
818 Returns
819 -------
820 pandas.DatetimeIndex
821 Time edges defining the boundaries of the time bins. Has one more element
822 than number_of_bins.
824 Raises
825 ------
826 ValueError
827 If tedges has incorrect length (not number_of_bins + 1).
828 If tstart has incorrect length (not equal to number_of_bins).
829 If tend has incorrect length (not equal to number_of_bins).
830 If none of tedges, tstart, or tend are provided.
832 Notes
833 -----
834 - When using tstart, the function assumes uniform spacing and extrapolates
835 the final edge based on the spacing between the last two start times.
836 - When using tend, the function assumes uniform spacing and extrapolates
837 the first edge based on the spacing between the first two end times.
838 - When ``tstart`` or ``tend`` are provided with non-uniformly-spaced bins,
839 the extrapolated edge uses only the very first or very last interval and
840 may be physically incorrect: the missing edge is implicitly assigned a
841 bin width equal to that single neighbouring interval, which is unrelated
842 to any other interval in the series. In such cases, supply ``tedges``
843 directly so that all bin widths are explicit.
844 - All input time data is converted to pandas.DatetimeIndex for consistency.
845 """
846 if tedges is not None:
847 if number_of_bins != len(tedges) - 1:
848 msg = "tedges must have one more element than number_of_bins"
849 raise ValueError(msg)
850 tedges = pd.DatetimeIndex(tedges)
851 # Ensure nanosecond precision while preserving timezone
852 return tedges.as_unit("ns")
854 if tstart is not None:
855 # Assume the index refers to the time at the start of the measurement interval
856 tstart = pd.DatetimeIndex(tstart).as_unit("ns")
857 if number_of_bins != len(tstart):
858 msg = "tstart must have the same number of elements as number_of_bins"
859 raise ValueError(msg)
860 if len(tstart) < 2: # noqa: PLR2004
861 msg = "tstart must have at least 2 elements to infer the bin width; pass tedges for a single bin"
862 raise ValueError(msg)
864 # Extrapolate final edge using uniform spacing
865 final_edge = tstart[-1] + (tstart[-1] - tstart[-2])
866 return pd.DatetimeIndex([*list(tstart), final_edge], dtype=tstart.dtype)
868 if tend is not None:
869 # Assume the index refers to the time at the end of the measurement interval
870 tend = pd.DatetimeIndex(tend).as_unit("ns")
871 if number_of_bins != len(tend):
872 msg = "tend must have the same number of elements as number_of_bins"
873 raise ValueError(msg)
874 if len(tend) < 2: # noqa: PLR2004
875 msg = "tend must have at least 2 elements to infer the bin width; pass tedges for a single bin"
876 raise ValueError(msg)
878 # Extrapolate initial edge using uniform spacing
879 initial_edge = tend[0] - (tend[1] - tend[0])
880 return pd.DatetimeIndex([initial_edge, *list(tend)], dtype=tend.dtype)
882 msg = "Either provide tedges, tstart, or tend"
883 raise ValueError(msg)
886def get_soil_temperature(*, station_number: int = 260, interpolate_missing_values: bool = True) -> pd.DataFrame:
887 """
888 Download soil temperature data from the KNMI and return it as a pandas DataFrame.
890 The data is available for the following KNMI weather stations:
891 - 260: De Bilt, the Netherlands (vanaf 1981)
892 - 273: Marknesse, the Netherlands (vanaf 1989)
893 - 286: Nieuw Beerta, the Netherlands (vanaf 1990)
894 - 323: Wilhelminadorp, the Netherlands (vanaf 1989)
896 TB1 = grondtemperatuur op 5 cm diepte (graden Celsius) tijdens de waarneming
897 TB2 = grondtemperatuur op 10 cm diepte (graden Celsius) tijdens de waarneming
898 TB3 = grondtemperatuur op 20 cm diepte (graden Celsius) tijdens de waarneming
899 TB4 = grondtemperatuur op 50 cm diepte (graden Celsius) tijdens de waarneming
900 TB5 = grondtemperatuur op 100 cm diepte (graden Celsius) tijdens de waarneming
901 TNB2 = minimum grondtemperatuur op 10 cm diepte in de afgelopen 6 uur (graden Celsius)
902 TNB1 = minimum grondtemperatuur op 5 cm diepte in de afgelopen 6 uur (graden Celsius)
903 TXB1 = maximum grondtemperatuur op 5 cm diepte in de afgelopen 6 uur (graden Celsius)
904 TXB2 = maximum grondtemperatuur op 10 cm diepte in de afgelopen 6 uur (graden Celsius)
906 Parameters
907 ----------
908 station_number : int, {260, 273, 286, 323}
909 The KNMI station number for which to download soil temperature data.
910 Default is 260 (De Bilt).
911 interpolate_missing_values : bool, optional
912 If True, missing values are interpolated and recent NaN values are extrapolated with the previous value.
913 If False, missing values remain as NaN. Default is True.
915 Returns
916 -------
917 pandas.DataFrame
918 DataFrame containing soil temperature data in Celsius with a DatetimeIndex.
919 Columns include TB1, TB2, TB3, TB4, TB5, TNB1, TNB2, TXB1, TXB2.
921 Notes
922 -----
923 - KNMI: Royal Netherlands Meteorological Institute
924 - The timeseries may contain NaN values for missing data.
925 """
926 # File-based daily cache
927 cache_dir.mkdir(exist_ok=True)
929 today = date.today().isoformat() # noqa: DTZ011
930 cache_path = cache_dir / f"soil_temp_{station_number}_{interpolate_missing_values}_{today}.pkl"
932 # Check if cached file exists and is from today
933 if cache_path.exists():
934 cached = pd.read_pickle(cache_path) # noqa: S301
935 assert isinstance(cached, pd.DataFrame) # noqa: S101 -- the cache only ever stores DataFrames
936 return cached
938 # Clean up old cache files to prevent disk bloat
939 for old_file in cache_dir.glob(f"soil_temp_{station_number}_{interpolate_missing_values}_*.pkl"):
940 old_file.unlink(missing_ok=True)
942 url = f"https://cdn.knmi.nl/knmi/map/page/klimatologie/gegevens/bodemtemps/bodemtemps_{station_number}.zip"
944 dtypes = {
945 "YYYYMMDD": "int32",
946 "HH": "int8",
947 " TB1": "float32",
948 " TB3": "float32",
949 " TB2": "float32",
950 " TB4": "float32",
951 " TB5": "float32",
952 " TNB1": "float32",
953 " TNB2": "float32",
954 " TXB1": "float32",
955 " TXB2": "float32",
956 }
958 # Imported lazily so the rest of the module remains importable in environments
959 # without ``requests`` (e.g. Pyodide/JupyterLite, where this KNMI download is the
960 # only feature that cannot run client-side).
961 import requests # noqa: PLC0415
963 # Download the ZIP file
964 with requests.get(url, params={"download": "zip"}, timeout=10) as response:
965 response.raise_for_status()
967 df = pd.read_csv( # type: ignore[call-overload] # ty: ignore[no-matching-overload]
968 io.BytesIO(response.content),
969 compression="zip",
970 dtype=dtypes, # pyright: ignore[reportArgumentType]
971 usecols=list(dtypes.keys()), # pyright: ignore[reportArgumentType]
972 skiprows=16,
973 sep=",",
974 na_values=[" "],
975 engine="c",
976 parse_dates=False,
977 )
979 df.index = pd.to_datetime(df["YYYYMMDD"].values, format=r"%Y%m%d").tz_localize("UTC") + pd.to_timedelta(
980 df["HH"].values, unit="h"
981 )
983 df.drop(columns=["YYYYMMDD", "HH"], inplace=True)
984 df.columns = df.columns.str.strip()
985 df /= 10.0
987 if interpolate_missing_values:
988 # Fill NaN values with interpolate linearly and then forward fill
989 df.interpolate(method="linear", inplace=True)
990 df.ffill(inplace=True)
992 # Save to cache for future use
993 df.to_pickle(cache_path)
994 return df
997def solve_underdetermined_system(
998 *,
999 coefficient_matrix: npt.ArrayLike,
1000 rhs_vector: npt.ArrayLike,
1001 nullspace_objective: str
1002 | Callable[
1003 [npt.NDArray[np.floating], npt.NDArray[np.floating], npt.NDArray[np.floating]], float
1004 ] = "squared_differences",
1005 optimization_method: str = "BFGS",
1006 rcond: float | None = None,
1007) -> npt.NDArray[np.floating]:
1008 """
1009 Solve an underdetermined linear system with nullspace regularization.
1011 For an underdetermined system Ax = b where A has more columns than rows,
1012 multiple solutions exist. This function computes a least-squares solution
1013 and then selects a specific solution from the nullspace based on a
1014 regularization objective.
1016 Parameters
1017 ----------
1018 coefficient_matrix : array-like
1019 Coefficient matrix of shape (m, n) where m < n (underdetermined).
1020 May contain NaN values in some rows, which will be excluded from the system.
1021 rhs_vector : array-like
1022 Right-hand side vector of length m. May contain NaN values corresponding
1023 to NaN rows in coefficient_matrix, which will be excluded from the system.
1024 nullspace_objective : str or callable, optional
1025 Objective function to minimize in the nullspace. Options:
1027 * "squared_differences" : Minimize sum of squared differences between
1028 adjacent elements: ``sum((x[i+1] - x[i])**2)``
1029 * "summed_differences" : Minimize sum of absolute differences between
1030 adjacent elements: ``sum(|x[i+1] - x[i]|)``
1031 * callable : Custom objective function with signature
1032 ``objective(coeffs, x_ls, nullspace_basis)`` where:
1034 - coeffs : optimization variables (nullspace coefficients)
1035 - x_ls : least-squares solution
1036 - nullspace_basis : nullspace basis matrix
1038 Default is "squared_differences".
1039 optimization_method : str, optional
1040 Optimization method passed to scipy.optimize.minimize.
1041 Default is "BFGS".
1042 rcond : float or None, optional
1043 Cutoff ratio for small singular values in both ``numpy.linalg.lstsq``
1044 and ``scipy.linalg.null_space``. Singular values smaller than
1045 ``rcond * largest_singular_value`` are treated as zero.
1046 Default is None, which uses the default of each function.
1047 Increasing rcond truncates more modes, expanding the nullspace
1048 available for smoothness optimization. Useful for noisy data.
1050 Returns
1051 -------
1052 ndarray
1053 Solution vector that minimizes the specified nullspace objective.
1054 Has length n (number of columns in coefficient_matrix).
1056 Raises
1057 ------
1058 ValueError
1059 If optimization fails, if coefficient_matrix and rhs_vector have incompatible shapes,
1060 or if an unknown nullspace objective is specified.
1062 Notes
1063 -----
1064 The algorithm follows these steps:
1066 1. Remove rows with NaN values from both coefficient_matrix and rhs_vector
1067 2. Compute least-squares solution: x_ls = pinv(valid_matrix) @ valid_rhs
1068 3. Compute nullspace basis: N = null_space(valid_matrix)
1069 4. Find nullspace coefficients: coeffs = argmin objective(x_ls + N @ coeffs)
1070 5. Return final solution: x = x_ls + N @ coeffs
1072 For the built-in objectives:
1074 * "squared_differences" provides smooth solutions, minimizing rapid changes
1075 * "summed_differences" provides sparse solutions, promoting piecewise constant behavior
1077 Examples
1078 --------
1079 Basic usage with default squared differences objective:
1081 >>> import numpy as np
1082 >>> from gwtransport.utils import solve_underdetermined_system
1083 >>>
1084 >>> # Create underdetermined system (2 equations, 4 unknowns)
1085 >>> matrix = np.array([[1, 2, 1, 0], [0, 1, 2, 1]])
1086 >>> rhs = np.array([3, 4])
1087 >>>
1088 >>> # Solve with squared differences regularization
1089 >>> x = solve_underdetermined_system(coefficient_matrix=matrix, rhs_vector=rhs)
1090 >>> print(f"Solution: {x}") # doctest: +SKIP
1091 >>> print(f"Residual: {np.linalg.norm(matrix @ x - rhs):.2e}") # doctest: +SKIP
1093 With summed differences objective:
1095 >>> x_sparse = solve_underdetermined_system( # doctest: +SKIP
1096 ... coefficient_matrix=matrix,
1097 ... rhs_vector=rhs,
1098 ... nullspace_objective="summed_differences",
1099 ... )
1101 With custom objective function:
1103 >>> def custom_objective(coeffs, x_ls, nullspace_basis):
1104 ... x = x_ls + nullspace_basis @ coeffs
1105 ... return np.sum(x**2) # Minimize L2 norm
1106 >>>
1107 >>> x_custom = solve_underdetermined_system( # doctest: +SKIP
1108 ... coefficient_matrix=matrix,
1109 ... rhs_vector=rhs,
1110 ... nullspace_objective=custom_objective,
1111 ... )
1113 Handling NaN values:
1115 >>> # System with missing data
1116 >>> matrix_nan = np.array([
1117 ... [1, 2, 1, 0],
1118 ... [np.nan, np.nan, np.nan, np.nan],
1119 ... [0, 1, 2, 1],
1120 ... ])
1121 >>> rhs_nan = np.array([3, np.nan, 4])
1122 >>>
1123 >>> x_nan = solve_underdetermined_system(
1124 ... coefficient_matrix=matrix_nan, rhs_vector=rhs_nan
1125 ... ) # doctest: +SKIP
1126 """
1127 matrix = np.asarray(coefficient_matrix)
1128 rhs = np.asarray(rhs_vector)
1130 if matrix.shape[0] != len(rhs):
1131 msg = f"coefficient_matrix has {matrix.shape[0]} rows but rhs_vector has {len(rhs)} elements"
1132 raise ValueError(msg)
1134 # Identify valid rows (no NaN values in either matrix or rhs)
1135 valid_rows = ~np.isnan(matrix).any(axis=1) & ~np.isnan(rhs)
1137 if not np.any(valid_rows):
1138 msg = "No valid rows found (all contain NaN values)"
1139 raise ValueError(msg)
1141 valid_matrix = matrix[valid_rows]
1142 valid_rhs = rhs[valid_rows]
1144 # Compute least-squares solution
1145 x_ls, *_ = np.linalg.lstsq(valid_matrix, valid_rhs, rcond=rcond)
1147 # Compute nullspace
1148 nullspace_basis = null_space(valid_matrix, rcond=rcond)
1149 nullrank = nullspace_basis.shape[1]
1151 if nullrank == 0:
1152 # System is determined, return least-squares solution
1153 return x_ls
1155 # Optimize in nullspace
1156 coeffs = _optimize_nullspace_coefficients(
1157 x_ls=x_ls,
1158 nullspace_basis=nullspace_basis,
1159 nullspace_objective=nullspace_objective,
1160 optimization_method=optimization_method,
1161 )
1163 return x_ls + nullspace_basis @ coeffs
1166def _optimize_nullspace_coefficients(
1167 *,
1168 x_ls: npt.NDArray[np.floating],
1169 nullspace_basis: npt.NDArray[np.floating],
1170 nullspace_objective: str
1171 | Callable[[npt.NDArray[np.floating], npt.NDArray[np.floating], npt.NDArray[np.floating]], float],
1172 optimization_method: str,
1173) -> npt.NDArray[np.floating]:
1174 """Optimize coefficients in the nullspace to minimize the objective.
1176 Parameters
1177 ----------
1178 x_ls : ndarray
1179 Least-squares solution vector.
1180 nullspace_basis : ndarray
1181 Nullspace basis matrix of shape (n, nullrank).
1182 nullspace_objective : str or callable
1183 Objective to minimize. Supported string values are
1184 ``'squared_differences'`` and ``'summed_differences'``. A callable
1185 with signature ``objective(coeffs, x_ls, nullspace_basis)`` is also
1186 accepted.
1187 optimization_method : str
1188 Optimization method passed to ``scipy.optimize.minimize``.
1190 Returns
1191 -------
1192 ndarray
1193 Optimal nullspace coefficient vector of length nullrank.
1195 Raises
1196 ------
1197 ValueError
1198 If iterative optimization fails to converge.
1199 """
1200 # For squared_differences, solve the quadratic form analytically:
1201 # min ||D(x_ls + N c)||^2 => (N'D'DN) c = -N'D'D x_ls
1202 coeffs_sq = _solve_squared_differences_analytical(x_ls=x_ls, nullspace_basis=nullspace_basis)
1204 if nullspace_objective == "squared_differences":
1205 return coeffs_sq
1207 # For other objectives, use iterative optimization starting from the
1208 # squared_differences solution for stability
1209 objective_func = _get_nullspace_objective_function(nullspace_objective=nullspace_objective)
1210 coeffs_0 = coeffs_sq
1212 res = minimize(
1213 objective_func,
1214 x0=coeffs_0,
1215 args=(x_ls, nullspace_basis),
1216 method=optimization_method,
1217 )
1219 if not res.success:
1220 msg = f"Optimization failed: {res.message}"
1221 raise ValueError(msg)
1223 return res.x
1226def _solve_squared_differences_analytical(
1227 *,
1228 x_ls: npt.NDArray[np.floating],
1229 nullspace_basis: npt.NDArray[np.floating],
1230) -> npt.NDArray[np.floating]:
1231 """Solve the squared-differences nullspace problem analytically.
1233 Minimizes ``sum((x[i+1] - x[i])^2)`` where ``x = x_ls + N @ c`` by
1234 solving the normal equations ``(N^T D^T D N) c = -N^T D^T D x_ls``.
1236 Parameters
1237 ----------
1238 x_ls : ndarray
1239 Least-squares solution vector of length n.
1240 nullspace_basis : ndarray
1241 Nullspace basis matrix of shape (n, nullrank).
1243 Returns
1244 -------
1245 ndarray
1246 Optimal nullspace coefficient vector of length nullrank.
1248 Raises
1249 ------
1250 numpy.linalg.LinAlgError
1251 If the normal equations matrix ``(DN)^T(DN)`` is ill-conditioned
1252 (condition number exceeds 1e12).
1253 """
1254 # D is the (n-1, n) first-difference matrix; D @ x = x[1:] - x[:-1]
1255 # D^T D is the tridiagonal matrix with 2 on diagonal, -1 on off-diagonals
1256 # (except corners which have 1 on diagonal)
1257 # Instead of forming D explicitly, compute D @ N and D @ x_ls directly
1258 dn = nullspace_basis[1:, :] - nullspace_basis[:-1, :] # (n-1, nullrank)
1259 dx = x_ls[1:] - x_ls[:-1] # (n-1,)
1261 # Normal equations: (DN)^T (DN) c = -(DN)^T (D x_ls)
1262 dntdn = dn.T @ dn # (nullrank, nullrank)
1263 rhs = -(dn.T @ dx) # (nullrank,)
1265 cond = np.linalg.cond(dntdn)
1266 cond_threshold = 1e12
1267 if cond > cond_threshold:
1268 msg = (
1269 f"The normal equations matrix (DN)^T(DN) is ill-conditioned "
1270 f"(condition number: {cond:.2e}). This typically means the "
1271 f"nullspace contains a near-constant vector, so the "
1272 f"squared-differences objective cannot distinguish between "
1273 f"nullspace directions. Consider using a different "
1274 f"nullspace_objective (e.g., 'summed_differences'), reducing "
1275 f"the problem's degrees of freedom, or lowering rcond to "
1276 f"shrink the nullspace (if the near-constant vector has a "
1277 f"small but non-zero singular value)."
1278 )
1279 raise np.linalg.LinAlgError(msg)
1281 return np.linalg.solve(dntdn, rhs)
1284def compute_reverse_target(
1285 *,
1286 coeff_matrix: npt.NDArray[np.floating],
1287 rhs_vector: npt.NDArray[np.floating],
1288) -> npt.NDArray[np.floating]:
1289 """Compute reverse matrix target from forward coefficient matrix.
1291 Constructs a target solution for the inverse problem by transposing the
1292 forward coefficient matrix and normalizing rows. For ``W_forward[i,j]``
1293 representing the fraction of ``cin[j]`` arriving in ``cout[i]``, the
1294 transpose-and-normalize approach reconstructs ``cin[j]`` as a weighted
1295 average of ``cout`` bins, weighted by how much ``cin[j]`` contributed
1296 to each ``cout`` bin.
1298 Parameters
1299 ----------
1300 coeff_matrix : ndarray
1301 Forward coefficient matrix of shape (n_cout, n_cin).
1302 rhs_vector : ndarray
1303 Right-hand side vector of length n_cout (e.g., cout values).
1305 Returns
1306 -------
1307 ndarray
1308 Target solution vector of length n_cin. Entries with near-zero
1309 column sums in the forward matrix are set to NaN.
1311 See Also
1312 --------
1313 solve_tikhonov : Consumes this target as the regularization reference.
1314 """
1315 min_row_sum = 1e-10
1316 wt = coeff_matrix.T # (n_cin, n_cout)
1317 row_sums = wt.sum(axis=1)
1318 valid = row_sums > min_row_sum
1319 w_reverse = np.zeros_like(wt)
1320 w_reverse[valid] = wt[valid] / row_sums[valid, None]
1321 x_target = w_reverse @ rhs_vector
1322 x_target[~valid] = np.nan
1323 return x_target
1326def solve_tikhonov(
1327 *,
1328 coefficient_matrix: npt.ArrayLike,
1329 rhs_vector: npt.ArrayLike,
1330 x_target: npt.NDArray[np.floating],
1331 regularization_strength: float = 1e-10,
1332 return_resolution: bool = False,
1333) -> npt.NDArray[np.floating] | tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
1334 """Solve a linear system with Tikhonov regularization toward a target.
1336 Minimizes ``||A x - b||² + λ ||x - x_target||²`` by solving the
1337 equivalent augmented least-squares problem::
1339 [A; √λ I_v] x = [b; √λ x_target_v]
1341 where ``I_v`` selects only entries where ``x_target`` is not NaN.
1343 Well-determined modes (large singular values relative to √λ) are
1344 dominated by the data; poorly-determined modes are pulled toward
1345 ``x_target``. The solution varies continuously with λ, unlike the
1346 hard singular-value cutoff of ``rcond`` in truncated SVD.
1348 Parameters
1349 ----------
1350 coefficient_matrix : array-like
1351 Coefficient matrix of shape (m, n). May contain NaN rows, which
1352 are excluded from the system.
1353 rhs_vector : array-like
1354 Right-hand side vector of length m. May contain NaN values
1355 corresponding to NaN rows in coefficient_matrix.
1356 x_target : ndarray
1357 Target solution of length n, typically from
1358 :func:`compute_reverse_target`. NaN entries are excluded from the
1359 regularization term.
1360 regularization_strength : float, optional
1361 Tikhonov parameter λ. Controls the tradeoff between fitting the
1362 data and staying close to ``x_target``. Larger values trust the
1363 target more; smaller values trust the data more. Default is 1e-10.
1365 A good starting value for noisy data is
1366 ``λ ≈ (noise_std / signal_amplitude)²``. For noiseless synthetic
1367 data, the default 1e-10 preserves machine precision.
1368 return_resolution : bool, optional
1369 If True, also return the per-element fraction of the solution that
1370 comes from data (vs from the regularization target). Default is
1371 False.
1373 Returns
1374 -------
1375 ndarray or tuple of ndarray
1376 If ``return_resolution`` is False (default), returns the solution
1377 vector of length n.
1379 If ``return_resolution`` is True, returns ``(x, fraction_data)``
1380 where ``fraction_data[j]`` is the diagonal of the model resolution
1381 matrix ``R = (A^T A + λ D)^{-1} A^T A``:
1383 - ``fraction_data[j] ≈ 1``: element *j* is data-driven
1384 - ``fraction_data[j] ≈ 0``: element *j* is target-driven
1385 - Non-regularized entries (NaN in ``x_target``):
1386 ``fraction_data[j] = 1.0``
1388 Raises
1389 ------
1390 ValueError
1391 If ``coefficient_matrix`` and ``rhs_vector`` have incompatible shapes, or if
1392 all rows contain NaN values.
1394 See Also
1395 --------
1396 compute_reverse_target : Compute the regularization target from the
1397 forward matrix.
1398 solve_underdetermined_system : Alternative solver using nullspace
1399 optimization.
1400 """
1401 matrix = np.asarray(coefficient_matrix)
1402 rhs = np.asarray(rhs_vector)
1404 if matrix.shape[0] != len(rhs):
1405 msg = f"coefficient_matrix has {matrix.shape[0]} rows but rhs_vector has {len(rhs)} elements"
1406 raise ValueError(msg)
1408 # Filter NaN rows
1409 valid_rows = ~np.isnan(matrix).any(axis=1) & ~np.isnan(rhs)
1411 if not np.any(valid_rows):
1412 msg = "No valid rows found (all contain NaN values)"
1413 raise ValueError(msg)
1415 valid_matrix = matrix[valid_rows]
1416 valid_rhs = rhs[valid_rows]
1418 n_cin = valid_matrix.shape[1]
1419 sqrt_lam = np.sqrt(regularization_strength)
1421 # Only regularize entries where x_target is valid
1422 valid_target = ~np.isnan(x_target)
1423 target_indices = np.where(valid_target)[0]
1425 # Build augmented system: [A; √λ I_v] x = [b; √λ x_target_v]
1426 n_reg = len(target_indices)
1427 reg_matrix = np.zeros((n_reg, n_cin))
1428 reg_matrix[np.arange(n_reg), target_indices] = sqrt_lam
1429 reg_rhs = sqrt_lam * x_target[target_indices]
1431 augmented_matrix = np.vstack([valid_matrix, reg_matrix])
1432 augmented_rhs = np.concatenate([valid_rhs, reg_rhs])
1434 x, *_ = np.linalg.lstsq(augmented_matrix, augmented_rhs, rcond=None)
1436 if return_resolution:
1437 # Compute fraction_data from model resolution matrix diagonal:
1438 # R = G^{-1} A^T A where G = A^T A + λ diag(d)
1439 # fraction_data[j] = R[j,j] = 1 - λ d[j] G_inv[j,j]
1440 d_reg = np.zeros(n_cin)
1441 d_reg[target_indices] = 1.0
1442 gram = valid_matrix.T @ valid_matrix
1443 gram[np.arange(n_cin), np.arange(n_cin)] += regularization_strength * d_reg
1444 gram_inv_diag = np.diag(np.linalg.inv(gram))
1445 fraction_data = 1.0 - regularization_strength * gram_inv_diag * d_reg
1446 return x, fraction_data
1448 return x
1451# Numerical tolerance for coefficient sum to determine valid output bins
1452_EPSILON_COEFF_SUM = 1e-10
1454# Corrected semi-normal-equation refinement steps in solve_inverse_transport_banded. One
1455# step reaches the QR-accurate solution; a second is a cheap, stable safety margin.
1456_BANDED_REFINEMENT_STEPS = 2
1459def solve_inverse_transport(
1460 *,
1461 w_forward: npt.NDArray[np.floating],
1462 observed: npt.NDArray[np.floating],
1463 n_output: int,
1464 regularization_strength: float,
1465 valid_rows: npt.NDArray[np.bool_] | None = None,
1466 warn_rank_deficient: bool = False,
1467) -> npt.NDArray[np.floating]:
1468 """Solve the inverse transport problem via Tikhonov regularization.
1470 Given the forward model ``w_forward @ x = observed``, recovers ``x`` by
1471 building the regularization target from the transpose of ``w_forward`` and
1472 solving the regularized least-squares problem.
1474 Parameters
1475 ----------
1476 w_forward : ndarray
1477 Forward coefficient matrix with shape ``(n_obs, n_output)``.
1478 observed : ndarray
1479 Observed values with shape ``(n_obs,)`` (e.g., extraction
1480 concentrations).
1481 n_output : int
1482 Length of the output vector (e.g., number of cin bins).
1483 regularization_strength : float
1484 Tikhonov regularization parameter.
1485 valid_rows : ndarray of bool, optional
1486 Which observation rows are valid, with shape ``(n_obs,)``. If None,
1487 rows with ``row_sum > 1e-10`` are considered valid.
1488 warn_rank_deficient : bool, optional
1489 If True, emit a warning when the forward matrix has rank
1490 deficiency among its active columns. Default is False.
1492 Returns
1493 -------
1494 ndarray
1495 Recovered signal with shape ``(n_output,)``. NaN for bins with no
1496 active columns.
1498 Warns
1499 -----
1500 UserWarning
1501 When ``warn_rank_deficient=True`` and the matrix is rank-deficient.
1503 See Also
1504 --------
1505 solve_inverse_transport_banded : Memory-light banded equivalent.
1506 """
1507 row_sums = w_forward.sum(axis=1)
1508 col_active: npt.NDArray[np.bool_] = w_forward.sum(axis=0) > 0
1510 if not np.any(col_active):
1511 return np.full(n_output, np.nan)
1513 if warn_rank_deficient:
1514 n_active = int(col_active.sum())
1515 rank = np.linalg.matrix_rank(w_forward[:, col_active])
1516 if rank < n_active:
1517 warnings.warn(
1518 f"Forward matrix is rank-deficient (rank {rank} < {n_active} active "
1519 f"columns). This occurs with constant flow when the residence time "
1520 f"is an integer multiple of the time step width. The "
1521 f"underdetermined modes will be pulled toward the regularization "
1522 f"target instead of being determined by data. To achieve full rank, "
1523 f"adjust aquifer_pore_volumes slightly (e.g., multiply by 1.001).",
1524 stacklevel=2,
1525 )
1527 valid: npt.NDArray[np.bool_] = row_sums > _EPSILON_COEFF_SUM if valid_rows is None else valid_rows
1529 rhs = np.where(valid, row_sums * observed, np.nan)
1530 w_solve = w_forward.copy()
1531 w_solve[~valid, :] = np.nan
1533 x_target = compute_reverse_target(coeff_matrix=w_forward, rhs_vector=observed)
1535 x_solved = solve_tikhonov(
1536 coefficient_matrix=w_solve,
1537 rhs_vector=rhs,
1538 x_target=x_target,
1539 regularization_strength=regularization_strength,
1540 )
1542 out = np.full(n_output, np.nan)
1543 idx = np.flatnonzero(col_active)
1544 out[idx] = x_solved[idx]
1545 return out
1548def solve_inverse_transport_banded(
1549 *,
1550 band_vals: npt.NDArray[np.floating],
1551 col_start: npt.NDArray[np.intp],
1552 observed: npt.NDArray[np.floating],
1553 n_output: int,
1554 regularization_strength: float,
1555) -> npt.NDArray[np.floating]:
1556 """Solve the inverse transport problem from a banded forward operator.
1558 Memory-light equivalent of :func:`solve_inverse_transport` for a forward
1559 weight matrix stored in banded layout: row ``k`` of the dense operator
1560 ``W`` is ``band_vals[k]`` placed at columns
1561 ``[col_start[k], col_start[k] + full_band)``. The Tikhonov normal
1562 equations ``(WᵀW + λ D) x = Wᵀ observed + λ D x_target`` are stored **in
1563 banded form** -- ``WᵀW`` is symmetric with half-bandwidth ``full_band - 1``
1564 -- and Cholesky-factored with :func:`scipy.linalg.cholesky_banded`. The Gram
1565 matrix ``WᵀW`` is built with a single dense BLAS matmul (``~24x`` a
1566 per-diagonal scatter) before its sub-diagonals are read into the banded
1567 layout. Forming ``WᵀW`` squares the condition number, so the bare Cholesky
1568 solve loses accuracy in the under-determined (spin-up nullspace) directions;
1569 **corrected semi-normal equations** restore it by refining with the residual
1570 evaluated through ``W`` itself rather than ``WᵀW`` (matching the dense
1571 least-squares solution to ~1e-10). The banded Cholesky factor, solve, and
1572 refinement stay at ``O(n_output * full_band)``; only the one-shot Gram
1573 assembly transiently materializes ``W`` and ``WᵀW`` densely.
1575 The regularization target ``x_target`` is the transpose-and-normalize of
1576 ``W`` applied to ``observed`` (the banded form of
1577 :func:`compute_reverse_target`), matching the dense solver. Columns with no
1578 forward contribution are decoupled (unit diagonal) so the system stays
1579 symmetric positive definite, and are returned as NaN.
1581 Parameters
1582 ----------
1583 band_vals : ndarray
1584 Banded forward weights of shape ``(n_obs, full_band)``. Rows the caller
1585 considers invalid must already be zeroed (as ``_resolve_spinup_mask``
1586 does); zero rows contribute nothing to the normal equations.
1587 col_start : ndarray of int
1588 First output-column index of each row's band, shape ``(n_obs,)``.
1589 observed : ndarray
1590 Observed values of shape ``(n_obs,)`` (e.g. extraction concentrations).
1591 Must not contain NaN.
1592 n_output : int
1593 Length of the output vector (number of cin bins).
1594 regularization_strength : float
1595 Tikhonov parameter λ. See :func:`solve_inverse_transport`. Must be
1596 strictly positive: deconvolution is generically rank-deficient, and λ
1597 is what makes the banded Cholesky factor positive definite (unlike the
1598 dense least-squares path, this solver cannot return a λ=0 min-norm
1599 solution).
1601 Returns
1602 -------
1603 ndarray
1604 Recovered signal of shape ``(n_output,)``. NaN for output bins with no
1605 forward contribution (zero column).
1607 Raises
1608 ------
1609 ValueError
1610 If ``regularization_strength`` is not strictly positive.
1612 See Also
1613 --------
1614 solve_inverse_transport : Dense-matrix equivalent.
1615 ``gwtransport.advection_utils._infiltration_to_extraction_weights`` : Banded builder.
1616 """
1617 if regularization_strength <= 0:
1618 msg = "regularization_strength must be > 0 for the banded inverse (Tikhonov positive-definiteness)"
1619 raise ValueError(msg)
1620 # Precondition: the caller's valid rows sum to 1 (guaranteed by
1621 # _resolve_spinup_mask), so the data equation is W x ≈ observed and the RHS
1622 # needs no row_sums scaling -- matching the dense solve_inverse_transport.
1623 band_vals = np.asarray(band_vals, dtype=float)
1624 observed = np.asarray(observed, dtype=float)
1625 full_band = band_vals.shape[1]
1626 n_cin = n_output
1627 cols = col_start[:, None] + np.arange(full_band)[None, :] # (n_obs, full_band) output-column index
1628 in_range = cols < n_cin
1629 cols_clipped = np.clip(cols, 0, n_cin - 1)
1631 # Column sums and Wᵀ observed (the reverse-target numerator) by scattering the band.
1632 col_sum = np.zeros(n_cin)
1633 wt_observed = np.zeros(n_cin)
1634 np.add.at(col_sum, cols_clipped[in_range], band_vals[in_range])
1635 np.add.at(wt_observed, cols_clipped[in_range], (band_vals * observed[:, None])[in_range])
1637 col_active = col_sum > 0
1638 if not np.any(col_active):
1639 return np.full(n_output, np.nan)
1641 # Reverse-target: transpose-and-normalize W applied to observed (banded form of
1642 # compute_reverse_target). The sliver 0 < col_sum <= _EPSILON_COEFF_SUM is left
1643 # untargeted (filled with 0) as in the dense path.
1644 with np.errstate(invalid="ignore", divide="ignore"):
1645 x_target = np.where(col_sum > _EPSILON_COEFF_SUM, wt_observed / col_sum, 0.0)
1647 # Lower-banded WᵀW via a dense BLAS matmul. Materialize the forward operator W densely
1648 # (row k is band_vals[k] at columns [col_start[k], col_start[k] + full_band)), form the
1649 # symmetric Gram matrix WᵀW with a single optimized matmul, then read its lower sub-diagonals
1650 # into the banded layout (band row d is the d-th sub-diagonal, WᵀW[j + d, j]). Each row's
1651 # in-range band columns are distinct, so the scatter into W needs no accumulation. This is
1652 # ~24x the per-diagonal np.add.at scatter; the matmul reorders the summation, so ab matches
1653 # the scatter to ~1e-13 -- well inside the Tikhonov + refinement tolerance.
1654 n_obs = band_vals.shape[0]
1655 w_dense = np.zeros((n_obs, n_cin))
1656 obs_idx = np.broadcast_to(np.arange(n_obs)[:, None], cols.shape)
1657 w_dense[obs_idx[in_range], cols_clipped[in_range]] = band_vals[in_range]
1658 gram = w_dense.T @ w_dense
1659 ab = np.zeros((full_band, n_cin))
1660 for d in range(full_band):
1661 ab[d, : n_cin - d] = np.diagonal(gram, offset=-d)
1663 lam = regularization_strength
1664 d_reg = lam * col_active
1665 ab[0] += d_reg
1666 # d_reg is zero off the active columns, so x_target needs no masking here or in
1667 # the refinement loop: the product d_reg * x_target vanishes wherever col_active is False.
1668 rhs = wt_observed + d_reg * x_target
1670 # Decouple zero (inactive, unregularized) diagonals so the matrix is SPD.
1671 dead = ab[0] <= 0.0
1672 ab[0, dead] = 1.0
1673 rhs[dead] = 0.0
1675 factor = cholesky_banded(ab, lower=True)
1676 x = cho_solve_banded((factor, True), rhs)
1678 # Forming WᵀW squares the condition number, so the bare Cholesky solution loses
1679 # accuracy in the under-determined (spin-up nullspace) directions. Corrected
1680 # semi-normal equations recover it: the residual is evaluated through W itself
1681 # (in observation space) rather than through WᵀW, avoiding the cancellation that
1682 # makes plain normal-equation refinement stall. One step reaches the QR-accurate
1683 # solution; the rest are a safety margin (the iteration's fixed point is stable).
1684 for _ in range(_BANDED_REFINEMENT_STEPS):
1685 gathered = x[cols_clipped]
1686 gathered[~in_range] = 0.0
1687 residual = observed - (band_vals * gathered).sum(axis=1) # b - W x (n_obs,)
1688 gradient = np.zeros(n_cin)
1689 np.add.at(gradient, cols_clipped[in_range], (band_vals * residual[:, None])[in_range]) # Wᵀ (b - W x)
1690 gradient += d_reg * (x_target - x)
1691 gradient[dead] = 0.0
1692 x += cho_solve_banded((factor, True), gradient)
1694 out = np.full(n_output, np.nan)
1695 out[col_active] = x[col_active]
1696 return out
1699def _squared_differences_objective(
1700 coeffs: npt.NDArray[np.floating], x_ls: npt.NDArray[np.floating], nullspace_basis: npt.NDArray[np.floating]
1701) -> float:
1702 """Minimize sum of squared differences between adjacent elements.
1704 Parameters
1705 ----------
1706 coeffs : ndarray
1707 Nullspace coefficient vector.
1708 x_ls : ndarray
1709 Least-squares solution vector.
1710 nullspace_basis : ndarray
1711 Nullspace basis matrix.
1713 Returns
1714 -------
1715 float
1716 Sum of squared differences between adjacent elements of the solution.
1717 """
1718 x = x_ls + nullspace_basis @ coeffs
1719 return np.sum(np.square(x[1:] - x[:-1]))
1722def _summed_differences_objective(
1723 coeffs: npt.NDArray[np.floating], x_ls: npt.NDArray[np.floating], nullspace_basis: npt.NDArray[np.floating]
1724) -> float:
1725 """Minimize sum of absolute differences between adjacent elements.
1727 Parameters
1728 ----------
1729 coeffs : ndarray
1730 Nullspace coefficient vector.
1731 x_ls : ndarray
1732 Least-squares solution vector.
1733 nullspace_basis : ndarray
1734 Nullspace basis matrix.
1736 Returns
1737 -------
1738 float
1739 Sum of absolute differences between adjacent elements of the solution.
1740 """
1741 x = x_ls + nullspace_basis @ coeffs
1742 return np.sum(np.abs(x[1:] - x[:-1]))
1745def _get_nullspace_objective_function(
1746 *,
1747 nullspace_objective: str
1748 | Callable[[npt.NDArray[np.floating], npt.NDArray[np.floating], npt.NDArray[np.floating]], float],
1749) -> Callable[[npt.NDArray[np.floating], npt.NDArray[np.floating], npt.NDArray[np.floating]], float]:
1750 """Get the objective function for nullspace optimization.
1752 Parameters
1753 ----------
1754 nullspace_objective : str or callable
1755 Objective identifier. Supported string values are
1756 ``'squared_differences'`` and ``'summed_differences'``. A callable
1757 with signature ``objective(coeffs, x_ls, nullspace_basis)`` is also
1758 accepted and returned as-is.
1760 Returns
1761 -------
1762 callable
1763 Objective function with signature
1764 ``(coeffs, x_ls, nullspace_basis) -> float``.
1766 Raises
1767 ------
1768 ValueError
1769 If ``nullspace_objective`` is an unrecognized string.
1770 """
1771 if nullspace_objective == "squared_differences":
1772 return _squared_differences_objective
1773 if nullspace_objective == "summed_differences":
1774 return _summed_differences_objective
1775 if callable(nullspace_objective):
1776 return nullspace_objective # type: ignore[return-value] # ty: ignore[invalid-return-type]
1777 msg = f"Unknown nullspace objective: {nullspace_objective}"
1778 raise ValueError(msg)