import numba
import numpy as np
import xarray as xr
from imod.prepare import common
# LayerRegrid does not support conductance method, nearest, or linear
METHODS = common.METHODS.copy()
METHODS.pop("conductance")
METHODS.pop("nearest")
METHODS.pop("multilinear")
@numba.njit(cache=True)
def _regrid_layers(src, dst, src_top, dst_top, src_bot, dst_bot, method):
"""
Maps one set of layers unto the other.
"""
nlayer_src, nrow, ncol = src.shape
nlayer_dst = dst.shape[0]
values = np.zeros(nlayer_src)
weights = np.zeros(nlayer_src)
for i in range(nrow):
for j in range(ncol):
src_t = src_top[:, i, j]
dst_t = dst_top[:, i, j]
src_b = src_bot[:, i, j]
dst_b = dst_bot[:, i, j]
# ii is index of dst
for ii in range(nlayer_dst):
dt = dst_t[ii]
db = dst_b[ii]
if np.isnan(dt) or np.isnan(db):
continue
count = 0
has_value = False
# jj is index of src
for jj in range(nlayer_src):
st = src_t[jj]
sb = src_b[jj]
if np.isnan(st) or np.isnan(sb):
continue
overlap = common._overlap((db, dt), (sb, st))
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 LayerRegridder:
"""
Object to repeatedly regrid layers of 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"}``
"""
[docs]
def __init__(self, method):
_method = common._get_method(method, METHODS)
self.method = _method
self._first_call = True
def _make_regrid(self):
"""
Use closure to avoid numba overhead
"""
jit_method = numba.njit(self.method)
@numba.njit
def regrid(src, dst, src_top, dst_top, src_bot, dst_bot):
return _regrid_layers(
src, dst, src_top, dst_top, src_bot, dst_bot, jit_method
)
self._regrid = regrid
def regrid(
self, source, source_top, source_bottom, destination_top, destination_bottom
):
"""
Parameters
----------
source : xr.DataArray
The values of the layered model.
source_top : xr.DataArray
The vertical location of the layer tops.
destination_top : xr.DataArray
The vertical location of the layer tops.
source_bottom : xr.DataArray
The vertical location of the layer bottoms.
destination_bottom : xr.DataArray
The vertical location of the layer bottoms.
Returns
-------
regridded : xr.DataArray
"""
def dim_format(dims):
return ", ".join(dim for dim in dims)
# Checks on inputs
for da in [
source_top,
source_bottom,
source,
destination_bottom,
destination_top,
]:
if not isinstance(da, xr.DataArray):
raise TypeError("All arguments must be DataArrays.")
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 [source_bottom, source]:
for dim in ["layer", "y", "x"]:
if not source_top[dim].equals(da[dim]):
raise ValueError(f"Input coordinates do not match along {dim}")
if self._first_call:
self._make_regrid()
self._first_call = False
dst = xr.full_like(destination_top, np.nan, dtype=source.dtype)
dst.values = self._regrid(
source.values,
dst.values,
source_top.values,
destination_top.values,
source_bottom.values,
destination_bottom.values,
)
return dst
