Coverage for src/gwtransport/deposition_utils.py: 100%

1""" 

2Utility Functions for the Deposition Module. 

3 

4This module provides the clipped-trapezoid integral helpers (``_clipped_linear_integral`` 

5and ``_positive_part_integral``) used by the deposition module's banded weight builder to 

6integrate ``clip(y(x), y_lower, y_upper)`` over each cin bin of a streamtube's residence window. 

7 

8This file is part of gwtransport which is released under AGPL-3.0 license. 

9See the ./LICENSE file or go to https://github.com/gwtransport/gwtransport/blob/main/LICENSE for full license details. 

10""" 

11 

12import numpy as np 

13import numpy.typing as npt 

14 

15 

16def _positive_part_integral( 

17 a: npt.NDArray[np.floating], b: npt.NDArray[np.floating], w: npt.NDArray[np.floating] 

18) -> npt.NDArray[np.floating]: 

19 """ 

20 Integrate max(f(x), 0) from x=0 to x=w where f is linear from a to b. 

21 

22 Parameters 

23 ---------- 

24 a : ndarray 

25 Function values at x=0. 

26 b : ndarray 

27 Function values at x=w. 

28 w : ndarray 

29 Integration width. 

30 

31 Returns 

32 ------- 

33 ndarray 

34 Integral values. 

35 """ 

36 both_pos = (a > 0) & (b > 0) 

37 

38 abs_diff = np.abs(a - b) 

39 # Sentinel ``1.0`` avoids division by zero in the ``excess**2 / (2*safe_diff)`` 

40 # branch when a == b; the surrounding ``np.where`` discards this branch 

41 # whenever both endpoints have the same sign (where the trapezoid formula 

42 # is used instead), so the sentinel value is never observed in the output. 

43 safe_diff = np.where(abs_diff > 0, abs_diff, 1.0) 

44 

45 # When exactly one endpoint is positive, ``excess`` is that endpoint; 

46 # when neither is positive it is 0. ``max(max(a, b), 0)`` yields both: 

47 # the positive endpoint when only one is positive, 0 otherwise. The 

48 # both-positive case is discarded by the ``np.where`` below. 

49 excess = np.maximum(np.maximum(a, b), 0.0) 

50 

51 return np.where( 

52 both_pos, 

53 w * (a + b) / 2, 

54 w * excess**2 / (2 * safe_diff), 

55 ) 

56 

57 

58def _clipped_linear_integral( 

59 a: npt.NDArray[np.floating], 

60 b: npt.NDArray[np.floating], 

61 w: npt.NDArray[np.floating], 

62 lo: float, 

63 hi: float, 

64) -> npt.NDArray[np.floating]: 

65 """ 

66 Integrate clip(f(x), lo, hi) from x=0 to x=w where f is linear from a to b. 

67 

68 Uses the identity ``clip(f) = f - max(f - hi, 0) + max(lo - f, 0)`` to 

69 compute the exact integral analytically. 

70 

71 Parameters 

72 ---------- 

73 a : ndarray 

74 Function values at x=0. 

75 b : ndarray 

76 Function values at x=w. 

77 w : ndarray 

78 Integration width. 

79 lo : float 

80 Lower clipping bound. 

81 hi : float 

82 Upper clipping bound. 

83 

84 Returns 

85 ------- 

86 ndarray 

87 Integral values. 

88 """ 

89 raw = w * (a + b) / 2 

90 excess_above = _positive_part_integral(a - hi, b - hi, w) 

91 deficit_below = _positive_part_integral(lo - a, lo - b, w) 

92 return raw - excess_above + deficit_below