from typing import Any
import numba
import numpy as np
import xarray as xr
from imod.prepare import common
# Voxelize does not support conductance method, nearest, or linear
METHODS: dict[str, Any] = common.METHODS.copy()
METHODS.pop("conductance")
METHODS.pop("nearest")
METHODS.pop("multilinear")
@numba.njit(cache=True)
def _voxelize(src, dst, src_top, src_bot, dst_z, method):
nlayer, nrow, ncol = src.shape
nz = dst_z.size - 1
values = np.zeros(nlayer)
weights = np.zeros(nlayer)
for i in range(nrow):
for j in range(ncol):
tops = src_top[:, i, j]
bots = src_bot[:, i, j]
# ii is index of dst
for ii in range(nz):
z0 = dst_z[ii]
z1 = dst_z[ii + 1]
if np.isnan(z0) or np.isnan(z1):
continue
zb = min(z0, z1)
zt = max(z0, z1)
count = 0
has_value = False
# jj is index of src
for jj in range(nlayer):
top = tops[jj]
bot = bots[jj]
overlap = common._overlap((bot, top), (zb, zt))
if overlap == 0:
continue
has_value = True
values[count] = src[jj, i, j]
weights[count] = overlap
count += 1
else:
if has_value:
dst[ii, i, j] = method(values, weights)
# Reset
values[:count] = 0
weights[:count] = 0
return dst
[docs]
class Voxelizer:
"""
Object to repeatedly voxelize similar objects. Compiles once on first call,
can then be repeatedly called without JIT compilation overhead.
Attributes
----------
method : str, function
The method to use for regridding. Default available methods are:
``{"mean", "harmonic_mean", "geometric_mean", "sum", "minimum",
"maximum", "mode", "median", "max_overlap"}``
Examples
--------
Usage is similar to the regridding. Initialize the Voxelizer object:
>>> mean_voxelizer = imod.prepare.Voxelizer(method="mean")
Then call the ``voxelize`` method to transform a layered dataset into a
voxel based one. The vertical coordinates of the layers must be provided
by ``top`` and ``bottom``.
>>> mean_voxelizer.voxelize(source, top, bottom, like)
If your data is already voxel based, i.e. the layers have tops and bottoms
that do not differ with x or y, you should use a ``Regridder`` instead.
It's possible to provide your own methods to the ``Regridder``, provided that
numba can compile them. They need to take the arguments ``values`` and
``weights``. Make sure they deal with ``nan`` values gracefully!
>>> def p30(values, weights):
>>> return np.nanpercentile(values, 30)
>>> p30_voxelizer = imod.prepare.Voxelizer(method=p30)
>>> p30_result = p30_voxelizer.regrid(source, top, bottom, like)
The Numba developers maintain a list of support Numpy features here:
https://numba.pydata.org/numba-doc/dev/reference/numpysupported.html
In general, however, the provided methods should be adequate for your
voxelizing needs.
"""
[docs]
def __init__(self, method, use_relative_weights=False):
_method = common._get_method(method, METHODS)
self.method = _method
self._first_call = True
def _make_voxelize(self):
"""
Use closure to avoid numba overhead
"""
jit_method = numba.njit(self.method)
@numba.njit
def voxelize(src, dst, src_top, src_bot, dst_z):
return _voxelize(src, dst, src_top, src_bot, dst_z, jit_method)
self._voxelize = voxelize
def voxelize(self, source, top, bottom, like):
"""
Parameters
----------
source : xr.DataArray
The values of the layered model.
top : xr.DataArray
The vertical location of the layer tops.
bottom : xr.DataArray
The vertical location of the layer bottoms.
like : xr.DataArray
An example DataArray providing the coordinates of the voxelized
results; what it should look like in terms of dimensions, data type,
and coordinates.
Returns
-------
voxelized : xr.DataArray
"""
def dim_format(dims):
return ", ".join(dim for dim in dims)
# Checks on inputs
if "z" not in like.dims:
# might be a coordinate
if "layer" in like.dims:
if not like.coords["z"].dims == ("layer",):
raise ValueError('"z" has to be given in ``like`` coordinates')
if "dz" not in like.coords:
dzs = np.diff(like.coords["z"].values)
dz = dzs[0]
if not np.allclose(dzs, dz):
raise ValueError(
'"dz" has to be given as a coordinate in case of'
' non-equidistant "z" coordinate.'
)
like["dz"] = dz
for da in [top, bottom, source]:
if not da.dims == ("layer", "y", "x"):
raise ValueError(
"Dimensions for top, bottom, and source have to be exactly"
f' ("layer", "y", "x"). Got instead {dim_format(da.dims)}.'
)
for da in [bottom, source]:
for dim in ["layer", "y", "x"]:
if not top[dim].equals(da[dim]):
raise ValueError(f"Input coordinates do not match along {dim}")
if self._first_call:
self._make_voxelize()
self._first_call = False
like_z = like["z"]
if not like_z.indexes["z"].is_monotonic_increasing:
like_z = like_z.isel(z=slice(None, None, -1))
dst_z = common._coord(like_z, "z")[::-1]
else:
dst_z = common._coord(like_z, "z")
dst_nlayer = like["z"].size
_, nrow, ncol = source.shape
dst_coords = {
"z": like.coords["z"],
"y": source.coords["y"],
"x": source.coords["x"],
}
dst_dims = ("z", "y", "x")
dst_shape = (dst_nlayer, nrow, ncol)
dst = xr.DataArray(np.full(dst_shape, np.nan), dst_coords, dst_dims)
dst.values = self._voxelize(
source.values, dst.values, top.values, bottom.values, dst_z
)
return dst
