Issue #344 part1 imod refactor primod update (#348)

* Temporarily install imod python from feature branch for feature development

* Use Mf6Wel object instead of WellDisStructuredDis

* 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.

pixi.toml

@@ -20,7 +20,6 @@ build-imod-coupler = "rm -rf dist && pyinstaller imod_coupler/__main__.py --name
 netCDF4 = "*"
 loguru = "*"
 numpy = "<2.0"
-imod = "*"
 pip = "*"
 pydantic = "2.*"
 pyinstaller = "*"
@@ -29,6 +28,13 @@ scipy = "*"
 tomli = "*"
 tomli-w = "*"
 xmipy = "*"
+filelock = "*"
+gdal = "3.9.1"
+geopandas = "*"
+pandamesh = "*"
+rasterio = "1.3.10"
+rioxarray = "*"
+xugrid = ">=0.11.0"
 
 #[pypi-dependencies]
 #ribasim = {git = "https://github.com/Deltares/Ribasim.git/#subdirectory=python/ribasim"}
@@ -39,10 +45,12 @@ install-ribasim-python = "pip install git+https://github.com/Deltares/Ribasim.gi
 install-ribasim-testmodels = "pip install git+https://github.com/Deltares/Ribasim.git/#subdirectory=python/ribasim_testmodels"
 install-primod = "pip install --no-deps --editable pre-processing"
 install-metaswap-testmodels = "svn checkout https://repos.deltares.nl/repos/DSCTestbench/trunk/cases/e150_metaswap/f00_common/c00_common/LHM2016_v01vrz .imod_collector/e150_metaswap"
+install-imod = "pip install git+https://github.com/Deltares/imod-python.git@msw_refactoring_feature_branch"
 install-imod-collector = "python scripts/download_imod_collector.py"
 install-imod-collector-regression = "python scripts/download_imod_collector.py regression"
 generate-env-file = "python scripts/generate_env_file.py"
 install = { depends_on = [
+    "install-imod",
     "install-minimal",
     "install-ribasim-testmodels",
     "install-primod",

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

pre-processing/primod/driver_coupling/ribamod.py

@@ -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}

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)

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.mf6_wel_adapter import Mf6Wel, cellid_from_arrays__structured
 from numpy import float64, int_
 from numpy.typing import NDArray
 
