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

1 import numba

2 import numpy as np

3 import xarray as xr

4

5

6@numba.njit

7 def _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

35 def _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

65 def _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

95 def 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