Issue #1779 Implementation of MetaSWAP model splitter

docs/api/changelog.rst

@@ -17,6 +17,7 @@ Added
   :meth:`imod.mf6.Modflow6Simulation.clip_box` to consider a package empty if
   its first times step is all nodata. This can save a lot of clipping or masking
   transient models with many timesteps.
+- Added :meth:`imod.msw.MetaSwapModel.split` to split MetaSWAP models.
 
 Fixed
 ~~~~~
@@ -53,7 +54,7 @@ Changed
 - ``proportion_depth`` and ``proportion_rate`` in
   :class:`imod.mf6.Evapotranspiration` are now optional variables. If provided,
   now require ``"segment"`` dimension when ``proportion_depth`` and
-  ``proportion_rate``. 
+  ``proportion_rate``.
 
 
 [1.0.0] - 2025-11-11
@@ -132,7 +133,7 @@ Fixed
 - Fixed bug where :meth:`imod.mf6.Modflow6Simulation.split`,
   :meth:`imod.mf6.Modflow6Simulation.regrid_like`, and
   :meth:`imod.mf6.Modflow6Simulation.clip_box` would not copy
-  :class:`imod.mf6.ValidationSettings`. 
+  :class:`imod.mf6.ValidationSettings`.
 - ``landuse``, ``soil_physical_unit``, ``active`` for :class:`imod.msw.GridData`
   are now properly regridded with the ``mode`` statistic when using
   :meth:`imod.msw.GridData.regrid_like`.
@@ -189,7 +190,7 @@ Fixed
   ``states_for_boundary`` argument would place these bc at the
   incorrect places with unstructured grids.
 - Fixed bug where :meth:`imod.mf6.SourceSinkMixing.from_flow_model` would return
-  an error upon adding a package which cannot have a ``concentration``, such as 
+  an error upon adding a package which cannot have a ``concentration``, such as
   :class:`imod.mf6.HorizontalFlowBarrierResistance`.
 - Broken names for ``outer_csvfile`` and ``inner_csvfile`` in the
   :class:`imod.mf6.Solution` MODFLOW 6 template file.
@@ -343,8 +344,8 @@ Changed
   :meth:`imod.mf6.Recharge.from_imod5_data` got extra arguments for
   ``period_data``, ``time_min`` and ``time_max``.
 - :func:`imod.visualize.read_imod_legend` now also returns the labels as an extra
-  argument. Update your code by changing 
-  ``colors, levels = read_imod_legend(...)`` to 
+  argument. Update your code by changing
+  ``colors, levels = read_imod_legend(...)`` to
   ``colors, levels, labels = read_imod_legend(...)``.
 
 

imod/common/utilities/partitioninfo.py

@@ -0,0 +1,53 @@
+from typing import List, NamedTuple
+
+import numpy as np
+
+from imod.typing import GridDataArray
+from imod.typing.grid import ones_like
+
+
+class PartitionInfo(NamedTuple):
+    active_domain: GridDataArray
+    id: int
+
+
+def create_partition_info(submodel_labels: GridDataArray) -> List[PartitionInfo]:
+    """
+    A PartitionInfo is used to partition a model or package. The partition
+    info's of a domain are created using a submodel_labels array. The
+    submodel_labels provided as input should have the same shape as a single
+    layer of the model grid (all layers are split the same way), and contains an
+    integer value in each cell. Each cell in the model grid will end up in the
+    submodel with the index specified by the corresponding label of that cell.
+    The labels should be numbers between 0 and the number of partitions.
+    """
+    _validate_submodel_label_array(submodel_labels)
+
+    unique_labels = np.unique(submodel_labels.values)
+
+    partition_infos = []
+    for label_id in unique_labels:
+        active_domain = submodel_labels.where(submodel_labels.values == label_id)
+        active_domain = ones_like(active_domain).where(active_domain.notnull(), 0)
+        active_domain = active_domain.astype(submodel_labels.dtype)
+
+        submodel_partition_info = PartitionInfo(
+            id=label_id, active_domain=active_domain
+        )
+        partition_infos.append(submodel_partition_info)
+
+    return partition_infos
+
+
+def _validate_submodel_label_array(submodel_labels: GridDataArray) -> None:
+    unique_labels = np.unique(submodel_labels.values)
+
+    if not (
+        len(unique_labels) == unique_labels.max() + 1
+        and unique_labels.min() == 0
+        and np.issubdtype(submodel_labels.dtype, np.integer)
+    ):
+        raise ValueError(
+            "The submodel_label  array should be integer and contain all the numbers between 0 and the number of "
+            "partitions minus 1."
+        )

imod/mf6/multimodel/modelsplitter.py

@@ -1,69 +1,18 @@
-from typing import List, NamedTuple, cast
-
-import numpy as np
+from typing import List, cast
 
 from imod.common.interfaces.iagnosticpackage import IAgnosticPackage
 from imod.common.interfaces.imodel import IModel
 from imod.common.interfaces.ipackage import IPackage
 from imod.common.utilities.clip import clip_by_grid
+from imod.common.utilities.partitioninfo import PartitionInfo
 from imod.mf6.boundary_condition import BoundaryCondition
 from imod.mf6.ims import Solution
 from imod.mf6.model_gwf import GroundwaterFlowModel
 from imod.mf6.model_gwt import GroundwaterTransportModel
 from imod.mf6.ssm import SourceSinkMixing
 from imod.mf6.validation_settings import ValidationSettings, trim_time_dimension
 from imod.typing import GridDataArray
