Issue #728 metaswap grid agnostic well

docs/api/changelog.rst

@@ -7,6 +7,17 @@ The format is based on `Keep a Changelog`_, and this project adheres to
 `Semantic Versioning`_.
 
 
+[Unreleased - feature branch]
+-----------------------------
+
+Changed
+~~~~~~~
+
+- :class:`imod.msw.CouplingMapping` and :class:`imod.msw.Sprinkling` now take
+  the :class:`imod.mf6.mf6_wel_adapter.Mf6Wel` as well argument instead of the
+  deprecated :class:`imod.mf6.WellDisStructured`.
+
+
 [Unreleased]
 ------------
 

imod/mf6/mf6_wel_adapter.py

@@ -17,14 +17,15 @@
 import numpy as np
 
 from imod.mf6.boundary_condition import BoundaryCondition
+from imod.mf6.interfaces.ipackage import IPackage
 from imod.schemata import DTypeSchema
 
 # FUTURE: There was an idea to autogenerate modflow 6 adapters.
 # This was relevant:
 # https://github.com/Deltares/xugrid/blob/main/xugrid/core/wrap.py#L90
 
 
-class Mf6Wel(BoundaryCondition):
+class Mf6Wel(BoundaryCondition, IPackage):
     """
     Package resembling input for Modflow 6 List Input. This class has
     methods for the modflow 6 wel packages with time component.
@@ -86,7 +87,7 @@ def _ds_to_arrdict(self, ds):
             dsvar[var] = ds[var]
         arrdict["var_values"] = dsvar
 
-        arrdict["cellid_names"] = ds.coords["nmax_cellid"].values
+        arrdict["cellid_names"] = ds.coords["dim_cellid"].values
         arrdict["nrow"] = ds.coords["ncellid"].size
         arrdict["cellid"] = ds["cellid"]
 
@@ -100,7 +101,7 @@ def _to_struct_array(self, arrdict, _):
         # Initialize the structured array
         recarr = np.empty(arrdict["nrow"], dtype=sparse_dtype)
         for cellid_name in arrdict["cellid_names"]:
-            recarr[cellid_name] = arrdict["cellid"].sel(nmax_cellid=cellid_name).values
+            recarr[cellid_name] = arrdict["cellid"].sel(dim_cellid=cellid_name).values
 
         for var in arrdict["data_vars"]:
             recarr[var] = arrdict["var_values"][var]

imod/mf6/wel.py

@@ -218,6 +218,78 @@ def get_all_imod5_prj_well_times(imod5_data: dict) -> list[datetime]:
     return sorted(set(wel_times_flat))
 
 
+def derive_cellid_from_points(
+    dst_grid: GridDataArray,
+    x: list,
+    y: list,
+    layer: list,
+) -> GridDataArray:
+    """
+    Create DataArray with Modflow6 cell identifiers based on x, y coordinates
+    in a dataframe. For structured grid this DataArray contains 3 columns:
+    ``layer, row, column``. For unstructured grids, this contains 2 columns:
+    ``layer, cell2d``.
+    See also: https://water.usgs.gov/water-resources/software/MODFLOW-6/mf6io_6.4.0.pdf#page=35
+
+    Note
+    ----
+    The "layer" coordinate should already be provided in the dataframe.
+    To determine the layer coordinate based on screen depts, look at
+    :func:`imod.prepare.wells.assign_wells`.
+
+    Parameters
+    ----------
+    dst_grid: {xr.DataArray, xu.UgridDataArray}
+        Destination grid to map the points to based on their x and y coordinates.
+    x: {list, np.array}
+        array-like with x-coordinates
+    y: {list, np.array}
+        array-like with y-coordinates
+    layer: {list, np.array}
+        array-like with layer-coordinates
+
+    Returns
+    -------
+    cellid : xr.DataArray
+        2D DataArray with a ``ncellid`` rows and 3 to 2 columns, depending
+        on whether on a structured or unstructured grid."""
+
+    # Find indices belonging to x, y coordinates
+    indices_cell2d = points_indices(dst_grid, out_of_bounds="ignore", x=x, y=y)
+    # Convert cell2d indices from 0-based to 1-based.
+    indices_cell2d = {dim: index + 1 for dim, index in indices_cell2d.items()}
+    # Prepare layer indices, for later concatenation
+
+    if isinstance(dst_grid, xu.UgridDataArray):
+        indices_layer = xr.DataArray(
+            layer, coords=indices_cell2d["mesh2d_nFaces"].coords
+        )
+        face_dim = dst_grid.ugrid.grid.face_dimension
+        indices_cell2d_dims = [face_dim]
+        cell2d_coords = ["cell2d"]
+    else:
+        indices_layer = xr.DataArray(layer, coords=indices_cell2d["x"].coords)
+        indices_cell2d_dims = ["y", "x"]
+        cell2d_coords = ["row", "column"]
+
+    # Prepare cellid array of the right shape.
+    cellid_ls = [indices_layer] + [indices_cell2d[dim] for dim in indices_cell2d_dims]
+    cellid = xr.concat(cellid_ls, dim="dim_cellid")
+    # Rename generic dimension name "index" to ncellid.
+    cellid = cellid.rename(index="ncellid")
+    # Put dimensions in right order after concatenation.
+    cellid = cellid.transpose("ncellid", "dim_cellid")
+    # Assign extra coordinate names.
+    coords = {
+        "dim_cellid": ["layer"] + cell2d_coords,
+        "x": ("ncellid", x),
+        "y": ("ncellid", y),
+    }
+    cellid = cellid.assign_coords(coords=coords)
+
+    return cellid
+
+
 class GridAgnosticWell(BoundaryCondition, IPointDataPackage, abc.ABC):
     """
     Abstract base class for grid agnostic wells
