diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst index cf97902cb..a5b40a49a 100644 --- a/docs/api/changelog.rst +++ b/docs/api/changelog.rst @@ -9,6 +9,16 @@ The format is based on `Keep a Changelog`_, and this project adheres to [Unreleased] ------------ +Added +~~~~~ + +- Experimental class :class:`imod.msw.SprinklingPoints` to specify sprinkling + from points for MetaSWAP models, instead of from grid. You can use this to + specify sprinkling wells from IPF files in an iMOD5 CAP dataset with + :meth:`imod.msw.SprinklingPoints.from_imod5_data`. +- :class:`imod.mf6.LayeredWell.from_imod5_cap_data` now also supports loading + wells from IPF files in an iMOD5 CAP dataset. + Fixed ~~~~~ @@ -16,6 +26,13 @@ Fixed :meth:`imod.mf6.LayeredWell.from_imod5_data` when simulation timesteps precede the first well timestep. +Changed +~~~~~~~ + +- Deprecated :class:`imod.msw.Sprinkling` in favor of + :class:`imod.msw.SprinklingGrid`. Call :class:`imod.msw.SprinklingGrid` to get + the same behavior as you were used to. + [1.1.0] - 2026-08-03 -------------------- diff --git a/docs/api/msw.rst b/docs/api/msw.rst index a14f60f59..b2142b9e9 100644 --- a/docs/api/msw.rst +++ b/docs/api/msw.rst @@ -51,11 +51,18 @@ Grid packages ScalingFactors.get_regrid_methods ScalingFactors.write Sprinkling - Sprinkling.regrid_like - Sprinkling.clip_box - Sprinkling.from_imod5_data - Sprinkling.get_regrid_methods - Sprinkling.write + SprinklingGrid + SprinklingGrid.regrid_like + SprinklingGrid.clip_box + SprinklingGrid.from_imod5_data + SprinklingGrid.get_regrid_methods + SprinklingGrid.write + SprinklingPoints + SprinklingPoints.regrid_like + SprinklingPoints.clip_box + SprinklingPoints.from_imod5_data + SprinklingPoints.get_regrid_methods + SprinklingPoints.write Initial conditions ================== diff --git a/imod/mf6/mf6_wel_adapter.py b/imod/mf6/mf6_wel_adapter.py index c536e6612..b044e899c 100644 --- a/imod/mf6/mf6_wel_adapter.py +++ b/imod/mf6/mf6_wel_adapter.py @@ -128,6 +128,9 @@ class Mf6Wel(BoundaryCondition, IPackage): _pkg_id = "wel" _period_data = ("cellid", "rate") + # Workaround that this needs to be ignored in get_period_varnames. This can + # become an _auxiliary_data in the future with some extra work. + _optional_data = ("id",) _keyword_map = {} _template = BoundaryCondition._initialize_template(_pkg_id) _auxiliary_data = {"concentration": "species"} @@ -143,6 +146,7 @@ def __init__( self, cellid, rate, + id, concentration=None, concentration_boundary_type="aux", save_flows: Optional[bool] = None, @@ -153,6 +157,7 @@ def __init__( dict_dataset = { "cellid": cellid, "rate": rate, + "id": id, "concentration": concentration, "concentration_boundary_type": concentration_boundary_type, "save_flows": save_flows, @@ -169,7 +174,7 @@ def _ds_to_arrdict(self, ds): arrdict: Dict[str, Any] = {} arrdict["data_vars"] = [ - var_name for var_name in ds.data_vars if var_name != "cellid" + var_name for var_name in ds.data_vars if var_name not in ("cellid", "id") ] dsvar = {} diff --git a/imod/mf6/utilities/imod5_converter.py b/imod/mf6/utilities/imod5_converter.py index adee28a6e..e27c21581 100644 --- a/imod/mf6/utilities/imod5_converter.py +++ b/imod/mf6/utilities/imod5_converter.py @@ -58,9 +58,16 @@ def fill_missing_layers( def _well_from_imod5_cap_point_data(cap_data: GridDataDict) -> dict[str, np.ndarray]: - raise NotImplementedError( - "Assigning sprinkling wells with an IPF file is not supported, please specify them as IDF." - ) + df_points = cap_data["artificial_recharge_layer"] + data = {} + # Order of columns is x, y, layer, the other columns are irrelevant here. + data["x"] = df_points.iloc[:, 0].to_numpy().astype(float) + data["y"] = df_points.iloc[:, 1].to_numpy().astype(float) + data["layer"] = df_points.iloc[:, 2].to_numpy().astype(int) + data["rate"] = np.zeros_like(data["x"], dtype=float) + data["id"] = df_points.index.to_numpy() + + return data def _well_from_imod5_cap_grid_data(cap_data: GridDataDict) -> dict[str, np.ndarray]: @@ -85,7 +92,7 @@ def _well_from_imod5_cap_grid_data(cap_data: GridDataDict) -> dict[str, np.ndarr def well_from_imod5_cap_data( imod5_data: Imod5DataDict, - target_dis: IRegridPackage, + target_dis: Optional[IRegridPackage], regridder_types: DataclassType, regrid_cache: RegridderWeightsCache, ) -> dict[str, np.ndarray]: @@ -121,6 +128,11 @@ def well_from_imod5_cap_data( if has_ipf_well: return _well_from_imod5_cap_point_data(cap_data) else: + if target_dis is None: + raise ValueError( + "target_dis must be provided when converting iMOD5 cap data " + "from grids (IDF)" + ) cap_data_regridded = regrid_imod5_cap_data( imod5_data, target_dis, regridder_types, regrid_cache )["cap"] diff --git a/imod/mf6/wel.py b/imod/mf6/wel.py index 3319c5744..75beb25bb 100644 --- a/imod/mf6/wel.py +++ b/imod/mf6/wel.py @@ -61,11 +61,11 @@ def _assign_dims(arg: Any) -> tuple[Any, ...] | xr.DataArray: if arg.dims[0] != "time": arg = arg.transpose() da = xr.DataArray( - data=arg.values, coords={"time": arg["time"]}, dims=["time", "index"] + data=arg.to_numpy(), coords={"time": arg["time"]}, dims=["time", "index"] ) return da elif is_da: - return "index", arg.values + return "index", arg.to_numpy() else: return "index", arg @@ -347,11 +347,11 @@ class GridAgnosticWell(BoundaryCondition, IPointDataPackage, abc.ABC): @property def x(self) -> npt.NDArray[np.float64]: - return self.dataset["x"].values + return self.dataset["x"].to_numpy() @property def y(self) -> npt.NDArray[np.float64]: - return self.dataset["y"].values + return self.dataset["y"].to_numpy() @classmethod def _is_grid_agnostic_package(cls) -> bool: @@ -394,7 +394,9 @@ def _create_dataset_vars( # Carefully rename the dimension and set coordinates d_rename = {"index": "ncellid"} ds_vars = ds_vars.rename_dims(**d_rename).rename_vars(**d_rename) - ds_vars = ds_vars.assign_coords(**{"ncellid": cellid.coords["ncellid"].values}) + ds_vars = ds_vars.assign_coords( + **{"ncellid": cellid.coords["ncellid"].to_numpy()} + ) return ds_vars @@ -534,9 +536,9 @@ def _to_mf6_pkg( ds = ds.assign(**data_vars_dict) # type: ignore[arg-type] ds = remove_inactive(ds, idomain) - ds["save_flows"] = self["save_flows"].values[()] - ds["print_flows"] = self["print_flows"].values[()] - ds["print_input"] = self["print_input"].values[()] + ds["save_flows"] = enforce_scalar(self["save_flows"]) + ds["print_flows"] = enforce_scalar(self["print_flows"]) + ds["print_input"] = enforce_scalar(self["print_input"]) filtered_final_well_ids = self._gather_filtered_well_ids(ds, wells_df) if len(filtered_final_well_ids) > 0: @@ -546,8 +548,6 @@ def _to_mf6_pkg( ) logger.log(loglevel=LogLevel.WARNING, message=message_end) - ds = ds.drop_vars("id") - data_vars_dict = {str(k): v for k, v in ds.data_vars.items()} return Mf6Wel(**data_vars_dict) # type: ignore[arg-type] @@ -1080,8 +1080,8 @@ def _find_well_value_at_layer( if (value is not None) and is_spatial_grid(value): value = imod.select.points_values( value, - x=well_dataset["x"].values, - y=well_dataset["y"].values, + x=well_dataset["x"].to_numpy(), + y=well_dataset["y"].to_numpy(), out_of_bounds="ignore", ) in_bounds = np.full(well_dataset.sizes["index"], False) @@ -1477,7 +1477,7 @@ def _validate_imod5_depth_information( def from_imod5_cap_data( cls, imod5_data: Imod5DataDict, - target_dis: StructuredDiscretization, + target_dis: Optional[StructuredDiscretization] = None, regridder_types: CapDataWellRegridMethod = CapDataWellRegridMethod(), regrid_cache: RegridderWeightsCache = RegridderWeightsCache(), ): @@ -1517,6 +1517,17 @@ def from_imod5_cap_data( xarray datasets, under the key of the package type to which it belongs, as returned by :func:`imod.formats.prj.open_projectfile_data`. + target_dis: Optional[StructuredDiscretization] + The target discretization to which the data should be regridded. + Only necessary when "artificial_recharge_layer" is an IDF grid, + otherwise ignored. + regridder_types: CapDataWellRegridMethod + The regridder type to use for the regridding of the "artificial_recharge_layer" + and "artificial_recharge_capacity" grids. Only necessary when + "artificial_recharge_layer" is an IDF grid, otherwise ignored. + regrid_cache: RegridderWeightsCache + Cache for storing intermediate regridding results. Only necessary when + "artificial_recharge_layer" is an IDF grid, otherwise ignored. """ data = well_from_imod5_cap_data( imod5_data, target_dis, regridder_types, regrid_cache diff --git a/imod/msw/__init__.py b/imod/msw/__init__.py index d8d053c49..313521ba0 100644 --- a/imod/msw/__init__.py +++ b/imod/msw/__init__.py @@ -16,5 +16,7 @@ from imod.msw.output_control import TimeOutputControl, VariableOutputControl from imod.msw.ponding import Ponding from imod.msw.scaling_factors import ScalingFactors -from imod.msw.sprinkling import Sprinkling + +# Import deprecated class Sprinkling to keep public API stable +from imod.msw.sprinkling import Sprinkling, SprinklingGrid, SprinklingPoints from imod.msw.vegetation import AnnualCropFactors diff --git a/imod/msw/model.py b/imod/msw/model.py index f295b76d2..2a4081296 100644 --- a/imod/msw/model.py +++ b/imod/msw/model.py @@ -46,11 +46,12 @@ from imod.msw.ponding import Ponding from imod.msw.regrid.regrid_schemes import CapDataRegridMethod from imod.msw.scaling_factors import ScalingFactors -from imod.msw.sprinkling import Sprinkling +from imod.msw.sprinkling import SprinklingGrid, SprinklingPoints from imod.msw.timeutil import to_metaswap_timeformat from imod.msw.utilities.common import find_in_file_list from imod.msw.utilities.imod5_converter import ( has_active_scaling_factor, + is_sprinkling_from_points, ) from imod.msw.utilities.mask import ( MetaSwapActive, @@ -302,8 +303,21 @@ def get_pkgkey( str The key of the package of type ``pkg_type``. """ + # Loop over all packages in the model and match based on filename + # attached to the package class. This is a temporary solution to make + # primod stable, where get_pkgkey(Sprinkling) was used and it should + # return a SprinklingGrid or SprinklingPoints package depending on the + # data in the model, while still making Sprinkling behave like + # SprinklingGrid. This can be reverted once primod is updated and had a + # few releases. + # + # The reason why this works for MetaSWAP as each package maps to one + # single unique file, so there can never be a model with both a + # SprinklingGrid and a SprinklingPoints package. The first occurrence of + # the filename is always the correct one. + for pkg_key, pkg in self.items(): - if isinstance(pkg, pkg_type): + if pkg._file_name == pkg_type._file_name: return pkg_key if not optional_package: @@ -830,7 +844,10 @@ def from_imod5_data( } model["infiltration"] = Infiltration.from_imod5_data(imod5_masked) model["ponding"] = Ponding.from_imod5_data(imod5_masked) - model["sprinkling"] = Sprinkling.from_imod5_data(imod5_masked) + if is_sprinkling_from_points(imod5_masked): + model["sprinkling"] = SprinklingPoints.from_imod5_data(imod5_masked) + else: + model["sprinkling"] = SprinklingGrid.from_imod5_data(imod5_masked) model["meteo_grid"] = MeteoGridCopy.from_imod5_data(imod5_masked) model["prec_mapping"] = PrecipitationMapping.from_imod5_data(imod5_masked) model["evt_mapping"] = EvapotranspirationMapping.from_imod5_data(imod5_masked) diff --git a/imod/msw/regrid/regrid_schemes.py b/imod/msw/regrid/regrid_schemes.py index cf82018b1..30547abfc 100644 --- a/imod/msw/regrid/regrid_schemes.py +++ b/imod/msw/regrid/regrid_schemes.py @@ -104,6 +104,28 @@ class SprinklingRegridMethod(DataclassType): max_abstraction_surfacewater: RegridVarType = (RegridderType.OVERLAP, "mean") +@dataclass(config=_CONFIG) +class SprinklingPointsRegridMethod(DataclassType): + """ + Object containing regridder methods for the + :class:`imod.msw.SprinklingPoints` package. This can be provided to the + ``regrid_like`` method to regrid with custom settings. + + Parameters + ---------- + art_grid: tuple, default (RegridderType.OVERLAP, "mode") + + Examples + -------- + Regrid with custom settings: + + >>> regrid_method = SprinklingPointsRegridMethod(art_grid=(RegridderType.OVERLAP,"min")) + >>> sprinkling.regrid_like(target_grid, RegridderWeightsCache(), regrid_method) + """ + + art_grid: RegridVarType = (RegridderType.OVERLAP, "mode") + + @dataclass(config=_CONFIG) class MeteoGridRegridMethod(DataclassType): """ diff --git a/imod/msw/sprinkling.py b/imod/msw/sprinkling.py index 802aec550..31957b6f3 100644 --- a/imod/msw/sprinkling.py +++ b/imod/msw/sprinkling.py @@ -1,4 +1,7 @@ -from typing import TextIO +import abc +import textwrap +import warnings +from typing import TextIO, cast import numpy as np import pandas as pd @@ -9,13 +12,17 @@ from imod.mf6.mf6_wel_adapter import Mf6Wel from imod.msw.fixed_format import VariableMetaData from imod.msw.pkgbase import MetaSwapPackage -from imod.msw.regrid.regrid_schemes import SprinklingRegridMethod -from imod.msw.utilities.common import concat_imod5 +from imod.msw.regrid.regrid_schemes import ( + SprinklingPointsRegridMethod, + SprinklingRegridMethod, +) from imod.msw.utilities.imod5_converter import ( - get_cell_area_from_imod5_data, + CapSprinklingDataDict, + is_sprinkling_from_points, + sprinkling_data_from_imod5_grid, + sprinkling_data_from_imod5_ipf, ) -from imod.typing import GridDataDict, Imod5DataDict, IntArray -from imod.typing.grid import zeros_like +from imod.typing import Imod5DataDict, IntArray def _ravel_per_subunit(da: xr.DataArray) -> np.ndarray: @@ -25,61 +32,127 @@ def _ravel_per_subunit(da: xr.DataArray) -> np.ndarray: return array_out[np.isfinite(array_out)] -def _sprinkling_data_from_imod5_ipf(cap_data: GridDataDict) -> GridDataDict: - raise NotImplementedError( - "Assigning sprinkling wells with an IPF file is not supported, please specify them as IDF." - ) - +def _extract_indexer_for_svat(df: pd.DataFrame, columns: list[str]): + """ + Get the indexer for a dataframe of wells to select the SVAT subunit for each + well based on its row/col location in the model grid. -def _sprinkling_data_from_imod5_grid(cap_data: GridDataDict) -> GridDataDict: - # Convert units from mm/d to m3/d - msw_area = get_cell_area_from_imod5_data(cap_data) - capacity_mmd = cap_data["artificial_recharge_capacity"] - capacity_m3d = capacity_mmd * 1e-3 * msw_area.sel(subunit=0, drop=True) + Parameters + ---------- + df : pd.DataFrame + DataFrame containing the wells with columns "subunit", "row", + and "column". "row" and "column" are 1-based indices. + columns : list[str] + List of column names to use for indexing. Must include "row" and "column". + + Returns + ------- + np.ndarray + Indexer array for selecting SVAT subunits from svat_da. + """ + if not {"row", "column"}.issubset(columns): + raise ValueError("columns must contain 'row' and 'column'") + df.loc[:, ["row", "column"]] -= 1 # Convert to 0-based indexing for xarray - artificial_rch_type = cap_data["artificial_recharge"] - from_groundwater = artificial_rch_type == 1 - from_surfacewater = artificial_rch_type == 2 - is_active = artificial_rch_type != 0 + indexer = df.loc[:, columns].to_numpy() + return indexer.T - zero_where_active = zeros_like(artificial_rch_type).where(is_active) - # Add zero where active, to have active cells set to 0.0. - max_abstraction_groundwater_rural = zero_where_active.where( - ~from_groundwater, capacity_m3d +def _get_mf6_cellid_dataframe(mf6_well: Mf6Wel) -> pd.DataFrame: + """ + Get cellids from the Mf6Wel objects dataset and convert to a dataframe for + easy merging with sprinkling data. + """ + # Promote id to dim to join datasets + mf6_well_ds = mf6_well.dataset.set_coords("id").swap_dims({"ncellid": "id"}) + # Convert the cellid DataArray to a broad table for easier manipulation. + mf6_cellid_df = mf6_well_ds["cellid"].to_dataset("dim_cellid").to_dataframe() + # Select only the cellid columns we need and reset index to promote id to column + # for merging + dim_cellid = ["layer", "row", "column"] + mf6_cellid_df = mf6_cellid_df.loc[:, dim_cellid].reset_index() + return mf6_cellid_df + + +def _make_sprinkling_well_points_dataframe( + sprinkling_dataset: xr.Dataset, mf6_cellid_df: pd.DataFrame +) -> pd.DataFrame: + """ + Create a dataframe of sprinkling well points from the sprinkling dataset and + merge it with the mf6_cellid_df to get the row/col of each well. + """ + # Get point data from sprinkling dataset and convert to dataframe for easy merging + points_keys = [ + key for key, da in sprinkling_dataset.data_vars.items() if "id" in da.dims + ] + sprinkling_points_df = ( + sprinkling_dataset[points_keys].drop_vars(["dx", "dy"]).to_dataframe() ) - max_abstraction_surfacewater_rural = zero_where_active.where( - ~from_surfacewater, capacity_m3d + # Merge again to confine to wells actually used in the modflow6 model. + # This drops points that are outside model domain. + return sprinkling_points_df.reset_index().merge( + mf6_cellid_df, on="id", how="right", validate="many_to_one" ) - # No sprinkling for urban environments - max_abstraction_urban = zero_where_active - data = {} - data["max_abstraction_groundwater"] = concat_imod5( - max_abstraction_groundwater_rural, max_abstraction_urban +def _merge_sprinkling_points_with_grids( + points_df: pd.DataFrame, svat: xr.DataArray, sprinkling_id_grid: xr.DataArray +) -> pd.DataFrame: + """ + Merge sprinkling points with SVAT grids. + """ + + # Flatten id_sprinkling grid → (y, x, id_sprinkling) table, drop cells with no well + grids = xr.merge([sprinkling_id_grid, svat]) + art_df = ( + grids.to_dataframe().reset_index().query("(id_sprinkling > 0) & (svat > 0)") ) - data["max_abstraction_surfacewater"] = concat_imod5( - max_abstraction_surfacewater_rural, max_abstraction_urban + # Drop unnecessary columns. We preserve the x, y coords as they might + # prove useful for debugging. + art_df = art_df.drop(["dx", "dy"], axis=1) + + # Join: each SVAT cell gets the matching well row(s) from arl_points + return art_df.merge( + points_df, # brings id back as a column + left_on="id_sprinkling", + right_on="id_sprinkling_p", + how="inner", + validate="many_to_one", ) - return data -class Sprinkling(MetaSwapPackage, IRegridPackage): +def align_svat_with_dis( + svat: xr.DataArray, dis_pkg: StructuredDiscretization +) -> xr.DataArray: """ - This contains the sprinkling capacities of links between SVAT units and - groundwater/surface water locations. + Align the SVAT grid with the dis_pkg grid as the SVAT grid might be smaller. + """ + idomain_flat = dis_pkg.dataset["idomain"].isel(layer=0, drop=True) + # Assign to _dummy instead of _ to avoid MyPy 2.3 crashing on the next line. + # See https://github.com/python/mypy/issues/21824 + _dummy, svat_aligned = xr.align(idomain_flat, svat, join="left") + return svat_aligned - This class is responsible for the file `scap_svat.inp` - Parameters - ---------- - max_abstraction_groundwater: array of floats (xr.DataArray) - Describes the maximum abstraction of groundwater to SVAT units in m3 per - day. This array must not have a subunit coordinate. - max_abstraction_surfacewater: array of floats (xr.DataArray) - Describes the maximum abstraction of surfacewater to SVAT units in m3 - per day. This array must not have a subunit coordinate. +def _get_svat_groundwater_for_wells( + msw_mf6_sprinkling_df: pd.DataFrame, svat_aligned: xr.DataArray +) -> np.ndarray: + """ + Get the SVAT subunit for each well from the SVAT grid based on the + well's row/col location. + """ + indexer = _extract_indexer_for_svat( + msw_mf6_sprinkling_df, columns=["subunit", "row", "column"] + ) + svat_groundwater = svat_aligned.data[*indexer] + return svat_groundwater.astype(int) + + +class SprinklingBase(MetaSwapPackage, IRegridPackage): + """ + Base class for sprinkling packages. This class is not meant to be + instantiated directly, but rather through the subclasses + :class:`imod.msw.SprinklingGrid` and :class:`imod.msw.SprinklingPoints`. """ _file_name = "scap_svat.inp" @@ -94,6 +167,61 @@ class Sprinkling(MetaSwapPackage, IRegridPackage): "trajectory": VariableMetaData(10, None, None, str), } + @abc.abstractmethod + def _render( + self, + file: TextIO, + index: IntArray, + svat: xr.DataArray, + mf6_dis: StructuredDiscretization, + mf6_well: Mf6Wel, + ) -> None: + raise NotImplementedError( + "method _render() must be implemented in subclasses of SprinklingBase." + ) + + @classmethod + @abc.abstractmethod + def from_imod5_data(cls, imod5_data: Imod5DataDict): + raise NotImplementedError( + "method from_imod5_data() must be implemented in subclasses of SprinklingBase." + ) + + +class SprinklingGrid(SprinklingBase): + """ + This contains the sprinkling capacities of links between SVAT units and + groundwater/surface water locations. Input is provided as grids for the + maximum abstraction of groundwater and surfacewater to SVAT units. To + specify the sprinkling capacity as points, see + :class:`imod.msw.SprinklingPoints`. + + This class is responsible for the file `scap_svat.inp` + + Parameters + ---------- + max_abstraction_groundwater: array of floats (xr.DataArray) + Describes the maximum abstraction of groundwater to SVAT units in m3 per + day. This array must have a subunit coordinate. + max_abstraction_surfacewater: array of floats (xr.DataArray) + Describes the maximum abstraction of surfacewater to SVAT units in m3 + per day. This array must have a subunit coordinate. + + + Examples + -------- + + >>> import xarray as xr + >>> import imod + >>> grid = imod.util.empty_2d(dx=25.0, dy=25.0, xmin=0.0, xmax=50.0, ymin=0.0, ymax=50.0) + >>> max_abstraction_groundwater = xr.concat([grid.fillna(25.0), grid.fillna(0.0)], dim="subunit").assign_coords(subunit=[0,1]) + >>> max_abstraction_surfacewater = xr.concat([grid.fillna(0.0), grid.fillna(25.0)], dim="subunit").assign_coords(subunit=[0,1]) + >>> sprinkling_grid = imod.msw.SprinklingGrid( + ... max_abstraction_groundwater=max_abstraction_groundwater, + ... max_abstraction_surfacewater=max_abstraction_surfacewater, + ... ) + """ + _with_subunit = ( "max_abstraction_groundwater", "max_abstraction_surfacewater", @@ -126,7 +254,7 @@ def _render( svat: xr.DataArray, mf6_dis: StructuredDiscretization, mf6_well: Mf6Wel, - ): + ) -> None: if not isinstance(mf6_well, Mf6Wel): raise TypeError(rf"well not of type 'Mf6Wel', got '{type(mf6_well)}'") @@ -173,34 +301,30 @@ def _render( return self._write_dataframe_fixed_width(file, dataframe) @classmethod - def from_imod5_data(cls, imod5_data: Imod5DataDict) -> "Sprinkling": + def from_imod5_data(cls, imod5_data: Imod5DataDict) -> "SprinklingGrid": """ - Import sprinkling data from imod5 data. Abstraction data for sprinkling - is defined in iMOD5 either with grids (IDF) or points (IPF) combined - with a grid. Depending on the type, the method does different conversions: - - - grids (IDF) - The ``"artifical_recharge_layer"`` variable was defined as grid - (IDF), this grid defines in which layer a groundwater abstraction - well should be placed. The ``"artificial_recharge"`` grid contains - types which point to the type of abstraction: - - * 0: no abstraction - * 1: groundwater abstraction - * 2: surfacewater abstraction - - The ``"artificial_recharge_capacity"`` grid/constant defines the - capacity of each groundwater or surfacewater abstraction. This is an - ``1:1`` mapping: Each grid cell maps to a separate well. - - - points with grid (IPF & IDF) - The ``"artifical_recharge_layer"`` variable was defined as point - data (IPF), this table contains wellids with an abstraction capacity - and layer. The ``"artificial_recharge"`` grid contains a mapping of - grid cells to wellids in the point data. The - ``"artificial_recharge_capacity"`` is ignored as the abstraction - capacity is already defined in the point data. This is an ``n:1`` - mapping: multiple grid cells can map to one well. + Import sprinkling data from imod5 data artificial recharge grids. + Abstraction data for sprinkling is defined in iMOD5 either with grids + (IDF) or points (IPF) combined with a grid. This class can handle only + the purely grid (IDF) variant. For point data (IPF), use + :class:`imod.msw.SprinklingPoints.from_imod5_data()` instead. + + The iMOD5 data is expected to contain three grids for sprinkling: + + 1. The ``"artificial_recharge"`` grid contains types which point to the + type of abstraction: + + * **0**: no abstraction + * **1**: groundwater abstraction + * **2**: surfacewater abstraction + + 2. The ``"artificial_recharge_layer"`` defines in which layer a groundwater + abstraction well should be placed. + 3. The ``"artificial_recharge_capacity"`` grid/constant defines the + capacity of each groundwater or surfacewater abstraction. This is + converted from mm/d to m3/d using the cell area of the SVAT grid. + + This is an ``1:1`` mapping: Each grid cell maps to a separate well. Parameters ---------- @@ -214,10 +338,249 @@ def from_imod5_data(cls, imod5_data: Imod5DataDict) -> "Sprinkling": ------- Sprinkling package """ + if is_sprinkling_from_points(imod5_data): + msg = textwrap.dedent( + """ + Unsupported format for artificial_recharge_layer: expected a + grid (IDF) got a DataFrame for point data (IPF). Call + imod.msw.SprinklingPoints.from_imod5_data() instead. + """ + ) + raise TypeError(msg) cap_data = imod5_data["cap"] - if isinstance(cap_data["artificial_recharge_layer"], pd.DataFrame): - data = _sprinkling_data_from_imod5_ipf(cap_data) - else: - data = _sprinkling_data_from_imod5_grid(cap_data) + data = sprinkling_data_from_imod5_grid(cap_data) return cls(**data) + + +class Sprinkling(SprinklingGrid): + """ + Deprecated class for sprinkling. Use :class:`imod.msw.SprinklingGrid` + instead for the same behavior. This class is kept for backwards + compatibility and will be removed in a future version. + """ + + def __init__(self, *args, **kwargs): + warnings.warn( + "Sprinkling is deprecated and will be removed in a future version. " + "Use SprinklingGrid for the same behavior instead.", + DeprecationWarning, + ) + super().__init__(*args, **kwargs) + + +class SprinklingPoints(SprinklingBase): + """ + This contains the sprinkling capacities of links between SVAT units and + groundwater/surface water locations. This class is capable of handling point + data (IPF) for sprinkling wells, which is a mapping of grid cells to well + locations. To specify the sprinkling capacity as grid, see + :class:`imod.msw.Sprinkling`. + + This class is responsible for the file `scap_svat.inp` + + .. note:: + This class is still in an experimental state, and might change in future + versions. It is not yet fully tested and validated. + + Parameters + ---------- + art_grid: xr.DataArray + Grid of the artificial recharge ids, with subunit coordinate. These will + be used to map the sprinkling points with the id provided in + ``id_sprinkling_p``. + x_p: np.ndarray | list[float] + x-coordinates of the artificial recharge locations. + y_p: np.ndarray | list[float] + y-coordinates of the artificial recharge locations. + layer_p: np.ndarray | list[int] + layer indices of the artificial recharge locations. + id_sprinkling_p: np.ndarray | list[int] + ids mapping of the artificial recharge locations to the grid cells in + ``art_grid``. + capacity_p: np.ndarray | list[float] + abstraction capacities of the artificial recharge locations. + + Examples + -------- + + Map a single sprinkling point to a grid with two subunits. The well will + sprinkle to all cells in the first subunit. We will use the arbitrary + number 43 as sprinkling point id. + + >>> import xarray as xr + >>> import imod + >>> grid = imod.util.empty_2d(dx=25.0, dy=25.0, xmin=0.0, xmax=50.0, ymin=0.0, ymax=50.0) + >>> art_grid = xr.concat( + ... [grid.fillna(43), grid.fillna(0)], dim="subunit" + ... ).assign_coords(subunit=[0,1]).astype(int) + >>> sprinkling_points = imod.msw.SprinklingPoints( + ... art_grid=art_grid, + ... x_p=[12.5], + ... y_p=[12.5], + ... layer_p=[2], + ... id_sprinkling_p=[43], + ... capacity_p=[25.0], + ... ) + """ + + _with_subunit = ("id_sprinkling",) + _without_subunit = () + + _to_fill = ( + "max_abstraction_groundwater_mm_d", + "max_abstraction_surfacewater_mm_d", + "trajectory", + ) + + _regrid_method = SprinklingPointsRegridMethod() + + def __init__( + self, + art_grid: xr.DataArray, + x_p: np.ndarray | list[float], + y_p: np.ndarray | list[float], + layer_p: np.ndarray | list[int], + id_sprinkling_p: np.ndarray | list[int], + capacity_p: np.ndarray | list[float], + ): + super().__init__() + # Replicate well ids as they were also created in + # imod.mf6.LayeredWell.from_imod5_cap_data() + id_index = pd.Index(range(len(x_p)), name="id").astype(str) + points_ds = xr.Dataset( + { + "x_p": (("id",), x_p), + "y_p": (("id",), y_p), + "layer_p": (("id",), layer_p), + "id_sprinkling_p": (("id",), id_sprinkling_p), + "capacity_p": (("id",), capacity_p), + }, + coords={"id": id_index}, + ) + art_grid = art_grid.rename("id_sprinkling") + self.dataset = xr.merge([art_grid, points_ds]) + + self._pkgcheck() + + @classmethod + def from_imod5_data(cls, imod5_data: Imod5DataDict) -> "SprinklingPoints": + """ + Import sprinkling data from imod5 data artificial recharge grids. + Abstraction data for sprinkling is defined in iMOD5 either with grids + (IDF) or points (IPF) combined with a grid. This class can handle only + the point (IPF) variant. For grid data (IDF), use + :class:`imod.msw.SprinklingGrid.from_imod5_data()` instead. + + The iMOD5 data is expected to contain one grid (IDF) and one table with + points for sprinkling (IPF): + + 1. The ``"artificial_recharge"`` grid contains a mapping of + grid cells to wellids in the point data. + 2. The ``"artificial_recharge_layer"`` variable was defined as point + data (IPF), this table contains wellids with an abstraction capacity + and layer. + 3. The ``"artificial_recharge_capacity"`` is ignored as the abstraction + capacity is already defined in the point data. + + This is an ``n:1`` mapping: multiple grid cells can map to one well. + + Parameters + ---------- + imod5_data: dict[str, dict[str, GridDataArray]] + dictionary containing the arrays mentioned in the project file as + xarray datasets, under the key of the package type to which it + belongs, as returned by + :func:`imod.formats.prj.open_projectfile_data`. + + Returns + ------- + SprinklingPoints package + """ + if is_sprinkling_from_points(imod5_data): + cap_data = cast(CapSprinklingDataDict, imod5_data["cap"]) + data = sprinkling_data_from_imod5_ipf(cap_data) + return cls(**data) + else: + msg = textwrap.dedent( + """ + Unsupported format for artificial_recharge_layer: expected a + DataFrame for point data (IPF), got a grid (IDF). Call + imod.msw.Sprinkling.from_imod5_data() instead. + """ + ) + raise TypeError(msg) + + def _render(self, file, index, svat, mf6_dis, mf6_well): + """ + Render the sprinkling points to the scap_svat.inp file. + + This method first merges the sprinkling points with the mf6_well cellids + to get the row/col of each well, then merges the sprinkling points with + the svat and id_sprinkling grid. It then selects the columns that need to be + written to scap_svat.inp and sets wells with layer > 0 to groundwater + abstraction, and wells with layer = 0 to surfacewater abstraction. + Finally, it deals with edge cases for wells that are outside art_grid + but in the model domain, and writes the dataframe to the file. + """ + # Merge the sprinkling points with the mf6_well cellids to get the + # row/col of each well. + mf6_cellid_df = _get_mf6_cellid_dataframe(mf6_well) + points_df = _make_sprinkling_well_points_dataframe(self.dataset, mf6_cellid_df) + # Merge the sprinkling points with the svat and id_sprinkling grid + msw_mf6_sprinkling_df = _merge_sprinkling_points_with_grids( + points_df, svat, self.dataset["id_sprinkling"] + ) + svat_aligned = align_svat_with_dis(svat, mf6_dis) + msw_mf6_sprinkling_df["svat_groundwater"] = _get_svat_groundwater_for_wells( + msw_mf6_sprinkling_df, svat_aligned + ) + is_point_inside = msw_mf6_sprinkling_df["svat_groundwater"] > 0 + # Select columns that need to be written to scap_svat.inp + inside_df = msw_mf6_sprinkling_df.loc[ + is_point_inside, ["svat", "layer", "svat_groundwater"] + ] + inside_df["svat"] = inside_df["svat"].astype(int) + # Set wells with layer > 0 to groundwater abstraction, and wells with layer = 0 + # to surfacewater abstraction. + capacity = msw_mf6_sprinkling_df.loc[is_point_inside, "capacity_p"] + is_gw_extraction = inside_df["layer"] > 0 + inside_df["max_abstraction_groundwater"] = capacity.where(is_gw_extraction, 0.0) + inside_df["max_abstraction_surfacewater"] = capacity.where( + ~is_gw_extraction, 0.0 + ) + ############## + # EDGE CASES # + ############## + # 1. Wells that are outside art_grid, but in model domain. + # These will be assigned to surfacewater abstraction. + outside_df = msw_mf6_sprinkling_df.loc[ + ~is_point_inside, ["svat", "layer", "svat_groundwater", "capacity_p"] + ] + # Set capacity to surfacewater abstraction, and set groundwater abstraction to 0. + outside_df = outside_df.rename( + columns={"capacity_p": "max_abstraction_surfacewater"} + ) + outside_df["max_abstraction_groundwater"] = 0.0 + # Set svat_groundwater to svat, as these wells are outside art_grid and + # will be assigned to surfacewater abstraction. + outside_df["svat_groundwater"] = outside_df["svat"] + ############ + # FINALIZE # + ############ + # Prepare the final dataframe to be written to scap_svat.inp + dataframe_out = pd.concat([inside_df, outside_df], axis=0, ignore_index=True) + # Order rows by SVAT number to ensure consistent output for testing and + # debugging. + dataframe_out = dataframe_out.sort_values(by=["svat"]).reset_index(drop=True) + # Fill last columns with empty strings, as they are not used in the + # iMOD5 implementation but required by MetaSWAP. + for var in self._to_fill: + dataframe_out[var] = "" + # Order columns to match the metadata dict, which defines the order of + # columns in scap_svat.inp. + dataframe_out = dataframe_out[list(self._metadata_dict.keys())] + + self._check_range(dataframe_out) + + return self._write_dataframe_fixed_width(file, dataframe_out) diff --git a/imod/msw/utilities/imod5_converter.py b/imod/msw/utilities/imod5_converter.py index c6a3cd32d..6656488e8 100644 --- a/imod/msw/utilities/imod5_converter.py +++ b/imod/msw/utilities/imod5_converter.py @@ -1,3 +1,8 @@ +from typing import TypedDict, cast + +import numpy as np +import pandas as pd +import xarray as xr from xarray.core.utils import is_scalar from imod.common.constants import MaskValues @@ -5,11 +10,31 @@ from imod.mf6 import StructuredDiscretization from imod.msw.utilities.common import concat_imod5 from imod.msw.utilities.mask import MetaSwapActive -from imod.typing import GridDataArray, GridDataDict -from imod.typing.grid import ones_like +from imod.typing import GridDataArray, GridDataDict, Imod5DataDict +from imod.typing.grid import ones_like, zeros_like from imod.util.spatial import get_cell_area +# Some additional type aliases for sprinkling data, which is a bit more complex +# than other packages. +class CapSprinklingDataDict(TypedDict, total=False): + artificial_recharge: GridDataArray + artificial_recharge_layer: pd.DataFrame + artificial_recharge_capacity: GridDataArray + + +class SprinklingPointsDataDict(TypedDict, total=False): + x_p: np.ndarray | list[float] + y_p: np.ndarray | list[float] + layer_p: np.ndarray | list[int] + id_sprinkling_p: np.ndarray | list[int] + capacity_p: np.ndarray | list[float] + + +class SprinklingPointsGridDataDict(SprinklingPointsDataDict, total=False): + art_grid: GridDataArray + + def get_cell_area_from_imod5_data( imod5_cap: GridDataDict, ) -> GridDataArray: @@ -116,3 +141,85 @@ def has_active_scaling_factor(imod5_cap: GridDataDict): ) return not scaling_factor_inactive + + +def is_sprinkling_from_points(imod5_data: Imod5DataDict) -> bool: + """ + Check if sprinkling is specified from points, based on the presence of + sprinkling layer and sprinkling points data in the iMOD5 CAP dataset. + """ + cap_data = cast(CapSprinklingDataDict, imod5_data["cap"]) + if isinstance(cap_data.get("artificial_recharge_layer"), pd.DataFrame): + return True + return False + + +def sprinkling_data_from_imod5_ipf( + cap_data: CapSprinklingDataDict, +) -> SprinklingPointsGridDataDict: + art_grid = cap_data["artificial_recharge"] + # Set urban landuse irrigation to 0, as sprinkling is not allowed for urban landuse. + subunit_template = xr.DataArray( + np.array([1, 0], dtype=int), dims="subunit", coords={"subunit": [0, 1]} + ) + art_grid = subunit_template * art_grid + + df_points = cap_data["artificial_recharge_layer"] + # Select first 5 columns and enforce column names, iMOD5 expects columns in + # this order. The additional columns are metadata for the user and can be + # ignored. + arl_points = df_points.iloc[:, :5] + arl_points.columns = ["x_p", "y_p", "layer_p", "id_sprinkling_p", "capacity_p"] + # Enforce dtypes + dtype_dict = { + "x_p": float, + "y_p": float, + "layer_p": int, + "id_sprinkling_p": int, + "capacity_p": float, + } + + arl_points = arl_points.astype(dtype_dict) + arl_point_dict = cast( + SprinklingPointsDataDict, + {key: arl_points[key].to_numpy() for key in dtype_dict.keys()}, + ) + + return { + "art_grid": art_grid, + **arl_point_dict, + } + + +def sprinkling_data_from_imod5_grid(cap_data: GridDataDict) -> GridDataDict: + # Convert units from mm/d to m3/d + msw_area = get_cell_area_from_imod5_data(cap_data) + capacity_mmd = cap_data["artificial_recharge_capacity"] + capacity_m3d = capacity_mmd * 1e-3 * msw_area.sel(subunit=0, drop=True) + + artificial_rch_type = cap_data["artificial_recharge"] + from_groundwater = artificial_rch_type == 1 + from_surfacewater = artificial_rch_type == 2 + is_active = artificial_rch_type != 0 + + zero_where_active = zeros_like(artificial_rch_type).where(is_active) + + # Add zero where active, to have active cells set to 0.0. + max_abstraction_groundwater_rural = zero_where_active.where( + ~from_groundwater, capacity_m3d + ) + max_abstraction_surfacewater_rural = zero_where_active.where( + ~from_surfacewater, capacity_m3d + ) + + # No sprinkling for urban environments + max_abstraction_urban = zero_where_active + + data = {} + data["max_abstraction_groundwater"] = concat_imod5( + max_abstraction_groundwater_rural, max_abstraction_urban + ) + data["max_abstraction_surfacewater"] = concat_imod5( + max_abstraction_surfacewater_rural, max_abstraction_urban + ) + return data diff --git a/imod/tests/fixtures/imod5_cap_data.py b/imod/tests/fixtures/imod5_cap_data.py index 88a2a5d7d..8e9b29a9b 100644 --- a/imod/tests/fixtures/imod5_cap_data.py +++ b/imod/tests/fixtures/imod5_cap_data.py @@ -143,23 +143,23 @@ def cap_data_sprinkling_points() -> Imod5DataDict: artificial_rch_type = zeros_grid(n) artificial_rch_type[:, 1] = 3000 - artificial_rch_type[:, 2] = 4000 + artificial_rch_type[2, 1] = 4000 data = { - "id": [3000, 4000], + "x": [2.0, 2.0], + "y": [3.0, 2.0], "layer": [2, 3], + "id": [3000, 4000], "capacity": [15.0, 30.0], - "y": [1.0, 2.0], - "x": [1.0, 2.0], } - layer = pd.DataFrame(data=data) + dataframe = pd.DataFrame(data=data) cap_data = { "boundary": boundary, "wetted_area": wetted_area, "urban_area": urban_area, "artificial_recharge": artificial_rch_type, - "artificial_recharge_layer": layer, + "artificial_recharge_layer": dataframe, "artificial_recharge_capacity": xr.DataArray(25.0), } diff --git a/imod/tests/fixtures/msw_model_fixture.py b/imod/tests/fixtures/msw_model_fixture.py index 06f7db8b7..2d0ebe124 100644 --- a/imod/tests/fixtures/msw_model_fixture.py +++ b/imod/tests/fixtures/msw_model_fixture.py @@ -276,7 +276,7 @@ def msw_add_sprinkling(msw_model): # %% Sprinkling area = msw_model["grid"].dataset["area"] - msw_model["sprinkling"] = msw.Sprinkling( + msw_model["sprinkling"] = msw.SprinklingGrid( max_abstraction_groundwater=xr.full_like(area, 100.0), max_abstraction_surfacewater=xr.full_like(area, 100.0), ) @@ -307,9 +307,10 @@ def coupled_mf6wel(): well_y = np.repeat(y, ncol) well_rate = np.zeros(well_x.shape) well_layer = np.full_like(well_x, layer, dtype=int) + well_id = [str(i) for i in range(len(well_rate))] cellids = derive_cellid_from_points(idomain, well_x, well_y, well_layer) - well_msw = Mf6Wel(cellids, well_rate) + well_msw = Mf6Wel(cellids, well_rate, well_id) return well_msw diff --git a/imod/tests/test_mf6/test_mf6_wel.py b/imod/tests/test_mf6/test_mf6_wel.py index 436e862c7..dff1dd6f3 100644 --- a/imod/tests/test_mf6/test_mf6_wel.py +++ b/imod/tests/test_mf6/test_mf6_wel.py @@ -1168,7 +1168,14 @@ def test_from_imod5_cap_data__big_grid( @pytest.mark.unittest_jit def test_from_imod5_cap_data__points(cap_data_sprinkling_points, cap_coupled_dis_grid): - with pytest.raises(NotImplementedError): - LayeredWell.from_imod5_cap_data( - cap_data_sprinkling_points, cap_coupled_dis_grid - ) + # Act + well = LayeredWell.from_imod5_cap_data( + cap_data_sprinkling_points, cap_coupled_dis_grid + ) + # Assert + ds = well.dataset + np.testing.assert_allclose(ds["x"].to_numpy(), np.array([2.0, 2.0])) + np.testing.assert_allclose(ds["y"].to_numpy(), np.array([3.0, 2.0])) + np.testing.assert_equal(ds["layer"].to_numpy(), np.array([2, 3])) + np.testing.assert_allclose(ds["rate"].to_numpy(), np.array([0.0, 0.0])) + np.testing.assert_equal(ds["id"].to_numpy(), np.array(["0", "1"])) diff --git a/imod/tests/test_mf6/test_mf6_wel_lowlvl.py b/imod/tests/test_mf6/test_mf6_wel_lowlvl.py index 7e4ea7957..b312a98b1 100644 --- a/imod/tests/test_mf6/test_mf6_wel_lowlvl.py +++ b/imod/tests/test_mf6/test_mf6_wel_lowlvl.py @@ -40,7 +40,8 @@ def test_mf6wel_to_struct_array__stationary( ): # Arrange cellid, rate = mf6wel_test_data_stationary - mf6wel = imod.mf6.mf6_wel_adapter.Mf6Wel(cellid=cellid, rate=rate) + id_names = [str(i) for i in range(len(rate))] + mf6wel = imod.mf6.mf6_wel_adapter.Mf6Wel(cellid=cellid, rate=rate, id=id_names) # Act bin_ds = mf6wel._get_bin_ds() @@ -56,7 +57,8 @@ def test_mf6wel_to_struct_array__transient( ): # Arrange cellid, rate = mf6wel_test_data_transient - mf6wel = imod.mf6.mf6_wel_adapter.Mf6Wel(cellid=cellid, rate=rate) + id_names = [str(i) for i in range(len(rate))] + mf6wel = imod.mf6.mf6_wel_adapter.Mf6Wel(cellid=cellid, rate=rate, id=id_names) ds = mf6wel._get_bin_ds().isel(time=0) # Act @@ -72,7 +74,8 @@ def test_mf6wel_write_datafile__stationary( ): # Arrange cellid, rate = mf6wel_test_data_stationary - mf6wel = imod.mf6.mf6_wel_adapter.Mf6Wel(cellid=cellid, rate=rate) + id_names = [str(i) for i in range(len(rate))] + mf6wel = imod.mf6.mf6_wel_adapter.Mf6Wel(cellid=cellid, rate=rate, id=id_names) ds = mf6wel._get_bin_ds() file_path = Path(tmp_path) / "mf6wel.bin" @@ -90,7 +93,8 @@ def test_mf6wel_write__stationary( ): # Arrange cellid, rate = mf6wel_test_data_stationary - mf6wel = imod.mf6.mf6_wel_adapter.Mf6Wel(cellid=cellid, rate=rate) + id_names = [str(i) for i in range(len(rate))] + mf6wel = imod.mf6.mf6_wel_adapter.Mf6Wel(cellid=cellid, rate=rate, id=id_names) globaltimes = pd.date_range("2000-01-01", "2000-01-06") pkgname = "wel" directory = Path(tmp_path) / "mf6wel" @@ -113,7 +117,8 @@ def test_mf6wel_write__transient( ): # Arrange cellid, rate = mf6wel_test_data_transient - mf6wel = imod.mf6.mf6_wel_adapter.Mf6Wel(cellid=cellid, rate=rate) + id_names = [str(i) for i in range(len(rate))] + mf6wel = imod.mf6.mf6_wel_adapter.Mf6Wel(cellid=cellid, rate=rate, id=id_names) globaltimes = pd.date_range("2000-01-01", "2000-01-06") pkgname = "wel" directory = Path(tmp_path) / "mf6wel" @@ -135,7 +140,8 @@ def test_mf6wel_render__transient( ): # Arrange cellid, rate = mf6wel_test_data_transient - mf6wel = imod.mf6.mf6_wel_adapter.Mf6Wel(cellid=cellid, rate=rate) + id_names = [str(i) for i in range(len(rate))] + mf6wel = imod.mf6.mf6_wel_adapter.Mf6Wel(cellid=cellid, rate=rate, id=id_names) globaltimes = pd.date_range("2000-01-01", "2000-01-06") pkgname = "wel" directory = Path(tmp_path) / "mf6wel" diff --git a/imod/tests/test_msw/test_coupler_mapping.py b/imod/tests/test_msw/test_coupler_mapping.py index 373cdd222..38779c4dd 100644 --- a/imod/tests/test_msw/test_coupler_mapping.py +++ b/imod/tests/test_msw/test_coupler_mapping.py @@ -47,7 +47,8 @@ def get_mf6_wel(svat_data): well_x = [2.0, 2.0, 2.0] well_rate = [-5.0] * 3 cellids = derive_cellid_from_points(svat_data, well_x, well_y, well_layer) - return Mf6Wel(cellids, well_rate) + well_id = [str(i) for i in range(len(well_rate))] + return Mf6Wel(cellids, well_rate, well_id) def get_mf6_dis(svat_data): diff --git a/imod/tests/test_msw/test_sprinkling.py b/imod/tests/test_msw/test_sprinkling.py index 63812df9f..6f75d3633 100644 --- a/imod/tests/test_msw/test_sprinkling.py +++ b/imod/tests/test_msw/test_sprinkling.py @@ -1,26 +1,86 @@ import tempfile +from dataclasses import dataclass from pathlib import Path +from typing import Callable, Optional import numpy as np import pytest import xarray as xr from numpy import nan from numpy.testing import assert_almost_equal, assert_equal +from pytest_cases import parametrize_with_cases from imod import msw +from imod.mf6.dis import StructuredDiscretization from imod.mf6.mf6_wel_adapter import Mf6Wel from imod.mf6.wel import derive_cellid_from_points -def test_simple_model_all_svats(fixed_format_parser): +@pytest.fixture(scope="function") +def sprinkling_svat_index(): 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 + dy = -1.0 # fmt: off - max_abstraction_groundwater = xr.DataArray( + svat = 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}, + name="svat", + ) + # fmt: on + index = (svat != 0).values.ravel() + return svat, index + + +@dataclass +class ExpectedCaseData: + xfail: Optional[str] = None + abs_gw: Optional[np.ndarray] = None + abs_sw: Optional[np.ndarray] = None + layer: Optional[np.ndarray] = None + svat: Optional[np.ndarray] = None + svat_gw: Optional[np.ndarray] = None + + +@dataclass +class SprinklingGridCaseData: + max_abstraction_groundwater: Optional[xr.DataArray] = None + max_abstraction_surfacewater: Optional[xr.DataArray] = None + + +@dataclass +class SprinklingPointsCaseData: + art_grid: Optional[xr.DataArray] = None + x_p: Optional[np.ndarray] = None + y_p: Optional[np.ndarray] = None + layer_p: Optional[np.ndarray] = None + id_sprinkling_p: Optional[np.ndarray] = None + capacity_p: Optional[np.ndarray] = None + + +class SprinklingGridCases: + def case_all_svats( + self, sprinkling_svat_index + ) -> tuple[SprinklingGridCaseData, ExpectedCaseData]: + svat, _ = sprinkling_svat_index + case_data = SprinklingGridCaseData() + case_data.max_abstraction_groundwater = xr.full_like(svat, 0.0) + case_data.max_abstraction_surfacewater = xr.full_like(svat, 0.0) + # fmt: off + case_data.max_abstraction_groundwater.data = np.array( [ [[nan, 100.0, nan], [nan, 200.0, nan], @@ -29,13 +89,8 @@ def test_simple_model_all_svats(fixed_format_parser): [nan, 200.0, nan], [nan, 300.0, nan]] ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} - ) - - max_abstraction_surfacewater = xr.DataArray( - np.array( + ) + case_data.max_abstraction_surfacewater.data = np.array( [ [[nan, 100.0, nan], [nan, 200.0, nan], @@ -44,73 +99,36 @@ def test_simple_model_all_svats(fixed_format_parser): [nan, 200.0, nan], [nan, 300.0, nan]] ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} - ) - - svat = xr.DataArray( - np.array( + ) + # fmt: on + expected_data = ExpectedCaseData() + expected_data.abs_gw = np.array([100.0, 300.0, 100.0, 200.0]) + expected_data.abs_sw = np.array([100.0, 300.0, 100.0, 200.0]) + expected_data.layer = np.array([3, 1, 3, 2]) + expected_data.svat = np.array([1, 2, 3, 4]) + expected_data.svat_gw = np.array([1, 2, 3, 4]) + case_data.expected_data = expected_data + return case_data, expected_data + + def case_some_svats( + self, sprinkling_svat_index + ) -> tuple[SprinklingGridCaseData, ExpectedCaseData]: + svat, _ = sprinkling_svat_index + case_data = SprinklingGridCaseData() + case_data.max_abstraction_groundwater = xr.full_like(svat, 0.0) + case_data.max_abstraction_surfacewater = xr.full_like(svat, 0.0) + # fmt: off + case_data.max_abstraction_groundwater.data = np.array( [ - [[0, 1, 0], - [0, 0, 0], - [0, 2, 0]], - - [[0, 3, 0], - [0, 4, 0], - [0, 0, 0]], + [[nan, 100.0, nan], + [nan, 200.0, nan], + [nan, 300.0, nan]], + [[nan, nan, nan], + [nan, 200.0, nan], + [nan, nan, nan]] ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} - ) - # fmt: on - index = (svat != 0).values.ravel() - - # Well - well_layer = [3, 2, 1] - well_y = y - well_x = [2.0, 2.0, 2.0] - well_rate = [-5.0] * 3 - cellids = derive_cellid_from_points(svat, well_x, well_y, well_layer) - well = Mf6Wel(cellids, well_rate) - - sprinkling = msw.Sprinkling( - max_abstraction_groundwater, - max_abstraction_surfacewater, - ) - - with tempfile.TemporaryDirectory() as output_dir: - output_dir = Path(output_dir) - sprinkling.write(output_dir, index, svat, None, well) - - results = fixed_format_parser( - output_dir / msw.Sprinkling._file_name, - msw.Sprinkling._metadata_dict, ) - - assert_equal(results["svat"], np.array([1, 2, 3, 4])) - assert_almost_equal( - results["max_abstraction_groundwater"], - np.array([100.0, 300.0, 100.0, 200.0]), - ) - assert_almost_equal( - results["max_abstraction_surfacewater"], - np.array([100.0, 300.0, 100.0, 200.0]), - ) - assert_equal(results["layer"], np.array([3, 1, 3, 2])) - assert_equal(results["svat_groundwater"], np.array([1, 2, 3, 4])) - - -def test_simple_model_some_svats(fixed_format_parser): - x = [1.0, 2.0, 3.0] - y = [1.0, 2.0, 3.0] - subunit = [0, 1] - dx = 1.0 - dy = 1.0 - # fmt: off - max_abstraction_groundwater = xr.DataArray( - np.array( + case_data.max_abstraction_surfacewater.data = np.array( [ [[nan, 100.0, nan], [nan, 200.0, nan], @@ -119,253 +137,439 @@ def test_simple_model_some_svats(fixed_format_parser): [nan, 200.0, nan], [nan, nan, nan]] ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} - ) - - max_abstraction_surfacewater = xr.DataArray( - np.array( + ) + # fmt: on + expected_data = ExpectedCaseData() + expected_data.abs_gw = np.array([100.0, 300.0, 200.0]) + expected_data.abs_sw = np.array([100.0, 300.0, 200.0]) + expected_data.layer = np.array([3, 1, 2]) + expected_data.svat = np.array([1, 2, 4]) + expected_data.svat_gw = np.array([1, 2, 4]) + + return case_data, expected_data + + def case_inconsistent_active_capacity( + self, sprinkling_svat_index + ) -> tuple[SprinklingGridCaseData, ExpectedCaseData]: + svat, _ = sprinkling_svat_index + case_data = SprinklingGridCaseData() + case_data.max_abstraction_groundwater = xr.full_like(svat, 0.0) + case_data.max_abstraction_surfacewater = xr.full_like(svat, 0.0) + # fmt: off + case_data.max_abstraction_groundwater.data = np.array( [ [[nan, 100.0, nan], + [nan, 0.0, nan], + [nan, 0.0, nan]], + [[nan, nan, nan], + [nan, 200.0, nan], + [nan, nan, nan]] + ] + ) + case_data.max_abstraction_surfacewater.data = np.array( + [ + [[nan, 0.0, nan], [nan, 200.0, nan], [nan, 300.0, nan]], [[nan, nan, nan], [nan, 200.0, nan], [nan, nan, nan]] ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} - ) - - svat = xr.DataArray( - np.array( - [ - [[0, 1, 0], - [0, 0, 0], - [0, 2, 0]], + ) + # fmt: on + expected_data = ExpectedCaseData() + expected_data.abs_gw = np.array([100.0, 0.0, 200.0]) + expected_data.abs_sw = np.array([0.0, 300.0, 200.0]) + expected_data.layer = np.array([3, 1, 2]) + expected_data.svat = np.array([1, 2, 4]) + expected_data.svat_gw = np.array([1, 2, 4]) + + return case_data, expected_data + + +class SprinklingPointsCases: + def case_one_point_one_art_cell( + self, sprinkling_svat_index + ) -> tuple[SprinklingPointsCaseData, ExpectedCaseData]: + """ + Simple test case for sprinkling points. Each point is mapped to one svat. + """ + svat, _ = sprinkling_svat_index + case_data = SprinklingPointsCaseData() + case_data.art_grid = xr.full_like(svat, 0, dtype=int) + # fmt: off + case_data.art_grid.data = np.array( + [[[0, 1, 0], + [0, 2, 0], + [0, 3, 0],], + [[0, 1, 0], + [0, 2, 0], + [0, 3, 0]]] + ) + # fmt: on + case_data.x_p = [2.0, 2.0, 2.0] + case_data.y_p = [3.0, 2.0, 1.0] + case_data.layer_p = [1, 2, 3] + case_data.id_sprinkling_p = [1, 2, 3] + case_data.capacity_p = [10.0, 20.0, 30.0] + + expected_data = ExpectedCaseData() + expected_data.svat = np.array([1, 2, 3, 4]) + expected_data.svat_gw = np.array([1, 2, 3, 4]) + expected_data.layer = np.array([1, 3, 1, 2]) + expected_data.abs_gw = np.array([10.0, 30.0, 10.0, 20.0]) + expected_data.abs_sw = np.array([0.0, 0.0, 0.0, 0.0]) + + return case_data, expected_data + + def case_one_point_one_art_cell__one_subunit( + self, sprinkling_svat_index + ) -> tuple[SprinklingPointsCaseData, ExpectedCaseData]: + """ + Simple test case for sprinkling points. Each point is mapped to one + svat. Only one subunit is used, similar to when imported from iMOD5 + DBASE + """ + svat, _ = sprinkling_svat_index + case_data = SprinklingPointsCaseData() + case_data.art_grid = xr.full_like(svat, 0, dtype=int) + # fmt: off + case_data.art_grid.data = np.array( + [[[0, 1, 0], + [0, 2, 0], + [0, 3, 0],], + [[0, 0, 0], + [0, 0, 0], + [0, 0, 0]]] + ) + # fmt: on + case_data.x_p = [2.0, 2.0, 2.0] + case_data.y_p = [3.0, 2.0, 1.0] + case_data.layer_p = [1, 2, 3] + case_data.id_sprinkling_p = [1, 2, 3] + case_data.capacity_p = [10.0, 20.0, 30.0] + + expected_data = ExpectedCaseData() + expected_data.svat = np.array([1, 2]) + expected_data.svat_gw = np.array([1, 2]) + expected_data.layer = np.array([1, 3]) + expected_data.abs_gw = np.array([10.0, 30.0]) + expected_data.abs_sw = np.array([0.0, 0.0]) + + return case_data, expected_data + + def case_multi_point_one_art_cell( + self, sprinkling_svat_index + ) -> tuple[SprinklingPointsCaseData, ExpectedCaseData]: + """ + Case where multiple points are assigned to the same SVAT. Not a common + usecase. Usually multiple cells coupled to one point. + """ + svat, _ = sprinkling_svat_index + case_data = SprinklingPointsCaseData() + case_data.art_grid = xr.full_like(svat, 0, dtype=int) + # fmt: off + case_data.art_grid.data = np.array( + [[[0, 0, 0], + [0, 1, 0], + [0, 0, 0]], + [[0, 0, 0], + [0, 1, 0], + [0, 0, 0]]] + ) + # fmt: on + case_data.x_p = [2.0, 2.0, 2.0] + case_data.y_p = [3.0, 2.0, 1.0] + case_data.layer_p = [1, 2, 3] + case_data.id_sprinkling_p = [1, 1, 1] + case_data.capacity_p = [10.0, 20.0, 30.0] + + expected_data = ExpectedCaseData() + expected_data.xfail = "Multiple points cannot be connected to one grid cell" + return case_data, expected_data + + def case_one_point_multi_art_cell( + self, sprinkling_svat_index + ) -> tuple[SprinklingPointsCaseData, ExpectedCaseData]: + """ + Case where one point is assigned to multiple art_grid cells. Quite a + common usecase. The point is located in the centre of the grid, where + there is only an svat in subunit 1. In subunit 0 this cell is not + active, therefore sprinkling capacity is assigned to surface water. + """ + svat, _ = sprinkling_svat_index + case_data = SprinklingPointsCaseData() + case_data.art_grid = xr.full_like(svat, 0, dtype=int) + # fmt: off + case_data.art_grid.data = np.array( + [[[0, 1, 0], + [0, 1, 0], + [0, 1, 0]], + [[0, 1, 0], + [0, 1, 0], + [0, 1, 0]]] + ) + # fmt: on + case_data.x_p = [2.0] + case_data.y_p = [2.0] + case_data.layer_p = [2] + case_data.id_sprinkling_p = [1] + case_data.capacity_p = [10.0] + + expected_data = ExpectedCaseData() + expected_data.svat = np.array([1, 2, 3, 4]) + expected_data.svat_gw = np.array([1, 2, 4, 4]) + expected_data.layer = np.array([2, 2, 2, 2]) + expected_data.abs_gw = np.array([0.0, 0.0, 10.0, 10.0]) + expected_data.abs_sw = np.array([10.0, 10.0, 0.0, 0.0]) + + return case_data, expected_data + + def case_art_grid_outside( + self, sprinkling_svat_index + ) -> tuple[SprinklingPointsCaseData, ExpectedCaseData]: + """ + Case where art grid is located outside the active SVAT area, but still + in the model domain. The well is inside the model domain. Sprinkling + should not be assigned. + """ + svat, _ = sprinkling_svat_index + case_data = SprinklingPointsCaseData() + case_data.art_grid = xr.full_like(svat, 0, dtype=int) + # fmt: off + case_data.art_grid.data = np.array( + [[[0, 0, 4], + [0, 0, 0], + [0, 0, 0]], + [[0, 0, 4], + [0, 0, 0], + [0, 0, 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} - ) - # fmt: on - index = (svat != 0).values.ravel() + # fmt: on + case_data.x_p = [2.0] + case_data.y_p = [2.0] + case_data.layer_p = [3] + case_data.id_sprinkling_p = [4] + case_data.capacity_p = [40.0] + + expected_data = ExpectedCaseData() + expected_data.svat = np.array([]) + expected_data.svat_gw = np.array([]) + expected_data.layer = np.array([]) + expected_data.abs_gw = np.array([]) + expected_data.abs_sw = np.array([]) + + return case_data, expected_data + + def case_point_outside( + self, sprinkling_svat_index + ) -> tuple[SprinklingPointsCaseData, ExpectedCaseData]: + """ + Case where one point is located outside the active SVAT area, but still in + the model domain. The well should be assigned as surface water extraction. + """ + svat, _ = sprinkling_svat_index + case_data = SprinklingPointsCaseData() + case_data.art_grid = xr.full_like(svat, 0, dtype=int) + # fmt: off + case_data.art_grid.data = np.array( + [[[0, 0, 0], + [0, 0, 0], + [0, 5, 0]], + [[0, 0, 0], + [0, 0, 0], + [0, 5, 0]]] + ) + # fmt: on + case_data.x_p = [3.0] + case_data.y_p = [1.0] + case_data.layer_p = [3] + case_data.id_sprinkling_p = [5] + case_data.capacity_p = [40.0] + + expected_data = ExpectedCaseData() + expected_data.svat = np.array([2]) + expected_data.svat_gw = np.array([2]) + expected_data.layer = np.array([3]) + expected_data.abs_gw = np.array([0.0]) + expected_data.abs_sw = np.array([40.0]) + return case_data, expected_data + + +@parametrize_with_cases("case_data, expected_data", cases=SprinklingGridCases) +def test_grid_simple_model( + fixed_format_parser: Callable, + sprinkling_svat_index: tuple[xr.DataArray, np.ndarray], + case_data: SprinklingGridCaseData, + expected_data: ExpectedCaseData, +): + svat, index = sprinkling_svat_index # Well well_layer = [3, 2, 1] - well_y = [1.0, 2.0, 3.0] + well_y = [3.0, 2.0, 1.0] well_x = [2.0, 2.0, 2.0] - well_rate = [-5.0] * 3 + well_rate_values = [-5.0] * 3 + well_rate = xr.DataArray(well_rate_values, dims=("ncellid",)) + well_id_values = ["0", "1", "2"] + well_id = xr.DataArray(well_id_values, dims=("ncellid",)) cellids = derive_cellid_from_points(svat, well_x, well_y, well_layer) - well = Mf6Wel(cellids, well_rate) + well = Mf6Wel(cellids, well_rate, well_id) - coupler_mapping = msw.Sprinkling( - max_abstraction_groundwater, - max_abstraction_surfacewater, + sprinkling = msw.SprinklingGrid( + case_data.max_abstraction_groundwater, + case_data.max_abstraction_surfacewater, ) with tempfile.TemporaryDirectory() as output_dir: output_dir = Path(output_dir) - coupler_mapping.write(output_dir, index, svat, None, well) + sprinkling.write(output_dir, index, svat, None, well) results = fixed_format_parser( - output_dir / msw.Sprinkling._file_name, - msw.Sprinkling._metadata_dict, + output_dir / msw.SprinklingGrid._file_name, + msw.SprinklingGrid._metadata_dict, ) - assert_equal(results["svat"], np.array([1, 2, 4])) + assert_equal(results["svat"], expected_data.svat) assert_almost_equal( results["max_abstraction_groundwater"], - np.array([100.0, 300.0, 200.0]), + expected_data.abs_gw, ) assert_almost_equal( results["max_abstraction_surfacewater"], - np.array([100.0, 300.0, 200.0]), + expected_data.abs_sw, ) - assert_equal(results["layer"], np.array([3, 1, 2])) - assert_equal(results["svat_groundwater"], np.array([1, 2, 4])) + assert_equal(results["layer"], expected_data.layer) + assert_equal(results["svat_groundwater"], expected_data.svat_gw) -def test_simple_model_inconsistent_active_capacity(fixed_format_parser): - x = [1.0, 2.0, 3.0] - y = [1.0, 2.0, 3.0] - subunit = [0, 1] - dx = 1.0 - dy = 1.0 - # fmt: off - max_abstraction_groundwater = xr.DataArray( - np.array( - [ - [[nan, 100.0, nan], - [nan, 0.0, nan], - [nan, 0.0, nan]], - [[nan, nan, nan], - [nan, 200.0, nan], - [nan, nan, nan]] - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} - ) +@parametrize_with_cases("case_data, expected_data", cases=SprinklingGridCases) +def test_grid_simple_model_1_subunit( + fixed_format_parser: Callable, + sprinkling_svat_index: tuple[xr.DataArray, np.ndarray], + case_data: SprinklingGridCaseData, + expected_data: ExpectedCaseData, +): + svat, index = sprinkling_svat_index - max_abstraction_surfacewater = xr.DataArray( - np.array( - [ - [[nan, 0.0, nan], - [nan, 200.0, nan], - [nan, 300.0, nan]], - [[nan, nan, nan], - [nan, 200.0, nan], - [nan, nan, nan]] - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} - ) - - svat = 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} - ) - # fmt: on - index = (svat != 0).values.ravel() + svat = svat.isel(subunit=[0]) + index = index[:9] # Only the first subunit # Well - well_layer = [3, 2, 1] - well_y = [1.0, 2.0, 3.0] - well_x = [2.0, 2.0, 2.0] - well_rate = [-5.0] * 3 + well_layer = [3, 1] + well_y = [3.0, 1.0] + well_x = [2.0, 2.0] + well_rate_values = [-5.0] * 2 + well_rate = xr.DataArray(well_rate_values, dims=("ncellid",)) + well_id_values = ["0", "2"] + well_id = xr.DataArray(well_id_values, dims=("ncellid",)) cellids = derive_cellid_from_points(svat, well_x, well_y, well_layer) - well = Mf6Wel(cellids, well_rate) + well = Mf6Wel(cellids, well_rate, well_id) - coupler_mapping = msw.Sprinkling( - max_abstraction_groundwater, - max_abstraction_surfacewater, + sprinkling = msw.SprinklingGrid( + case_data.max_abstraction_groundwater.isel(subunit=[0]), + case_data.max_abstraction_surfacewater.isel(subunit=[0]), ) with tempfile.TemporaryDirectory() as output_dir: output_dir = Path(output_dir) - coupler_mapping.write(output_dir, index, svat, None, well) + sprinkling.write(output_dir, index, svat, None, well) results = fixed_format_parser( - output_dir / msw.Sprinkling._file_name, - msw.Sprinkling._metadata_dict, + output_dir / msw.SprinklingGrid._file_name, + msw.SprinklingGrid._metadata_dict, ) - assert_equal(results["svat"], np.array([1, 2, 4])) + assert_equal(results["svat"], expected_data.svat[:2]) assert_almost_equal( results["max_abstraction_groundwater"], - np.array([100.0, 0.0, 200.0]), + expected_data.abs_gw[:2], ) assert_almost_equal( results["max_abstraction_surfacewater"], - np.array([0.0, 300.0, 200.0]), + expected_data.abs_sw[:2], ) - assert_equal(results["layer"], np.array([3, 1, 2])) - assert_equal(results["svat_groundwater"], np.array([1, 2, 4])) + assert_equal(results["layer"], expected_data.layer[:2]) + assert_equal(results["svat_groundwater"], expected_data.svat_gw[:2]) -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 - max_abstraction_groundwater = xr.DataArray( - np.array( - [ - [[nan, 100.0, nan], - [nan, 200.0, nan], - [nan, 300.0, nan]] - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} - ) +@parametrize_with_cases("case_data, expected_data", cases=SprinklingPointsCases) +def test_points_simple_model( + fixed_format_parser: Callable, + sprinkling_svat_index: tuple[xr.DataArray, np.ndarray], + case_data: SprinklingPointsCaseData, + expected_data: ExpectedCaseData, +): + if expected_data.xfail: + pytest.xfail(expected_data.xfail) + svat, index = sprinkling_svat_index - max_abstraction_surfacewater = xr.DataArray( - np.array( - [ - [[nan, 100.0, nan], - [nan, 200.0, nan], - [nan, 300.0, nan]] - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} + # Well + n_wells = len(case_data.x_p) + well_rate_values = [-5.0] * n_wells + well_rate = xr.DataArray(well_rate_values, dims=("ncellid",)) + well_id_values = [str(i) for i in np.arange(n_wells)] + well_id = xr.DataArray(well_id_values, dims=("ncellid",)) + + cellids = derive_cellid_from_points( + svat, case_data.x_p, case_data.y_p, case_data.layer_p ) + well = Mf6Wel(cellids, well_rate, well_id) - svat = xr.DataArray( - np.array( - [ - [[0, 1, 0], - [0, 0, 0], - [0, 2, 0]], - ] - ), - dims=("subunit", "y", "x"), - coords={"subunit": subunit, "y": y, "x": x, "dx": dx, "dy": dy} + layer_template = xr.DataArray( + [1.0, 2.0, 3.0], coords={"layer": [1, 2, 3]}, dims=("layer",) ) - # fmt: on - index = (svat != 0).values.ravel() + grid_2d_template = xr.ones_like(svat.isel(subunit=0, drop=True), dtype=float) + mf6_dis_template = layer_template * grid_2d_template - # Well - well_layer = [3, 2] - well_y = [1.0, 3.0] - well_x = [2.0, 2.0] - well_rate = [-5.0] * 2 - cellids = derive_cellid_from_points(svat, well_x, well_y, well_layer) - well = Mf6Wel(cellids, well_rate) + dis = StructuredDiscretization( + top=grid_2d_template, + bottom=-mf6_dis_template, + idomain=mf6_dis_template.astype(int), + ) - sprinkling = msw.Sprinkling( - max_abstraction_groundwater, - max_abstraction_surfacewater, + sprinkling = msw.SprinklingPoints( + case_data.art_grid, + case_data.x_p, + case_data.y_p, + case_data.layer_p, + case_data.id_sprinkling_p, + case_data.capacity_p, ) with tempfile.TemporaryDirectory() as output_dir: output_dir = Path(output_dir) - sprinkling.write(output_dir, index, svat, None, well) + sprinkling.write(output_dir, index, svat, dis, well) results = fixed_format_parser( - output_dir / msw.Sprinkling._file_name, - msw.Sprinkling._metadata_dict, + output_dir / msw.SprinklingPoints._file_name, + msw.SprinklingPoints._metadata_dict, ) - assert_equal(results["svat"], np.array([1, 2])) + assert_equal(results["svat"], expected_data.svat) assert_almost_equal( results["max_abstraction_groundwater"], - np.array([100.0, 300.0]), + expected_data.abs_gw, ) assert_almost_equal( results["max_abstraction_surfacewater"], - np.array([100.0, 300.0]), + expected_data.abs_sw, ) - assert_equal(results["layer"], np.array([3, 2])) - assert_equal(results["svat_groundwater"], np.array([1, 2])) + assert_equal(results["layer"], expected_data.layer) + assert_equal(results["svat_groundwater"], expected_data.svat_gw) @pytest.mark.unittest_jit def test_sprinkling_from_imod5_data__points(cap_data_sprinkling_points): - with pytest.raises(NotImplementedError): - msw.Sprinkling.from_imod5_data(cap_data_sprinkling_points) + with pytest.raises(TypeError): + msw.SprinklingGrid.from_imod5_data(cap_data_sprinkling_points) + + +@pytest.mark.unittest_jit +def test_sprinklingpoints_from_imod5_data__grid(cap_data_sprinkling_grid): + with pytest.raises(TypeError): + msw.SprinklingPoints.from_imod5_data(cap_data_sprinkling_grid) @pytest.mark.unittest_jit @@ -387,10 +591,10 @@ def test_sprinkling_from_imod5_data__grid(cap_data_sprinkling_grid): # fmt: on # Act - sprinkling = msw.Sprinkling.from_imod5_data(cap_data_sprinkling_grid) + sprinkling = msw.SprinklingGrid.from_imod5_data(cap_data_sprinkling_grid) # Assert - assert isinstance(sprinkling, msw.Sprinkling) + assert isinstance(sprinkling, msw.SprinklingGrid) ds = sprinkling.dataset assert (ds.sel(subunit=1) == 0).all() rural_ds = ds.sel(subunit=0) @@ -400,3 +604,73 @@ def test_sprinkling_from_imod5_data__grid(cap_data_sprinkling_grid): np.testing.assert_array_equal( rural_ds["max_abstraction_surfacewater"].to_numpy(), expected_sw_abstraction ) + + +@pytest.mark.unittest_jit +def test_sprinklingpoints_from_imod5_data__points(cap_data_sprinkling_points): + # Arrange + expected_vars = { + "id_sprinkling_p", + "capacity_p", + "layer_p", + "y_p", + "x_p", + "id_sprinkling", + } + + # Act + sprinkling = msw.SprinklingPoints.from_imod5_data(cap_data_sprinkling_points) + + # Assert + assert sprinkling.dataset.sizes == {"subunit": 2, "id": 2, "x": 3, "y": 3} + assert set(sprinkling.dataset.keys()) == expected_vars + # No unit conversion is done in SprinklingPoints, as the capacity is already + # in m3/d + np.testing.assert_almost_equal(sprinkling.dataset["capacity_p"], [15.0, 30.0]) + + +@pytest.mark.unittest_jit +def test_sprinklingpoints_from_imod5_data_write__points( + sprinkling_svat_index, + fixed_format_parser, + cap_data_sprinkling_points, + cap_coupled_dis_grid, + tmp_path, +): + """ + Test with two wells: one inside the active SVAT area, and one outside the + active SVAT area but still in the model domain. Well nr 2. is not assigned + to anything. Well nr 1. is assigned and is located in the centre cell. In + subunit 1 this cell is inactive and the svats coupled to this well are + assigned to surface water, in subunit 2 this cell is active and this well is + coupled to the groundwater svat. + """ + # Arrange + svat, index = sprinkling_svat_index + df = cap_data_sprinkling_points["cap"]["artificial_recharge_layer"] + well_x = df["x"].to_numpy() + well_y = df["y"].to_numpy() + well_layer = df["layer"].to_numpy() + well_rate = xr.DataArray([0.0, 0.0], dims=("ncellid",)) + well_id = xr.DataArray(["0", "1"], dims=("ncellid",)) + cellids = derive_cellid_from_points(svat, well_x, well_y, well_layer) + mf6_well = Mf6Wel(cellids, well_rate, well_id) + mf6_dis = cap_coupled_dis_grid + directory = tmp_path / "sprinkling_points" + directory.mkdir(parents=True, exist_ok=True) + + # Act + sprinkling = msw.SprinklingPoints.from_imod5_data(cap_data_sprinkling_points) + sprinkling.write(directory, index, svat, mf6_dis, mf6_well) + + results = fixed_format_parser( + directory / msw.SprinklingGrid._file_name, + msw.SprinklingGrid._metadata_dict, + ) + + # Assert + np.testing.assert_equal(results["svat"], [1, 2]) + np.testing.assert_equal(results["svat_groundwater"], [1, 2]) + np.testing.assert_equal(results["layer"], [2, 3]) + np.testing.assert_equal(results["max_abstraction_surfacewater"], [0.0, 30.0]) + np.testing.assert_equal(results["max_abstraction_groundwater"], [15.0, 0.0])