From 66800d939409be723e34b62609b2950c06760480 Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Wed, 8 Jul 2026 16:32:23 +0200 Subject: [PATCH 01/16] Boyscouting: Get rid of calls to .values --- imod/mf6/wel.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/imod/mf6/wel.py b/imod/mf6/wel.py index e39af906a..5ebec41de 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: @@ -1060,8 +1062,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) From 0473920565aac87aab777b236e82d5dd2b8a51a0 Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Thu, 9 Jul 2026 15:17:59 +0200 Subject: [PATCH 02/16] Pass id on to Mf6Wel --- imod/mf6/mf6_wel_adapter.py | 4 +++- imod/mf6/wel.py | 2 -- 2 files changed, 3 insertions(+), 3 deletions(-) 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/wel.py b/imod/mf6/wel.py index 5ebec41de..30a201233 100644 --- a/imod/mf6/wel.py +++ b/imod/mf6/wel.py @@ -539,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] From 7040e132f38d1386008f8a5ded7e0494c083434f Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Mon, 13 Jul 2026 14:29:18 +0200 Subject: [PATCH 03/16] Support IPF wells for LayeredWell.from_imod5_cap_data --- imod/mf6/utilities/imod5_converter.py | 20 ++++++++++++++++---- imod/mf6/wel.py | 13 ++++++++++++- 2 files changed, 28 insertions(+), 5 deletions(-) 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 30a201233..abd280628 100644 --- a/imod/mf6/wel.py +++ b/imod/mf6/wel.py @@ -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 From aa1ab088e0acff2014e0056f61db48488610355d Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Wed, 15 Jul 2026 16:12:59 +0200 Subject: [PATCH 04/16] Temp add scratch script to try out data reworking --- imod/tests/_scratch.py | 120 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 120 insertions(+) create mode 100644 imod/tests/_scratch.py diff --git a/imod/tests/_scratch.py b/imod/tests/_scratch.py new file mode 100644 index 000000000..8d3e986b4 --- /dev/null +++ b/imod/tests/_scratch.py @@ -0,0 +1,120 @@ +# %% +import imod +import xarray as xr +from imod.util.dims import drop_layer_dim_cap_data + + +# %% +SUBUNIT = 0 + +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 +) +# %% +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() +art_grid = art_grid.rename("id_msw") +dataset = xr.merge([art_grid, points_ds]) + +# %% +# Flatten id_msw grid → (y, x, id_msw) table, drop cells with no well +# TODO: Verify with Peter that only the first subunit (landuse: agriculture) is +# relevant for the SVAT mapping. +grids = xr.merge([art_grid, svat.sel(subunit=SUBUNIT, drop=True)]) + +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_merged_df = art_df.merge( + arl_points.reset_index(), # brings id back as a column + left_on="id_msw", + right_on="id2grid_p", + how="inner", + validate="many_to_one" +) + +# %% +# 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. +msw_mf6_merged_df = msw_merged_df.merge( + mf6_cellid_df, + on="id", + 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 = msw_mf6_merged_df.loc[:, dim_cellid].to_numpy() - 1 +# 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[SUBUNIT, indexer[:, 1], indexer[:, 2]] + +msw_mf6_merged_df["svat_groundwater"] = svat_groundwater.astype(int) + +# %% +# Deal with edge case: wells that are outside art_grid, but in model domain +outside_art = ~mf6_cellid_df["id"].isin(msw_mf6_merged_df["id"].unique()) +indexer_outside = mf6_cellid_df.loc[outside_art, dim_cellid].to_numpy() - 1 +svat_outside = svat_aligned.data[SUBUNIT, indexer_outside[:, 1], indexer_outside[:, 2]] + +# %% +# 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"] +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) + +# Cases to catch: +# 1. Well is outside the model domain → surfacewater +# 2. Well is in layer 0 → surfacewater + + From 6405be447b34fcd41e1184809ada799f8383b0ce Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Thu, 16 Jul 2026 15:15:02 +0200 Subject: [PATCH 05/16] Change order of merging to first merge tables, then to grid and deal with edge case --- imod/tests/_scratch.py | 77 ++++++++++++++++++++++++------------------ 1 file changed, 45 insertions(+), 32 deletions(-) diff --git a/imod/tests/_scratch.py b/imod/tests/_scratch.py index 8d3e986b4..56b3dba8e 100644 --- a/imod/tests/_scratch.py +++ b/imod/tests/_scratch.py @@ -2,7 +2,7 @@ import imod import xarray as xr from imod.util.dims import drop_layer_dim_cap_data - +import pandas as pd # %% SUBUNIT = 0 @@ -16,16 +16,15 @@ 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() @@ -43,6 +42,23 @@ art_grid = art_grid.rename("id_msw") dataset = xr.merge([art_grid, points_ds]) +# %% +# 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 # TODO: Verify with Peter that only the first subunit (landuse: agriculture) is @@ -60,32 +76,14 @@ art_df = art_df.drop(["dx", "dy"], axis=1) # Join: each SVAT cell gets the matching well row(s) from arl_points -msw_merged_df = art_df.merge( - arl_points.reset_index(), # brings id back as a column +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" ) -# %% -# 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. -msw_mf6_merged_df = msw_merged_df.merge( - mf6_cellid_df, - on="id", - 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. @@ -98,23 +96,38 @@ msw_mf6_merged_df["svat_groundwater"] = svat_groundwater.astype(int) -# %% -# Deal with edge case: wells that are outside art_grid, but in model domain -outside_art = ~mf6_cellid_df["id"].isin(msw_mf6_merged_df["id"].unique()) -indexer_outside = mf6_cellid_df.loc[outside_art, dim_cellid].to_numpy() - 1 -svat_outside = svat_aligned.data[SUBUNIT, indexer_outside[:, 1], indexer_outside[:, 2]] - # %% # 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) -# Cases to catch: -# 1. Well is outside the model domain → surfacewater -# 2. Well is in layer 0 → surfacewater +# %% +# +# 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"]] +# Select the SVAT subunit for these wells based on their row/col location. +indexer_outside = mf6_cellid_df.loc[is_outside_art, dim_cellid].to_numpy() - 1 +svat_outside = svat_aligned.data[SUBUNIT, 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 +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) +# %% From fe251dd0163580a135326c7c5f7ae0f3b3967c2b Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Thu, 16 Jul 2026 16:57:42 +0200 Subject: [PATCH 06/16] Make script work for both subunits --- imod/tests/_scratch.py | 69 ++++++++++++++++++++++++++++++++++-------- 1 file changed, 57 insertions(+), 12 deletions(-) diff --git a/imod/tests/_scratch.py b/imod/tests/_scratch.py index 56b3dba8e..048f3fc4a 100644 --- a/imod/tests/_scratch.py +++ b/imod/tests/_scratch.py @@ -4,9 +4,50 @@ from imod.util.dims import drop_layer_dim_cap_data import pandas as pd -# %% -SUBUNIT = 0 +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) @@ -61,10 +102,7 @@ ) # %% # Flatten id_msw grid → (y, x, id_msw) table, drop cells with no well -# TODO: Verify with Peter that only the first subunit (landuse: agriculture) is -# relevant for the SVAT mapping. -grids = xr.merge([art_grid, svat.sel(subunit=SUBUNIT, drop=True)]) - +grids = xr.merge([art_grid, svat]) art_df = ( grids .to_dataframe() @@ -84,15 +122,15 @@ validate="many_to_one" ) -# %% +# %% # Derive the SVAT subunit for each well from the SVAT grid based on the # well's row/col location. -indexer = msw_mf6_merged_df.loc[:, dim_cellid].to_numpy() - 1 +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[SUBUNIT, indexer[:, 1], indexer[:, 2]] +svat_groundwater = svat_aligned.data[*indexer] msw_mf6_merged_df["svat_groundwater"] = svat_groundwater.astype(int) @@ -114,15 +152,20 @@ # 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. -indexer_outside = mf6_cellid_df.loc[is_outside_art, dim_cellid].to_numpy() - 1 -svat_outside = svat_aligned.data[SUBUNIT, indexer_outside[:, 1], indexer_outside[:, 2]] +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) # %% @@ -131,3 +174,5 @@ 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. \ No newline at end of file From 42948e5545e1a0735d8bf323fcfc3eddac10616c Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Wed, 22 Jul 2026 09:43:48 +0200 Subject: [PATCH 07/16] Prototype SprinklingPoints class --- imod/msw/regrid/regrid_schemes.py | 22 ++ imod/msw/sprinkling.py | 342 +++++++++++++++++++++++++++++- 2 files changed, 362 insertions(+), 2 deletions(-) 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) From 87a5e02a2a91efc4a6419972f2cbfdb42a5fed8b Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Wed, 22 Jul 2026 09:43:58 +0200 Subject: [PATCH 08/16] Format --- imod/tests/_scratch.py | 82 ++++++++++++++++++++++++++---------------- 1 file changed, 52 insertions(+), 30 deletions(-) diff --git a/imod/tests/_scratch.py b/imod/tests/_scratch.py index 048f3fc4a..be0f2dba2 100644 --- a/imod/tests/_scratch.py +++ b/imod/tests/_scratch.py @@ -1,8 +1,10 @@ # %% -import imod +import pandas as pd import xarray as xr + +import imod from imod.util.dims import drop_layer_dim_cap_data -import pandas as pd + def get_indexer(df: pd.DataFrame, columns: list[str]): """ @@ -20,7 +22,7 @@ def get_indexer(df: pd.DataFrame, columns: list[str]): np.ndarray Indexer array for selecting SVAT subunits from svat_da. """ - if not set(["row", "column"]).issubset(columns): + 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 @@ -28,7 +30,9 @@ def get_indexer(df: pd.DataFrame, columns: list[str]): return indexer.T -def double_length_df_subunit(df: pd.DataFrame, subunit_col: str = "subunit") -> pd.DataFrame: +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). @@ -44,30 +48,49 @@ def double_length_df_subunit(df: pd.DataFrame, subunit_col: str = "subunit") -> 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) + 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) +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}} +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") +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"]) +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"]) +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() # %% @@ -75,15 +98,20 @@ def double_length_df_subunit(df: pd.DataFrame, subunit_col: str = "subunit") -> 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 = 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. @@ -95,36 +123,28 @@ def double_length_df_subunit(df: pd.DataFrame, subunit_col: str = "subunit") -> # 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" + 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)") -) +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 + points_mf6_merged_df, # brings id back as a column left_on="id_msw", right_on="id2grid_p", how="inner", - validate="many_to_one" + validate="many_to_one", ) # %% # Derive the SVAT subunit for each well from the SVAT grid based on the -# well's row/col location. +# 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. @@ -141,12 +161,12 @@ def double_length_df_subunit(df: pd.DataFrame, subunit_col: str = "subunit") -> 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 +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. @@ -157,7 +177,9 @@ def double_length_df_subunit(df: pd.DataFrame, subunit_col: str = "subunit") -> 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]] +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. @@ -175,4 +197,4 @@ def double_length_df_subunit(df: pd.DataFrame, subunit_col: str = "subunit") -> 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. \ No newline at end of file +# TODO: Verify if iMOD5 SVAT grid the same as the one derived in this script. From f088ad279b4582936a709ccc404ffde61d0b4127 Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Wed, 12 Aug 2026 13:41:40 +0200 Subject: [PATCH 09/16] Fix mypy errors --- imod/msw/sprinkling.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/imod/msw/sprinkling.py b/imod/msw/sprinkling.py index 6be532f94..148f16ebc 100644 --- a/imod/msw/sprinkling.py +++ b/imod/msw/sprinkling.py @@ -86,7 +86,7 @@ def _extract_indexer_for_svat(df: pd.DataFrame, columns: list[str]): np.ndarray Indexer array for selecting SVAT subunits from svat_da. """ - if not set("row", "column").issubset(columns): + 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 @@ -186,7 +186,9 @@ def align_svat_with_dis( 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") + # 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 From eb8a9e2ae20effc3cba563694c4d8ce0a0e4a139 Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Wed, 12 Aug 2026 13:45:48 +0200 Subject: [PATCH 10/16] Convert to set literal --- imod/msw/sprinkling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/imod/msw/sprinkling.py b/imod/msw/sprinkling.py index 148f16ebc..6f065d99c 100644 --- a/imod/msw/sprinkling.py +++ b/imod/msw/sprinkling.py @@ -86,7 +86,7 @@ def _extract_indexer_for_svat(df: pd.DataFrame, columns: list[str]): np.ndarray Indexer array for selecting SVAT subunits from svat_da. """ - if not set(["row", "column"]).issubset(columns): + 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 From 308b1577727043744fb9ff173b609c6e22c2e910 Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Wed, 12 Aug 2026 14:29:51 +0200 Subject: [PATCH 11/16] Clearer varname for inside_df --- imod/msw/sprinkling.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/imod/msw/sprinkling.py b/imod/msw/sprinkling.py index 6f065d99c..516100b94 100644 --- a/imod/msw/sprinkling.py +++ b/imod/msw/sprinkling.py @@ -532,18 +532,18 @@ def _render(self, file, index, svat, mf6_dis, mf6_well): ) # 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) + inside_df = msw_mf6_sprinkling_df[["svat", "layer", "svat_groundwater"]] + inside_df["svat"] = inside_df["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( + inside_df["max_abstraction_groundwater"] = capacity.where(is_gw_extraction, 0.0) + inside_df["max_abstraction_surfacewater"] = capacity.where( ~is_gw_extraction, 0.0 ) - # TODO: Make sure wells in svats are all present in the dataframe. If + # TODO: Make sure wells in svats are all present in the inside_df. 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. @@ -552,7 +552,7 @@ def _render(self, file, index, svat, mf6_dis, mf6_well): 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 = pd.concat([inside_df, 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: From 99dba8f5bbdda76a3aa239bb1829355f2084f91f Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Wed, 12 Aug 2026 15:27:03 +0200 Subject: [PATCH 12/16] Add comment --- imod/msw/sprinkling.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/imod/msw/sprinkling.py b/imod/msw/sprinkling.py index 516100b94..d4b6f6056 100644 --- a/imod/msw/sprinkling.py +++ b/imod/msw/sprinkling.py @@ -551,7 +551,7 @@ def _render(self, file, index, svat, mf6_dis, mf6_well): outside_df = _get_wells_outside_art_grid_dataframe( mf6_cellid_df, points_df, msw_mf6_sprinkling_df, svat_aligned ) - + # Prepare the final dataframe to be written to scap_svat.inp dataframe_out = pd.concat([inside_df, outside_df], axis=0, ignore_index=True) dataframe_out = dataframe_out.sort_values(by=["svat"]).reset_index(drop=True) From 157e91d98197524be3b6c7fd5782d696d69b4ca0 Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Wed, 12 Aug 2026 15:27:23 +0200 Subject: [PATCH 13/16] Add SprinklingPoints to msw namespace --- imod/msw/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 From ab4815d7dcdbe11ccaa2d01cd257e0dd657aa1a9 Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Wed, 12 Aug 2026 15:39:26 +0200 Subject: [PATCH 14/16] Put columns in right order --- imod/tests/fixtures/imod5_cap_data.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/imod/tests/fixtures/imod5_cap_data.py b/imod/tests/fixtures/imod5_cap_data.py index 88a2a5d7d..38c8e814d 100644 --- a/imod/tests/fixtures/imod5_cap_data.py +++ b/imod/tests/fixtures/imod5_cap_data.py @@ -146,11 +146,11 @@ 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) From e5af555565b155e1aed55029fd1fce002ac8eaed Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Wed, 12 Aug 2026 15:39:39 +0200 Subject: [PATCH 15/16] Start adding unittest --- imod/tests/test_msw/test_sprinkling.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/imod/tests/test_msw/test_sprinkling.py b/imod/tests/test_msw/test_sprinkling.py index 63812df9f..2350e9efc 100644 --- a/imod/tests/test_msw/test_sprinkling.py +++ b/imod/tests/test_msw/test_sprinkling.py @@ -400,3 +400,18 @@ 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'} + rate = 25.0 * 0.25 * 1.0e-3 + + # 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 + + From 01f4a5e535e4e587a23c7e1e592965054f0154cf Mon Sep 17 00:00:00 2001 From: JoerivanEngelen Date: Wed, 12 Aug 2026 17:08:25 +0200 Subject: [PATCH 16/16] Start refactoring test_module: Separate cases from tests --- imod/tests/test_msw/test_sprinkling.py | 317 ++++++++++--------------- 1 file changed, 127 insertions(+), 190 deletions(-) diff --git a/imod/tests/test_msw/test_sprinkling.py b/imod/tests/test_msw/test_sprinkling.py index 2350e9efc..f217038b8 100644 --- a/imod/tests/test_msw/test_sprinkling.py +++ b/imod/tests/test_msw/test_sprinkling.py @@ -1,54 +1,28 @@ 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.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] subunit = [0, 1] 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]], - [[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} - ) - - max_abstraction_surfacewater = xr.DataArray( - np.array( - [ - [[nan, 100.0, nan], - [nan, 200.0, nan], - [nan, 300.0, nan]], - [[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} - ) - svat = xr.DataArray( np.array( [ @@ -66,66 +40,60 @@ def test_simple_model_all_svats(fixed_format_parser): ) # 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, + return svat, index + + +@dataclass +class AbstractionCaseData: + max_abstraction_groundwater: Optional[xr.DataArray] = None + max_abstraction_surfacewater: Optional[xr.DataArray] = None + expected_abs_gw: Optional[np.ndarray] = None + expected_abs_sw: Optional[np.ndarray] = None + expected_layer: Optional[np.ndarray] = None + expected_svat_gw: Optional[np.ndarray] = None + + +class AbstractionGrids: + def case_all_svats(self, sprinkling_svat_index) -> AbstractionCaseData: + svat, _ = sprinkling_svat_index + case_data = AbstractionCaseData() + 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], + [nan, 300.0, nan]], + [[nan, 100.0, nan], + [nan, 200.0, nan], + [nan, 300.0, nan]] + ] ) - - 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], [nan, 300.0, nan]], - [[nan, nan, nan], + [[nan, 100.0, nan], [nan, 200.0, nan], - [nan, nan, 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( + ) + # fmt: on + case_data.expected_abs_gw = np.array([100.0, 300.0, 100.0, 200.0]) + case_data.expected_abs_sw = np.array([100.0, 300.0, 100.0, 200.0]) + case_data.expected_layer = np.array([3, 1, 3, 2]) + case_data.expected_svat_gw = np.array([1, 2, 3, 4]) + return case_data + + def case_some_svats(self, sprinkling_svat_index): + svat, _ = sprinkling_svat_index + case_data = AbstractionCaseData() + 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], @@ -134,73 +102,32 @@ 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( + ) + case_data.max_abstraction_surfacewater.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 = [1.0, 2.0, 3.0] - 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) - - coupler_mapping = msw.Sprinkling( - max_abstraction_groundwater, - max_abstraction_surfacewater, - ) - - with tempfile.TemporaryDirectory() as output_dir: - output_dir = Path(output_dir) - coupler_mapping.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_almost_equal( - results["max_abstraction_groundwater"], - np.array([100.0, 300.0, 200.0]), - ) - assert_almost_equal( - results["max_abstraction_surfacewater"], - np.array([100.0, 300.0, 200.0]), - ) - assert_equal(results["layer"], np.array([3, 1, 2])) - assert_equal(results["svat_groundwater"], np.array([1, 2, 4])) - - -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( + # fmt: on + case_data.expected_abs_gw = np.array([100.0, 300.0, 200.0]) + case_data.expected_abs_sw = np.array([100.0, 300.0, 200.0]) + case_data.expected_layer = np.array([3, 1, 2]) + case_data.expected_svat_gw = np.array([1, 2, 4]) + + return case_data + + def case_inconsistent_active_capacity(self, sprinkling_svat_index): + svat, _ = sprinkling_svat_index + case_data = AbstractionCaseData() + 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], @@ -209,13 +136,8 @@ def test_simple_model_inconsistent_active_capacity(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, 0.0, nan], [nan, 200.0, nan], @@ -224,62 +146,58 @@ def test_simple_model_inconsistent_active_capacity(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} - ) + ) + # fmt: on + case_data.expected_abs_gw = np.array([100.0, 0.0, 200.0]) + case_data.expected_abs_sw = np.array([0.0, 300.0, 200.0]) + case_data.expected_layer = np.array([3, 1, 2]) + case_data.expected_svat_gw = np.array([1, 2, 4]) - svat = xr.DataArray( - np.array( - [ - [[0, 1, 0], - [0, 0, 0], - [0, 2, 0]], + return case_data - [[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() + +@parametrize_with_cases("case_data", cases=AbstractionGrids) +def test_simple_model( + fixed_format_parser: Callable, + sprinkling_svat_index: tuple[xr.DataArray, np.ndarray], + case_data: AbstractionCaseData, +): + 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_id = ["a", "b", "c"] 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"], case_data.expected_svat_gw) assert_almost_equal( results["max_abstraction_groundwater"], - np.array([100.0, 0.0, 200.0]), + case_data.expected_abs_gw, ) assert_almost_equal( results["max_abstraction_surfacewater"], - np.array([0.0, 300.0, 200.0]), + case_data.expected_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"], case_data.expected_layer) + assert_equal(results["svat_groundwater"], case_data.expected_svat_gw) def test_simple_model_1_subunit(fixed_format_parser): @@ -332,8 +250,9 @@ def test_simple_model_1_subunit(fixed_format_parser): well_y = [1.0, 3.0] well_x = [2.0, 2.0] well_rate = [-5.0] * 2 + well_id = ["a", "c"] cellids = derive_cellid_from_points(svat, well_x, well_y, well_layer) - well = Mf6Wel(cellids, well_rate) + well = Mf6Wel(cellids, well_rate, well_id) sprinkling = msw.Sprinkling( max_abstraction_groundwater, @@ -401,11 +320,11 @@ def test_sprinkling_from_imod5_data__grid(cap_data_sprinkling_grid): 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'} - rate = 25.0 * 0.25 * 1.0e-3 + 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) @@ -413,5 +332,23 @@ def test_sprinklingpoints_from_imod5_data__points(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"].max(), 100.0) + + +# @pytest.mark.unittest_jit +# def test_sprinklingpoints_write__points(cap_data_sprinkling_points, tmp_path): +# well_x = cap_data_sprinkling_points["cap"]["x"].values +# well_y = cap_data_sprinkling_points["cap"]["y"].values +# well_layer = cap_data_sprinkling_points["cap"]["layer"].values +# +# # Arrange +# # cellids = derive_cellid_from_points(svat, well_x, well_y, well_layer) +# # well = Mf6Wel(cellids, well_rate) +# +# # Act +# sprinkling = msw.SprinklingPoints.from_imod5_data(cap_data_sprinkling_points) +# +# # sprinkling.write() +#