Issue #1255 msw refactor mf6 obj at write

docs/api/changelog.rst

@@ -13,10 +13,21 @@ The format is based on `Keep a Changelog`_, and this project adheres to
 Changed
 ~~~~~~~
 
+- :class:`imod.msw.CouplingMapping`, :class:`imod.msw.Sprinkling`,
+  `imod.msw.Sprinkling.MetaSwapModel`, now take the
+  :class:`imod.mf6.mf6_wel_adapter.Mf6Wel` and the
+  :class:`imod.mf6.StructuredDiscretization` packages as arguments at their
+  respective ``write`` method, instead of upon initializing these MetaSWAP
+  objects.
 - :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`.
 
+Added
+~~~~~
+
+- :meth:`imod.msw.MetaSwapModel.regrid_like` to regrid MetaSWAP models.
+
 
 [Unreleased]
 ------------
@@ -97,9 +108,7 @@ Added
   :meth:`imod.mf6.GeneralHeadboundary.cleanup`, :meth:`imod.mf6.River.cleanup`,
   :meth:`imod.mf6.Well.cleanup` convenience methods to call the corresponding
   cleanup utility functions with the appropriate arguments.
-- :meth:`imod.msw.MetaSwapModel.regrid_like` to regrid MetaSWAP models. This is
-  still experimental functionality, regridding the :class:`imod.msw.Sprinkling`
-  is not yet supported.
+
 
 Removed
 ~~~~~~~

examples/metaswap/read_metaswap_file.py

@@ -75,7 +75,7 @@
 directory = imod.util.temporary_directory()
 os.makedirs(directory)
 
-landuse_options.write(directory, None, None)
+landuse_options.write(directory, None, None, None, None)
 
 # %%
 # Reading the .inp file

imod/mf6/wel.py

@@ -283,7 +283,7 @@ def derive_cellid_from_points(
     }
     cellid = cellid.assign_coords(coords=xy_coords)
 
-    return cellid
+    return cellid.astype(int)
 
 
 class GridAgnosticWell(BoundaryCondition, IPointDataPackage, abc.ABC):

imod/msw/coupler_mapping.py

@@ -1,4 +1,4 @@
-from typing import Optional, cast
+from typing import Any, TextIO
 
 import numpy as np
 import pandas as pd
@@ -7,7 +7,7 @@
 from imod.mf6.dis import StructuredDiscretization
 from imod.mf6.mf6_wel_adapter import Mf6Wel
 from imod.msw.fixed_format import VariableMetaData
-from imod.msw.pkgbase import MetaSwapPackage
+from imod.msw.pkgbase import DataDictType, MetaSwapPackage
 from imod.typing import IntArray
 
 
@@ -18,13 +18,6 @@ class CouplerMapping(MetaSwapPackage):
     This class is responsible for the file `mod2svat.inp`. It also includes
     connection to wells.
 
-    Parameters
-    ----------
-    modflow_dis: StructuredDiscretization
-        Modflow 6 structured discretization
-    well: Mf6Wel (optional)
-        If given, this parameter describes sprinkling of SVAT units from MODFLOW
-        cells.
     """
 
     _file_name = "mod2svat.inp"
@@ -35,59 +28,62 @@ class CouplerMapping(MetaSwapPackage):
         "layer": VariableMetaData(5, 0, 9999, int),
     }
 
-    _with_subunit = ("mod_id",)
+    _with_subunit = ()
     _without_subunit = ()
     _to_fill = ("free",)
 
     def __init__(
         self,
-        modflow_dis: StructuredDiscretization,
-        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
-        # MetaSWAP only covers part of the Modflow grid domain.
-        self.idomain_active = modflow_dis["idomain"] >= 1
-
-    def _create_mod_id_rch(self, svat):
+    def _create_mod_id_rch(
+        self, svat: xr.DataArray, idomain_top_active: xr.DataArray
+    ) -> xr.DataArray:
         """
         Create modflow indices for the recharge layer, which is where
         infiltration will take place.
         """
-        self.dataset["mod_id"] = xr.full_like(svat, fill_value=0, dtype=np.int64)
+        mod_id = xr.full_like(svat, fill_value=0, dtype=np.int64)
         n_subunit = svat["subunit"].size
-        idomain_top_active = self.idomain_active.sel(layer=1, drop=True)
 
         n_mod_top = idomain_top_active.sum()
 
         # idomain does not have a subunit dimension, so tile for n_subunits
-        mod_id_1d = np.tile(np.arange(1, n_mod_top + 1), (n_subunit, 1))
+        mod_id_1d: IntArray = np.tile(np.arange(1, n_mod_top + 1), (n_subunit, 1))
 
-        self.dataset["mod_id"].values[:, idomain_top_active.values] = mod_id_1d
+        mod_id.data[:, idomain_top_active.data] = mod_id_1d
+        return mod_id
 
-    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()
+    def _render(
+        self,
+        file: TextIO,
+        index: IntArray,
+        svat: xr.DataArray,
+        mf6_dis: StructuredDiscretization,
+        mf6_well: Mf6Wel,
+        *args: Any,
+    ):
+        if mf6_well and (not isinstance(mf6_well, Mf6Wel)):
+            raise TypeError(rf"mf6_well not of type 'Mf6Wel', got '{type(mf6_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
+        # MetaSWAP only covers part of the Modflow grid domain.
+        idomain_active = mf6_dis["idomain"] >= 1
+        idomain_top_active = idomain_active.sel(layer=1, drop=True)
+        mod_id = self._create_mod_id_rch(svat, idomain_top_active)
 
-        data_dict = {"svat": svat.values.ravel()[index]}
+        self._pkgcheck()
 
+        data_dict: DataDictType = {"svat": svat.values.ravel()[index]}
         data_dict["layer"] = np.full_like(data_dict["svat"], 1)
-
-        for var in self._with_subunit:
-            data_dict[var] = self._index_da(self.dataset[var], index)
+        data_dict["mod_id"] = self._index_da(mod_id, index)
 
         # Get well values
-        if self.well:
-            mod_id_well, svat_well, layer_well = self._create_well_id(svat)
-            data_dict["mod_id"] = np.append(mod_id_well, data_dict["mod_id"])
-            data_dict["svat"] = np.append(svat_well, data_dict["svat"])
-            data_dict["layer"] = np.append(layer_well, data_dict["layer"])
+        if mf6_well:
+            well_data_dict = self._create_well_id(svat, idomain_active, mf6_well)
+            for key in data_dict.keys():
+                data_dict[key] = np.append(well_data_dict[key], data_dict[key])
 
         for var in self._to_fill:
             data_dict[var] = ""
@@ -101,25 +97,27 @@ def _render(self, file, index, svat: xr.DataArray):
         return self.write_dataframe_fixed_width(file, dataframe)
 
     def _create_well_id(
-        self, svat: xr.DataArray
-    ) -> tuple[IntArray, IntArray, IntArray]:
+        self,
+        svat: xr.DataArray,
+        idomain_active: xr.DataArray,
+        mf6_well: Mf6Wel,
+    ) -> DataDictType:
         """
         Get modflow indices, svats, and layer number for the wells
         """
         n_subunit = svat["subunit"].size
 
-        well = cast(Mf6Wel, self.well)
-        well_cellid = well.dataset["cellid"]
+        well_cellid = mf6_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.data[self.idomain_active.data] = np.arange(1, n_mod + 1)
+        n_mod = idomain_active.sum()
+        mod_id = xr.full_like(idomain_active, 0, dtype=np.int64)
+        mod_id.data[idomain_active.data] = np.arange(1, n_mod + 1)
 
         well_mod_id = mod_id.data[well_layer - 1, well_row, well_column]
         well_mod_id = np.tile(well_mod_id, (n_subunit, 1))
@@ -135,7 +133,7 @@ def _create_well_id(
         layer = np.tile(well_layer, (n_subunit, 1))
         layer_1d = layer[well_active]
 
-        return (well_mod_id_1d, well_svat_1d, layer_1d)
+        return {"mod_id": well_mod_id_1d, "svat": well_svat_1d, "layer": layer_1d}
 
     def is_regridding_supported(self) -> bool:
         return False

imod/msw/initial_conditions.py

@@ -1,5 +1,6 @@
 import pathlib
 import shutil
+from typing import Any, TextIO
 
 from imod.mf6.interfaces.iregridpackage import IRegridPackage
 from imod.msw.fixed_format import VariableMetaData
@@ -23,7 +24,7 @@ class InitialConditionsEquilibrium(MetaSwapPackage, IRegridPackage):
     def __init__(self):
         super().__init__()
 
-    def _render(self, file, *args):
+    def _render(self, file: TextIO, *args: Any):
         file.write(self._option + "\n")
 
 
@@ -50,7 +51,7 @@ def __init__(self, initial_pF=2.2):
         super().__init__()
         self.dataset["initial_pF"] = initial_pF
 
-    def _render(self, file, *args):
+    def _render(self, file: TextIO, *args: Any):
         file.write(self._option + "\n")
 
         dataframe = self.dataset.assign_coords(index=[0]).to_dataframe()
@@ -78,7 +79,7 @@ class InitialConditionsPercolation(MetaSwapPackage, IRegridPackage):
     def __init__(self):
         super().__init__()
 
-    def _render(self, file, *args):
+    def _render(self, file: TextIO, *args: Any):
         file.write(self._option + "\n")
 
 

imod/msw/landuse.py

@@ -1,3 +1,5 @@
+from typing import Any, TextIO
+
 from imod.mf6.interfaces.iregridpackage import IRegridPackage
 from imod.msw.fixed_format import VariableMetaData
 from imod.msw.pkgbase import MetaSwapPackage
@@ -193,7 +195,7 @@ def __init__(
 
         self._pkgcheck()
 
-    def _render(self, file, *args):
+    def _render(self, file: TextIO, *args: Any):
         dataframe = self.dataset.to_dataframe(
             dim_order=("landuse_index",)
         ).reset_index()

imod/msw/meteo_mapping.py

@@ -1,5 +1,5 @@
 from copy import deepcopy
-from typing import Optional
+from typing import Any, Optional, TextIO
 
 import numpy as np
 import pandas as pd
@@ -9,7 +9,7 @@
 from imod.msw.fixed_format import VariableMetaData
 from imod.msw.pkgbase import MetaSwapPackage
 from imod.prepare import common
-from imod.typing import GridDataArray
+from imod.typing import GridDataArray, IntArray
 from imod.util.regrid_method_type import RegridMethodType
 
 
@@ -21,10 +21,18 @@ class MeteoMapping(MetaSwapPackage):
     new packages.
     """
 
+    meteo: GridDataArray
+
     def __init__(self):
         super().__init__()
 
-    def _render(self, file, index, svat):
+    def _render(
+        self,
+        file: TextIO,
+        index: IntArray,
+        svat: xr.DataArray,
+        *args: Any,
+    ):
         data_dict = {"svat": svat.values.ravel()[index]}
 
         row, column = self.grid_mapping(svat, self.meteo)

imod/msw/model.py

@@ -6,7 +6,8 @@
 import jinja2
 import numpy as np
 
-from imod import mf6
+from imod.mf6.dis import StructuredDiscretization
+from imod.mf6.mf6_wel_adapter import Mf6Wel
 from imod.mf6.utilities.regrid import RegridderWeightsCache
 from imod.msw.coupler_mapping import CouplerMapping
 from imod.msw.grid_data import GridData
@@ -106,7 +107,7 @@ def __init__(self, unsaturated_database):
             self._render_unsaturated_database_path(unsaturated_database)
         )
 
-    def _render_unsaturated_database_path(self, unsaturated_database):
+    def _render_unsaturated_database_path(self, unsaturated_database: Union[str, Path]):
         # Force to Path object
         unsaturated_database = Path(unsaturated_database)
 
@@ -204,7 +205,12 @@ def _get_pkg_key(self, pkg_type: type, optional_package: bool = False):
         if not optional_package:
             raise KeyError(f"Could not find package of type: {pkg_type}")
 
-    def write(self, directory: Union[str, Path]):
+    def write(
+        self,
+        directory: Union[str, Path],
+        mf6_dis: StructuredDiscretization,
+        mf6_wel: Mf6Wel,
+    ):
         """
         Write packages and simulation settings (para_sim.inp).
 
@@ -244,11 +250,11 @@ def write(self, directory: Union[str, Path]):
 
         # write package contents
         for pkgname in self:
-            self[pkgname].write(directory, index, svat)
+            self[pkgname].write(directory, index, svat, mf6_dis, mf6_wel)
 
     def regrid_like(
         self,
-        mf6_regridded_dis: mf6.StructuredDiscretization,
+        mf6_regridded_dis: StructuredDiscretization,
         regrid_context: Optional[RegridderWeightsCache] = None,
         regridder_types: Optional[dict[str, Tuple[RegridderType, str]]] = None,
     ) -> "MetaSwapModel":
@@ -273,6 +279,6 @@ def regrid_like(
                 raise ValueError(f"package {pkgname} cannot be  regridded")
             regridded_model[pkgname] = regridded_package
         if mod2svat_name is not None:
-            regridded_model[mod2svat_name] = CouplerMapping(mf6_regridded_dis)
+            regridded_model[mod2svat_name] = CouplerMapping()
 
         return regridded_model

imod/msw/output_control.py

@@ -1,5 +1,5 @@
 from copy import deepcopy
-from typing import Optional
+from typing import Any, Optional, TextIO
 
 import pandas as pd
 
@@ -87,7 +87,7 @@ def __init__(
         for key, value in kwargs.items():
             self.dataset[key] = int(value)
 
-    def _render(self, file, *args):
+    def _render(self, file: TextIO, *args: Any):
         variable, option = zip(
             *[(var, self.dataset[var].values) for var in self.dataset.data_vars]
         )
@@ -125,7 +125,7 @@ def __init__(self, time):
 
         self.dataset["times"] = time
 
-    def _render(self, file, *args):
+    def _render(self, file: TextIO, *args: Any):
         year, time_since_start_year = to_metaswap_timeformat(self.dataset["times"])
 
         dataframe = pd.DataFrame(