Coverage for C:\src\imod-python\imod\prepare\voxelize.py: 90%

89 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-08 13:27 +0200

1import numba 

2import numpy as np 

3import xarray as xr 

4 

5from imod.prepare import common 

6 

7# Voxelize does not support conductance method, nearest, or linear 

8METHODS = common.METHODS.copy() 

9METHODS.pop("conductance") 

10METHODS.pop("nearest") 

11METHODS.pop("multilinear") 

12 

13 

14@numba.njit(cache=True) 

15def _voxelize(src, dst, src_top, src_bot, dst_z, method): 

16 nlayer, nrow, ncol = src.shape 

17 nz = dst_z.size - 1 

18 values = np.zeros(nlayer) 

19 weights = np.zeros(nlayer) 

20 

21 for i in range(nrow): 

22 for j in range(ncol): 

23 tops = src_top[:, i, j] 

24 bots = src_bot[:, i, j] 

25 

26 # ii is index of dst 

27 for ii in range(nz): 

28 z0 = dst_z[ii] 

29 z1 = dst_z[ii + 1] 

30 if np.isnan(z0) or np.isnan(z1): 

31 continue 

32 

33 zb = min(z0, z1) 

34 zt = max(z0, z1) 

35 count = 0 

36 has_value = False 

37 # jj is index of src 

38 for jj in range(nlayer): 

39 top = tops[jj] 

40 bot = bots[jj] 

41 

42 overlap = common._overlap((bot, top), (zb, zt)) 

43 if overlap == 0: 

44 continue 

45 

46 has_value = True 

47 values[count] = src[jj, i, j] 

48 weights[count] = overlap 

49 count += 1 

50 else: 

51 if has_value: 

52 dst[ii, i, j] = method(values, weights) 

53 # Reset 

54 values[:count] = 0 

55 weights[:count] = 0 

56 

57 return dst 

58 

59 

60class Voxelizer: 

61 """ 

62 Object to repeatedly voxelize similar objects. Compiles once on first call, 

63 can then be repeatedly called without JIT compilation overhead. 

64 

65 Attributes 

66 ---------- 

67 method : str, function 

68 The method to use for regridding. Default available methods are: 

69 ``{"mean", "harmonic_mean", "geometric_mean", "sum", "minimum", 

70 "maximum", "mode", "median", "max_overlap"}`` 

71 

72 Examples 

73 -------- 

74 Usage is similar to the regridding. Initialize the Voxelizer object: 

75 

76 >>> mean_voxelizer = imod.prepare.Voxelizer(method="mean") 

77 

78 Then call the ``voxelize`` method to transform a layered dataset into a 

79 voxel based one. The vertical coordinates of the layers must be provided 

80 by ``top`` and ``bottom``. 

81 

82 >>> mean_voxelizer.voxelize(source, top, bottom, like) 

83 

84 If your data is already voxel based, i.e. the layers have tops and bottoms 

85 that do not differ with x or y, you should use a ``Regridder`` instead. 

86 

87 It's possible to provide your own methods to the ``Regridder``, provided that 

88 numba can compile them. They need to take the arguments ``values`` and 

89 ``weights``. Make sure they deal with ``nan`` values gracefully! 

90 

91 >>> def p30(values, weights): 

92 >>> return np.nanpercentile(values, 30) 

93 

94 >>> p30_voxelizer = imod.prepare.Voxelizer(method=p30) 

95 >>> p30_result = p30_voxelizer.regrid(source, top, bottom, like) 

96 

97 The Numba developers maintain a list of support Numpy features here: 

98 https://numba.pydata.org/numba-doc/dev/reference/numpysupported.html 

99 

100 In general, however, the provided methods should be adequate for your 

101 voxelizing needs. 

102 """ 

103 

104 def __init__(self, method, use_relative_weights=False): 

105 _method = common._get_method(method, METHODS) 

106 self.method = _method 

107 self._first_call = True 

108 

109 def _make_voxelize(self): 

110 """ 

111 Use closure to avoid numba overhead 

112 """ 

113 jit_method = numba.njit(self.method) 

114 

115 @numba.njit 

116 def voxelize(src, dst, src_top, src_bot, dst_z): 

117 return _voxelize(src, dst, src_top, src_bot, dst_z, jit_method) 

118 

119 self._voxelize = voxelize 

120 

121 def voxelize(self, source, top, bottom, like): 

122 """ 

123 

124 Parameters 

125 ---------- 

126 source : xr.DataArray 

127 The values of the layered model. 

128 top : xr.DataArray 

129 The vertical location of the layer tops. 

130 bottom : xr.DataArray 

131 The vertical location of the layer bottoms. 

132 like : xr.DataArray 

133 An example DataArray providing the coordinates of the voxelized 

134 results; what it should look like in terms of dimensions, data type, 

135 and coordinates. 

136 

137 Returns 

138 ------- 

139 voxelized : xr.DataArray 

140 """ 

141 

142 def dim_format(dims): 

143 return ", ".join(dim for dim in dims) 

144 

145 # Checks on inputs 

146 if "z" not in like.dims: 

147 # might be a coordinate 

148 if "layer" in like.dims: 

149 if not like.coords["z"].dims == ("layer",): 

150 raise ValueError('"z" has to be given in ``like`` coordinates') 

151 if "dz" not in like.coords: 

152 dzs = np.diff(like.coords["z"].values) 

153 dz = dzs[0] 

154 if not np.allclose(dzs, dz): 

155 raise ValueError( 

156 '"dz" has to be given as a coordinate in case of' 

157 ' non-equidistant "z" coordinate.' 

158 ) 

159 like["dz"] = dz 

160 for da in [top, bottom, source]: 

161 if not da.dims == ("layer", "y", "x"): 

162 raise ValueError( 

163 "Dimensions for top, bottom, and source have to be exactly" 

164 f' ("layer", "y", "x"). Got instead {dim_format(da.dims)}.' 

165 ) 

166 for da in [bottom, source]: 

167 for dim in ["layer", "y", "x"]: 

168 if not top[dim].equals(da[dim]): 

169 raise ValueError(f"Input coordinates do not match along {dim}") 

170 

171 if self._first_call: 

172 self._make_voxelize() 

173 self._first_call = False 

174 

175 like_z = like["z"] 

176 if not like_z.indexes["z"].is_monotonic_increasing: 

177 like_z = like_z.isel(z=slice(None, None, -1)) 

178 dst_z = common._coord(like_z, "z")[::-1] 

179 else: 

180 dst_z = common._coord(like_z, "z") 

181 

182 dst_nlayer = like["z"].size 

183 _, nrow, ncol = source.shape 

184 

185 dst_coords = { 

186 "z": like.coords["z"], 

187 "y": source.coords["y"], 

188 "x": source.coords["x"], 

189 } 

190 dst_dims = ("z", "y", "x") 

191 dst_shape = (dst_nlayer, nrow, ncol) 

192 

193 dst = xr.DataArray(np.full(dst_shape, np.nan), dst_coords, dst_dims) 

194 dst.values = self._voxelize( 

195 source.values, dst.values, top.values, bottom.values, dst_z 

196 ) 

197 

198 return dst