Model split perf

imod/mf6/boundary_condition.py

@@ -202,12 +202,23 @@ def _get_unfiltered_pkg_options(
         options = copy(predefined_options)
 
         if not_options is None:
-            not_options = self._get_period_varnames()
+            not_options = []
+            if hasattr(self, "_period_data"):
+                not_options.extend(self._period_data)
+            if hasattr(self, "_auxiliary_data"):
+                not_options.extend(get_variable_names(self))
+                not_options.extend(self._auxiliary_data.keys())
 
         for varname in self.dataset.data_vars.keys():  # pylint:disable=no-member
             if varname in not_options:
                 continue
-            v = self.dataset[varname].values[()]
+            # TODO: can we easily avoid this try-except?
+            # On which keys does it fail?
+            try:
+                v = self.dataset[varname].item()
+            except ValueError:
+                # Apparently not a scalar, therefore not an option entry.
+                pass
             options[varname] = v
         return options
 

imod/mf6/hfb.py

@@ -33,6 +33,7 @@
 from imod.mf6.disv import VerticesDiscretization
 from imod.mf6.mf6_hfb_adapter import Mf6HorizontalFlowBarrier
 from imod.mf6.package import Package
+from imod.mf6.utilities.zarr_helper import to_zarr
 from imod.mf6.validation_settings import ValidationSettings
 from imod.prepare.cleanup import cleanup_hfb
 from imod.schemata import (
@@ -562,6 +563,11 @@ def to_netcdf(
         new.dataset["geometry"] = new.line_data.to_json()
         new.dataset.to_netcdf(*args, **kwargs)
 
+    def to_zarr(self, path, engine: str, **kwargs):
+        new = deepcopy(self)
+        new.dataset["geometry"] = new.line_data.to_json()
+        to_zarr(new.dataset, path, engine, **kwargs)
+
     def _netcdf_encoding(self):
         return {"geometry": {"dtype": "str"}}
 

imod/mf6/model.py

@@ -592,6 +592,7 @@ def dump(
         validate: bool = True,
         mdal_compliant: bool = False,
         crs: Optional[Any] = None,
+        engine="netCDF4",
     ):
         """
         Dump simulation to files. Writes a model definition as .TOML file, which
@@ -615,6 +616,8 @@ def dump(
         crs: Any, optional
             Anything accepted by rasterio.crs.CRS.from_user_input
             Requires ``rioxarray`` installed.
+        engine: str, optional
+            "netCDF4" or "zarr" or "zarr.zip". Defaults to "netCDF4".
         """
         modeldirectory = pathlib.Path(directory) / modelname
         modeldirectory.mkdir(exist_ok=True, parents=True)
@@ -624,13 +627,26 @@ def dump(
             if statusinfo.has_errors():
                 raise ValidationError(statusinfo.to_string())
 
+        match engine:
+            case "netCDF4":
+                ext = "nc"
+            case "zarr":
+                ext = "zarr"
+            case "zarr.zip":
+                ext = "zarr.zip"
+            case _:
+                raise ValueError(f"Unknown engine: {engine}")
+
         toml_content: dict = collections.defaultdict(dict)
         for pkgname, pkg in self.items():
-            pkg_path = f"{pkgname}.nc"
+            pkg_path = f"{pkgname}.{ext}"
             toml_content[type(pkg).__name__][pkgname] = pkg_path
-            pkg.to_netcdf(
-                modeldirectory / pkg_path, crs=crs, mdal_compliant=mdal_compliant
-            )
+            if engine == "netCDF4":
+                pkg.to_netcdf(
+                    modeldirectory / pkg_path, crs=crs, mdal_compliant=mdal_compliant
+                )
+            else:
+                pkg.to_zarr(modeldirectory / pkg_path, engine=engine)
 
         toml_path = modeldirectory / f"{modelname}.toml"
         with open(toml_path, "wb") as f:

imod/mf6/multimodel/exchange_creator.py

@@ -1,5 +1,5 @@
 import abc
-from typing import Dict
+from typing import Dict, NamedTuple
 
 import numpy as np
 import pandas as pd
@@ -8,10 +8,14 @@
 from imod.common.utilities.grid import get_active_domain_slice, to_cell_idx
 from imod.mf6.gwfgwf import GWFGWF
 from imod.mf6.gwtgwt import GWTGWT
-from imod.mf6.multimodel.modelsplitter import PartitionInfo
 from imod.typing import GridDataArray
 
 
+class PartitionInfo(NamedTuple):
+    active_domain: GridDataArray
+    partition_id: int
+
+
 def _adjust_gridblock_indexing(connected_cells: xr.Dataset) -> xr.Dataset:
     """
     adjusts the gridblock numbering from 0-based to 1-based.
@@ -25,8 +29,7 @@ class ExchangeCreator(abc.ABC):
     """
     Creates the GroundWaterFlow to GroundWaterFlow exchange package (gwfgwf) as a function of a submodel label array
     and a PartitionInfo object. This file contains the cell indices of coupled cells. With coupled cells we mean
-    cells that are adjacent but that are located in different subdomains.  At the moment only structured grids are
-    supported, for unstructured grids the geometric information is still set to default values.
+    cells that are adjacent but that are located in different subdomains.
 
     The submodel_labels array should have the same topology as the domain being partitioned. The array will be used
     to determine the connectivity of the submodels after the split operation has been performed.
@@ -248,7 +251,7 @@ def _create_global_cellidx_to_local_cellid_mapping(
 
         mapping = {}
         for submodel_partition_info in partition_info:
-            model_id = submodel_partition_info.id
+            model_id = submodel_partition_info.partition_id
             mapping[model_id] = pd.merge(
                 global_to_local_idx[model_id], local_cell_idx_to_id[model_id]
             )
@@ -268,7 +271,7 @@ def _get_local_cell_indices(
     def _local_cell_idx_to_id(cls, partition_info) -> Dict[int, pd.DataFrame]:
         local_cell_idx_to_id = {}
         for submodel_partition_info in partition_info:
-            model_id = submodel_partition_info.id
+            model_id = submodel_partition_info.partition_id
 
             local_cell_indices = cls._get_local_cell_indices(submodel_partition_info)
             local_cell_id = list(np.ndindex(local_cell_indices.shape))

imod/mf6/multimodel/exchange_creator_structured.py

@@ -1,17 +1,21 @@
-from typing import Dict, List
+from typing import Dict, List, NamedTuple
 
 import numpy as np
 import pandas as pd
 import xarray as xr
 
 from imod.common.utilities.grid import create_geometric_grid_info
 from imod.mf6.multimodel.exchange_creator import ExchangeCreator
-from imod.mf6.multimodel.modelsplitter import PartitionInfo
 from imod.typing import GridDataArray
 
 NOT_CONNECTED_VALUE = -999
 
 
+class PartitionInfo(NamedTuple):
+    active_domain: GridDataArray
+    partition_id: int
+
+
 class ExchangeCreator_Structured(ExchangeCreator):
     """
     Creates the GroundWaterFlow to GroundWaterFlow exchange package (gwfgwf) as
@@ -130,7 +134,7 @@ def _create_global_to_local_idx(
                 compat="override",
             )["label"]
 
-            model_id = submodel_partition_info.id
+            model_id = submodel_partition_info.partition_id
             global_to_local_idx[model_id] = pd.DataFrame(
                 {
                     "global_idx": overlap.values.flatten(),

imod/mf6/multimodel/exchange_creator_unstructured.py

@@ -4,8 +4,7 @@
 import pandas as pd
 import xarray as xr
 
-from imod.mf6.multimodel.exchange_creator import ExchangeCreator
-from imod.mf6.multimodel.modelsplitter import PartitionInfo
+from imod.mf6.multimodel.exchange_creator import ExchangeCreator, PartitionInfo
 from imod.typing import GridDataArray
 
 
@@ -122,7 +121,7 @@ def _create_global_to_local_idx(
                 compat="override",
             )["label"]
 
-            model_id = submodel_partition_info.id
+            model_id = submodel_partition_info.partition_id
             global_to_local_idx[model_id] = pd.DataFrame(
                 {
                     "global_idx": overlap.values.flatten(),