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/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..6be532f94 100644 --- a/imod/msw/sprinkling.py +++ b/imod/msw/sprinkling.py @@ -9,7 +9,10 @@ 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, @@ -65,10 +68,181 @@ 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 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 _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) + _, 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) + + +def _get_wells_outside_art_grid_dataframe( + mf6_cellid_df: pd.DataFrame, + points_df: pd.DataFrame, + msw_mf6_sprinkling_df: pd.DataFrame, + svat_aligned: xr.DataArray, +) -> pd.DataFrame: + """ + 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_sprinkling_df["id"].unique()) + outside_df = points_df.loc[is_outside_art, ["layer", "capacity"]] + outside_df = _replicate_dataframe_by_subunit(outside_df) + # Select the SVAT subunit for these wells based on their row/col location. + cellid_outside_df = mf6_cellid_df.loc[is_outside_art] + cellid_outside_df = _replicate_dataframe_by_subunit(cellid_outside_df) + indexer_outside = _extract_indexer_for_svat( + cellid_outside_df, columns=["subunit", "row", "column"] + ) + svat_outside = svat_aligned.data[*indexer_outside] + 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) + return outside_df.query("svat > 0").reset_index(drop=True) + + 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` @@ -221,3 +395,167 @@ def from_imod5_data(cls, imod5_data: Imod5DataDict) -> "Sprinkling": 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": + cap_data = imod5_data["cap"] + art_grid = cap_data["artificial_recharge"] + df_points = cap_data["artificial_recharge_layer"] + + arl_points = df_points.iloc[:, :5] + arl_points.columns = ["x_p", "y_p", "layer_p", "id2grid_p", "capacity_p"] + # Enforce dtypes + arl_points = arl_points.astype( + { + "x_p": float, + "y_p": float, + "layer_p": int, + "id2grid_p": int, + "capacity_p": float, + } + ) + + return cls( + art_grid=art_grid, + x_p=arl_points["x_p"].to_numpy(), + y_p=arl_points["y_p"].to_numpy(), + layer_p=arl_points["layer_p"].to_numpy(), + id2grid_p=arl_points["id2grid_p"].to_numpy(), + capacity_p=arl_points["capacity_p"].to_numpy(), + ) + + 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 + ) + + # Select columns that need to be written to scap_svat.inp + dataframe = msw_mf6_sprinkling_df[["svat", "layer", "svat_groundwater"]] + dataframe["svat"] = dataframe["svat"].astype(int) + capacity = msw_mf6_sprinkling_df["capacity"] + # Set wells with layer > 0 to groundwater abstraction, and wells with layer = 0 + # to surfacewater abstraction. + is_gw_extraction = msw_mf6_sprinkling_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 + ) + + # TODO: Make sure wells in svats are all present in the dataframe. If + # these svats are 0 in the art_grid, they should get a 0.0 capacity. + + # Deal with edge case: wells that are outside art_grid, but in model domain. + # These will be assigned to surfacewater abstraction. + outside_df = _get_wells_outside_art_grid_dataframe( + mf6_cellid_df, points_df, msw_mf6_sprinkling_df, svat_aligned + ) + + dataframe_out = pd.concat([dataframe, outside_df], axis=0, ignore_index=True) + dataframe_out = dataframe_out.sort_values(by=["svat"]).reset_index(drop=True) + + for var in self._to_fill: + dataframe_out[var] = "" + + 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.