Issue #1698 Model write optimization (#1700)

Fixes #1698


# Description
This is part 2 of fixing the performance issues with large model.
In part 1 #1693 the modelsplitter has been optimized. In this PR the
focus is on wiring the partitioned model.

As @Huite pointed out in #1686 the performance bottleneck had to do with
the fact that the same package had to be loaded from file multiple times
while only a part of the file is actually needed.

After digging around for a while i discovered that this had to do with
the fact how we open de the dataset.
`dataset = xr.open_dataset(path, **kwargs)`

In the line above we don't specify anything chunk related. That has as a
result that when you access the dataset the entire file has to be loaded
from disk. By simply adding `chunks="auto"` this is no longer the case
and a huge performance gain is achieved.

There are some other changes related to setting chunking to auto. There
are some parts of the code that don't expect to receive dask arrays. For
instance you can use .item() on a dask array. Instead i now use
.values[()].

I was also getting some errors when the to_netcdf method were called on
the package. All of them had something to do with wrong/unsupported
datatypes. In this PR you will find that an encoding is added for
float16 types. And that in some packages the from_file method has been
updated to ensure that he loaded type is converted to a supported type

An unrelated change but performance wise significant change has been
applied to the `_get_transport_models_per_flow_model` method. This
method is used to match gwf models to gwt models so that gwfgwt
exchanges can be created. This method was doing a full comparison of
domains, which is expensive. There is also a method available that does
the comparison on domain level. By switching to this method the matching
algorithm becomes almost instantaneously.

**NOTE**
This PR has issue #1699 as a base. The base needs to altered to master
once that PR is in

**NOTE**
This PR also improves the `dump` method

**NOTE**
some timmings:

<img width="833" height="739" alt="image"
src="https://github.com/user-attachments/assets/974c841c-0413-4433-8486-1abe98dc0715"
/>

<img width="843" height="215" alt="image"
src="https://github.com/user-attachments/assets/c7082975-af35-4143-a6f9-860557b3eb09"
/>

<img width="842" height="705" alt="image"
src="https://github.com/user-attachments/assets/383bf1a6-f028-4cb4-aa72-48ab95e84e3d"
/>



<!---
Before requesting review, please go through this checklist:
-->

- [x] Links to correct issue
- [ ] Update changelog, if changes affect users
- [x] PR title starts with ``Issue #nr``, e.g. ``Issue #737``
- [ ] Unit tests were added
- [ ] **If feature added**: Added/extended example
- [ ] **If feature added**: Added feature to API documentation
- [ ] **If pixi.lock was changed**: Ran `pixi run generate-sbom` and
committed changes

---------

Co-authored-by: JoerivanEngelen <joerivanengelen@hotmail.com>

Author: Manangka

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

@@ -94,6 +94,17 @@ def to_netcdf(
         kwargs.update({"encoding": self._netcdf_encoding()})
 
         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 +179,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)
@@ -183,4 +194,12 @@ def from_file(cls, path: str | Path, **kwargs) -> Self:
             if _is_scalar_nan(value):
                 dataset[key] = None
 
+        # 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 dataset:
+                dataset[var] = dataset[var].astype(str)
+
         return cls._from_dataset(dataset)

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,
 )
@@ -1622,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]
@@ -1073,9 +1073,9 @@ def _assign_wells_to_layers(
     ) -> pd.DataFrame:
         # Ensure top, bottom & k
         # are broadcasted to 3d grid
-        like = ones_like(active)
-        bottom = like * bottom
-        top_2d = (like * top).sel(layer=1)
+        like = ones_like(active.compute())
+        bottom = like * bottom.compute()
+        top_2d = (like * top.compute()).sel(layer=1)
         top_3d = bottom.shift(layer=1).fillna(top_2d)
         k = like * k
 

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)
 
 

pyproject.toml

@@ -80,7 +80,6 @@ ignore = [
     "E501", # line-too-long. This rule can't be fullfilled by the ruff formatter. The same behavior as black.
     "PD003",
     "PD004",
-    "PD901",
     "PD011",
     "PD013",
     "PD015",