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

90 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-08 14:15 +0200

1from typing import Any 

2 

3import numba 

4import numpy as np 

5import xarray as xr 

6 

7from imod.prepare import common 

8 

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

10METHODS: dict[str, Any] = common.METHODS.copy() 

11METHODS.pop("conductance") 

12METHODS.pop("nearest") 

13METHODS.pop("multilinear") 

14 

15 

16@numba.njit(cache=True) 

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

18 nlayer, nrow, ncol = src.shape 

19 nz = dst_z.size - 1 

20 values = np.zeros(nlayer) 

21 weights = np.zeros(nlayer) 

22 

23 for i in range(nrow): 

24 for j in range(ncol): 

25 tops = src_top[:, i, j] 

26 bots = src_bot[:, i, j] 

27 

28 # ii is index of dst 

29 for ii in range(nz): 

30 z0 = dst_z[ii] 

31 z1 = dst_z[ii + 1] 

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

33 continue 

34 

35 zb = min(z0, z1) 

36 zt = max(z0, z1) 

37 count = 0 

38 has_value = False 

39 # jj is index of src 

40 for jj in range(nlayer): 

41 top = tops[jj] 

42 bot = bots[jj] 

43 

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

45 if overlap == 0: 

46 continue 

47 

48 has_value = True 

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

50 weights[count] = overlap 

51 count += 1 

52 else: 

53 if has_value: 

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

55 # Reset 

56 values[:count] = 0 

57 weights[:count] = 0 

58 

59 return dst 

60 

61 

62class Voxelizer: 

63 """ 

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

65 can then be repeatedly called without JIT compilation overhead. 

66 

67 Attributes 

68 ---------- 

69 method : str, function 

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

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

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

73 

74 Examples 

75 -------- 

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

77 

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

79 

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

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

82 by ``top`` and ``bottom``. 

83 

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

85 

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

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

88 

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

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

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

92 

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

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

95 

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

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

98 

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

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

101 

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

103 voxelizing needs. 

104 """ 

105 

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

107 _method = common._get_method(method, METHODS) 

108 self.method = _method 

109 self._first_call = True 

110 

111 def _make_voxelize(self): 

112 """ 

113 Use closure to avoid numba overhead 

114 """ 

115 jit_method = numba.njit(self.method) 

116 

117 @numba.njit 

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

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

120 

121 self._voxelize = voxelize 

122 

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

124 """ 

125 

126 Parameters 

127 ---------- 

128 source : xr.DataArray 

129 The values of the layered model. 

130 top : xr.DataArray 

131 The vertical location of the layer tops. 

132 bottom : xr.DataArray 

133 The vertical location of the layer bottoms. 

134 like : xr.DataArray 

135 An example DataArray providing the coordinates of the voxelized 

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

137 and coordinates. 

138 

139 Returns 

140 ------- 

141 voxelized : xr.DataArray 

142 """ 

143 

144 def dim_format(dims): 

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

146 

147 # Checks on inputs 

148 if "z" not in like.dims: 

149 # might be a coordinate 

150 if "layer" in like.dims: 

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

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

153 if "dz" not in like.coords: 

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

155 dz = dzs[0] 

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

157 raise ValueError( 

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

159 ' non-equidistant "z" coordinate.' 

160 ) 

161 like["dz"] = dz 

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

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

164 raise ValueError( 

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

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

167 ) 

168 for da in [bottom, source]: 

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

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

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

172 

173 if self._first_call: 

174 self._make_voxelize() 

175 self._first_call = False 

176 

177 like_z = like["z"] 

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

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

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

181 else: 

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

183 

184 dst_nlayer = like["z"].size 

185 _, nrow, ncol = source.shape 

186 

187 dst_coords = { 

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

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

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

191 } 

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

193 dst_shape = (dst_nlayer, nrow, ncol) 

194 

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

196 dst.values = self._voxelize( 

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

198 ) 

199 

200 return dst