-from imod.typing.grid import (
-    get_non_spatial_dimension_names,
-    ones_like,
-)
-
-
-class PartitionInfo(NamedTuple):
-    active_domain: GridDataArray
-    id: int
-
-
-def create_partition_info(submodel_labels: GridDataArray) -> List[PartitionInfo]:
-    """
-    A PartitionInfo is used to partition a model or package. The partition
-    info's of a domain are created using a submodel_labels array. The
-    submodel_labels provided as input should have the same shape as a single
-    layer of the model grid (all layers are split the same way), and contains an
-    integer value in each cell. Each cell in the model grid will end up in the
-    submodel with the index specified by the corresponding label of that cell.
-    The labels should be numbers between 0 and the number of partitions.
-    """
-    _validate_submodel_label_array(submodel_labels)
-
-    unique_labels = np.unique(submodel_labels.values)
-
-    partition_infos = []
-    for label_id in unique_labels:
-        active_domain = submodel_labels.where(submodel_labels.values == label_id)
-        active_domain = ones_like(active_domain).where(active_domain.notnull(), 0)
-        active_domain = active_domain.astype(submodel_labels.dtype)
-
-        submodel_partition_info = PartitionInfo(
-            id=label_id, active_domain=active_domain
-        )
-        partition_infos.append(submodel_partition_info)
-
-    return partition_infos
-
-
-def _validate_submodel_label_array(submodel_labels: GridDataArray) -> None:
-    unique_labels = np.unique(submodel_labels.values)
-
-    if not (
-        len(unique_labels) == unique_labels.max() + 1
-        and unique_labels.min() == 0
-        and np.issubdtype(submodel_labels.dtype, np.integer)
-    ):
-        raise ValueError(
-            "The submodel_label  array should be integer and contain all the numbers between 0 and the number of "
-            "partitions minus 1."
-        )
+from imod.typing.grid import get_non_spatial_dimension_names
 
 
 class ModelSplitter:

imod/mf6/simulation.py

@@ -25,6 +25,7 @@
 from imod.common.statusinfo import NestedStatusInfo
 from imod.common.utilities.dataclass_type import DataclassType
 from imod.common.utilities.mask import mask_all_models
+from imod.common.utilities.partitioninfo import create_partition_info
 from imod.common.utilities.regrid import _regrid_like
 from imod.common.utilities.version import (
     get_version,
@@ -43,10 +44,7 @@
 from imod.mf6.multimodel.exchange_creator_unstructured import (
     ExchangeCreator_Unstructured,
 )
-from imod.mf6.multimodel.modelsplitter import (
-    ModelSplitter,
-    create_partition_info,
-)
+from imod.mf6.multimodel.modelsplitter import ModelSplitter
 from imod.mf6.out import open_cbc, open_conc, open_hds
 from imod.mf6.package import Package
 from imod.mf6.validation_settings import ValidationSettings

imod/msw/coupler_mapping.py

@@ -4,14 +4,15 @@
 import pandas as pd
 import xarray as xr
 
+from imod.common.interfaces.ipackagebase import IPackageBase
 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 DataDictType, MetaSwapPackage
 from imod.typing import IntArray
 
 
-class CouplerMapping(MetaSwapPackage):
+class CouplerMapping(MetaSwapPackage, IPackageBase):
     """
     This contains the data to connect MODFLOW 6 cells to MetaSWAP svats.
 

imod/msw/grid_data.py

@@ -83,7 +83,29 @@ def __init__(
 
         self._pkgcheck()
 
-    def generate_index_array(self) -> tuple[np.ndarray, xr.DataArray]:
+    def generate_index_array(self) -> np.ndarray:
+        """
+        Generate index array to be used on other packages.
+
+        Returns
+        -------
+        np.ndarray
+            Index array and svat grid.
+            The index array is a 1D array with the index of the active cells.
+        """
+        area = self.dataset["area"]
+        active = self.dataset["active"]
+
+        isactive = area.where(active).notnull()
+        # Load into memory to avoid dask issue
+        # https://github.com/dask/dask/issues/11753
+        isactive.load()
+
+        index = isactive.values.ravel()
+
+        return index
+
+    def generate_index_svat_array(self) -> tuple[np.ndarray, xr.DataArray]:
         """
         Generate index array and svat grid to be used on other packages.
 
@@ -94,17 +116,18 @@ def generate_index_array(self) -> tuple[np.ndarray, xr.DataArray]:
             The index array is a 1D array with the index of the active cells.
             The svat grid is a 2D array with the SVAT numbers for each cell.
         """
+        index = self.generate_index_array()
+
         area = self.dataset["area"]
         active = self.dataset["active"]
-
         isactive = area.where(active).notnull()
+
         svat = xr.full_like(area, fill_value=0, dtype=np.int64).rename("svat")
         # Load into memory to avoid dask issue
         # https://github.com/dask/dask/issues/11753
         isactive.load()
         svat.load()
 
-        index = isactive.values.ravel()
         svat.data[isactive.data] = np.arange(1, index.sum() + 1)
 
         return index, svat

imod/msw/meteo_mapping.py

@@ -9,6 +9,7 @@
 import xarray as xr
 
 import imod
+from imod.common.interfaces.ipackagebase import IPackageBase
 from imod.common.utilities.clip import clip_spatial_box, clip_time_slice
 from imod.common.utilities.dataclass_type import DataclassType
 from imod.msw.fixed_format import VariableMetaData
@@ -65,7 +66,7 @@ def open_first_meteo_grid_from_imod5_data(imod5_data: Imod5DataDict, column_nr:
     return open_first_meteo_grid(metegrid_path, column_nr=column_nr)
 
 
-class MeteoMapping(MetaSwapPackage):
+class MeteoMapping(MetaSwapPackage, IPackageBase):
     """
     This class provides common methods for creating mappings between
     meteorological data and MetaSWAP grids. It should not be instantiated