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

1""" 

2Utility Functions for the Deposition Module. 

3 

4This module provides the clipped-trapezoid integral helper ``_clipped_linear_integral``, which the 

5deposition module's banded weight builder uses to integrate ``clip(f(x), lo, hi)`` over each cin bin 

6inside an output bin's residence window, and ``_positive_part_integral``, the positive-part integral 

7it is built from. 

8 

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

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

11""" 

12 

13import numpy as np 

14import numpy.typing as npt 

15 

16 

17def _positive_part_integral( 

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

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

20 """ 

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

22 

23 Parameters 

24 ---------- 

25 a : ndarray 

26 Function values at x=0. 

27 b : ndarray 

28 Function values at x=w. 

29 w : ndarray 

30 Integration width. 

31 

32 Returns 

33 ------- 

34 ndarray 

35 Integral values. 

36 """ 

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

38 

39 abs_diff = np.abs(a - b) 

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

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

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

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

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

45 

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

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

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

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

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

51 

52 return np.where( 

53 both_pos, 

54 w * (a + b) / 2, 

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

56 ) 

57 

58 

59def _clipped_linear_integral( 

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

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

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

63 lo: float, 

64 hi: float, 

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

66 """ 

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

68 

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

70 compute the exact integral analytically. 

71 

72 Parameters 

73 ---------- 

74 a : ndarray 

75 Function values at x=0. 

76 b : ndarray 

77 Function values at x=w. 

78 w : ndarray 

79 Integration width. 

80 lo : float 

81 Lower clipping bound. 

82 hi : float 

83 Upper clipping bound. 

84 

85 Returns 

86 ------- 

87 ndarray 

88 Integral values. 

89 """ 

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

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

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

93 return raw - excess_above + deficit_below