from dataclasses import asdict
from typing import Optional
import numpy as np
import xarray as xr
from imod.common.interfaces.iregridpackage import IRegridPackage
from imod.common.utilities.dataclass_type import DataclassType
from imod.common.utilities.regrid import _regrid_array
from imod.msw.fixed_format import VariableMetaData
from imod.msw.pkgbase import MetaSwapPackage
from imod.msw.regrid.regrid_schemes import IdfMappingRegridMethod
from imod.typing import GridDataArray
from imod.util.regrid import RegridderWeightsCache
from imod.util.spatial import spatial_reference
[docs]
class IdfMapping(MetaSwapPackage, IRegridPackage):
"""
Describes svat location in the IDF grid.
Note that MetaSWAP can only write equidistant grids.
"""
_file_name = "idf_svat.inp"
_metadata_dict = {
"svat": VariableMetaData(10, 1, 9999999, int),
"rows": VariableMetaData(10, 1, 9999999, int),
"columns": VariableMetaData(10, 1, 9999999, int),
"x_grid": VariableMetaData(15, -9999999.0, 9999999.0, float),
"y_grid": VariableMetaData(15, -9999999.0, 9999999.0, float),
}
_with_subunit = ()
_without_subunit = ("rows", "columns", "x_grid", "y_grid")
_to_fill = ()
_regrid_method = IdfMappingRegridMethod()
# NOTE that it is stated in the IO manual: "The x- and y-coordinates should
# increase with increasing col, row." But the example works with decreasing
# y-coordinates.
[docs]
def __init__(self, area, nodata):
super().__init__()
self.dataset["area"] = area
self.dataset["nodata"] = nodata
nrow = self.dataset.coords["y"].size
ncol = self.dataset.coords["x"].size
y_index = xr.DataArray(
np.arange(1, nrow + 1), coords={"y": self.dataset.coords["y"]}, dims=("y",)
)
x_index = xr.DataArray(
np.arange(1, ncol + 1), coords={"x": self.dataset.coords["x"]}, dims=("x",)
)
rows, columns = xr.broadcast(y_index, x_index)
self.dataset["rows"] = rows
self.dataset["columns"] = columns
y_grid, x_grid = xr.broadcast(self.dataset["y"], self.dataset["x"])
self.dataset["x_grid"] = x_grid
self.dataset["y_grid"] = y_grid
def _get_output_settings(self):
grid = self.dataset["area"]
dx, xmin, _, dy, ymin, _ = spatial_reference(grid)
ncol = grid["x"].size
nrow = grid["y"].size
# If non-equidistant, spatial_reference returned a 1d array instead of
# float
if (not np.isscalar(dx)) or (not np.isscalar(dy)):
raise ValueError("MetaSWAP only supports equidistant grids")
nodata = self.dataset["nodata"].values
return {
"simgro_opt": -1,
"idf_per": 1,
"idf_dx": dx,
"idf_dy": np.abs(dy),
"idf_ncol": ncol,
"idf_nrow": nrow,
"idf_xmin": xmin,
"idf_ymin": ymin,
"idf_nodata": nodata,
}
[docs]
def regrid_like(
self,
target_grid: GridDataArray,
regrid_cache: RegridderWeightsCache,
regridder_types: Optional[DataclassType] = None,
) -> "MetaSwapPackage":
"""
Creates a package of the same type as this package, based on another
discretization. It regrids all the arrays in this package to the desired
discretization, and leaves the options unmodified. At the moment only
regridding to a different planar grid is supported, meaning
``target_grid`` has different ``"x"`` and ``"y"``.
The default regridding methods are obtained by calling
``.get_regrid_methods()`` on the package, which returns a dataclass with
the default regridding methods for each variable in the package.
Parameters
----------
target_grid: xr.DataArray or xu.UgridDataArray
a grid defined using the same discretization as the one we want to
regrid the package to.
regrid_cache: RegridderWeightsCache
stores regridder weights for different regridders. Can be used to
speed up regridding, if the same regridders are used several times
for regridding different arrays.
regridder_types: RegridMethodType, optional
dictionary mapping arraynames (str) to a tuple of regrid type (a
specialization class of BaseRegridder) and function name (str) this
dictionary can be used to override the default mapping method.
Examples
--------
To regrid the infiltration package with a non-default method for the
infiltration capacity, call ``regrid_like`` with these arguments:
>>> regridder_types = imod.msw.regrid.InfiltrationRegridMethod(infiltration_capacity=(imod.RegridderType.OVERLAP, "max"))
>>> regrid_cache = imod.util.regrid.RegridderWeightsCache()
>>> new_infiltration = infiltration.regrid_like(like, regrid_cache, regridder_types)
Returns
-------
A package with the same options as this package, and with all the
data-arrays regridded to another discretization, similar to the one used
in input argument "target_grid"
"""
if regridder_types is None:
regridder_settings = asdict(self.get_regrid_methods(), dict_factory=dict)
else:
regridder_settings = asdict(regridder_types, dict_factory=dict)
nodata = self.dataset["nodata"].values[()]
regridded_area = _regrid_array(
self.dataset["area"],
regrid_cache,
regridder_settings["area"][0],
regridder_settings["area"][1],
target_grid,
)
return type(self)(regridded_area, nodata)
