Coverage for src / gwtransport / deposition_utils.py: 100%
26 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-27 06:32 +0000
1"""
2Utility Functions for the Deposition Module.
4This module provides geometric utilities used by the deposition module for computing
5surface areas and average heights between streamlines.
7Available functions:
9- :func:`compute_average_heights` - Compute average heights of clipped trapezoids formed by
10 streamlines. Trapezoids have vertical sides defined by x_edges and top/bottom edges defined
11 by y_edges (2D array). Clipping bounds (y_lower, y_upper) restrict the integration domain.
12 Returns area/width ratios representing average heights for use in pore volume calculations
13 from streamline geometry.
15This file is part of gwtransport which is released under AGPL-3.0 license.
16See the ./LICENSE file or go to https://github.com/gwtransport/gwtransport/blob/main/LICENSE for full license details.
17"""
19import numpy as np
20import numpy.typing as npt
23def _positive_part_integral(
24 a: npt.NDArray[np.floating], b: npt.NDArray[np.floating], w: npt.NDArray[np.floating]
25) -> npt.NDArray[np.floating]:
26 """
27 Integrate max(f(x), 0) from x=0 to x=w where f is linear from a to b.
29 Parameters
30 ----------
31 a : ndarray
32 Function values at x=0.
33 b : ndarray
34 Function values at x=w.
35 w : ndarray
36 Integration width.
38 Returns
39 -------
40 ndarray
41 Integral values.
42 """
43 both_pos = (a > 0) & (b > 0)
44 only_a_pos = (a > 0) & ~both_pos
45 only_b_pos = (b > 0) & ~both_pos
47 abs_diff = np.abs(a - b)
48 # Sentinel ``1.0`` avoids division by zero in the ``excess**2 / (2*safe_diff)``
49 # branch when a == b; the surrounding ``np.where`` discards this branch
50 # whenever both endpoints have the same sign (where the trapezoid formula
51 # is used instead), so the sentinel value is never observed in the output.
52 safe_diff = np.where(abs_diff > 0, abs_diff, 1.0)
54 excess = np.where(only_a_pos, a, np.where(only_b_pos, b, 0.0))
56 return np.where(
57 both_pos,
58 w * (a + b) / 2,
59 w * excess**2 / (2 * safe_diff),
60 )
63def _clipped_linear_integral(
64 a: npt.NDArray[np.floating],
65 b: npt.NDArray[np.floating],
66 w: npt.NDArray[np.floating],
67 lo: float,
68 hi: float,
69) -> npt.NDArray[np.floating]:
70 """
71 Integrate clip(f(x), lo, hi) from x=0 to x=w where f is linear from a to b.
73 Uses the identity ``clip(f) = f - max(f - hi, 0) + max(lo - f, 0)`` to
74 compute the exact integral analytically.
76 Parameters
77 ----------
78 a : ndarray
79 Function values at x=0.
80 b : ndarray
81 Function values at x=w.
82 w : ndarray
83 Integration width.
84 lo : float
85 Lower clipping bound.
86 hi : float
87 Upper clipping bound.
89 Returns
90 -------
91 ndarray
92 Integral values.
93 """
94 raw = w * (a + b) / 2
95 excess_above = _positive_part_integral(a - hi, b - hi, w)
96 deficit_below = _positive_part_integral(lo - a, lo - b, w)
97 return raw - excess_above + deficit_below
100def compute_average_heights(
101 *, x_edges: npt.ArrayLike, y_edges: npt.ArrayLike, y_lower: float, y_upper: float
102) -> npt.NDArray[np.floating]:
103 """
104 Compute average heights of clipped trapezoids.
106 Trapezoids have vertical left and right sides, with corners at:
107 - top-left: (y_edges[i, j], x_edges[j])
108 - top-right: (y_edges[i, j+1], x_edges[j+1])
109 - bottom-left: (y_edges[i+1, j], x_edges[j])
110 - bottom-right: (y_edges[i+1, j+1], x_edges[j+1])
112 The area is computed as the exact integral of
113 ``clip(top(x), y_lower, y_upper) - clip(bottom(x), y_lower, y_upper)``
114 where ``top(x)`` and ``bottom(x)`` are linear interpolants of the corner
115 values, and ``clip`` restricts values to ``[y_lower, y_upper]``.
117 Parameters
118 ----------
119 x_edges : numpy.ndarray
120 1D array of x coordinates, shape (n_x,)
121 y_edges : numpy.ndarray
122 2D array of y coordinates, shape (n_y, n_x)
123 y_lower : float
124 Lower horizontal clipping bound
125 y_upper : float
126 Upper horizontal clipping bound
128 Returns
129 -------
130 avg_heights : numpy.ndarray
131 2D array of average heights (area/width) for each clipped trapezoid,
132 shape (n_y-1, n_x-1)
134 See Also
135 --------
136 gwtransport.advection.infiltration_to_extraction : Use computed areas in transport calculations
138 Examples
139 --------
140 >>> import numpy as np
141 >>> from gwtransport.deposition_utils import compute_average_heights
142 >>> # Create simple grid
143 >>> x_edges = np.array([0.0, 1.0, 2.0])
144 >>> y_edges = np.array([[0.0, 0.5, 1.0], [1.0, 1.5, 2.0], [2.0, 2.5, 3.0]])
145 >>> # Compute average heights with clipping bounds
146 >>> avg_heights = compute_average_heights(
147 ... x_edges=x_edges, y_edges=y_edges, y_lower=0.5, y_upper=2.5
148 ... )
149 >>> print(f"Shape: {avg_heights.shape}")
150 Shape: (2, 2)
151 """
152 x_edges = np.asarray(x_edges, dtype=float)
153 y_edges = np.asarray(y_edges, dtype=float)
155 y_tl, y_tr = y_edges[:-1, :-1], y_edges[:-1, 1:]
156 y_bl, y_br = y_edges[1:, :-1], y_edges[1:, 1:]
157 widths = np.diff(x_edges)
159 top_integral = _clipped_linear_integral(y_tl, y_tr, widths, y_lower, y_upper)
160 bottom_integral = _clipped_linear_integral(y_bl, y_br, widths, y_lower, y_upper)
162 areas = np.maximum(top_integral - bottom_integral, 0.0)
164 with np.errstate(invalid="ignore"):
165 return areas / widths