diff --git a/docs/api/msw.rst b/docs/api/msw.rst index a14f60f59..2f7a1194a 100644 --- a/docs/api/msw.rst +++ b/docs/api/msw.rst @@ -56,6 +56,12 @@ Grid packages Sprinkling.from_imod5_data Sprinkling.get_regrid_methods Sprinkling.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..23866bf0e 100644 --- a/imod/mf6/mf6_wel_adapter.py +++ b/imod/mf6/mf6_wel_adapter.py @@ -143,6 +143,7 @@ def __init__( self, cellid, rate, + id, concentration=None, concentration_boundary_type="aux", save_flows: Optional[bool] = None, @@ -153,6 +154,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 +171,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 e39af906a..abd280628 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 @@ -525,9 +527,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: @@ -537,8 +539,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] @@ -1060,8 +1060,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) @@ -1449,7 +1449,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(), ): @@ -1489,6 +1489,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..6169964e6 100644 --- a/imod/msw/__init__.py +++ b/imod/msw/__init__.py @@ -16,5 +16,5 @@ 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 +from imod.msw.sprinkling import Sprinkling, SprinklingPoints from imod.msw.vegetation import AnnualCropFactors diff --git a/imod/msw/regrid/regrid_schemes.py b/imod/msw/regrid/regrid_schemes.py index cf82018b1..084e78841 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.Sprinkling` 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")) + >>> sprinking.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..4b71c94c5 100644 --- a/imod/msw/sprinkling.py +++ b/imod/msw/sprinkling.py @@ -1,4 +1,5 @@ -from typing import TextIO +import textwrap +from typing import TextIO, TypedDict, cast import numpy as np import pandas as pd @@ -9,15 +10,38 @@ 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.regrid.regrid_schemes import ( + SprinklingPointsRegridMethod, + SprinklingRegridMethod, +) from imod.msw.utilities.common import concat_imod5 from imod.msw.utilities.imod5_converter import ( get_cell_area_from_imod5_data, ) -from imod.typing import GridDataDict, Imod5DataDict, IntArray +from imod.typing import GridDataArray, GridDataDict, Imod5DataDict, IntArray from imod.typing.grid import zeros_like +# 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] + id2grid_p: np.ndarray | list[int] + capacity_p: np.ndarray | list[float] + + +class SprinklingPointsGridDataDict(SprinklingPointsDataDict, total=False): + art_grid: GridDataArray + + def _ravel_per_subunit(da: xr.DataArray) -> np.ndarray: # per defined well element, all subunits array_out = da.to_numpy().ravel() @@ -25,11 +49,37 @@ 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 _sprinkling_data_from_imod5_ipf( + cap_data: CapSprinklingDataDict, +) -> SprinklingPointsGridDataDict: + art_grid = cap_data["artificial_recharge"] + 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", "id2grid_p", "capacity_p"] + # Enforce dtypes + dtype_dict = { + "x_p": float, + "y_p": float, + "layer_p": int, + "id2grid_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 @@ -65,10 +115,151 @@ def _sprinkling_data_from_imod5_grid(cap_data: GridDataDict) -> GridDataDict: return data +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. + + 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 + + indexer = df.loc[:, columns].to_numpy() + return indexer.T + + +def _replicate_dataframe_by_subunit( + df: pd.DataFrame, subunit_col: str = "subunit" +) -> pd.DataFrame: + """ + Duplicate the rows of a DataFrame for each subunit (0 and 1). + + Parameters + ---------- + df : pd.DataFrame + Input DataFrame to duplicate. + subunit_col : str, optional + Name of the column to assign subunit values, by default "subunit". + + Returns + ------- + pd.DataFrame + DataFrame with duplicated rows for each subunit. + """ + subunit_nrs = [0, 1] + df_ls = [df.assign(**{subunit_col: subunit_nr}) for subunit_nr in subunit_nrs] + return pd.concat(df_ls, ignore_index=True) + + +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() + ) + # 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" + ) + + +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. + """ + + # TODO: Rename "id_msw" and "id2grid_p" to something clearer like "id_sprinkling" + # Flatten id_msw grid → (y, x, id_msw) table, drop cells with no well + grids = xr.merge([sprinkling_id_grid, svat]) + art_df = grids.to_dataframe().reset_index().query("(id_msw > 0) & (svat > 0)") + # 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_msw", + right_on="id2grid_p", + how="inner", + validate="many_to_one", + ) + + +def align_svat_with_dis( + svat: xr.DataArray, dis_pkg: StructuredDiscretization +) -> xr.DataArray: + """ + 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 + + +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 Sprinkling(MetaSwapPackage, IRegridPackage): """ This contains the sprinkling capacities of links between SVAT units and - groundwater/surface water locations. + 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` @@ -175,32 +366,28 @@ def _render( @classmethod def from_imod5_data(cls, imod5_data: Imod5DataDict) -> "Sprinkling": """ - 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 ---------- @@ -215,9 +402,218 @@ def from_imod5_data(cls, imod5_data: Imod5DataDict) -> "Sprinkling": Sprinkling package """ cap_data = imod5_data["cap"] + if isinstance(cap_data["artificial_recharge_layer"], pd.DataFrame): + 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) + + data = _sprinkling_data_from_imod5_grid(cap_data) + + return cls(**data) + + +class SprinklingPoints(MetaSwapPackage, IRegridPackage): + """ + 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` + + Parameters + ---------- + art_grid: xr.DataArray + Grid of the artificial recharge types, with subunit coordinate. + 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. + id2grid_p: np.ndarray | list[int] + mapping of the artificial recharge locations to the grid cells. + capacity_p: np.ndarray | list[float] + abstraction capacities of the artificial recharge locations. + + """ + + _file_name = "scap_svat.inp" + _metadata_dict = { + "svat": VariableMetaData(10, 1, 99999999, int), + "max_abstraction_groundwater_mm_d": VariableMetaData(8, None, None, str), + "max_abstraction_surfacewater_mm_d": VariableMetaData(8, None, None, str), + "max_abstraction_groundwater": VariableMetaData(8, 0.0, 1e9, float), + "max_abstraction_surfacewater": VariableMetaData(8, 0.0, 1e9, float), + "svat_groundwater": VariableMetaData(10, 1, 99999999, int), + "layer": VariableMetaData(6, 1, 9999, int), + "trajectory": VariableMetaData(10, None, None, str), + } + + _with_subunit = ( + "max_abstraction_groundwater", + "max_abstraction_surfacewater", + ) + _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], + id2grid_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), + "id2grid_p": (("id",), id2grid_p), + "capacity_p": (("id",), capacity_p), + }, + coords={"id": id_index}, + ) + art_grid = art_grid.rename("id_msw") + self.dataset = xr.merge([art_grid, points_ds]) + + @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.Sprinkling.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 + """ + cap_data = cast(CapSprinklingDataDict, imod5_data["cap"]) if isinstance(cap_data["artificial_recharge_layer"], pd.DataFrame): data = _sprinkling_data_from_imod5_ipf(cap_data) + return cls(**data) else: - data = _sprinkling_data_from_imod5_grid(cap_data) + 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_msw 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_msw grid + msw_mf6_sprinkling_df = _merge_sprinkling_points_with_grids( + points_df, svat, self.dataset["id_msw"] + ) + 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())] - return cls(**data) + self._check_range(dataframe_out) + + return self._write_dataframe_fixed_width(file, dataframe_out) diff --git a/imod/tests/_scratch.py b/imod/tests/_scratch.py new file mode 100644 index 000000000..be0f2dba2 --- /dev/null +++ b/imod/tests/_scratch.py @@ -0,0 +1,200 @@ +# %% +import pandas as pd +import xarray as xr + +import imod +from imod.util.dims import drop_layer_dim_cap_data + + +def get_indexer(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. + + Parameters + ---------- + df : pd.DataFrame + DataFrame containing the wells with columns "subunit", "row", + and "column" (1-based). + + Returns + ------- + np.ndarray + Indexer array for selecting SVAT subunits from svat_da. + """ + if not set("row", "column").issubset(columns): + raise ValueError("columns must contain 'row' and 'column'") + df.loc[:, ["row", "column"]] -= 1 # Convert to 0-based indexing for xarray + + indexer = df.loc[:, columns].to_numpy() + return indexer.T + + +def double_length_df_subunit( + df: pd.DataFrame, subunit_col: str = "subunit" +) -> pd.DataFrame: + """ + Duplicate the rows of a DataFrame for each subunit (0 and 1). + + Parameters + ---------- + df : pd.DataFrame + Input DataFrame to duplicate. + subunit_col : str, optional + Name of the column to assign subunit values, by default "subunit". + + Returns + ------- + pd.DataFrame + DataFrame with duplicated rows for each subunit. + """ + return pd.concat( + [df.assign(**{subunit_col: 0}), df.assign(**{subunit_col: 1})], + ignore_index=True, + ) + + +# %% +df_points = imod.ipf.read( + r"c:\Users\engelen\projects_wdir\imod-python\imod5_converter\NHI_sprint\Peelvenen\BASIS7\METASWAP\grid\Sprinkling\BEREGEN_LOC.IPF" +) +art_grid = ( + imod.idf.open( + r"c:\Users\engelen\projects_wdir\imod-python\imod5_converter\NHI_sprint\Peelvenen\BASIS7\METASWAP\grid\Sprinkling\BEREGENINGS_LOCATIES.IDF" + ) + .compute() + .astype(int) +) + +imod5_data = { + "cap": {"artificial_recharge": art_grid, "artificial_recharge_layer": df_points} +} + +well = imod.mf6.LayeredWell.from_imod5_cap_data( + imod5_data, target_dis=None, regridder_types=None, regrid_cache=None +) +# %% +# Setup to get example args for _render() of SprinklingPoints +prj_data, period_data = imod.formats.prj.open_projectfile_data( + r"c:\Users\engelen\projects_wdir\imod-python\imod5_converter\NHI_sprint\Peelvenen\prjfiles\Peelvenen_relative_paths_dis_npf.PRJ" +) + +dis_pkg = imod.mf6.StructuredDiscretization.from_imod5_data(prj_data, validate=False) +dis_pkg["idomain"] = dis_pkg["idomain"].clip(min=0) +npf_pkg = imod.mf6.NodePropertyFlow.from_imod5_data( + prj_data, dis_pkg.dataset["idomain"] +) + +prj_data = drop_layer_dim_cap_data(prj_data) +griddata, msw_active = imod.msw.GridData.from_imod5_data(prj_data, dis_pkg) +# Convert to args of Sprinkling._render() +mf6_well = well.to_mf6_pkg( + dis_pkg["idomain"], dis_pkg["top"], dis_pkg["bottom"], npf_pkg["k"] +) +isactive_1d, svat = griddata.generate_isactive_svat_arrays() + +# %% +# In from_imod5_cap_data +arl_points = df_points.iloc[:, :5] +arl_points.columns = ["x_p", "y_p", "layer_p", "id2grid_p", "capacity"] +# Enforce dtypes +arl_points = arl_points.astype( + {"x_p": float, "y_p": float, "layer_p": int, "id2grid_p": int, "capacity": float} +) +arl_points["id"] = arl_points.index.astype(str) +arl_points = arl_points.set_index("id") + +points_ds = arl_points.to_xarray() + +# in def __init__ +art_grid = art_grid.rename("id_msw") +dataset = xr.merge([art_grid, points_ds]) + +# %% +# In render() +# 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() +# Merge again to confine to wells actually used in the modflow6 model. +# This drops points that are not in art_grid. +points_mf6_merged_df = arl_points.reset_index().merge( + mf6_cellid_df, on="id", how="right", validate="many_to_one" +) +# %% +# Flatten id_msw grid → (y, x, id_msw) table, drop cells with no well +grids = xr.merge([art_grid, svat]) +art_df = grids.to_dataframe().reset_index().query("(id_msw > 0) & (svat > 0)") +# 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 +msw_mf6_merged_df = art_df.merge( + points_mf6_merged_df, # brings id back as a column + left_on="id_msw", + right_on="id2grid_p", + how="inner", + validate="many_to_one", +) + +# %% +# Derive the SVAT subunit for each well from the SVAT grid based on the +# well's row/col location. +indexer = get_indexer(msw_mf6_merged_df, columns=["subunit", "row", "column"]) +# We need to 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) +_, svat_aligned = xr.align(idomain_flat, svat, join="left") +svat_groundwater = svat_aligned.data[*indexer] + +msw_mf6_merged_df["svat_groundwater"] = svat_groundwater.astype(int) + +# %% +# Select columns that need to be written to scap_svat.inp +dataframe = msw_mf6_merged_df[["svat", "layer", "svat_groundwater"]] +dataframe["svat"] = dataframe["svat"].astype(int) +capacity = msw_mf6_merged_df["capacity"] +# Set wells with layer > 0 to groundwater abstraction, and wells with layer = 0 +# to surfacewater abstraction. +is_gw_extraction = msw_mf6_merged_df["layer"] > 0 +dataframe["max_abstraction_groundwater"] = capacity.where(is_gw_extraction, 0.0) +dataframe["max_abstraction_surfacewater"] = capacity.where(~is_gw_extraction, 0.0) + +# %% +# +# Deal with edge case: wells that are outside art_grid, but in model domain. +# These will be assigned to surfacewater abstraction. We can identify these by +# checking which wells in mf6_cellid_df are not in msw_mf6_merged_df. +is_outside_art = ~mf6_cellid_df["id"].isin(msw_mf6_merged_df["id"].unique()) +outside_df = points_mf6_merged_df.loc[is_outside_art, ["layer", "capacity"]] +outside_df = double_length_df_subunit(outside_df) +# Select the SVAT subunit for these wells based on their row/col location. +df_cellid_outside = mf6_cellid_df.loc[is_outside_art] +df_cellid_outside = double_length_df_subunit(df_cellid_outside) +indexer_outside = get_indexer(df_cellid_outside, columns=["subunit", "row", "column"]) +svat_outside = svat_aligned.data[ + indexer_outside[0], indexer_outside[1], indexer_outside[2] +] +outside_df["svat_groundwater"] = svat_outside.astype(int) +outside_df["svat"] = svat_outside.astype(int) +# Set capacity to surfacewater abstraction, and set groundwater abstraction to 0. +outside_df = outside_df.rename(columns={"capacity": "max_abstraction_surfacewater"}) +outside_df["max_abstraction_groundwater"] = 0.0 +# drop subunit column as it is no longer needed +outside_df = outside_df.drop(columns=["subunit"]) +# drop wells that are outside the active metaswap model domain (svat = 0) +outside_df = outside_df.query("svat > 0").reset_index(drop=True) + +# %% +# +# Combine +dataframe_out = pd.concat([dataframe, outside_df], axis=0, ignore_index=True) +dataframe_out = dataframe_out.sort_values(by=["svat"]).reset_index(drop=True) +# %% +# +# TODO: Verify if iMOD5 SVAT grid the same as the one derived in this script. diff --git a/imod/tests/fixtures/imod5_cap_data.py b/imod/tests/fixtures/imod5_cap_data.py index 88a2a5d7d..b670abdc1 100644 --- a/imod/tests/fixtures/imod5_cap_data.py +++ b/imod/tests/fixtures/imod5_cap_data.py @@ -146,20 +146,20 @@ def cap_data_sprinkling_points() -> Imod5DataDict: artificial_rch_type[:, 2] = 4000 data = { - "id": [3000, 4000], + "x": [1.5, 2.5], + "y": [1.5, 2.5], "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/test_msw/test_sprinkling.py b/imod/tests/test_msw/test_sprinkling.py index 63812df9f..e10aa8e9e 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 # 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 + id2grid_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,26 @@ 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( - [ - [[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() - - # 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( + # 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( [ [[nan, 100.0, nan], [nan, 200.0, nan], @@ -119,13 +127,8 @@ 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( + ) + case_data.max_abstraction_surfacewater.data = np.array( [ [[nan, 100.0, nan], [nan, 200.0, nan], @@ -134,240 +137,389 @@ 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} - ) - - svat = 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( [ - [[0, 1, 0], - [0, 0, 0], - [0, 2, 0]], - - [[0, 3, 0], - [0, 4, 0], - [0, 0, 0]], + [[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} - ) - # fmt: on - index = (svat != 0).values.ravel() + ) + 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]] + ] + ) + # 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.isel(subunit=0, drop=True), 0, dtype=int) + # fmt: off + case_data.art_grid.data = np.array( + [[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.id2grid_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_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.isel(subunit=0, drop=True), 0, dtype=int) + # fmt: off + case_data.art_grid.data = np.array( + [[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.id2grid_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.isel(subunit=0, drop=True), 0, dtype=int) + # fmt: off + case_data.art_grid.data = np.array( + [[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.id2grid_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.isel(subunit=0, drop=True), 0, dtype=int) + # fmt: off + case_data.art_grid.data = np.array( + [[0, 0, 4], + [0, 0, 0], + [0, 0, 0]] + ) + # fmt: on + case_data.x_p = [2.0] + case_data.y_p = [2.0] + case_data.layer_p = [3] + case_data.id2grid_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.isel(subunit=0, drop=True), 0, dtype=int) + # fmt: off + case_data.art_grid.data = np.array( + [[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.id2grid_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_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.Sprinkling( + 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, ) - 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 = [1.0, 3.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.Sprinkling( + 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, ) - 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.id2grid_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): + with pytest.raises(TypeError): msw.Sprinkling.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 def test_sprinkling_from_imod5_data__grid(cap_data_sprinkling_grid): # Arrange @@ -400,3 +552,71 @@ 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 = {"id2grid_p", "capacity_p", "layer_p", "y_p", "x_p", "id_msw"} + + # Act + sprinkling = msw.SprinklingPoints.from_imod5_data(cap_data_sprinkling_points) + + # Assert + assert sprinkling.dataset.sizes == {"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.Sprinkling._file_name, + msw.Sprinkling._metadata_dict, + ) + + # Assert + # TODO: Check with Hendrik whether this is the appropriate behaviour. + np.testing.assert_equal(results["svat"], [1, 2, 3, 4]) + np.testing.assert_equal(results["svat_groundwater"], [1, 2, 4, 4]) + np.testing.assert_equal(results["layer"], [2, 2, 2, 2]) + np.testing.assert_equal( + results["max_abstraction_surfacewater"], [15.0, 15.0, 0.0, 0.0] + ) + np.testing.assert_equal( + results["max_abstraction_groundwater"], [0.0, 0.0, 15.0, 15.0] + )