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/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="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
+
+
 class GridAgnosticWell(BoundaryCondition, IPointDataPackage, abc.ABC):
     """
     Abstract base class for grid agnostic wells
@@ -248,7 +320,7 @@ def _create_cellid(
         df_for_cellid = assigned_wells.groupby("index").first()
         d_for_cellid = df_for_cellid[["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

@@ -5,7 +5,7 @@
 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
 
@@ -21,7 +21,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 +41,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
@@ -103,27 +105,30 @@ def _create_well_id(self, svat):
         """
         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
+        cellid = self.well["cellid"]
+        if len(cellid.coords["nmax_cellid"]) != 3:
+            raise TypeError("Coupling to unstructured grids is not supported.")
+
+        well_layer = cellid.sel(nmax_cellid="layer").data
+        well_row = cellid.sel(nmax_cellid="row").data - 1
+        well_column = cellid.sel(nmax_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"]
+        cellid = self.well["cellid"]
+        if len(cellid.coords["nmax_cellid"]) != 3:
+            raise TypeError("Coupling to unstructured grids is not supported.")
+
+        well_layer = cellid.sel(nmax_cellid="layer").data
+        well_row = cellid.sel(nmax_cellid="row").data - 1
+        well_column = cellid.sel(nmax_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]
 

imod/tests/fixtures/msw_model_fixture.py

@@ -5,6 +5,8 @@
 from numpy import nan
 
 from imod import mf6, msw
+from imod.mf6.mf6_wel_adapter import Mf6Wel
+from imod.mf6.wel import derive_cellid_from_points
 
 
 def make_coupled_mf6_model():
@@ -75,14 +77,13 @@ def make_coupled_mf6_model():
     # Create wells
     wel_layer = 3
 
-    ix = np.tile(np.arange(ncol) + 1, nrow)
-    iy = np.repeat(np.arange(nrow) + 1, ncol)
-    rate = np.zeros(ix.shape)
-    layer = np.full_like(ix, wel_layer)
+    well_x = np.tile(x, nrow)
+    well_y = np.repeat(y, ncol)
+    well_rate = np.zeros(well_x.shape)
+    well_layer = np.full_like(well_x, wel_layer)
 
-    gwf_model["wells_msw"] = mf6.WellDisStructured(
-        layer=layer, row=iy, column=ix, rate=rate
-    )
+    cellids = derive_cellid_from_points(idomain, well_x, well_y, well_layer)
+    gwf_model["well_msw"] = Mf6Wel(cellids, well_rate)
 
     # Attach it to a simulation
     simulation = mf6.Modflow6Simulation("test")
@@ -192,12 +193,13 @@ def make_msw_model():
 
     wel_layer = 3
 
-    ix = np.tile(np.arange(ncol) + 1, nrow)
-    iy = np.repeat(np.arange(nrow) + 1, ncol)
-    rate = np.zeros(ix.shape)
-    layer = np.full_like(ix, wel_layer)
+    well_x = np.tile(x, nrow)
+    well_y = np.repeat(y, ncol)
+    well_rate = np.zeros(well_x.shape)
+    well_layer = np.full_like(well_x, wel_layer)
 
-    well = mf6.WellDisStructured(layer=layer, row=iy, column=ix, rate=rate)
+    cellids = derive_cellid_from_points(active, well_x, well_y, well_layer).astype(int)
+    well = Mf6Wel(cellids, well_rate)
 
     # %% Modflow 6
     dz = np.array([1, 10, 100])

imod/tests/test_mf6/test_mf6_wel.py

@@ -19,7 +19,7 @@
 from imod.mf6.dis import StructuredDiscretization
 from imod.mf6.npf import NodePropertyFlow
 from imod.mf6.utilities.grid import broadcast_to_full_domain
-from imod.mf6.wel import LayeredWell, Well
+from imod.mf6.wel import LayeredWell, Well, derive_cellid_from_points
 from imod.mf6.write_context import WriteContext
 from imod.schemata import ValidationError
 from imod.tests.fixtures.flow_basic_fixture import BasicDisSettings
@@ -708,7 +708,7 @@ def test_derive_cellid_from_points(basic_dis, well_high_lvl_test_data_stationary
     )
 
     # Act
-    cellid = imod.mf6.wel.Well._derive_cellid_from_points(idomain, x, y, layer)
+    cellid = derive_cellid_from_points(idomain, x, y, layer)
 
     # Assert
     np.testing.assert_array_equal(cellid, cellid_expected)