Issue #342 imod refactor primod update (#352)

Make iMOD Coupler compatible with iMOD Python 0.18.0:

* Use Mf6Wel object instead of WellDisStructuredDis

* Provide WEL at write instead of at init.

* Derive Mf6Wel package from coupler dict

* Drop unused "layer" coord that was causing trouble

* Set fill value to -1 for rasterio >=1.3.10

* Set dtype explicitly after imod-python update. Otherwise defaults to using the dtype of like, which is an integer. This causes problem for the fill value "nan" which the ribamod coupling depends on.

* Use special method to get diskey, for future support of unstructured grids

* Make sprinkling optional and improve function name.

* Add unittest for get_mf6_pkgs_with_coupling_dict

* Update remove sprinkling function, as initialization of Coupling object is not necessary anymore.

* Add pandera dependency

* Use python-build instead of build, as build-feedstock has been archived

* Force pandas > 2.0, otherwise ribasim-test-models doesn't work

* fix mypy error: unused type: ignore comment

* Cast to geodataframe to work around mypy error.

pixi.toml

@@ -20,7 +20,7 @@ build-imod-coupler = "rm -rf dist && pyinstaller imod_coupler/__main__.py --name
 netCDF4 = "*"
 loguru = "*"
 numpy = "<2.0"
-imod = "<0.18.0"
+imod = ">=0.18.0"
 pip = "*"
 pydantic = "2.*"
 pyinstaller = "*"
@@ -58,10 +58,12 @@ test-primod = "pytest --junitxml=report.xml tests/test_primod"
 
 
 [feature.common.dependencies]
-build = "*"
+python-build = "*"
+pandera = "*"
+pandas = ">=2"
 geopandas = "*"
 httpx = "*"
-imod = "<0.18.0"
+imod = ">=0.18.0"
 ipython = "*"
 jupyterlab = "*"
 mypy = "*"

pre-processing/primod/driver_coupling/metamod.py

@@ -68,7 +68,7 @@ def derive_mapping(
             pkgname for pkgname, pkg in msw_model.items() if isinstance(pkg, GridData)
         ][0]
 
-        dis = gwf_model[gwf_model._get_pkgkey("dis")]
+        dis = gwf_model[gwf_model._get_diskey()]
 
         index, svat = msw_model[grid_data_key].generate_index_array()
         grid_mapping = NodeSvatMapping(svat=svat, modflow_dis=dis, index=index)
