diff --git a/imod/msw/grid_data.py b/imod/msw/grid_data.py index 27c08829c..3e08a69b0 100644 --- a/imod/msw/grid_data.py +++ b/imod/msw/grid_data.py @@ -13,15 +13,12 @@ from imod.msw.fixed_format import VariableMetaData from imod.msw.pkgbase import MetaSwapPackage from imod.msw.regrid.regrid_schemes import GridDataRegridMethod +from imod.msw.utilities.common import concat_imod5 from imod.typing import GridDataArray, GridDataDict -from imod.typing.grid import concat, ones_like +from imod.typing.grid import ones_like from imod.util.spatial import get_cell_area, spatial_reference -def _concat_subunits(arg1: GridDataArray, arg2: GridDataArray): - return concat([arg1, arg2], dim="subunit").assign_coords(subunit=[0, 1]) - - def get_cell_area_from_imod5_data( imod5_cap: GridDataDict, ) -> GridDataArray: @@ -45,7 +42,7 @@ def get_cell_area_from_imod5_data( ) urban_area = urban_area.where(rural_area > 0.0, other=0.0) rural_area = mf6_area - (wetted_area + urban_area) - return _concat_subunits(rural_area, urban_area) + return concat_imod5(rural_area, urban_area) def get_landuse_from_imod5_data( @@ -59,7 +56,7 @@ def get_landuse_from_imod5_data( rural_landuse = imod5_cap["landuse"] # Urban landuse = 18 urban_landuse = ones_like(rural_landuse) * 18 - return _concat_subunits(rural_landuse, urban_landuse) + return concat_imod5(rural_landuse, urban_landuse) def get_rootzone_depth_from_imod5_data( @@ -72,7 +69,7 @@ def get_rootzone_depth_from_imod5_data( """ rootzone_thickness = imod5_cap["rootzone_thickness"] * 0.01 # rootzone depth is equal for both svats. - return _concat_subunits(rootzone_thickness, rootzone_thickness) + return concat_imod5(rootzone_thickness, rootzone_thickness) class GridData(MetaSwapPackage, IRegridPackage): diff --git a/imod/msw/ponding.py b/imod/msw/ponding.py index afc1629af..89ab6af44 100644 --- a/imod/msw/ponding.py +++ b/imod/msw/ponding.py @@ -7,7 +7,8 @@ from imod.msw.fixed_format import VariableMetaData from imod.msw.pkgbase import DataDictType, MetaSwapPackage from imod.msw.regrid.regrid_schemes import PondingRegridMethod -from imod.typing import IntArray +from imod.msw.utilities.common import concat_imod5 +from imod.typing import GridDataDict, IntArray class Ponding(MetaSwapPackage, IRegridPackage): @@ -70,3 +71,16 @@ def _render(self, file: TextIO, index: IntArray, svat: xr.DataArray, *args: Any) self._check_range(dataframe) return self.write_dataframe_fixed_width(file, dataframe) + + @classmethod + def from_imod5_data(cls, imod5_data: dict[str, GridDataDict]) -> "Ponding": + """ + Concatenate ponding depths along subunits + """ + cap_data = imod5_data["cap"] + data = {} + for key in cls._with_subunit: + data_ls = [cap_data[f"{landuse}_{key}"] for landuse in ["rural", "urban"]] + data[key] = concat_imod5(*data_ls) + + return cls(**data) diff --git a/imod/msw/utilities/common.py b/imod/msw/utilities/common.py new file mode 100644 index 000000000..836902662 --- /dev/null +++ b/imod/msw/utilities/common.py @@ -0,0 +1,6 @@ +from imod.typing import GridDataArray +from imod.typing.grid import concat + + +def concat_imod5(arg1: GridDataArray, arg2: GridDataArray) -> GridDataArray: + return concat([arg1, arg2], dim="subunit").assign_coords(subunit=[0, 1]) diff --git a/imod/tests/test_msw/test_ponding.py b/imod/tests/test_msw/test_ponding.py index 626cd6412..e7bb47e06 100644 --- a/imod/tests/test_msw/test_ponding.py +++ b/imod/tests/test_msw/test_ponding.py @@ -10,11 +10,12 @@ RegridderWeightsCache, ) from imod.msw import Ponding +from imod.typing import GridDataArray, GridDataDict -def setup_ponding(): +def setup_ponding() -> tuple[GridDataDict, np.ndarray, GridDataArray]: x = [1.0, 2.0, 3.0] - y = [1.0, 2.0, 3.0] + y = [3.0, 2.0, 1.0] subunit = [0, 1] dx = 1.0 dy = 1.0 @@ -69,16 +70,17 @@ def setup_ponding(): # fmt: on index = (svat != 0).values.ravel() - ponding = Ponding( - ponding_depth=ponding_depth, - runoff_resistance=runoff_resistance, - runon_resistance=runoff_resistance, - ) - return ponding, index, svat + data_ponding = { + "ponding_depth": ponding_depth, + "runoff_resistance": runoff_resistance, + "runon_resistance": runoff_resistance, + } + return data_ponding, index, svat def test_simple_model(fixed_format_parser): - ponding, index, svat = setup_ponding() + data_ponding, index, svat = setup_ponding() + ponding = Ponding(**data_ponding) with tempfile.TemporaryDirectory() as output_dir: output_dir = Path(output_dir) ponding.write(output_dir, index, svat, None, None) @@ -95,7 +97,8 @@ def test_simple_model(fixed_format_parser): def test_regrid_ponding(simple_2d_grid_with_subunits): - ponding, _, _ = setup_ponding() + data_ponding, _, _ = setup_ponding() + ponding = Ponding(**data_ponding) new_grid = simple_2d_grid_with_subunits regrid_context = RegridderWeightsCache() @@ -104,3 +107,28 @@ def test_regrid_ponding(simple_2d_grid_with_subunits): assert np.all(regridded_ponding.dataset["x"].values == new_grid["x"].values) assert np.all(regridded_ponding.dataset["y"].values == new_grid["y"].values) + + +def test_from_imod5_data(): + data_ponding, _, _ = setup_ponding() + expected_ponding = Ponding(**data_ponding) + + # Create cap data + cap_data = {} + mapping_ls = [ + ("rural_runoff_resistance", "runoff_resistance", 0), + ("rural_runoff_resistance", "runoff_resistance", 0), + ("urban_runoff_resistance", "runoff_resistance", 1), + ("rural_runon_resistance", "runon_resistance", 0), + ("urban_runon_resistance", "runon_resistance", 1), + ("rural_ponding_depth", "ponding_depth", 0), + ("urban_ponding_depth", "ponding_depth", 1), + ] + for cap_key, pkg_key, subunit_nr in mapping_ls: + cap_data[cap_key] = data_ponding[pkg_key].sel(subunit=subunit_nr, drop=True) + + imod5_data = {"cap": cap_data} + + actual_ponding = Ponding.from_imod5_data(imod5_data) + + xr.testing.assert_equal(expected_ponding.dataset, actual_ponding.dataset)