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
« 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
5from imod.prepare import common
7# Voxelize does not support conductance method, nearest, or linear
8METHODS = common.METHODS.copy()
9METHODS.pop("conductance")
10METHODS.pop("nearest")
11METHODS.pop("multilinear")
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)
21 for i in range(nrow):
22 for j in range(ncol):
23 tops = src_top[:, i, j]
24 bots = src_bot[:, i, j]
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
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]
42 overlap = common._overlap((bot, top), (zb, zt))
43 if overlap == 0:
44 continue
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
57 return dst
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.
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"}``
72 Examples
73 --------
74 Usage is similar to the regridding. Initialize the Voxelizer object:
76 >>> mean_voxelizer = imod.prepare.Voxelizer(method="mean")
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``.
82 >>> mean_voxelizer.voxelize(source, top, bottom, like)
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.
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!
91 >>> def p30(values, weights):
92 >>> return np.nanpercentile(values, 30)
94 >>> p30_voxelizer = imod.prepare.Voxelizer(method=p30)
95 >>> p30_result = p30_voxelizer.regrid(source, top, bottom, like)
97 The Numba developers maintain a list of support Numpy features here:
98 https://numba.pydata.org/numba-doc/dev/reference/numpysupported.html
100 In general, however, the provided methods should be adequate for your
101 voxelizing needs.
102 """
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
109 def _make_voxelize(self):
110 """
111 Use closure to avoid numba overhead
112 """
113 jit_method = numba.njit(self.method)
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)
119 self._voxelize = voxelize
121 def voxelize(self, source, top, bottom, like):
122 """
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.
137 Returns
138 -------
139 voxelized : xr.DataArray
140 """
142 def dim_format(dims):
143 return ", ".join(dim for dim in dims)
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}")
171 if self._first_call:
172 self._make_voxelize()
173 self._first_call = False
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")
182 dst_nlayer = like["z"].size
183 _, nrow, ncol = source.shape
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)
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 )
198 return dst