Coverage for src/gwtransport/utils.py: 94%
272 statements
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-04 21:13 +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, cumulative flow volumes, time-edge
7construction, linear-system solvers, and external data retrieval.
9The inverse solvers below are two intentionally coexisting families: a Tikhonov family (the dense
10:func:`solve_inverse_transport` and its banded equivalent :func:`solve_inverse_transport_banded`,
11both fed by :func:`compute_reverse_target` and built on :func:`solve_tikhonov`) for the
12overdetermined deconvolution in advection/diffusion, and a separate nullspace solver
13(:func:`solve_underdetermined_system`) for the underdetermined deposition inverse.
15Available functions:
17- :func:`step_plot_coords` - Expand bin edges (n+1) and bin-averaged values (n) into paired x/y arrays of 2n
18 points each, so that ``ax.plot(x, y)`` draws the piecewise-constant series as a step function. Edges may be
19 numeric or datetime; each output keeps the dtype of its input.
21- :func:`cumulative_flow_volume` - Accumulate per-bin flow rates times bin widths into the cumulative volume at
22 every bin edge (n+1 values, starting at zero). With ``strictly_monotone=True`` the plateaus left by zero-flow
23 bins are bumped by a few ulps, which is required before inverting the sequence from volume back to time.
25- :func:`linear_interpolate` - Interpolate ``y_ref`` at ``x_query`` linearly; ``x_ref`` must be ascending and the
26 result has the shape of ``x_query``. Query points outside the reference range clamp to the end values unless
27 ``left`` / ``right`` supply a fill value such as NaN.
29- :func:`simplify_bins` - Merge adjacent bins of a piecewise-constant series until the peak-to-peak range within
30 each merged group is at most ``tol``, returning the merged edges, values and flow. Values are volume-weighted
31 (flow times width) when ``flow`` is given and width-weighted otherwise, while the merged flow is always
32 width-weighted. Splitting at the largest value jump makes the result independent of scan direction.
34- :func:`compute_time_edges` - Build the n+1 bin edges as a nanosecond-precision DatetimeIndex from exactly one of
35 explicit edges, per-bin start times, or per-bin end times, validating the length against ``number_of_bins``.
36 From ``tstart`` or ``tend`` the single missing outer edge is extrapolated from the adjacent interval alone, so
37 pass ``tedges`` directly when the bins are not uniformly spaced.
39- :func:`get_soil_temperature` - Download the KNMI soil-temperature record of one of four Dutch weather stations
40 and return it as a DataFrame in degrees Celsius on a UTC DatetimeIndex, with columns for the temperature at 5,
41 10, 20, 50 and 100 cm depth and the six-hourly minima and maxima at 5 and 10 cm. The download is cached on disk
42 for the calendar day; missing values are interpolated and forward-filled unless disabled. This is the only
43 function here that needs ``requests``, which is therefore imported lazily.
45- :func:`solve_underdetermined_system` - Solve ``A x = b`` for a wide system (more unknowns than equations) by
46 taking the least-squares solution and adding the nullspace component that minimizes a roughness objective;
47 rows holding NaN are dropped first. The closed-form squared-differences optimum is always computed and also
48 seeds the iterative optimization of the other objectives, so a nullspace containing a near-constant vector
49 raises :class:`numpy.linalg.LinAlgError` whichever objective is requested.
51- :func:`compute_reverse_target` - Transpose the forward coefficient matrix, normalize its rows and apply it to
52 the observations, giving each input bin the contribution-weighted average of the output bins it fed. This is
53 the reference solution the Tikhonov solvers pull poorly-determined modes toward; input bins with negligible
54 forward weight are returned as NaN.
56- :func:`solve_tikhonov` - Solve ``min ||A x - b||² + λ ||x - x_target||²`` as a single augmented least-squares
57 problem. Rows of the system holding NaN are excluded, and NaN entries of ``x_target`` are left unregularized.
59- :func:`solve_inverse_transport` - Recover the input signal of the dense forward model
60 ``w_forward @ x = observed`` by building the target with :func:`compute_reverse_target` and solving it with
61 :func:`solve_tikhonov`. Observation rows that are NaN or carry negligible weight drop out (an explicit
62 ``valid_rows`` mask overrides the weight test), and output bins left without forward contribution come back
63 as NaN.
65- :func:`solve_inverse_transport_banded` - Solve the same inverse problem for a forward operator stored in banded
66 layout (row ``k`` is ``band_vals[k]`` placed at column ``col_start[k]``), through banded Cholesky normal
67 equations plus corrected semi-normal refinement. The factorization, solve and refinement stay at
68 ``O(n_output * full_band)``. The regularization strength must be strictly positive, since it is what makes the
69 banded factor positive definite.
71This file is part of gwtransport which is released under AGPL-3.0 license.
72See the ./LICENSE file or go to https://github.com/gwtransport/gwtransport/blob/main/LICENSE for full license details.
73"""
75from __future__ import annotations
77import io
78from collections.abc import Callable
79from datetime import date
80from pathlib import Path
82import numpy as np
83import numpy.typing as npt
84import pandas as pd
85from scipy.linalg import cho_solve_banded, cholesky_banded, null_space
86from scipy.optimize import minimize
88cache_dir = Path(__file__).parent.parent.parent / "cache"
91def step_plot_coords(edges: npt.ArrayLike, values: npt.ArrayLike) -> tuple[npt.NDArray, npt.NDArray]:
92 """Compute step-plot coordinates from bin edges and bin-averaged values.
94 Converts bin edges (n+1) and bin values (n) into paired x/y arrays
95 suitable for plotting piecewise-constant (step) functions with
96 ``ax.plot(x, y)``.
98 Parameters
99 ----------
100 edges : array-like
101 Bin edges (n+1 elements for n bins). Can be numeric, datetime, or
102 any type accepted by :func:`numpy.repeat`.
103 values : array-like
104 Bin-averaged values (n elements), one per bin.
106 Returns
107 -------
108 x : ndarray
109 Step x-coordinates (2n elements). Same dtype as *edges*.
110 y : ndarray
111 Step y-coordinates (2n elements). Same dtype as *values*.
113 Examples
114 --------
115 >>> import numpy as np
116 >>> edges = np.array([0.0, 1.0, 3.0, 6.0])
117 >>> values = np.array([2.0, 5.0, 1.0])
118 >>> x, y = step_plot_coords(edges, values)
119 >>> x
120 array([0., 1., 1., 3., 3., 6.])
121 >>> y
122 array([2., 2., 5., 5., 1., 1.])
123 """
124 x = np.repeat(edges, 2)[1:-1]
125 y = np.repeat(values, 2)
126 return x, y
129_DUP_BUMP_ULPS = 16 # safety factor in ulps; see _make_strictly_monotone docstring
132def _make_strictly_monotone(arr: npt.ArrayLike) -> npt.NDArray[np.floating]:
133 """Bump consecutive duplicates so a non-decreasing array becomes strictly monotone.
135 Returns the input unchanged if no consecutive duplicates are present. Otherwise returns a
136 new array with each duplicate bumped up by ``k * step``, where ``k`` is its 1-based
137 position within the consecutive duplicate run and ``step`` is ``16 * ulp(max(arr))``
138 capped per run so the largest bump stays strictly below the next genuine value above the
139 plateau (``step = min(16 * ulp(max(arr)), gap / (run_len + 1))``). The cap prevents a long
140 run from overshooting a closely-spaced successor; a gap narrower than the run length in
141 ulps is unrepresentable and cannot be separated.
143 The factor of 16 is a safety margin against IEEE 754 rounding noise in ``np.interp``'s
144 linear-interpolation arithmetic, which differs subtly between Linux x86_64 (with FMA)
145 and ARM macOS. A 1-ulp gap, while strictly monotone, can place a downstream query value
146 on the wrong side of a bracket boundary if the intermediate arithmetic rounds 1 ulp away
147 from the exact value. 16 ulps ensures the bracket selection is unambiguous on every
148 platform we support. The perturbation is relative to the array scale:
149 ``bump ≈ 16 * ulp(max(arr)) ≈ 3.5e-15 * max(arr)``, i.e. about 15 significant digits
150 below the data scale and well below physical relevance. The absolute size therefore grows
151 with the cumulative-volume magnitude (e.g. ``~1e-13`` only for ``max(arr) ~ 30``).
153 Parameters
154 ----------
155 arr : array-like
156 1D non-decreasing array (e.g., a cumulative volume sequence ``flow_cum`` that contains
157 plateaus from ``Q = 0`` bins).
159 Returns
160 -------
161 ndarray
162 Strictly monotone array of the same length.
164 Notes
165 -----
166 Use this before passing ``arr`` as ``x_ref`` to a ``V → t`` inversion via
167 :func:`linear_interpolate` or :func:`numpy.interp`. Plateaus in ``arr`` make ``arr⁻¹``
168 multi-valued, and ``np.interp`` would silently pick one of the two limits, biasing
169 integrals over output bins that span the kink.
170 """
171 arr = np.asarray(arr, dtype=float)
172 diffs = np.diff(arr)
173 if not np.any(diffs == 0):
174 return arr
175 ulp_max = np.nextafter(arr.max(), np.inf) - arr.max()
176 n = len(arr)
177 idx = np.arange(n)
178 is_dup = np.concatenate(([False], diffs == 0))
179 # 1-based position of each duplicate within its consecutive run.
180 last_nondup = np.maximum.accumulate(np.where(is_dup, -1, idx))
181 cumcount = np.where(is_dup, idx - last_nondup, 0)
183 # Per-run headroom: each bumped value must stay strictly below the next genuine
184 # (non-duplicate) value above the plateau, otherwise a long run can overshoot a
185 # closely-spaced next value and break monotonicity. ``next_nondup_idx`` is the first
186 # non-duplicate index after each position (``n`` when the run reaches the array end, where
187 # there is no successor and hence no overshoot risk). The gap to that successor caps the
188 # bump step so the last (largest) bump in a run of length L is at most ``L/(L+1)`` of the
189 # gap. A gap narrower than the run length in ulps is unrepresentable and cannot be split.
190 next_nondup_idx = np.minimum.accumulate(np.where(is_dup, n, idx)[::-1])[::-1]
191 has_successor = next_nondup_idx < n
192 gap_to_next = arr[np.clip(next_nondup_idx, 0, n - 1)] - arr[idx]
193 run_len = next_nondup_idx - last_nondup - 1
194 full_step = _DUP_BUMP_ULPS * ulp_max
195 with np.errstate(invalid="ignore", divide="ignore"):
196 capped_step = np.where(has_successor, np.minimum(full_step, gap_to_next / (run_len + 1.0)), full_step)
197 bump = np.where(is_dup, cumcount * capped_step, 0.0)
198 return arr + bump
201def cumulative_flow_volume(
202 flow: npt.ArrayLike, dt_days: npt.ArrayLike, *, strictly_monotone: bool = False
203) -> npt.NDArray[np.floating]:
204 """Cumulative infiltrated/extracted volume from per-bin flow rates.
206 Multiplies each per-bin flow rate by its bin width and accumulates, with a
207 leading zero prepended so the result has one entry per bin edge (n+1 values
208 for n bins). The result is the cumulative volume ``V`` at each time edge.
210 Parameters
211 ----------
212 flow : array-like
213 Flow rate per bin (m³/day), length n.
214 dt_days : array-like
215 Bin widths in days, length n (e.g. ``numpy.diff`` of edge days).
216 strictly_monotone : bool, optional
217 When ``True``, bump consecutive duplicates (plateaus from ``Q = 0``
218 bins) via ``_make_strictly_monotone`` so the cumulative volume is
219 strictly increasing. Required before a V → t inversion; leave ``False``
220 when the plateaus must be preserved. Default is ``False``.
222 Returns
223 -------
224 ndarray
225 Cumulative volume at each edge (length ``len(flow) + 1``), starting at
226 zero.
228 See Also
229 --------
230 ``_make_strictly_monotone`` : Bump duplicates before V → t inversion.
231 """
232 flow_cum = np.concatenate(([0.0], np.cumsum(np.asarray(flow) * np.asarray(dt_days))))
233 return _make_strictly_monotone(flow_cum) if strictly_monotone else flow_cum
236def linear_interpolate(
237 *,
238 x_ref: npt.ArrayLike,
239 y_ref: npt.ArrayLike,
240 x_query: npt.ArrayLike,
241 left: float | None = None,
242 right: float | None = None,
243) -> npt.NDArray[np.floating]:
244 """
245 Linear interpolation using numpy's optimized interp function.
247 Parameters
248 ----------
249 x_ref : array-like
250 Reference x-values, in ascending order.
251 y_ref : array-like
252 Reference y-values corresponding to x_ref.
253 x_query : array-like
254 Query x-values where interpolation is needed. Array may have any shape.
255 left : float, optional
256 Value to return for x_query < x_ref[0].
258 - If ``left=None``: clamp to y_ref[0] (default)
259 - If ``left=float``: use specified value (e.g., ``np.nan``)
261 right : float, optional
262 Value to return for x_query > x_ref[-1].
264 - If ``right=None``: clamp to y_ref[-1] (default)
265 - If ``right=float``: use specified value (e.g., ``np.nan``)
267 Returns
268 -------
269 ndarray
270 Interpolated y-values with the same shape as x_query.
272 Examples
273 --------
274 Basic interpolation with clamping (default):
276 >>> import numpy as np
277 >>> from gwtransport.utils import linear_interpolate
278 >>> x_ref = np.array([1.0, 2.0, 3.0, 4.0])
279 >>> y_ref = np.array([10.0, 20.0, 30.0, 40.0])
280 >>> x_query = np.array([0.5, 1.5, 2.5, 3.5, 4.5])
281 >>> linear_interpolate(x_ref=x_ref, y_ref=y_ref, x_query=x_query)
282 array([10., 15., 25., 35., 40.])
284 Using NaN for extrapolation:
286 >>> linear_interpolate(
287 ... x_ref=x_ref, y_ref=y_ref, x_query=x_query, left=np.nan, right=np.nan
288 ... )
289 array([nan, 15., 25., 35., nan])
290 """
291 return np.interp(np.asarray(x_query), np.asarray(x_ref), np.asarray(y_ref), left=left, right=right)
294def simplify_bins(
295 *,
296 edges: npt.ArrayLike,
297 values: npt.ArrayLike,
298 flow: npt.ArrayLike | None = None,
299 tol: float = 0.0,
300) -> tuple[
301 npt.NDArray[np.floating] | pd.DatetimeIndex,
302 npt.NDArray[np.floating],
303 npt.NDArray[np.floating] | None,
304]:
305 """Simplify a piecewise-constant time series by merging adjacent bins.
307 Splits at the largest value jump until the peak-to-peak range within
308 every group does not exceed `tol`. The result is independent of scan
309 direction.
311 Parameters
312 ----------
313 edges : array-like
314 Bin edges with shape ``(n+1,)``. May be numeric or pandas Timestamps.
315 values : array-like
316 Bin-averaged values with shape ``(n,)`` (e.g., concentrations).
317 flow : array-like, optional
318 Flow rate per bin with shape ``(n,)`` (e.g., m³/day). When provided,
319 merged-bin values are weighted by volume (flow x bin width) instead of
320 bin width alone.
321 tol : float, optional
322 Maximum peak-to-peak range within a merged group.
323 Default is 0.0, which merges only runs of identical values.
325 Returns
326 -------
327 new_edges : ndarray or DatetimeIndex
328 Simplified bin edges with shape ``(m+1,)``, preserving the type of
329 `edges`.
330 new_values : ndarray of float
331 Volume-weighted (or width-weighted) average values per simplified
332 bin, with shape ``(m,)``.
333 new_flow : ndarray of float or None
334 Time-weighted (width-weighted) average flow per simplified bin, with
335 shape ``(m,)``. None when `flow` is not provided.
336 """
337 edges = np.asarray(edges) if not isinstance(edges, pd.DatetimeIndex) else edges
338 values = np.asarray(values, dtype=float)
339 if len(values) == 0:
340 flow_out = np.asarray(flow, dtype=float) if flow is not None else None
341 return edges, values, flow_out
343 widths = np.asarray(np.diff(edges), dtype=float)
344 if flow is not None:
345 flow = np.asarray(flow, dtype=float)
346 weights = widths * flow
347 else:
348 weights = widths
350 # Iteratively split each segment at its largest value jump until every group's peak-to-peak
351 # range is within tol. An explicit LIFO stack replaces the natural recursion, which peels one
352 # element per level on smooth monotone data (argmax|diff| sits at a segment edge) and overflows
353 # the interpreter stack for a few thousand points. Every split index is interior to its
354 # (disjoint) segment, so sorting the collected splits reproduces the recursion's in-order
355 # output exactly -- the merged bins are identical.
356 splits: list[int] = []
357 stack: list[tuple[int, int]] = [(0, len(values))]
358 while stack:
359 lo, hi = stack.pop()
360 if np.ptp(values[lo:hi]) <= tol:
361 continue
362 i = lo + int(np.argmax(np.abs(np.diff(values[lo:hi])))) + 1
363 splits.append(i)
364 stack.extend(((lo, i), (i, hi)))
365 splits.sort()
366 s = np.array([0, *splits])
367 idx = np.append(s, len(values))
368 new_edges = edges[idx]
369 new_widths = np.add.reduceat(widths, s)
370 weight_sums = np.add.reduceat(weights, s)
371 # A merged group of all-zero-flow bins has zero volume weight (0/0 -> NaN); its average is
372 # still well defined, so fall back to width weighting there.
373 zero_weight = weight_sums == 0.0
374 new_values = np.where(
375 zero_weight, np.add.reduceat(widths * values, s), np.add.reduceat(weights * values, s)
376 ) / np.where(zero_weight, new_widths, weight_sums)
377 # When flow is given, weights == flow * widths, so weight_sums == reduceat(flow * widths, s) exactly.
378 new_flow = weight_sums / new_widths if flow is not None else None
380 return new_edges, new_values, new_flow
383def compute_time_edges(
384 *,
385 tedges: pd.DatetimeIndex | None,
386 tstart: pd.DatetimeIndex | None,
387 tend: pd.DatetimeIndex | None,
388 number_of_bins: int,
389) -> pd.DatetimeIndex:
390 """
391 Compute time edges for binning data based on provided time parameters.
393 This function creates a DatetimeIndex of time bin edges from one of three possible
394 input formats: explicit edges, start times, or end times. The resulting edges
395 define the boundaries of time intervals for data binning.
397 Define either explicit time edges, or start and end times for each bin and leave the others at None.
399 Parameters
400 ----------
401 tedges : pandas.DatetimeIndex or None
402 Explicit time edges for the bins. If provided, must have one more element
403 than the number of bins (n_bins + 1). Takes precedence over tstart and tend.
404 tstart : pandas.DatetimeIndex or None
405 Start times for each bin. Must have the same number of elements as the
406 number of bins. Used when tedges is None.
407 tend : pandas.DatetimeIndex or None
408 End times for each bin. Must have the same number of elements as the
409 number of bins. Used when both tedges and tstart are None.
410 number_of_bins : int
411 The expected number of time bins. Used for validation against the provided
412 time parameters.
414 Returns
415 -------
416 pandas.DatetimeIndex
417 Time edges defining the boundaries of the time bins. Has one more element
418 than number_of_bins.
420 Raises
421 ------
422 ValueError
423 If tedges has incorrect length (not number_of_bins + 1).
424 If tstart has incorrect length (not equal to number_of_bins).
425 If tend has incorrect length (not equal to number_of_bins).
426 If none of tedges, tstart, or tend are provided.
428 Notes
429 -----
430 - When using tstart, the function assumes uniform spacing and extrapolates
431 the final edge based on the spacing between the last two start times.
432 - When using tend, the function assumes uniform spacing and extrapolates
433 the first edge based on the spacing between the first two end times.
434 - When ``tstart`` or ``tend`` are provided with non-uniformly-spaced bins,
435 the extrapolated edge uses only the very first or very last interval and
436 may be physically incorrect: the missing edge is implicitly assigned a
437 bin width equal to that single neighbouring interval, which is unrelated
438 to any other interval in the series. In such cases, supply ``tedges``
439 directly so that all bin widths are explicit.
440 - All input time data is converted to pandas.DatetimeIndex for consistency.
441 """
442 if tedges is not None:
443 if number_of_bins != len(tedges) - 1:
444 msg = "tedges must have one more element than number_of_bins"
445 raise ValueError(msg)
446 tedges = pd.DatetimeIndex(tedges)
447 # Ensure nanosecond precision while preserving timezone
448 return tedges.as_unit("ns")
450 if tstart is not None:
451 # Assume the index refers to the time at the start of the measurement interval
452 tstart = pd.DatetimeIndex(tstart).as_unit("ns")
453 if number_of_bins != len(tstart):
454 msg = "tstart must have the same number of elements as number_of_bins"
455 raise ValueError(msg)
456 if len(tstart) < 2: # noqa: PLR2004
457 msg = "tstart must have at least 2 elements to infer the bin width; pass tedges for a single bin"
458 raise ValueError(msg)
460 # Extrapolate final edge using uniform spacing
461 final_edge = tstart[-1] + (tstart[-1] - tstart[-2])
462 return pd.DatetimeIndex([*list(tstart), final_edge], dtype=tstart.dtype)
464 if tend is not None:
465 # Assume the index refers to the time at the end of the measurement interval
466 tend = pd.DatetimeIndex(tend).as_unit("ns")
467 if number_of_bins != len(tend):
468 msg = "tend must have the same number of elements as number_of_bins"
469 raise ValueError(msg)
470 if len(tend) < 2: # noqa: PLR2004
471 msg = "tend must have at least 2 elements to infer the bin width; pass tedges for a single bin"
472 raise ValueError(msg)
474 # Extrapolate initial edge using uniform spacing
475 initial_edge = tend[0] - (tend[1] - tend[0])
476 return pd.DatetimeIndex([initial_edge, *list(tend)], dtype=tend.dtype)
478 msg = "Either provide tedges, tstart, or tend"
479 raise ValueError(msg)
482def get_soil_temperature(*, station_number: int = 260, interpolate_missing_values: bool = True) -> pd.DataFrame:
483 """
484 Download soil temperature data from the KNMI and return it as a pandas DataFrame.
486 The data is available for the following KNMI weather stations:
487 - 260: De Bilt, the Netherlands (vanaf 1981)
488 - 273: Marknesse, the Netherlands (vanaf 1989)
489 - 286: Nieuw Beerta, the Netherlands (vanaf 1990)
490 - 323: Wilhelminadorp, the Netherlands (vanaf 1989)
492 TB1 = grondtemperatuur op 5 cm diepte (graden Celsius) tijdens de waarneming
493 TB2 = grondtemperatuur op 10 cm diepte (graden Celsius) tijdens de waarneming
494 TB3 = grondtemperatuur op 20 cm diepte (graden Celsius) tijdens de waarneming
495 TB4 = grondtemperatuur op 50 cm diepte (graden Celsius) tijdens de waarneming
496 TB5 = grondtemperatuur op 100 cm diepte (graden Celsius) tijdens de waarneming
497 TNB2 = minimum grondtemperatuur op 10 cm diepte in de afgelopen 6 uur (graden Celsius)
498 TNB1 = minimum grondtemperatuur op 5 cm diepte in de afgelopen 6 uur (graden Celsius)
499 TXB1 = maximum grondtemperatuur op 5 cm diepte in de afgelopen 6 uur (graden Celsius)
500 TXB2 = maximum grondtemperatuur op 10 cm diepte in de afgelopen 6 uur (graden Celsius)
502 Parameters
503 ----------
504 station_number : int, {260, 273, 286, 323}
505 The KNMI station number for which to download soil temperature data.
506 Default is 260 (De Bilt).
507 interpolate_missing_values : bool, optional
508 If True, missing values are interpolated and recent NaN values are extrapolated with the previous value.
509 If False, missing values remain as NaN. Default is True.
511 Returns
512 -------
513 pandas.DataFrame
514 DataFrame containing soil temperature data in Celsius with a DatetimeIndex.
515 Columns include TB1, TB2, TB3, TB4, TB5, TNB1, TNB2, TXB1, TXB2.
517 Notes
518 -----
519 - KNMI: Royal Netherlands Meteorological Institute
520 - The timeseries may contain NaN values for missing data.
521 """
522 # File-based daily cache
523 cache_dir.mkdir(exist_ok=True)
525 today = date.today().isoformat() # noqa: DTZ011
526 cache_path = cache_dir / f"soil_temp_{station_number}_{interpolate_missing_values}_{today}.pkl"
528 # Check if cached file exists and is from today
529 if cache_path.exists():
530 cached = pd.read_pickle(cache_path) # noqa: S301
531 assert isinstance(cached, pd.DataFrame) # noqa: S101 -- the cache only ever stores DataFrames
532 return cached
534 # Clean up old cache files to prevent disk bloat
535 for old_file in cache_dir.glob(f"soil_temp_{station_number}_{interpolate_missing_values}_*.pkl"):
536 old_file.unlink(missing_ok=True)
538 url = f"https://cdn.knmi.nl/knmi/map/page/klimatologie/gegevens/bodemtemps/bodemtemps_{station_number}.zip"
540 dtypes = {
541 "YYYYMMDD": "int32",
542 "HH": "int8",
543 " TB1": "float32",
544 " TB3": "float32",
545 " TB2": "float32",
546 " TB4": "float32",
547 " TB5": "float32",
548 " TNB1": "float32",
549 " TNB2": "float32",
550 " TXB1": "float32",
551 " TXB2": "float32",
552 }
554 # Imported lazily so the rest of the module remains importable in environments
555 # without ``requests`` (e.g. Pyodide/JupyterLite, where this KNMI download is the
556 # only feature that cannot run client-side).
557 import requests # noqa: PLC0415
559 # Download the ZIP file
560 with requests.get(url, params={"download": "zip"}, timeout=10) as response:
561 response.raise_for_status()
563 df = pd.read_csv( # type: ignore[call-overload] # ty: ignore[no-matching-overload]
564 io.BytesIO(response.content),
565 compression="zip",
566 dtype=dtypes, # pyright: ignore[reportArgumentType]
567 usecols=list(dtypes.keys()), # pyright: ignore[reportArgumentType]
568 skiprows=16,
569 sep=",",
570 na_values=[" "],
571 engine="c",
572 parse_dates=False,
573 )
575 df.index = pd.to_datetime(df["YYYYMMDD"].values, format=r"%Y%m%d").tz_localize("UTC") + pd.to_timedelta(
576 df["HH"].values, unit="h"
577 )
579 df.drop(columns=["YYYYMMDD", "HH"], inplace=True)
580 df.columns = df.columns.str.strip()
581 df /= 10.0
583 if interpolate_missing_values:
584 # Fill NaN values with interpolate linearly and then forward fill
585 df.interpolate(method="linear", inplace=True)
586 df.ffill(inplace=True)
588 # Save to cache for future use
589 df.to_pickle(cache_path)
590 return df
593def solve_underdetermined_system(
594 *,
595 coefficient_matrix: npt.ArrayLike,
596 rhs_vector: npt.ArrayLike,
597 nullspace_objective: str
598 | Callable[
599 [npt.NDArray[np.floating], npt.NDArray[np.floating], npt.NDArray[np.floating]], float
600 ] = "squared_differences",
601 optimization_method: str = "BFGS",
602 rcond: float | None = None,
603) -> npt.NDArray[np.floating]:
604 """
605 Solve an underdetermined linear system with nullspace regularization.
607 For an underdetermined system Ax = b where A has more columns than rows,
608 multiple solutions exist. This function computes a least-squares solution
609 and then selects a specific solution from the nullspace based on a
610 regularization objective.
612 Parameters
613 ----------
614 coefficient_matrix : array-like
615 Coefficient matrix of shape (m, n) where m < n (underdetermined).
616 May contain NaN values in some rows, which will be excluded from the system.
617 rhs_vector : array-like
618 Right-hand side vector of length m. May contain NaN values corresponding
619 to NaN rows in coefficient_matrix, which will be excluded from the system.
620 nullspace_objective : str or callable, optional
621 Objective function to minimize in the nullspace. Options:
623 * "squared_differences" : Minimize sum of squared differences between
624 adjacent elements: ``sum((x[i+1] - x[i])**2)``
625 * "summed_differences" : Minimize sum of absolute differences between
626 adjacent elements: ``sum(|x[i+1] - x[i]|)``
627 * callable : Custom objective function with signature
628 ``objective(coeffs, x_ls, nullspace_basis)`` where:
630 - coeffs : optimization variables (nullspace coefficients)
631 - x_ls : least-squares solution
632 - nullspace_basis : nullspace basis matrix
634 Default is "squared_differences".
635 optimization_method : str, optional
636 Optimization method passed to scipy.optimize.minimize.
637 Default is "BFGS".
638 rcond : float or None, optional
639 Cutoff ratio for small singular values in both ``numpy.linalg.lstsq``
640 and ``scipy.linalg.null_space``. Singular values smaller than
641 ``rcond * largest_singular_value`` are treated as zero.
642 Default is None, which uses the default of each function.
643 Increasing rcond truncates more modes, expanding the nullspace
644 available for smoothness optimization. Useful for noisy data.
646 Returns
647 -------
648 ndarray
649 Solution vector that minimizes the specified nullspace objective.
650 Has length n (number of columns in coefficient_matrix).
652 Raises
653 ------
654 ValueError
655 If optimization fails, if coefficient_matrix and rhs_vector have incompatible shapes,
656 or if an unknown nullspace objective is specified.
657 numpy.linalg.LinAlgError
658 If the squared-differences normal equations ``(DN)^T(DN)`` are ill-conditioned
659 (condition number above 1e12), which happens when the nullspace contains a
660 near-constant vector.
662 Notes
663 -----
664 The algorithm follows these steps:
666 1. Remove rows with NaN values from both coefficient_matrix and rhs_vector
667 2. Compute least-squares solution: x_ls = pinv(valid_matrix) @ valid_rhs
668 3. Compute nullspace basis: N = null_space(valid_matrix)
669 4. Find nullspace coefficients: coeffs = argmin objective(x_ls + N @ coeffs)
670 5. Return final solution: x = x_ls + N @ coeffs
672 For the built-in objectives:
674 * "squared_differences" provides smooth solutions, minimizing rapid changes
675 * "summed_differences" provides sparse solutions, promoting piecewise constant behavior
677 Examples
678 --------
679 Rows containing NaN are dropped; the solution satisfies the remaining equations:
681 >>> import numpy as np
682 >>> from gwtransport.utils import solve_underdetermined_system
683 >>> matrix = np.array([
684 ... [1.0, 2.0, 1.0, 0.0],
685 ... [np.nan, np.nan, np.nan, np.nan],
686 ... [0.0, 1.0, 2.0, 1.0],
687 ... ])
688 >>> rhs = np.array([3.0, np.nan, 4.0])
689 >>> x = solve_underdetermined_system(coefficient_matrix=matrix, rhs_vector=rhs)
690 >>> bool(np.allclose(matrix[[0, 2]] @ x, [3.0, 4.0]))
691 True
692 """
693 matrix = np.asarray(coefficient_matrix)
694 rhs = np.asarray(rhs_vector)
696 if matrix.shape[0] != len(rhs):
697 msg = f"coefficient_matrix has {matrix.shape[0]} rows but rhs_vector has {len(rhs)} elements"
698 raise ValueError(msg)
700 # Identify valid rows (no NaN values in either matrix or rhs)
701 valid_rows = ~np.isnan(matrix).any(axis=1) & ~np.isnan(rhs)
703 if not np.any(valid_rows):
704 msg = "No valid rows found (all contain NaN values)"
705 raise ValueError(msg)
707 valid_matrix = matrix[valid_rows]
708 valid_rhs = rhs[valid_rows]
710 # Compute least-squares solution
711 x_ls, *_ = np.linalg.lstsq(valid_matrix, valid_rhs, rcond=rcond)
713 # Compute nullspace
714 nullspace_basis = null_space(valid_matrix, rcond=rcond)
716 if nullspace_basis.shape[1] == 0:
717 # System is determined, return least-squares solution
718 return x_ls
720 # Squared-differences optimum in closed form: minimizing ||D(x_ls + N c)||^2 gives the normal
721 # equations (DN)^T(DN) c = -(DN)^T(D x_ls), where D is the (n-1, n) first-difference matrix.
722 # D @ N and D @ x_ls are formed directly instead of materializing D.
723 dn = nullspace_basis[1:, :] - nullspace_basis[:-1, :] # (n-1, nullrank)
724 dx = x_ls[1:] - x_ls[:-1] # (n-1,)
725 dntdn = dn.T @ dn # (nullrank, nullrank)
727 cond = np.linalg.cond(dntdn)
728 cond_threshold = 1e12
729 if cond > cond_threshold:
730 msg = (
731 f"The normal equations matrix (DN)^T(DN) is ill-conditioned "
732 f"(condition number: {cond:.2e}). This typically means the "
733 f"nullspace contains a near-constant vector, so the "
734 f"squared-differences objective cannot distinguish between "
735 f"nullspace directions. Consider using a different "
736 f"nullspace_objective (e.g., 'summed_differences'), reducing "
737 f"the problem's degrees of freedom, or lowering rcond to "
738 f"shrink the nullspace (if the near-constant vector has a "
739 f"small but non-zero singular value)."
740 )
741 raise np.linalg.LinAlgError(msg)
743 coeffs = np.linalg.solve(dntdn, -(dn.T @ dx))
745 if nullspace_objective != "squared_differences":
746 # Other objectives are optimized iteratively, started from the squared-differences
747 # solution for stability.
748 if nullspace_objective == "summed_differences":
749 objective_func = _summed_differences_objective
750 elif callable(nullspace_objective):
751 objective_func = nullspace_objective
752 else:
753 msg = f"Unknown nullspace objective: {nullspace_objective}"
754 raise ValueError(msg)
756 res = minimize(objective_func, x0=coeffs, args=(x_ls, nullspace_basis), method=optimization_method)
757 if not res.success:
758 msg = f"Optimization failed: {res.message}"
759 raise ValueError(msg)
760 coeffs = res.x
762 return x_ls + nullspace_basis @ coeffs
765def compute_reverse_target(
766 *,
767 coeff_matrix: npt.NDArray[np.floating],
768 rhs_vector: npt.NDArray[np.floating],
769) -> npt.NDArray[np.floating]:
770 """Compute reverse matrix target from forward coefficient matrix.
772 Constructs a target solution for the inverse problem by transposing the
773 forward coefficient matrix and normalizing rows. For ``W_forward[i,j]``
774 representing the fraction of ``cin[j]`` arriving in ``cout[i]``, the
775 transpose-and-normalize approach reconstructs ``cin[j]`` as a weighted
776 average of ``cout`` bins, weighted by how much ``cin[j]`` contributed
777 to each ``cout`` bin.
779 Parameters
780 ----------
781 coeff_matrix : ndarray
782 Forward coefficient matrix of shape (n_cout, n_cin).
783 rhs_vector : ndarray
784 Right-hand side vector of length n_cout (e.g., cout values).
786 Returns
787 -------
788 ndarray
789 Target solution vector of length n_cin. Entries with near-zero
790 column sums in the forward matrix are set to NaN.
792 See Also
793 --------
794 solve_tikhonov : Consumes this target as the regularization reference.
795 """
796 min_row_sum = 1e-10
797 wt = coeff_matrix.T # (n_cin, n_cout)
798 row_sums = wt.sum(axis=1)
799 valid = row_sums > min_row_sum
800 w_reverse = np.zeros_like(wt)
801 w_reverse[valid] = wt[valid] / row_sums[valid, None]
802 x_target = w_reverse @ rhs_vector
803 x_target[~valid] = np.nan
804 return x_target
807def solve_tikhonov(
808 *,
809 coefficient_matrix: npt.ArrayLike,
810 rhs_vector: npt.ArrayLike,
811 x_target: npt.NDArray[np.floating],
812 regularization_strength: float = 1e-10,
813) -> npt.NDArray[np.floating]:
814 """Solve a linear system with Tikhonov regularization toward a target.
816 Minimizes ``||A x - b||² + λ ||x - x_target||²`` by solving the
817 equivalent augmented least-squares problem::
819 [A; √λ I_v] x = [b; √λ x_target_v]
821 where ``I_v`` selects only entries where ``x_target`` is not NaN.
823 Well-determined modes (large singular values relative to √λ) are
824 dominated by the data; poorly-determined modes are pulled toward
825 ``x_target``. The solution varies continuously with λ, unlike the
826 hard singular-value cutoff of ``rcond`` in truncated SVD.
828 Parameters
829 ----------
830 coefficient_matrix : array-like
831 Coefficient matrix of shape (m, n). May contain NaN rows, which
832 are excluded from the system.
833 rhs_vector : array-like
834 Right-hand side vector of length m. May contain NaN values
835 corresponding to NaN rows in coefficient_matrix.
836 x_target : ndarray
837 Target solution of length n, typically from
838 :func:`compute_reverse_target`. NaN entries are excluded from the
839 regularization term.
840 regularization_strength : float, optional
841 Tikhonov parameter λ. Controls the tradeoff between fitting the
842 data and staying close to ``x_target``. Larger values trust the
843 target more; smaller values trust the data more. Default is 1e-10.
845 A good starting value for noisy data is
846 ``λ ≈ (noise_std / signal_amplitude)²``. For noiseless synthetic
847 data, the default 1e-10 preserves machine precision.
849 Returns
850 -------
851 ndarray
852 Solution vector of length n.
854 Raises
855 ------
856 ValueError
857 If ``coefficient_matrix`` and ``rhs_vector`` have incompatible shapes, or if
858 all rows contain NaN values.
860 See Also
861 --------
862 compute_reverse_target : Compute the regularization target from the
863 forward matrix.
864 solve_underdetermined_system : Alternative solver using nullspace
865 optimization.
866 """
867 matrix = np.asarray(coefficient_matrix)
868 rhs = np.asarray(rhs_vector)
870 if matrix.shape[0] != len(rhs):
871 msg = f"coefficient_matrix has {matrix.shape[0]} rows but rhs_vector has {len(rhs)} elements"
872 raise ValueError(msg)
874 # Filter NaN rows
875 valid_rows = ~np.isnan(matrix).any(axis=1) & ~np.isnan(rhs)
877 if not np.any(valid_rows):
878 msg = "No valid rows found (all contain NaN values)"
879 raise ValueError(msg)
881 valid_matrix = matrix[valid_rows]
882 valid_rhs = rhs[valid_rows]
884 n_cin = valid_matrix.shape[1]
885 sqrt_lam = np.sqrt(regularization_strength)
887 # Only regularize entries where x_target is valid
888 valid_target = ~np.isnan(x_target)
889 target_indices = np.where(valid_target)[0]
891 # Build augmented system: [A; √λ I_v] x = [b; √λ x_target_v]
892 n_reg = len(target_indices)
893 reg_matrix = np.zeros((n_reg, n_cin))
894 reg_matrix[np.arange(n_reg), target_indices] = sqrt_lam
895 reg_rhs = sqrt_lam * x_target[target_indices]
897 augmented_matrix = np.vstack([valid_matrix, reg_matrix])
898 augmented_rhs = np.concatenate([valid_rhs, reg_rhs])
900 x, *_ = np.linalg.lstsq(augmented_matrix, augmented_rhs, rcond=None)
901 return x
904# Numerical tolerance for coefficient sum to determine valid output bins
905_EPSILON_COEFF_SUM = 1e-10
907# Corrected semi-normal-equation refinement steps in solve_inverse_transport_banded. One
908# step reaches the QR-accurate solution; a second is a cheap, stable safety margin.
909_BANDED_REFINEMENT_STEPS = 2
912def solve_inverse_transport(
913 *,
914 w_forward: npt.NDArray[np.floating],
915 observed: npt.NDArray[np.floating],
916 n_output: int,
917 regularization_strength: float,
918 valid_rows: npt.NDArray[np.bool_] | None = None,
919) -> npt.NDArray[np.floating]:
920 """Solve the inverse transport problem via Tikhonov regularization.
922 Given the forward model ``w_forward @ x = observed``, recovers ``x`` by
923 building the regularization target from the transpose of ``w_forward`` and
924 solving the regularized least-squares problem.
926 Parameters
927 ----------
928 w_forward : ndarray
929 Forward coefficient matrix with shape ``(n_obs, n_output)``.
930 observed : ndarray
931 Observed values with shape ``(n_obs,)`` (e.g., extraction
932 concentrations). NaN entries mark measurement gaps; their rows are
933 excluded from the solve and the regularization target.
934 n_output : int
935 Length of the output vector (e.g., number of cin bins).
936 regularization_strength : float
937 Tikhonov regularization parameter.
938 valid_rows : ndarray of bool, optional
939 Which observation rows are valid, with shape ``(n_obs,)``. If None,
940 rows with ``row_sum > 1e-10`` are considered valid.
942 Returns
943 -------
944 ndarray
945 Recovered signal with shape ``(n_output,)``. NaN for bins with no
946 active columns.
948 See Also
949 --------
950 solve_inverse_transport_banded : Memory-light banded equivalent.
951 """
952 row_sums = w_forward.sum(axis=1)
953 nan_obs = np.isnan(observed)
954 # Aliases w_forward when there are no gaps; otherwise one masked copy, shared with the
955 # regularization target below.
956 w_masked = np.where(nan_obs[:, None], 0.0, w_forward) if nan_obs.any() else w_forward
957 # A column is active when its weight over the surviving rows exceeds the regularization
958 # epsilon; sliver-support and gap-only columns emit NaN instead of a min-norm value.
959 col_active: npt.NDArray[np.bool_] = w_masked.sum(axis=0) > _EPSILON_COEFF_SUM
961 if not np.any(col_active):
962 return np.full(n_output, np.nan)
964 # Gapped rows drop out of the data equations and the regularization target.
965 valid: npt.NDArray[np.bool_] = (row_sums > _EPSILON_COEFF_SUM if valid_rows is None else valid_rows) & ~nan_obs
967 rhs = np.where(valid, row_sums * observed, np.nan)
968 w_solve = w_forward.copy()
969 w_solve[~valid, :] = np.nan
971 x_target = compute_reverse_target(
972 coeff_matrix=w_masked,
973 rhs_vector=np.where(nan_obs, 0.0, observed),
974 )
976 x_solved = solve_tikhonov(
977 coefficient_matrix=w_solve,
978 rhs_vector=rhs,
979 x_target=x_target,
980 regularization_strength=regularization_strength,
981 )
983 out = np.full(n_output, np.nan)
984 idx = np.flatnonzero(col_active)
985 out[idx] = x_solved[idx]
986 return out
989def solve_inverse_transport_banded(
990 *,
991 band_vals: npt.NDArray[np.floating],
992 col_start: npt.NDArray[np.intp],
993 observed: npt.NDArray[np.floating],
994 n_output: int,
995 regularization_strength: float,
996) -> npt.NDArray[np.floating]:
997 """Solve the inverse transport problem from a banded forward operator.
999 Memory-light equivalent of :func:`solve_inverse_transport` for a forward
1000 weight matrix stored in banded layout: row ``k`` of the dense operator
1001 ``W`` is ``band_vals[k]`` placed at columns
1002 ``[col_start[k], col_start[k] + full_band)``. The Tikhonov normal
1003 equations ``(WᵀW + λ D) x = Wᵀ observed + λ D x_target`` are stored **in
1004 banded form** -- ``WᵀW`` is symmetric with half-bandwidth ``full_band - 1``
1005 -- and Cholesky-factored with :func:`scipy.linalg.cholesky_banded`. The Gram
1006 matrix ``WᵀW`` is built with a single dense BLAS matmul (``~24x`` a
1007 per-diagonal scatter) before its sub-diagonals are read into the banded
1008 layout. Forming ``WᵀW`` squares the condition number, so the bare Cholesky
1009 solve loses accuracy in the under-determined (spin-up nullspace) directions;
1010 **corrected semi-normal equations** restore it by refining with the residual
1011 evaluated through ``W`` itself rather than ``WᵀW`` (matching the dense
1012 least-squares solution to ~1e-7 at the default regularization, degrading to
1013 ~1e-6 only at very small regularization with an ill-conditioned Gram). The
1014 banded Cholesky factor, solve, and refinement stay at
1015 ``O(n_output * full_band)``; only the one-shot Gram assembly transiently
1016 materializes ``W`` and ``WᵀW`` densely.
1018 The regularization target ``x_target`` is the transpose-and-normalize of
1019 ``W`` applied to ``observed`` (the banded form of
1020 :func:`compute_reverse_target`), matching the dense solver. Columns with no
1021 forward contribution are decoupled (unit diagonal) so the system stays
1022 symmetric positive definite, and are returned as NaN.
1024 Parameters
1025 ----------
1026 band_vals : ndarray
1027 Banded forward weights of shape ``(n_obs, full_band)``. Rows the caller
1028 considers invalid must already be zeroed (as ``_resolve_spinup_mask``
1029 does); zero rows contribute nothing to the normal equations.
1030 col_start : ndarray of int
1031 First output-column index of each row's band, shape ``(n_obs,)``.
1032 observed : ndarray
1033 Observed values of shape ``(n_obs,)`` (e.g. extraction concentrations).
1034 NaN entries mark measurement gaps; their rows are excluded from the
1035 normal equations (band row and observed value zeroed).
1036 n_output : int
1037 Length of the output vector (number of cin bins).
1038 regularization_strength : float
1039 Tikhonov parameter λ. See :func:`solve_inverse_transport`. Must be
1040 strictly positive: deconvolution is generically rank-deficient, and λ
1041 is what makes the banded Cholesky factor positive definite (unlike the
1042 dense least-squares path, this solver cannot return a λ=0 min-norm
1043 solution).
1045 Returns
1046 -------
1047 ndarray
1048 Recovered signal of shape ``(n_output,)``. NaN for output bins with no
1049 forward contribution (zero column).
1051 Raises
1052 ------
1053 ValueError
1054 If ``regularization_strength`` is not strictly positive.
1056 See Also
1057 --------
1058 solve_inverse_transport : Dense-matrix equivalent.
1059 ``gwtransport.advection_utils._infiltration_to_extraction_weights`` : Banded builder.
1060 """
1061 if regularization_strength <= 0:
1062 msg = "regularization_strength must be > 0 for the banded inverse (Tikhonov positive-definiteness)"
1063 raise ValueError(msg)
1064 # Precondition: the caller's valid rows sum to 1 (guaranteed by
1065 # _resolve_spinup_mask), so the data equation is W x ≈ observed and the RHS
1066 # needs no row_sums scaling -- matching the dense solve_inverse_transport.
1067 band_vals = np.asarray(band_vals, dtype=float)
1068 observed = np.asarray(observed, dtype=float)
1069 # Zeroed gapped rows drop out of the normal equations, and a zeroed observed value keeps
1070 # 0 * NaN out of Wᵀ·observed and the refinement residual.
1071 nan_obs = np.isnan(observed)
1072 if nan_obs.any():
1073 band_vals = np.where(nan_obs[:, None], 0.0, band_vals)
1074 observed = np.where(nan_obs, 0.0, observed)
1075 full_band = band_vals.shape[1]
1076 n_cin = n_output
1077 cols = col_start[:, None] + np.arange(full_band)[None, :] # (n_obs, full_band) output-column index
1078 in_range = cols < n_cin
1079 cols_clipped = np.clip(cols, 0, n_cin - 1)
1081 # Column sums and Wᵀ observed (the reverse-target numerator) by scattering the band.
1082 col_sum = np.zeros(n_cin)
1083 wt_observed = np.zeros(n_cin)
1084 np.add.at(col_sum, cols_clipped[in_range], band_vals[in_range])
1085 np.add.at(wt_observed, cols_clipped[in_range], (band_vals * observed[:, None])[in_range])
1087 col_active = col_sum > 0
1088 if not np.any(col_active):
1089 return np.full(n_output, np.nan)
1091 # Reverse-target: transpose-and-normalize W applied to observed (banded form of
1092 # compute_reverse_target). The sliver 0 < col_sum <= _EPSILON_COEFF_SUM is left
1093 # untargeted (filled with 0) as in the dense path.
1094 with np.errstate(invalid="ignore", divide="ignore"):
1095 x_target = np.where(col_sum > _EPSILON_COEFF_SUM, wt_observed / col_sum, 0.0)
1097 # Lower-banded WᵀW via a dense BLAS matmul. Materialize the forward operator W densely
1098 # (row k is band_vals[k] at columns [col_start[k], col_start[k] + full_band)), form the
1099 # symmetric Gram matrix WᵀW with a single optimized matmul, then read its lower sub-diagonals
1100 # into the banded layout (band row d is the d-th sub-diagonal, WᵀW[j + d, j]). Each row's
1101 # in-range band columns are distinct, so the scatter into W needs no accumulation. This is
1102 # ~24x the per-diagonal np.add.at scatter; the matmul reorders the summation, so ab matches
1103 # the scatter to ~1e-13 -- well inside the Tikhonov + refinement tolerance.
1104 n_obs = band_vals.shape[0]
1105 w_dense = np.zeros((n_obs, n_cin))
1106 obs_idx = np.broadcast_to(np.arange(n_obs)[:, None], cols.shape)
1107 w_dense[obs_idx[in_range], cols_clipped[in_range]] = band_vals[in_range]
1108 gram = w_dense.T @ w_dense
1109 ab = np.zeros((full_band, n_cin))
1110 for d in range(full_band):
1111 ab[d, : n_cin - d] = np.diagonal(gram, offset=-d)
1113 lam = regularization_strength
1114 d_reg = lam * col_active
1115 ab[0] += d_reg
1116 # d_reg is zero off the active columns, so x_target needs no masking here or in
1117 # the refinement loop: the product d_reg * x_target vanishes wherever col_active is False.
1118 rhs = wt_observed + d_reg * x_target
1120 # Decouple zero (inactive, unregularized) diagonals so the matrix is SPD.
1121 dead = ab[0] <= 0.0
1122 ab[0, dead] = 1.0
1123 rhs[dead] = 0.0
1125 factor = cholesky_banded(ab, lower=True)
1126 x = cho_solve_banded((factor, True), rhs)
1128 # Forming WᵀW squares the condition number, so the bare Cholesky solution loses
1129 # accuracy in the under-determined (spin-up nullspace) directions. Corrected
1130 # semi-normal equations recover it: the residual is evaluated through W itself
1131 # (in observation space) rather than through WᵀW, avoiding the cancellation that
1132 # makes plain normal-equation refinement stall. One step reaches the QR-accurate
1133 # solution; the rest are a safety margin (the iteration's fixed point is stable).
1134 for _ in range(_BANDED_REFINEMENT_STEPS):
1135 gathered = x[cols_clipped]
1136 gathered[~in_range] = 0.0
1137 residual = observed - (band_vals * gathered).sum(axis=1) # b - W x (n_obs,)
1138 gradient = np.zeros(n_cin)
1139 np.add.at(gradient, cols_clipped[in_range], (band_vals * residual[:, None])[in_range]) # Wᵀ (b - W x)
1140 gradient += d_reg * (x_target - x)
1141 gradient[dead] = 0.0
1142 x += cho_solve_banded((factor, True), gradient)
1144 out = np.full(n_output, np.nan)
1145 out[col_active] = x[col_active]
1146 return out
1149def _summed_differences_objective(
1150 coeffs: npt.NDArray[np.floating], x_ls: npt.NDArray[np.floating], nullspace_basis: npt.NDArray[np.floating]
1151) -> float:
1152 """Minimize sum of absolute differences between adjacent elements.
1154 Parameters
1155 ----------
1156 coeffs : ndarray
1157 Nullspace coefficient vector.
1158 x_ls : ndarray
1159 Least-squares solution vector.
1160 nullspace_basis : ndarray
1161 Nullspace basis matrix.
1163 Returns
1164 -------
1165 float
1166 Sum of absolute differences between adjacent elements of the solution.
1167 """
1168 x = x_ls + nullspace_basis @ coeffs
1169 return np.sum(np.abs(x[1:] - x[:-1]))