diff --git a/imod/msw/grid_data.py b/imod/msw/grid_data.py index 25003716b..27c08829c 100644 --- a/imod/msw/grid_data.py +++ b/imod/msw/grid_data.py @@ -1,13 +1,80 @@ +from typing import Optional + import numpy as np import xarray as xr +from imod.logging import LogLevel, logger +from imod.mf6 import StructuredDiscretization from imod.mf6.interfaces.iregridpackage import IRegridPackage +from imod.mf6.regrid.regrid_schemes import ( + RegridMethodType, +) +from imod.mf6.utilities.regrid import RegridderWeightsCache from imod.msw.fixed_format import VariableMetaData from imod.msw.pkgbase import MetaSwapPackage from imod.msw.regrid.regrid_schemes import GridDataRegridMethod +from imod.typing import GridDataArray, GridDataDict +from imod.typing.grid import concat, 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: + # area's per type of svats + mf6_area = get_cell_area(imod5_cap["wetted_area"]) + wetted_area = imod5_cap["wetted_area"] + urban_area = imod5_cap["urban_area"] + rural_area = mf6_area - (wetted_area + urban_area) + if (wetted_area > mf6_area).any(): + logger.log( + loglevel=LogLevel.WARNING, + message=f"wetted area was set to the max cell area of {mf6_area}", + additional_depth=0, + ) + wetted_area = wetted_area.where(wetted_area <= mf6_area, other=mf6_area) + if (rural_area < 0.0).any(): + logger.log( + loglevel=LogLevel.WARNING, + message="found urban area > than (cel-area - wetted area). Urban area was set to 0", + additional_depth=0, + ) + 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) + + +def get_landuse_from_imod5_data( + imod5_cap: GridDataDict, +) -> GridDataArray: + """ + Get landuse from imod5 capillary zone data. This adds two subunits, one + based on the landuse grid, which specifies rural landuse. The other + specifies urban landuse, which is coded to value 18. + """ + rural_landuse = imod5_cap["landuse"] + # Urban landuse = 18 + urban_landuse = ones_like(rural_landuse) * 18 + return _concat_subunits(rural_landuse, urban_landuse) + + +def get_rootzone_depth_from_imod5_data( + imod5_cap: GridDataDict, +) -> GridDataArray: + """ + Get rootzone depth from imod5 capillary zone data. Also does a unit + conversion: iMOD5 specifies rootzone thickness in centimeters, whereas + MetaSWAP requires rootzone depth in meters. + """ + rootzone_thickness = imod5_cap["rootzone_thickness"] * 0.01 + # rootzone depth is equal for both svats. + return _concat_subunits(rootzone_thickness, rootzone_thickness) + + class GridData(MetaSwapPackage, IRegridPackage): """ This contains the grid data of MetaSWAP. @@ -111,3 +178,37 @@ def _pkgcheck(self): raise ValueError( "Provided area grid with total areas larger than cell area" ) + + @classmethod + def from_imod5_data( + cls, + imod5_data: dict[str, GridDataDict], + target_dis: StructuredDiscretization, + regridder_types: Optional[RegridMethodType] = None, + regrid_cache: RegridderWeightsCache = RegridderWeightsCache(), + ) -> "GridData": + # Get iMOD5 capillary zone data + imod5_cap = imod5_data["cap"] + + data = {} + data["area"] = get_cell_area_from_imod5_data(imod5_cap) + data["landuse"] = get_landuse_from_imod5_data(imod5_cap) + data["rootzone_depth"] = get_rootzone_depth_from_imod5_data(imod5_cap) + data["surface_elevation"] = imod5_cap["surface_elevation"] + data["soil_physical_unit"] = imod5_cap["soil_physical_unit"] + + mf6_top_active = target_dis["idomain"].isel(layer=0, drop=True) + subunit_active = ( + (imod5_cap["boundary"] == 1) & (data["area"] > 0) & (mf6_top_active >= 1) + ) + active = subunit_active.all(dim="subunit") + data_active = { + key: ( + griddata.where(subunit_active) + if key in cls._with_subunit + else griddata.where(active) + ) + for key, griddata in data.items() + } + data_active["active"] = active + return cls(**data_active) diff --git a/imod/msw/pkgbase.py b/imod/msw/pkgbase.py index 5c99c1053..08b306e6d 100644 --- a/imod/msw/pkgbase.py +++ b/imod/msw/pkgbase.py @@ -240,3 +240,6 @@ def regrid_like( def get_regrid_methods(self) -> RegridMethodType: return deepcopy(self._regrid_method) + + def from_imod5_data(self, *args, **kwargs): + raise NotImplementedError("Method not implemented for this package.") diff --git a/imod/tests/test_msw/test_grid_data.py b/imod/tests/test_msw/test_grid_data.py index a3b62a64b..0d0fbf0c3 100644 --- a/imod/tests/test_msw/test_grid_data.py +++ b/imod/tests/test_msw/test_grid_data.py @@ -1,4 +1,5 @@ import tempfile +from copy import deepcopy from pathlib import Path import numpy as np @@ -8,7 +9,9 @@ from hypothesis.strategies import floats from numpy import nan from numpy.testing import assert_almost_equal, assert_equal +from pytest_cases import case, parametrize_with_cases +from imod.mf6.dis import StructuredDiscretization from imod.mf6.utilities.regrid import RegridderWeightsCache from imod.msw import GridData from imod.msw.fixed_format import format_fixed_width @@ -103,147 +106,107 @@ def test_write( assert_equal(results["soil_physical_unit"][0], int(soil_physical_unit)) -def test_generate_index_array(): +@pytest.fixture(scope="function") +def coords_planar() -> dict: x = [1.0, 2.0, 3.0] - y = [1.0, 2.0, 3.0] - subunit = [0, 1] + y = [3.0, 2.0, 1.0] dx = 1.0 dy = 1.0 + return {"y": y, "x": x, "dx": dx, "dy": dy} + + +@pytest.fixture(scope="function") +def coords_one_subunit(coords_planar: dict) -> dict: + coords_subunit = deepcopy(coords_planar) + coords_subunit["subunit"] = [0] + return coords_subunit + + +@pytest.fixture(scope="function") +def coords_two_subunit(coords_planar: dict) -> dict: + coords_subunit = deepcopy(coords_planar) + coords_subunit["subunit"] = [0, 1] + return coords_subunit + + +@case(tags="one_subunit") +def case_grid_data_one_subunits( + coords_one_subunit: dict, coords_planar: dict +) -> dict[str, xr.DataArray]: + data = {} # fmt: off - area = xr.DataArray( + data["area"] = xr.DataArray( np.array( [ [[0.5, 0.5, 0.5], - [nan, nan, nan], + [0.7, 0.7, 0.7], [1.0, 1.0, 1.0]], - [[0.5, 0.5, 0.5], - [1.0, 1.0, 1.0], - [nan, nan, nan]], ] ), dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} + coords=coords_one_subunit ) - landuse = xr.DataArray( + data["landuse"] = xr.DataArray( np.array( [ [[1.0, 1.0, 1.0], - [nan, nan, nan], + [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], - - [[2.0, 2.0, 2.0], - [2.0, 2.0, 2.0], - [nan, nan, nan]], ] ), dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} + coords=coords_one_subunit ) - rootzone_depth = xr.DataArray( + data["rootzone_depth"] = xr.DataArray( np.array( [ - [[1.0, 1.0, 1.0], - [nan, nan, nan], - [1.0, 1.0, 1.0]], - [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0], - [nan, nan, nan]], + [1.0, 1.0, 1.0]], ] ), dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} + coords=coords_one_subunit ) - surface_elevation = xr.DataArray( + data["surface_elevation"] = xr.DataArray( np.array( [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), dims=("y", "x"), - coords={"y": y, "x": x, "dx": dx, "dy": dy} + coords=coords_planar ) - soil_physical_unit = xr.DataArray( + data["soil_physical_unit"] = xr.DataArray( np.array( [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), dims=("y", "x"), - coords={"y": y, "x": x, "dx": dx, "dy": dy} + coords=coords_planar ) - active = xr.DataArray( + data["active"] = xr.DataArray( np.array( [[False, True, False], [False, True, False], [False, True, False]]), dims=("y", "x"), - coords={"y": y, "x": x} - ) - # fmt: on - grid_data = GridData( - area, - landuse, - rootzone_depth, - surface_elevation, - soil_physical_unit, - active, - ) - - index, svat = grid_data.generate_index_array() - - index_expected = [ - False, - True, - False, - False, - False, - False, - False, - True, - False, - False, - True, - False, - False, - True, - False, - False, - False, - False, - ] - - # fmt: off - svat_expected = xr.DataArray( - np.array( - [ - [[0, 1, 0], - [0, 0, 0], - [0, 2, 0]], - - [[0, 3, 0], - [0, 4, 0], - [0, 0, 0]], - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} + coords=coords_planar ) # fmt: on - assert_equal(index, np.array(index_expected)) - assert_equal(svat.values, svat_expected.values) + return data -def simple_model(): - x = [1.0, 2.0, 3.0] - y = [3.0, 2.0, 1.0] - subunit = [0, 1] - dx = 1.0 - dy = 1.0 +@case(tags="two_subunit") +def case_grid_data_two_subunits( + coords_two_subunit: dict, coords_planar: dict +) -> dict[str, xr.DataArray]: + data = {} # fmt: off - area = xr.DataArray( + data["area"] = xr.DataArray( np.array( [ [[0.5, 0.5, 0.5], @@ -256,9 +219,9 @@ def simple_model(): ] ), dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} + coords=coords_two_subunit ) - landuse = xr.DataArray( + data["landuse"] = xr.DataArray( np.array( [ [[1.0, 1.0, 1.0], @@ -271,9 +234,9 @@ def simple_model(): ] ), dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} + coords=coords_two_subunit ) - rootzone_depth = xr.DataArray( + data["rootzone_depth"] = xr.DataArray( np.array( [ [[1.0, 1.0, 1.0], @@ -286,51 +249,92 @@ def simple_model(): ] ), dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} + coords=coords_two_subunit ) - surface_elevation = xr.DataArray( + data["surface_elevation"] = xr.DataArray( np.array( [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), dims=("y", "x"), - coords={"y": y, "x": x, "dx": dx, "dy": dy} + coords=coords_planar ) - soil_physical_unit = xr.DataArray( + data["soil_physical_unit"] = xr.DataArray( np.array( [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]), dims=("y", "x"), - coords={"y": y, "x": x, "dx": dx, "dy": dy} + coords=coords_planar ) - active = xr.DataArray( + data["active"] = xr.DataArray( np.array( [[False, True, False], [False, True, False], [False, True, False]]), dims=("y", "x"), - coords={"y": y, "x": x} + coords=coords_planar ) # fmt: on + return data - grid_data = GridData( - area, - landuse, - rootzone_depth, - surface_elevation, - soil_physical_unit, - active, - ) - return grid_data +@parametrize_with_cases("grid_data_dict", cases=".", has_tag="two_subunit") +def test_generate_index_array( + grid_data_dict: dict[str, xr.DataArray], coords_two_subunit: dict +): + grid_data = GridData(**grid_data_dict) + + index, svat = grid_data.generate_index_array() + index_expected = [ + False, + True, + False, + False, + False, + False, + False, + True, + False, + False, + True, + False, + False, + True, + False, + False, + False, + False, + ] -def test_simple_model(fixed_format_parser): - grid_data = simple_model() + # fmt: off + svat_expected = xr.DataArray( + np.array( + [ + [[0, 1, 0], + [0, 0, 0], + [0, 2, 0]], + + [[0, 3, 0], + [0, 4, 0], + [0, 0, 0]], + ] + ), + dims=("subunit", "y", "x"), + coords=coords_two_subunit, + ) + # fmt: on + assert_equal(index, np.array(index_expected)) + assert_equal(svat.values, svat_expected.values) + + +@parametrize_with_cases("grid_data_dict", cases=".", has_tag="two_subunit") +def test_simple_model(fixed_format_parser, grid_data_dict: dict[str, xr.DataArray]): + grid_data = GridData(**grid_data_dict) index, svat = grid_data.generate_index_array() @@ -350,8 +354,13 @@ def test_simple_model(fixed_format_parser): assert_almost_equal(results["rootzone_depth"], np.array([1.0, 1.0, 1.0, 1.0])) -def test_simple_model_regrid(simple_2d_grid_with_subunits): - grid_data = simple_model() +# Only use two subunit case, as this one sums to total area whereas the one +# subunit case doesn't sum to 1 across subunit. +@parametrize_with_cases("grid_data_dict", cases=".", has_tag="two_subunit") +def test_simple_model_regrid( + simple_2d_grid_with_subunits, grid_data_dict: dict[str, xr.DataArray] +): + grid_data = GridData(**grid_data_dict) new_grid = simple_2d_grid_with_subunits regrid_context = RegridderWeightsCache() @@ -363,84 +372,11 @@ def test_simple_model_regrid(simple_2d_grid_with_subunits): assert np.sum(regridded_area.values) == regridded_total_area -def test_simple_model_1_subunit(fixed_format_parser): - x = [1.0, 2.0, 3.0] - y = [1.0, 2.0, 3.0] - subunit = [0] - dx = 1.0 - dy = 1.0 - # fmt: off - area = xr.DataArray( - np.array( - [ - [[0.5, 0.5, 0.5], - [0.7, 0.7, 0.7], - [1.0, 1.0, 1.0]], - - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} - ) - landuse = xr.DataArray( - np.array( - [ - [[1.0, 1.0, 1.0], - [1.0, 1.0, 1.0], - [1.0, 1.0, 1.0]], - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} - ) - rootzone_depth = xr.DataArray( - np.array( - [ - [[1.0, 1.0, 1.0], - [1.0, 1.0, 1.0], - [1.0, 1.0, 1.0]], - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} - ) - - surface_elevation = xr.DataArray( - np.array( - [[1.0, 2.0, 3.0], - [4.0, 5.0, 6.0], - [7.0, 8.0, 9.0]]), - dims=("y", "x"), - coords={"y": y, "x": x, "dx": dx, "dy": dy} - ) - - soil_physical_unit = xr.DataArray( - np.array( - [[1.0, 2.0, 3.0], - [4.0, 5.0, 6.0], - [7.0, 8.0, 9.0]]), - dims=("y", "x"), - coords={"y": y, "x": x, "dx": dx, "dy": dy} - ) - - active = xr.DataArray( - np.array( - [[False, True, False], - [False, True, False], - [False, True, False]]), - dims=("y", "x"), - coords={"y": y, "x": x} - ) - # fmt: on - - grid_data = GridData( - area, - landuse, - rootzone_depth, - surface_elevation, - soil_physical_unit, - active, - ) +@parametrize_with_cases("grid_data_dict", cases=".", has_tag="one_subunit") +def test_simple_model_1_subunit( + fixed_format_parser, grid_data_dict: dict[str, xr.DataArray] +): + grid_data = GridData(**grid_data_dict) index, svat = grid_data.generate_index_array() @@ -460,17 +396,15 @@ def test_simple_model_1_subunit(fixed_format_parser): assert_almost_equal(results["rootzone_depth"], np.array([1.0, 1.0, 1.0])) -def test_area_grid_exceeds_cell_area(): +@parametrize_with_cases("grid_data_dict", cases=".", has_tag="two_subunit") +def test_area_grid_exceeds_cell_area( + grid_data_dict: dict[str, xr.DataArray], coords_two_subunit: dict +): """ Test where provided area grid exceeds total cell area, should throw error. """ - x = [1.0, 3.0, 5.0] - y = [1.0, 2.0, 3.0] - subunit = [0, 1] - dx = 1.0 - dy = 1.0 # fmt: off - area = xr.DataArray( + grid_data_dict["area"] = xr.DataArray( np.array( [ [[0.5, 0.5, 0.5], @@ -483,167 +417,53 @@ def test_area_grid_exceeds_cell_area(): ] ), dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} - ) - landuse = xr.DataArray( - np.array( - [ - [[1.0, 1.0, 1.0], - [nan, nan, nan], - [1.0, 1.0, 1.0]], - - [[2.0, 2.0, 2.0], - [2.0, 2.0, 2.0], - [nan, nan, nan]], - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} - ) - rootzone_depth = xr.DataArray( - np.array( - [ - [[1.0, 1.0, 1.0], - [nan, nan, nan], - [1.0, 1.0, 1.0]], - - [[1.0, 1.0, 1.0], - [1.0, 1.0, 1.0], - [nan, nan, nan]], - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} - ) - - surface_elevation = xr.DataArray( - np.array( - [[1.0, 2.0, 3.0], - [4.0, 5.0, 6.0], - [7.0, 8.0, 9.0]]), - dims=("y", "x"), - coords={"y": y, "x": x, "dx": dx, "dy": dy} - ) - - soil_physical_unit = xr.DataArray( - np.array( - [[1.0, 2.0, 3.0], - [4.0, 5.0, 6.0], - [7.0, 8.0, 9.0]]), - dims=("y", "x"), - coords={"y": y, "x": x, "dx": dx, "dy": dy} - ) - - active = xr.DataArray( - np.array( - [[False, True, False], - [False, True, False], - [False, True, False]]), - dims=("y", "x"), - coords={"y": y, "x": x} + coords=coords_two_subunit ) # fmt: on with pytest.raises(ValueError): - GridData( - area, - landuse, - rootzone_depth, - surface_elevation, - soil_physical_unit, - active, - ) + GridData(**grid_data_dict) -def test_non_equidistant(): +@parametrize_with_cases("grid_data_dict", cases=".") +def test_non_equidistant(grid_data_dict: dict[str, xr.DataArray]): """ Test where provided grid is non-equidistant, should throw error. """ - x = [1.0, 2.0, 5.0] - y = [1.0, 2.0, 5.0] - subunit = [0, 1] dx = [1.0, 1.0, 5.0] dy = [1.0, 1.0, 5.0] - # fmt: off - area = xr.DataArray( - np.array( - [ - [[0.5, 0.5, 0.5], - [nan, nan, nan], - [1.0, 1.0, 1.0]], + for key, value in grid_data_dict.items(): + grid_data_dict[key] = value.assign_coords(dx=("x", dx), dy=("y", dy)) - [[0.5, 0.5, 0.5], - [1.0, 1.0, 1.0], - [nan, nan, nan]], - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": ("x", dx), "dy": ("y", dy)} - ) - landuse = xr.DataArray( - np.array( - [ - [[1.0, 1.0, 1.0], - [nan, nan, nan], - [1.0, 1.0, 1.0]], + with pytest.raises(ValueError): + GridData(**grid_data_dict) - [[2.0, 2.0, 2.0], - [2.0, 2.0, 2.0], - [nan, nan, nan]], - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": ("x", dx), "dy": ("y", dy)} - ) - rootzone_depth = xr.DataArray( - np.array( - [ - [[1.0, 1.0, 1.0], - [nan, nan, nan], - [1.0, 1.0, 1.0]], - [[1.0, 1.0, 1.0], - [1.0, 1.0, 1.0], - [nan, nan, nan]], - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": ("x", dx), "dy": ("y", dy)} - ) +@parametrize_with_cases("grid_data_dict", cases=".") +def test_from_imod5_data(grid_data_dict: dict[str, xr.DataArray]): + cap_data = {} + cap_data["wetted_area"] = 1 - grid_data_dict["area"].sum(dim="subunit") + like = grid_data_dict["area"].sel(subunit=0, drop=True) + top = xr.zeros_like(like, dtype=float) + cap_data["boundary"] = xr.ones_like(like, dtype=int) + cap_data["urban_area"] = xr.zeros_like(like, dtype=float) + 0.1 + cap_data["landuse"] = xr.ones_like(like, dtype=int) + cap_data["rootzone_thickness"] = xr.ones_like(like, dtype=int) + cap_data["surface_elevation"] = top + cap_data["soil_physical_unit"] = xr.ones_like(like, dtype=int) + cap_data["active"] = xr.ones_like(like, dtype=bool) - surface_elevation = xr.DataArray( - np.array( - [[1.0, 2.0, 3.0], - [4.0, 5.0, 6.0], - [7.0, 8.0, 9.0]]), - dims=("y", "x"), - coords={"y": y, "x": x, "dx": ("x", dx), "dy": ("y", dy)} - ) + imod5_data = {"cap": cap_data} - soil_physical_unit = xr.DataArray( - np.array( - [[1.0, 2.0, 3.0], - [4.0, 5.0, 6.0], - [7.0, 8.0, 9.0]]), - dims=("y", "x"), - coords={"y": y, "x": x, "dx": ("x", dx), "dy": ("y", dy)} - ) + layer = xr.DataArray([1, 1], coords={"layer": [1, 2]}, dims=("layer",)) + idomain = layer * xr.ones_like(like, dtype=int) - active = xr.DataArray( - np.array( - [[False, True, False], - [False, True, False], - [False, True, False]]), - dims=("y", "x"), - coords={"y": y, "x": x} + # Dis only needed for idomain + dis = StructuredDiscretization(top, top - 0.1, idomain, validate=False) + + griddata = GridData.from_imod5_data(imod5_data, target_dis=dis) + expected_rootzone_depth = cap_data["rootzone_thickness"] * 0.01 + xr.testing.assert_allclose( + expected_rootzone_depth, griddata["rootzone_depth"].sel(subunit=0, drop=True) ) - # fmt: on - with pytest.raises(ValueError): - GridData( - area, - landuse, - rootzone_depth, - surface_elevation, - soil_physical_unit, - active, - ) + assert (griddata["landuse"].sel(subunit=1, drop=True) == 18).all() diff --git a/imod/typing/__init__.py b/imod/typing/__init__.py index 0d17b1127..760346894 100644 --- a/imod/typing/__init__.py +++ b/imod/typing/__init__.py @@ -11,6 +11,7 @@ GridDataArray: TypeAlias = Union[xr.DataArray, xu.UgridDataArray] GridDataset: TypeAlias = Union[xr.Dataset, xu.UgridDataset] +GridDataDict: TypeAlias = dict[str, GridDataArray] ScalarAsDataArray: TypeAlias = Union[xr.DataArray, xu.UgridDataArray] ScalarAsDataset: TypeAlias = Union[xr.Dataset, xu.UgridDataset] UnstructuredData: TypeAlias = Union[xu.UgridDataset, xu.UgridDataArray]