Coverage for C:\src\imod-python\imod\evaluate\boundaries.py: 93%

75 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-16 11:41 +0200

1import numba 

2import numpy as np 

3import xarray as xr 

4 

5 

6@numba.njit 

7def _interpolate_value_boundaries_z1d(values, z, dz, threshold, out): 

8 nlay, nrow, ncol = values.shape 

9 for i in range(nrow): 

10 for j in range(ncol): 

11 n = 0 

12 if values[0, i, j] >= threshold: # top cell exceeds threshold 

13 out[n, i, j] = z[0] + 0.5 * dz[0] 

14 n += 1 

15 for k in range(1, nlay): 

16 # from top downward 

17 if (n % 2 == 0 and values[k, i, j] >= threshold) or ( 

18 n % 2 == 1 and values[k, i, j] <= threshold 

19 ): # exceeding (n even) or falling below threshold (n odd) 

20 if values[k, i, j] == values[k - 1, i, j]: 

21 # edge case: subsequent values equal threshold result in ZeroDivisionError 

22 continue 

23 if not np.isnan(values[k - 1, i, j]): 

24 # interpolate z coord of threshold 

25 out[n, i, j] = z[k - 1] + (threshold - values[k - 1, i, j]) * ( 

26 z[k] - z[k - 1] 

27 ) / (values[k, i, j] - values[k - 1, i, j]) 

28 else: 

29 # set to top of layer 

30 out[n, i, j] = z[k] + 0.5 * dz[k] 

31 n += 1 

32 

33 

34@numba.njit 

35def _interpolate_value_boundaries_z3d(values, z, dz, threshold, out): 

36 nlay, nrow, ncol = values.shape 

37 for i in range(nrow): 

38 for j in range(ncol): 

39 n = 0 

40 if values[0, i, j] >= threshold: # top cell exceeds threshold 

41 out[n, i, j] = z[0, i, j] + 0.5 * dz[0, i, j] 

42 n += 1 

43 for k in range(1, nlay): 

44 # from top downward 

45 if (n % 2 == 0 and values[k, i, j] >= threshold) or ( 

46 n % 2 == 1 and values[k, i, j] <= threshold 

47 ): # exceeding (n even) or falling below threshold (n odd) 

48 if values[k, i, j] == values[k - 1, i, j]: 

49 # edge case: subsequent values equal threshold result in ZeroDivisionError 

50 continue 

51 if not np.isnan(values[k - 1, i, j]): 

52 # interpolate z coord of threshold 

53 out[n, i, j] = z[k - 1, i, j] + ( 

54 threshold - values[k - 1, i, j] 

55 ) * (z[k, i, j] - z[k - 1, i, j]) / ( 

56 values[k, i, j] - values[k - 1, i, j] 

57 ) 

58 else: 

59 # set to top of layer 

60 out[n, i, j] = z[k, i, j] + 0.5 * dz[k, i, j] 

61 n += 1 

62 

63 

64@numba.njit 

65def _interpolate_value_boundaries_z3ddz1d(values, z, dz, threshold, out): 

66 nlay, nrow, ncol = values.shape 

67 for i in range(nrow): 

68 for j in range(ncol): 

69 n = 0 

70 if values[0, i, j] >= threshold: # top cell exceeds threshold 

71 out[n, i, j] = z[0, i, j] + 0.5 * dz[0] 

72 n += 1 

73 for k in range(1, nlay): 

74 # from top downward 

75 if (n % 2 == 0 and values[k, i, j] >= threshold) or ( 

76 n % 2 == 1 and values[k, i, j] <= threshold 

77 ): # exceeding (n even) or falling below threshold (n odd) 

78 # if (values[k, i, j] == threshold) and (values[k - 1, i, j] == threshold): 

79 if values[k, i, j] == values[k - 1, i, j]: 

80 # edge case: subsequent values equal threshold result in ZeroDivisionError 

81 continue 

82 if not np.isnan(values[k - 1, i, j]): 

83 # interpolate z coord of threshold 

84 out[n, i, j] = z[k - 1, i, j] + ( 

85 threshold - values[k - 1, i, j] 

86 ) * (z[k, i, j] - z[k - 1, i, j]) / ( 

87 values[k, i, j] - values[k - 1, i, j] 

88 ) 

89 else: 

90 # set to top of layer 

91 out[n, i, j] = z[k, i, j] + 0.5 * dz[k] 

92 n += 1 

93 

94 

95def interpolate_value_boundaries(values, z, threshold): 

96 """Function that returns all exceedance and non-exceedance boundaries for 

97 a given threshold in a 3D values DataArray. Returned z-coordinates are 

98 linearly interpolated between cell mids. As many boundaries are returned as are maximally 

99 present in the 3D values DataArray. Function returns xr.DataArray of exceedance boundaries 

100 and xr.DataArray of z-coordinates where values fall below the set treshold. 

101 

102 Parameters 

103 ---------- 

104 values : 3D xr.DataArray 

105 The datarray containing the values to search for boundaries. Dimensions ``layer``, ``y``, ``x`` 

106 z : 1D or 3D xr.DataArray 

107 Datarray containing z-coordinates of cell midpoints. Dimensions ``layer``, ``y``, ``x``. Should contain a dz coordinate. 

108 threshold : float 

109 Value threshold 

110 

111 Returns 

112 ------- 

113 xr.DataArray 

114 Z locations of successive exceedances of threshold from the top down. Dimensions ``boundary``, ``y``, ``x`` 

115 xr.DataArray 

116 Z locations of successive instances of falling below threshold from the top down. Dimensions ``boundary``, ``y``, ``x`` 

117 """ 

118 

119 if "dz" not in z.coords: 

120 raise ValueError('Dataarray "z" must contain a "dz" coordinate') 

121 

122 values = values.load() 

123 out = xr.full_like(values, np.nan) 

124 # if z.dz is a scalar, conform to z 

125 # TODO: broadcast dz to z, so _interpolate_value_boundaries_z3ddz1d is superfluous? 

126 if len(z.dz.dims) == 0: 

127 dz = np.ones_like(z) * z.dz.values 

128 else: 

129 dz = z.dz.values 

130 

131 if len(z.dims) == 1: 

132 _interpolate_value_boundaries_z1d( 

133 values.values, z.values, dz, threshold, out.values 

134 ) 

135 else: 

136 if len(z.dz.dims) <= 1: 

137 _interpolate_value_boundaries_z3ddz1d( 

138 values.values, z.values, dz, threshold, out.values 

139 ) 

140 else: 

141 _interpolate_value_boundaries_z3d( 

142 values.values, z.values, dz, threshold, out.values 

143 ) 

144 out = out.rename({"layer": "boundary"}) 

145 out = out.dropna(dim="boundary", how="all") 

146 

147 # _interpolate_value_boundaries returns exceedance / falling below as even- and odd-indexed boundaries 

148 # separate and renumber them for convenience 

149 out.coords["boundary"] = np.arange(len(out.coords["boundary"])) 

150 exceedance = out.sel(boundary=slice(0, len(out.coords["boundary"]), 2)) 

151 exceedance.coords["boundary"] = np.arange(len(exceedance.coords["boundary"])) 

152 fallbelow = out.sel(boundary=slice(1, len(out.coords["boundary"]), 2)) 

153 fallbelow.coords["boundary"] = np.arange(len(fallbelow.coords["boundary"])) 

154 

155 return exceedance, fallbelow