Add chunking when opening netcdf files. Handle errors for code that didn't expect to recieve dask objects. Optimize the flow-transport model matcher

Author: Manangka

imod/mf6/hfb.py

@@ -557,6 +557,7 @@ def to_netcdf(
 
         """
         kwargs.update({"encoding": self._netcdf_encoding()})
+        kwargs.update({"format": "NETCDF4"})
 
         new = deepcopy(self)
         new.dataset["geometry"] = new.line_data.to_json()

imod/mf6/oc.py

@@ -172,7 +172,7 @@ def _render(self, directory, pkgname, globaltimes, binary):
                 package_times = self.dataset[datavar].coords["time"].values
                 starts = np.searchsorted(globaltimes, package_times) + 1
                 for i, s in enumerate(starts):
-                    setting = self.dataset[datavar].isel(time=i).item()
+                    setting = self.dataset[datavar].isel(time=i).values[()]
                     periods[s][key] = self._get_ocsetting(setting)
 
             else:

imod/mf6/pkgbase.py

@@ -92,8 +92,20 @@ def to_netcdf(
 
         """
         kwargs.update({"encoding": self._netcdf_encoding()})
+        kwargs.update({"format": "NETCDF4"})
 
         dataset = self.dataset
+
+        # Create encoding dict for float16 variables
+        for var in dataset.data_vars:
+            if dataset[var].dtype == np.float16:
+                kwargs["encoding"][var] = {"dtype": "float32"}
+
+        # Also check coordinates
+        for coord in dataset.coords:
+            if dataset[coord].dtype == np.float16:
+                kwargs["encoding"][coord] = {"dtype": "float32"}
+
         if isinstance(dataset, xu.UgridDataset):
             if mdal_compliant:
                 dataset = dataset.ugrid.to_dataset()
@@ -168,7 +180,7 @@ def from_file(cls, path: str | Path, **kwargs) -> Self:
             # TODO: seems like a bug? Remove str() call if fixed in xarray/zarr
             dataset = xr.open_zarr(str(path), **kwargs)
         else:
-            dataset = xr.open_dataset(path, **kwargs)
+            dataset = xr.open_dataset(path, chunks="auto", **kwargs)
 
         if dataset.ugrid_roles.topology:
             dataset = xu.UgridDataset(dataset)

imod/mf6/rch.py

@@ -1,5 +1,6 @@
 from datetime import datetime
-from typing import Optional
+from pathlib import Path
+from typing import Optional, Self
 
 import numpy as np
 import xarray as xr
@@ -166,6 +167,22 @@ def __init__(
         super().__init__(dict_dataset)
         self._validate_init_schemata(validate)
 
+    @classmethod
+    def from_file(cls, path: str | Path, **kwargs) -> Self:
+        instance = super().from_file(path, **kwargs)
+
+        # to_netcdf converts strings into  NetCDF "variable‑length UTF‑8 strings"
+        # which are loaded as dtype=object arrays.
+        # This will convert them back to str.
+        vars = [
+            "species",
+        ]
+        for var in vars:
+            if var in instance.dataset:
+                instance.dataset[var] = instance.dataset[var].astype(str)
+
+        return instance
+
     def _validate(self, schemata, **kwargs):
         # Insert additional kwargs
         kwargs["rate"] = self["rate"]

imod/mf6/simulation.py

@@ -60,7 +60,7 @@
 from imod.typing import GridDataArray, GridDataset
 from imod.typing.grid import (
     concat,
-    is_equal,
+    is_same_domain,
     is_unstructured,
     merge_partitions,
 )
@@ -1037,12 +1037,14 @@ def dump(
                     _, filename, _, _ = exchange_package.get_specification()
                     exchange_class_short = type(exchange_package).__name__
                     path = f"{filename}.nc"
-                    exchange_package.dataset.to_netcdf(directory / path)
+                    exchange_package.dataset.to_netcdf(
+                        directory / path, format="NETCDF4"
+                    )
                     toml_content[key][exchange_class_short].append(path)
 
             else:
                 path = f"{key}.nc"
-                value.dataset.to_netcdf(directory / path)
+                value.dataset.to_netcdf(directory / path, format="NETCDF4")
                 toml_content[cls_name][key] = path
 
         with open(directory / f"{self.name}.toml", "wb") as f:
@@ -1620,10 +1622,16 @@ def _get_transport_models_per_flow_model(self) -> dict[str, list[str]]:
 
         for flow_model_name in flow_models:
             flow_model = self[flow_model_name]
+
+            matched_tsp_models = []
             for tpt_model_name in transport_models:
                 tpt_model = self[tpt_model_name]
-                if is_equal(tpt_model.domain, flow_model.domain):
+                if is_same_domain(tpt_model.domain, flow_model.domain):
                     result[flow_model_name].append(tpt_model_name)
+                    matched_tsp_models.append(tpt_model_name)
+            for tpt_model_name in matched_tsp_models:
+                transport_models.pop(tpt_model_name)
+
         return result
 
     def _generate_gwfgwt_exchanges(self) -> list[GWFGWT]:

imod/mf6/wel.py

@@ -540,7 +540,7 @@ def _to_mf6_package_information(
         else:
             message += " The first 10 unplaced wells are: \n"
 
-        is_filtered = self.dataset["id"].isin([filtered_wells])
+        is_filtered = self.dataset["id"].compute().isin(filtered_wells)
         for i in range(min(10, len(filtered_wells))):
             ids = filtered_wells[i]
             x = self.dataset["x"].data[is_filtered][i]
@@ -1076,7 +1076,7 @@ def _assign_wells_to_layers(
         like = ones_like(active)
         bottom = like * bottom
         top_2d = (like * top).sel(layer=1)
-        top_3d = bottom.shift(layer=1).fillna(top_2d)
+        top_3d = bottom.compute().shift(layer=1).fillna(top_2d)
         k = like * k
 
         index_names = wells_df.index.names

imod/typing/grid.py

@@ -316,11 +316,15 @@ def is_spatial_grid(_: Any) -> bool:  # noqa: F811
 
 @dispatch
 def is_equal(array1: xu.UgridDataArray, array2: xu.UgridDataArray) -> bool:
+    if not is_same_domain(array1, array2):
+        return False
     return array1.equals(array2) and array1.ugrid.grid.equals(array2.ugrid.grid)
 
 
 @dispatch  # type: ignore[no-redef]
 def is_equal(array1: xr.DataArray, array2: xr.DataArray) -> bool:  # noqa: F811
+    if not is_same_domain(array1, array2):
+        return False
     return array1.equals(array2)