@@ -245,10 +317,10 @@ def _create_cellid(
 
         # Groupby index and select first, to unset any duplicate records
         # introduced by the multi-indexed "time" dimension.
-        df_for_cellid = assigned_wells.groupby("index").first()
-        d_for_cellid = df_for_cellid[["x", "y", "layer"]].to_dict("list")
+        unique_assigned_wells = assigned_wells.groupby("index").first()
+        d_for_cellid = unique_assigned_wells[["x", "y", "layer"]].to_dict("list")
 
-        return self._derive_cellid_from_points(like, **d_for_cellid)
+        return derive_cellid_from_points(like, **d_for_cellid)
 
     def _create_dataset_vars(
         self, assigned_wells: pd.DataFrame, cellid: xr.DataArray
@@ -275,80 +347,6 @@ def _create_dataset_vars(
 
         return ds_vars
 
-    @staticmethod
-    def _derive_cellid_from_points(
-        dst_grid: GridDataArray,
-        x: list,
-        y: list,
-        layer: list,
-    ) -> GridDataArray:
-        """
-        Create DataArray with Modflow6 cell identifiers based on x, y coordinates
-        in a dataframe. For structured grid this DataArray contains 3 columns:
-        ``layer, row, column``. For unstructured grids, this contains 2 columns:
-        ``layer, cell2d``.
-        See also: https://water.usgs.gov/water-resources/software/MODFLOW-6/mf6io_6.4.0.pdf#page=35
-
-        Note
-        ----
-        The "layer" coordinate should already be provided in the dataframe.
-        To determine the layer coordinate based on screen depts, look at
-        :func:`imod.prepare.wells.assign_wells`.
-
-        Parameters
-        ----------
-        dst_grid: {xr.DataArray, xu.UgridDataArray}
-            Destination grid to map the points to based on their x and y coordinates.
-        x: {list, np.array}
-            array-like with x-coordinates
-        y: {list, np.array}
-            array-like with y-coordinates
-        layer: {list, np.array}
-            array-like with layer-coordinates
-
-        Returns
-        -------
-        cellid : xr.DataArray
-            2D DataArray with a ``ncellid`` rows and 3 to 2 columns, depending
-            on whether on a structured or unstructured grid."""
-
-        # Find indices belonging to x, y coordinates
-        indices_cell2d = points_indices(dst_grid, out_of_bounds="ignore", x=x, y=y)
-        # Convert cell2d indices from 0-based to 1-based.
-        indices_cell2d = {dim: index + 1 for dim, index in indices_cell2d.items()}
-        # Prepare layer indices, for later concatenation
-
-        if isinstance(dst_grid, xu.UgridDataArray):
-            indices_layer = xr.DataArray(
-                layer, coords=indices_cell2d["mesh2d_nFaces"].coords
-            )
-            face_dim = dst_grid.ugrid.grid.face_dimension
-            indices_cell2d_dims = [face_dim]
-            cell2d_coords = ["cell2d"]
-        else:
-            indices_layer = xr.DataArray(layer, coords=indices_cell2d["x"].coords)
-            indices_cell2d_dims = ["y", "x"]
-            cell2d_coords = ["row", "column"]
-
-        # Prepare cellid array of the right shape.
-        cellid_ls = [indices_layer] + [
-            indices_cell2d[dim] for dim in indices_cell2d_dims
-        ]
-        cellid = xr.concat(cellid_ls, dim="nmax_cellid")
-        # Rename generic dimension name "index" to ncellid.
-        cellid = cellid.rename(index="ncellid")
-        # Put dimensions in right order after concatenation.
-        cellid = cellid.transpose("ncellid", "nmax_cellid")
-        # Assign extra coordinate names.
-        coords = {
-            "nmax_cellid": ["layer"] + cell2d_coords,
-            "x": ("ncellid", x),
-            "y": ("ncellid", y),
-        }
-        cellid = cellid.assign_coords(coords=coords)
-
-        return cellid
-
     def render(self, directory, pkgname, globaltimes, binary):
         raise NotImplementedError(
             textwrap.dedent(

imod/msw/coupler_mapping.py

@@ -1,13 +1,14 @@
-from typing import Optional
+from typing import Optional, cast
 
 import numpy as np
 import pandas as pd
 import xarray as xr
 
 from imod.mf6.dis import StructuredDiscretization
-from imod.mf6.wel import WellDisStructured
+from imod.mf6.mf6_wel_adapter import Mf6Wel
 from imod.msw.fixed_format import VariableMetaData
 from imod.msw.pkgbase import MetaSwapPackage
+from imod.typing import IntArray
 
 
 class CouplerMapping(MetaSwapPackage):
@@ -21,7 +22,7 @@ class CouplerMapping(MetaSwapPackage):
     ----------
     modflow_dis: StructuredDiscretization
         Modflow 6 structured discretization
-    well: WellDisStructured (optional)
+    well: Mf6Wel (optional)
         If given, this parameter describes sprinkling of SVAT units from MODFLOW
         cells.
     """
@@ -41,10 +42,12 @@ class CouplerMapping(MetaSwapPackage):
     def __init__(
         self,
         modflow_dis: StructuredDiscretization,
-        well: Optional[WellDisStructured] = None,
+        well: Optional[Mf6Wel] = None,
     ):
         super().__init__()
 
+        if well and (not isinstance(well, Mf6Wel)):
+            raise TypeError(rf"well not of type 'Mf6Wel', got '{type(well)}'")
         self.well = well
         # Test if equal or larger than 1, to ignore idomain == -1 as well. Don't
         # assign to self.dataset, as grid extent might differ from svat when
@@ -67,7 +70,7 @@ def _create_mod_id_rch(self, svat):
 
         self.dataset["mod_id"].values[:, idomain_top_active.values] = mod_id_1d
 
-    def _render(self, file, index, svat):
+    def _render(self, file, index, svat: xr.DataArray):
         self._create_mod_id_rch(svat)
         # package check only possible after calling _create_mod_id_rch
         self._pkgcheck()
@@ -97,33 +100,39 @@ def _render(self, file, index, svat):
 
         return self.write_dataframe_fixed_width(file, dataframe)
 
-    def _create_well_id(self, svat):
+    def _create_well_id(
+        self, svat: xr.DataArray
+    ) -> tuple[IntArray, IntArray, IntArray]:
         """
         Get modflow indices, svats, and layer number for the wells
         """
         n_subunit = svat["subunit"].size
 
-        # Convert to Python's 0-based index
-        well_row = self.well["row"] - 1
-        well_column = self.well["column"] - 1
-        well_layer = self.well["layer"] - 1
+        well = cast(Mf6Wel, self.well)
+        well_cellid = well.dataset["cellid"]
+        if len(well_cellid.coords["dim_cellid"]) != 3:
+            raise TypeError("Coupling to unstructured grids is not supported.")
+
+        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_mod = self.idomain_active.sum()
         mod_id = xr.full_like(self.idomain_active, 0, dtype=np.int64)
-        mod_id.values[self.idomain_active.values] = np.arange(1, n_mod + 1)
+        mod_id.data[self.idomain_active.data] = np.arange(1, n_mod + 1)
 
-        well_mod_id = mod_id[well_layer, well_row, well_column]
+        well_mod_id = mod_id.data[well_layer - 1, well_row, well_column]
         well_mod_id = np.tile(well_mod_id, (n_subunit, 1))
 
-        well_svat = svat.values[:, well_row, well_column]
+        well_svat = svat.data[:, well_row, well_column]
 
         well_active = well_svat != 0
 
         well_svat_1d = well_svat[well_active]
         well_mod_id_1d = well_mod_id[well_active]
 
         # Tile well_layers for each subunit
-        layer = np.tile(well_layer + 1, (n_subunit, 1))
+        layer = np.tile(well_layer, (n_subunit, 1))
         layer_1d = layer[well_active]
 
         return (well_mod_id_1d, well_svat_1d, layer_1d)

imod/msw/sprinkling.py

@@ -2,7 +2,7 @@
 import pandas as pd
 import xarray as xr
 
-from imod.mf6.wel import WellDisStructured
+from imod.mf6.mf6_wel_adapter import Mf6Wel
 from imod.msw.fixed_format import VariableMetaData
 from imod.msw.pkgbase import MetaSwapPackage
 
@@ -22,7 +22,7 @@ class Sprinkling(MetaSwapPackage):
     max_abstraction_surfacewater: array of floats (xr.DataArray)
         Describes the maximum abstraction of surfacewater to SVAT units in m3
         per day. This array must not have a subunit coordinate.
-    well: WellDisStructured
+    well: Mf6Wel
         Describes the sprinkling of SVAT units coming groundwater.
     """
 
@@ -55,23 +55,29 @@ def __init__(
         self,
         max_abstraction_groundwater: xr.DataArray,
         max_abstraction_surfacewater: xr.DataArray,
-        well: WellDisStructured,
+        well: Mf6Wel,
     ):
         super().__init__()
         self.dataset["max_abstraction_groundwater_m3_d"] = max_abstraction_groundwater
         self.dataset["max_abstraction_surfacewater_m3_d"] = max_abstraction_surfacewater
+        if not isinstance(well, Mf6Wel):
+            raise TypeError(rf"well not of type 'Mf6Wel', got '{type(well)}'")
         self.well = well
 
         self._pkgcheck()
 
     def _render(self, file, index, svat):
-        well_row = self.well["row"] - 1
-        well_column = self.well["column"] - 1
-        well_layer = self.well["layer"]
+        well_cellid = self.well["cellid"]
+        if len(well_cellid.coords["dim_cellid"]) != 3:
+            raise TypeError("Coupling to unstructured grids is not supported.")
+
+        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
 
-        well_svat = svat.values[:, well_row, well_column]
+        well_svat = svat.data[:, well_row, well_column]
         well_active = well_svat != 0
 
         # Tile well_layers for each subunit
@@ -80,7 +86,7 @@ def _render(self, file, index, svat):
         data_dict = {"svat": well_svat[well_active], "layer": layer[well_active]}
 
         for var in self._without_subunit:
-            well_arr = self.dataset[var].values[well_row, well_column]
+            well_arr = self.dataset[var].data[well_row, well_column]
             well_arr = np.tile(well_arr, (n_subunit, 1))
             data_dict[var] = well_arr[well_active]