@@ -78,7 +78,7 @@ def derive_mapping(
         rch_mapping = RechargeSvatMapping(svat, recharge, index=index)
 
         if self._check_sprinkling(msw_model=msw_model, gwf_model=gwf_model):
-            well = gwf_model[self.mf6_wel_package]
+            well = gwf_model.prepare_wel_for_mf6(self.mf6_wel_package, True, True)
             well_mapping = WellSvatMapping(svat, well, index=index)
             return grid_mapping, rch_mapping, well_mapping
         else:

pre-processing/primod/driver_coupling/ribameta.py

@@ -66,6 +66,7 @@ def derive_mapping(
                 self.ribasim_basin_definition,
                 like=svat.isel(subunit=0, drop=True),
                 column="node_id",
+                dtype=np.float64,
             )
         elif isinstance(self.ribasim_basin_definition, xr.DataArray):
             gridded_basin = self.ribasim_basin_definition
@@ -91,6 +92,7 @@ def derive_mapping(
                     self.ribasim_user_demand_definition,
                     like=svat.isel(subunit=0, drop=True),
                     column="node_id",
+                    dtype=np.float64,
                 )
             elif isinstance(self.ribasim_user_demand_definition, xr.DataArray):
                 gridded_user_demand = self.ribasim_user_demand_definition
@@ -104,7 +106,7 @@ def derive_mapping(
             nsu = swspr_grid_data.dataset["area"].sizes["subunit"]
             swsprmax = msw_model["sprinkling"]
             swspr_grid_data.dataset["area"].values = np.tile(
-                swsprmax["max_abstraction_surfacewater_m3_d"].values,
+                swsprmax["max_abstraction_surfacewater"].values,
                 (nsu, 1, 1),
             )
             index_swspr, svat_swspr = swspr_grid_data.generate_index_array()

pre-processing/primod/driver_coupling/ribamod.py

@@ -1,6 +1,6 @@
 import abc
 from pathlib import Path
-from typing import Any
+from typing import Any, cast
 
 import imod
 import numpy as np
@@ -69,7 +69,7 @@ def write_exchanges(
         gwf_model = mf6_simulation[self.mf6_model]
         ribasim_model = coupled_model.ribasim_model
 
-        dis = gwf_model[gwf_model._get_pkgkey("dis")]
+        dis = gwf_model[gwf_model._get_diskey()]
         basin_definition = self.ribasim_basin_definition
         if isinstance(basin_definition, gpd.GeoDataFrame):
             # Validate and fetch
@@ -80,6 +80,7 @@ def write_exchanges(
                 basin_definition,
                 like=dis["idomain"].isel(layer=0, drop=True),
                 column="node_id",
+                dtype=np.float64,
             )
             basin_dataset = xr.Dataset(
                 {key: gridded_basin for key in self.mf6_packages}
@@ -161,11 +162,11 @@ def _write_exchange(
                 )
             # in active coupling, check subgrid levels versus modflow bottom elevation
             if isinstance(self, RibaModActiveDriverCoupling):
+                # Cast to geodataframe for mypy
+                subgrid_df = cast(gpd.GeoDataFrame, ribasim_model.basin.subgrid.df)
                 #  check on the bottom elevation and ribasim minimal subgrid level
                 minimum_subgrid_level = (
-                    ribasim_model.basin.subgrid.df.groupby("subgrid_id")  # type: ignore
-                    .min()["subgrid_level"]
-                    .to_numpy()
+                    subgrid_df.groupby("subgrid_id").min()["subgrid_level"].to_numpy()
                 )
                 subgrid_index = mapping.dataframe["subgrid_index"]
                 bound_index = mapping.dataframe["bound_index"]

pre-processing/primod/mapping/wel_svat_mapping.py

@@ -4,7 +4,7 @@
 import numpy as np
 import pandas as pd
 import xarray as xr
-from imod import mf6
+from imod.mf6.mf6_wel_adapter import Mf6Wel
 from imod.msw.fixed_format import VariableMetaData
 from numpy.typing import NDArray
 
@@ -38,9 +38,7 @@ class WellSvatMapping(MetaModMapping):
     _with_subunit = ("wel_id", "svat", "layer")
     _to_fill = ("free",)
 
-    def __init__(
-        self, svat: xr.DataArray, well: mf6.WellDisStructured, index: NDArray[Int]
-    ):
+    def __init__(self, svat: xr.DataArray, well: Mf6Wel, index: NDArray[Int]):
         super().__init__()
         self.index = index
         self.well = well
@@ -55,10 +53,14 @@ def _create_well_id(
         """
         Get modflow indices, svats, and layer number for the wells
         """
+        well_cellid = self.well["cellid"]
+        if len(well_cellid.coords["dim_cellid"]) != 3:
+            raise TypeError("Coupling to unstructured grids is not supported.")
+
         # Convert to Python's 0-based index
-        well_row = self.well["row"] - 1
-        well_column = self.well["column"] - 1
-        well_layer = self.well["layer"]
+        well_layer = well_cellid.sel(dim_cellid="layer").data
+        well_row = well_cellid.sel(dim_cellid="row").data - 1
+        well_column = well_cellid.sel(dim_cellid="column").data - 1
 
         n_subunit = svat["subunit"].size
 
@@ -71,7 +73,7 @@ def _create_well_id(
         layer = np.tile(well_layer, (n_subunit, 1))
         layer_1d = layer[well_active]
 
-        well_id = self.well.dataset.coords["index"] + 1
+        well_id = well_cellid.coords["ncellid"] + 1
         well_id_1d = np.tile(well_id, (n_subunit, 1))[well_active]
 
         return (well_id_1d, well_svat_1d, layer_1d)

pre-processing/primod/metamod.py

@@ -8,9 +8,10 @@
 
 from primod.coupled_model import CoupledModel
 from primod.driver_coupling.metamod import MetaModDriverCoupling
+from primod.model_mixin import MetaModMixin
 
 
-class MetaMod(CoupledModel):
+class MetaMod(CoupledModel, MetaModMixin):
     """Couple MetaSWAP and MODFLOW 6.
 
     Parameters
@@ -19,6 +20,8 @@ class MetaMod(CoupledModel):
         The MetaSWAP model that should be coupled.
     mf6_simulation : Modflow6Simulation
         The Modflow6 simulation that should be coupled.
+    coupling_list: list of DriverCoupling
+        One entry per MODFLOW 6 model that should be coupled
     """
 
     _toml_name = "imod_coupler.toml"
@@ -96,7 +99,16 @@ def write(
             directory / self._modflow6_model_dir,
             **modflow6_write_kwargs,
         )
-        self.msw_model.write(directory / self._metaswap_model_dir)
+
+        mf6_dis_pkg, mf6_wel_pkg = self.get_mf6_pkgs_for_metaswap(
+            coupling_dict, self.mf6_simulation
+        )
+
+        self.msw_model.write(
+            directory / self._metaswap_model_dir,
+            mf6_dis_pkg,
+            mf6_wel_pkg,
+        )
 
     def write_toml(
         self,

pre-processing/primod/model_mixin.py

@@ -0,0 +1,27 @@
+"""Module containing mixins for specific kernels, for example MODFLOW."""
+
+from imod.mf6 import Modflow6Simulation, StructuredDiscretization
+from imod.mf6.mf6_wel_adapter import Mf6Wel
+
+
+class MetaModMixin:
+    """MetaSWAP-Modflow coupling specific methods."""
+
+    @staticmethod
+    def get_mf6_pkgs_for_metaswap(
+        coupling_dict: dict[str, str], mf6_simulation: Modflow6Simulation
+    ) -> tuple[StructuredDiscretization, Mf6Wel | None]:
+        """
+        Get names of DIS and possibly WEL packages from coupling_dict then fetch
+        these MODFLOW 6 packages from simulation.
+        """
+        mf6_model_key = coupling_dict["mf6_model"]
+        gwf_model = mf6_simulation[mf6_model_key]
+        mf6_dis_key = gwf_model._get_diskey()
+        mf6_dis_pkg = gwf_model[mf6_dis_key]
+
+        mf6_wel_pkg = None
+        if "mf6_msw_well_pkg" in coupling_dict.keys():
+            mf6_well_key = coupling_dict["mf6_msw_well_pkg"]
+            mf6_wel_pkg = gwf_model.prepare_wel_for_mf6(mf6_well_key, True, True)
+        return mf6_dis_pkg, mf6_wel_pkg

pre-processing/primod/ribametamod.py

@@ -9,9 +9,10 @@
 
 from primod.coupled_model import CoupledModel
 from primod.driver_coupling.driver_coupling_base import DriverCoupling
+from primod.model_mixin import MetaModMixin
 
 
-class RibaMetaMod(CoupledModel):
+class RibaMetaMod(CoupledModel, MetaModMixin):
     """Couple Ribasim, MetaSWAP and MODFLOW 6.
     Parameters
     ----------
@@ -112,7 +113,12 @@ def write(
             directory / self._modflow6_model_dir,
             **modflow6_write_kwargs,
         )
-        self.msw_model.write(directory / self._metaswap_model_dir)
+        mf6_dis_pkg, mf6_wel_pkg = self.get_mf6_pkgs_for_metaswap(
+            coupling_dict, self.mf6_simulation
+        )
+        self.msw_model.write(
+            directory / self._metaswap_model_dir, mf6_dis_pkg, mf6_wel_pkg
+        )
         self.ribasim_model.write(directory / self._ribasim_model_dir / "ribasim.toml")
 
     def write_toml(

tests/fixtures/common.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pandas as pd
 import xarray as xr
-from imod import mf6
+from imod.mf6.wel import LayeredWell
 from numpy import float64, int_
 from numpy.typing import NDArray
 
@@ -32,9 +32,7 @@ def get_times() -> pd.DatetimeIndex:
     return pd.date_range(start="1/1/1971", end="8/1/1971", freq=freq)
 
 
-def create_wells(
-    nrow: int, ncol: int, idomain: xr.DataArray, wel_layer: int | None = None
-) -> mf6.WellDisStructured:
+def create_wells(idomain: xr.DataArray, wel_layer: int | None = None) -> LayeredWell:
     """
     Create wells, deactivate inactive cells. This function wouldn't be necessary
     if iMOD Python had a package to specify wells based on grids.
@@ -43,11 +41,16 @@ def create_wells(
     if wel_layer is None:
         wel_layer = 3
 
+    x = idomain.coords["x"].to_numpy()
+    y = idomain.coords["y"].to_numpy()
+
+    x_grid, y_grid = np.meshgrid(x, y)
+
     is_inactive = ~idomain.sel(layer=wel_layer).astype(bool)
     id_inactive = np.argwhere(is_inactive.values) + 1
 
-    ix = np.tile(np.arange(ncol) + 1, nrow)
-    iy = np.repeat(np.arange(nrow) + 1, ncol)
+    ix = np.ravel(x_grid)
+    iy = np.ravel(y_grid)
 
     to_deactivate = np.full_like(ix, False, dtype=bool)
     for i in id_inactive:
@@ -58,19 +61,15 @@ def create_wells(
     iy_active = iy[~to_deactivate]
 
     rate = np.zeros(ix_active.shape)
-    layer = np.full_like(ix_active, wel_layer)
+    layer = np.full_like(ix_active, wel_layer, dtype=int)
 
-    return mf6.WellDisStructured(
-        layer=layer, row=iy_active, column=ix_active, rate=rate
-    )
+    return LayeredWell(ix_active, iy_active, layer, rate)
 
 
-def create_wells_max_layer(
-    nrow: int, ncol: int, idomain: xr.DataArray
-) -> mf6.WellDisStructured:
+def create_wells_max_layer(idomain: xr.DataArray) -> LayeredWell:
     """
     Create wells in deepest layer of MODFLOW 6 model
     """
 
     wel_layer = idomain.layer.max().item()
-    return create_wells(nrow, ncol, idomain, wel_layer)
+    return create_wells(idomain, wel_layer)

tests/fixtures/fixture_metaswap.py

@@ -7,14 +7,13 @@
 from imod import mf6, msw
 from numpy import nan
 
-from .common import create_wells, get_times, grid_sizes
+from .common import get_times, grid_sizes
 
 
 def metaswap_model(
     times: list[pd.date_range],
     area: xr.DataArray,
     active: xr.DataArray,
-    well: mf6.WellDisStructured,
     dis: mf6.StructuredDiscretization,
     unsaturated_database: str,
 ) -> msw.MetaSwapModel:
@@ -85,7 +84,6 @@ def metaswap_model(
         max_abstraction_surfacewater=xr.full_like(
             msw_grid, 0.02 * (20 * 20)
         ),  # 20 mm/d
-        well=well,
     )
 
     # Ponding
@@ -149,7 +147,7 @@ def metaswap_model(
     )
 
     # Metaswap Mappings
-    msw_model["mod2svat"] = msw.CouplerMapping(modflow_dis=dis, well=well)
+    msw_model["mod2svat"] = msw.CouplerMapping()
 
     # Output Control
     msw_model["oc_idf"] = msw.IdfMapping(area, -9999.0)
@@ -166,9 +164,6 @@ def make_msw_model(idomain: xr.DataArray) -> msw.MetaSwapModel:
     x, y, layer, dx, dy, dz = grid_sizes()
     subunit = [0, 1]
 
-    nrow = len(y)
-    ncol = len(x)
-
     top = 0.0
     bottom = top - xr.DataArray(np.cumsum(dz), coords={"layer": layer}, dims="layer")
 
@@ -206,12 +201,9 @@ def make_msw_model(idomain: xr.DataArray) -> msw.MetaSwapModel:
     area = area.where(modflow_active)
     active = active & modflow_active
 
-    # Well
-    well = create_wells(nrow, ncol, idomain)
-
     dis = mf6.StructuredDiscretization(idomain=idomain, top=top, bottom=bottom)
 
-    return metaswap_model(times,area,active, well, dis,unsaturated_database)
+    return metaswap_model(times,area,active,dis,unsaturated_database)
 
 
 @pytest_cases.fixture(scope="function")

tests/fixtures/fixture_modflow.py

@@ -84,7 +84,7 @@ def make_coupled_mf6_model(idomain: xr.DataArray) -> mf6.Modflow6Simulation:
         head, print_input=True, print_flows=True, save_flows=True
     )
     gwf_model["rch_msw"] = make_recharge_pkg(idomain)
-    gwf_model["wells_msw"] = create_wells(nrow, ncol, idomain)
+    gwf_model["wells_msw"] = create_wells(idomain)
 
     simulation = make_mf6_simulation(gwf_model)
     return simulation