@@ -34,7 +34,7 @@ def get_times() -> pd.DatetimeIndex:
 
 def create_wells(
     nrow: int, ncol: int, idomain: xr.DataArray, wel_layer: int | None = None
-) -> mf6.WellDisStructured:
+) -> Mf6Wel:
     """
     Create wells, deactivate inactive cells. This function wouldn't be necessary
     if iMOD Python had a package to specify wells based on grids.
@@ -59,15 +59,17 @@ def create_wells(
 
     rate = np.zeros(ix_active.shape)
     layer = np.full_like(ix_active, wel_layer)
+    index = np.arange(ncol * nrow)[~to_deactivate]
 
-    return mf6.WellDisStructured(
-        layer=layer, row=iy_active, column=ix_active, rate=rate
+    cellid = cellid_from_arrays__structured(
+        layer=layer, row=iy_active, column=ix_active
     )
+    rate_da = xr.DataArray(rate, coords={"index": index}, dims=("index",))
 
+    return Mf6Wel(cellid=cellid, rate=rate_da)
 
-def create_wells_max_layer(
-    nrow: int, ncol: int, idomain: xr.DataArray
-) -> mf6.WellDisStructured:
+
+def create_wells_max_layer(nrow: int, ncol: int, idomain: xr.DataArray) -> Mf6Wel:
     """
     Create wells in deepest layer of MODFLOW 6 model
     """

tests/fixtures/fixture_metaswap.py

@@ -5,6 +5,7 @@
 import pytest_cases
 import xarray as xr
 from imod import mf6, msw
+from imod.mf6.mf6_wel_adapter import Mf6Wel
 from numpy import nan
 
 from .common import create_wells, get_times, grid_sizes
@@ -14,7 +15,7 @@ def metaswap_model(
     times: list[pd.date_range],
     area: xr.DataArray,
     active: xr.DataArray,
-    well: mf6.WellDisStructured,
+    well: Mf6Wel,
     dis: mf6.StructuredDiscretization,
     unsaturated_database: str,
 ) -> msw.MetaSwapModel:

tests/fixtures/fixture_ribametamod.py

@@ -17,7 +17,7 @@ def make_msw_model(
 
     idomain = gwf["GWF_1"]["dis"]["idomain"]
     if "layer" in idomain.dims:
-        idomain_layer1 = idomain.sel(layer=1)
+        idomain_layer1 = idomain.sel(layer=1, drop=True)
     if nsubunits is None:
         nsubunits = xr.ones_like(idomain)
 

tests/test_imod_coupler/test_metamod.py

@@ -209,6 +209,7 @@ def test_metamod_regression(
     subprocess.run(
         [imod_coupler_exec_regression, tmp_path_reg / metamod_model._toml_name],
         check=True,
+        capture_output=True,
     )
 
     # Read Modflow 6 output

tests/test_primod/test_wel_svat_mapping.py

@@ -3,7 +3,7 @@
 
 import numpy as np
 import xarray as xr
-from imod import mf6
+from imod.mf6.mf6_wel_adapter import Mf6Wel, cellid_from_arrays__structured
 from numpy.testing import assert_equal
 from primod.mapping.wel_svat_mapping import WellSvatMapping
 
@@ -34,14 +34,12 @@ def test_simple_model(fixed_format_parser):
     index = (svat != 0).to_numpy().ravel()
 
     # Well
-    well_layer = [3, 2, 1]
-    well_row = [1, 2, 3]
-    well_column = [2, 2, 2]
-    well_rate = [-5.0] * 3
-    well = mf6.WellDisStructured(
-        layer=well_layer,
-        row=well_row,
-        column=well_column,
+    cellid = cellid_from_arrays__structured(
+        layer=[3, 2, 1], row=[1, 2, 3], column=[2, 2, 2]
+    )
+    well_rate = xr.DataArray([-5.0] * 3, coords={"index": [0, 1, 2]}, dims=("index",))
+    well = Mf6Wel(
+        cellid=cellid,
         rate=well_rate,
     )
 
@@ -83,14 +81,10 @@ def test_simple_model_1_subunit(fixed_format_parser):
     index = (svat != 0).to_numpy().ravel()
 
     # Well
-    well_layer = [3, 2]
-    well_row = [1, 3]
-    well_column = [2, 2]
-    well_rate = [-5.0] * 2
-    well = mf6.WellDisStructured(
-        layer=well_layer,
-        row=well_row,
-        column=well_column,
+    cellid = cellid_from_arrays__structured(layer=[3, 2], row=[1, 3], column=[2, 2])
+    well_rate = xr.DataArray([-5.0] * 2, coords={"index": [0, 1]}, dims=("index",))
+    well = Mf6Wel(
+        cellid=cellid,
         rate=well_rate,
     )
 
@@ -141,14 +135,14 @@ def test_simple_model_inactive(fixed_format_parser):
     index = (svat != 0).to_numpy().ravel()
 
     # Well
-    well_layer = [1, 3, 2, 1]
-    well_row = [1, 1, 2, 3]
-    well_column = [1, 2, 2, 2]
-    well_rate = [-5.0] * 4
-    well = mf6.WellDisStructured(
-        layer=well_layer,
-        row=well_row,
-        column=well_column,
+    cellid = cellid_from_arrays__structured(
+        layer=[1, 3, 2, 1], row=[1, 1, 2, 3], column=[1, 2, 2, 2]
+    )
+    well_rate = xr.DataArray(
+        [-5.0] * 4, coords={"index": [0, 1, 2, 3]}, dims=("index",)
+    )
+    well = Mf6Wel(
+        cellid=cellid,
         rate=well_rate,
     )