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
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +0200
1from typing import Any
3import numba
4import numpy as np
5import xarray as xr
7from imod.prepare import common
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")
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)
23 for i in range(nrow):
24 for j in range(ncol):
25 tops = src_top[:, i, j]
26 bots = src_bot[:, i, j]
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
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]
44 overlap = common._overlap((bot, top), (zb, zt))
45 if overlap == 0:
46 continue
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
59 return dst
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.
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"}``
74 Examples
75 --------
76 Usage is similar to the regridding. Initialize the Voxelizer object:
78 >>> mean_voxelizer = imod.prepare.Voxelizer(method="mean")
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``.
84 >>> mean_voxelizer.voxelize(source, top, bottom, like)
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.
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!
93 >>> def p30(values, weights):
94 >>> return np.nanpercentile(values, 30)
96 >>> p30_voxelizer = imod.prepare.Voxelizer(method=p30)
97 >>> p30_result = p30_voxelizer.regrid(source, top, bottom, like)
99 The Numba developers maintain a list of support Numpy features here:
100 https://numba.pydata.org/numba-doc/dev/reference/numpysupported.html
102 In general, however, the provided methods should be adequate for your
103 voxelizing needs.
104 """
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
111 def _make_voxelize(self):
112 """
113 Use closure to avoid numba overhead
114 """
115 jit_method = numba.njit(self.method)
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)
121 self._voxelize = voxelize
123 def voxelize(self, source, top, bottom, like):
124 """
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.
139 Returns
140 -------
141 voxelized : xr.DataArray
142 """
144 def dim_format(dims):
145 return ", ".join(dim for dim in dims)
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}")
173 if self._first_call:
174 self._make_voxelize()
175 self._first_call = False
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")
184 dst_nlayer = like["z"].size
185 _, nrow, ncol = source.shape
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)
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 )
200 return dst