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

1 from typing import Any

2

3 import numba

4 import numpy as np

5 import xarray as xr

6

7 from imod.prepare import common

8

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

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

11 METHODS.pop("conductance")

12 METHODS.pop("nearest")

13 METHODS.pop("multilinear")

14

15

16@numba.njit(cache=True)

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

62 class 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