Issue #1343 hfb dump (#1346)

Fixes #1343

# Description
- Fix crash when calling ``GroundwaterFlowModel.dump`` when a HFB was
assigned to the model
- Really encode HFB as string this time when calling ``to_netcdf``...
- Move MDAL/GDAL logic to ``Package.to_netcdf``, which is a more logical
place for it then in ``Model.dump``. This has the advantage that this
stuff is automatically not done when we override the method for the HFB
package, as the MDAL/CRS stuff is irrelevant for HFBs.
- Drop node_dim from UGRID dataset if present, it still remains
available in the grid topology.
- Ensure non-scalar data not loaded into memory when calling
``from_file``.

Author: JoerivanEngelen

docs/api/changelog.rst

@@ -33,9 +33,13 @@ Added
 Fixed
 ~~~~~
 
-- Fixed bug were :class:`HorizontalFlowBarrierResistance`, 
+- Fixed bug where :class:`HorizontalFlowBarrierResistance`, 
   :class:`HorizontalFlowBarrierSingleLayerResistance` and other HFB packages could not 
-  be dumped, clipped, or copied with xarray >= 2024.10.0.
+  be clipped or copied with xarray >= 2024.10.0.
+- Fixed crash upon calling :meth:`imod.mf6.GroundwaterFlowModel.dump`, when a
+  :class:`HorizontalFlowBarrierResistance`,
+  :class:`HorizontalFlowBarrierSingleLayerResistance` or other HFB package was
+  assigned to the model.
 - :meth:`imod.mf6.Modflow6Simulation.regrid_like` can now regrid a structured
   model to an unstructured grid.
 - :meth:`imod.mf6.Modflow6Simulation.regrid_like` throws a

imod/mf6/hfb.py

@@ -4,7 +4,7 @@
 import typing
 from copy import deepcopy
 from enum import Enum
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
 
 import cftime
 import numpy as np
@@ -491,7 +491,11 @@ def render(self, directory, pkgname, globaltimes, binary):
             package, first convert to a Modflow6 package by calling pkg.to_mf6_pkg()"""
         )
 
-    def to_netcdf(self, *args, **kwargs):
+    def to_netcdf(
+        self, *args, mdal_compliant: bool = False, crs: Optional[Any] = None, **kwargs
+    ):
+        kwargs.update({"encoding": self._netcdf_encoding()})
+
         new = deepcopy(self)
         new.dataset["geometry"] = new.line_data.to_json()
         new.dataset.to_netcdf(*args, **kwargs)

imod/mf6/model.py

@@ -36,7 +36,6 @@
 from imod.mf6.write_context import WriteContext
 from imod.schemata import ValidationError
 from imod.typing import GridDataArray
-from imod.typing.grid import is_spatial_grid
 
 HFB_PKGNAME = "hfb_merged"
 SUGGESTION_TEXT = (
@@ -465,20 +464,9 @@ def dump(
         for pkgname, pkg in self.items():
             pkg_path = f"{pkgname}.nc"
             toml_content[type(pkg).__name__][pkgname] = pkg_path
-            dataset = pkg.dataset
-            if isinstance(dataset, xu.UgridDataset):
-                if mdal_compliant:
-                    dataset = dataset.ugrid.to_dataset()
-                    mdal_dataset = imod.util.spatial.mdal_compliant_ugrid2d(
-                        dataset, crs=crs
-                    )
-                    mdal_dataset.to_netcdf(modeldirectory / pkg_path)
-                else:
-                    dataset.ugrid.to_netcdf(modeldirectory / pkg_path)
-            else:
-                if is_spatial_grid(dataset):
-                    dataset = imod.util.spatial.gdal_compliant_grid(dataset, crs=crs)
-                dataset.to_netcdf(modeldirectory / pkg_path)
+            pkg.to_netcdf(
+                modeldirectory / pkg_path, crs=crs, mdal_compliant=mdal_compliant
+            )
 
         toml_path = modeldirectory / f"{modelname}.toml"
         with open(toml_path, "wb") as f:

imod/mf6/pkgbase.py

@@ -1,20 +1,37 @@
 import abc
 import numbers
 import pathlib
-from typing import Mapping
+from typing import Any, Mapping, Optional
 
 import numpy as np
 import xarray as xr
 import xugrid as xu
+from xarray.core.utils import is_scalar
 
 import imod
 from imod.mf6.interfaces.ipackagebase import IPackageBase
-from imod.typing.grid import GridDataArray, GridDataset, merge_with_dictionary
+from imod.typing.grid import (
+    GridDataArray,
+    GridDataset,
+    is_spatial_grid,
+    merge_with_dictionary,
+)
 
 TRANSPORT_PACKAGES = ("adv", "dsp", "ssm", "mst", "ist", "src")
 EXCHANGE_PACKAGES = ("gwfgwf", "gwfgwt", "gwtgwt")
 
 
+def _is_scalar_nan(da: GridDataArray):
+    """
+    Test if is_scalar_nan, carefully avoid loading grids in memory
+    """
+    scalar_data: bool = is_scalar(da)
+    if scalar_data:
+        stripped_value = da.to_numpy()[()]
+        return isinstance(stripped_value, numbers.Real) and np.isnan(stripped_value)  # type: ignore[call-overload]
+    return False
+
+
 class PackageBase(IPackageBase, abc.ABC):
     """
     This class is used for storing a collection of Xarray DataArrays or UgridDataArrays
@@ -48,15 +65,48 @@ def __getitem__(self, key):
     def __setitem__(self, key, value):
         self.dataset.__setitem__(key, value)
 
-    def to_netcdf(self, *args, **kwargs):
+    def to_netcdf(
+        self, *args, mdal_compliant: bool = False, crs: Optional[Any] = None, **kwargs
+    ):
         """
 
-        Write dataset contents to a netCDF file.
-        Custom encoding rules can be provided on package level by overriding the _netcdf_encoding in the package
+        Write dataset contents to a netCDF file. Custom encoding rules can be
+        provided on package level by overriding the _netcdf_encoding in the
+        package
+
+        Parameters
+        ----------
+        *args:
+            Will be passed on to ``xr.Dataset.to_netcdf`` or
+            ``xu.UgridDataset.to_netcdf``.
+        mdal_compliant: bool, optional
+            Convert data with
+            :func:`imod.prepare.spatial.mdal_compliant_ugrid2d` to MDAL
+            compliant unstructured grids. Defaults to False.
+        crs: Any, optional
+            Anything accepted by rasterio.crs.CRS.from_user_input
+            Requires ``rioxarray`` installed.
+        **kwargs:
+            Will be passed on to ``xr.Dataset.to_netcdf`` or
+            ``xu.UgridDataset.to_netcdf``.
 
         """
         kwargs.update({"encoding": self._netcdf_encoding()})
-        self.dataset.to_netcdf(*args, **kwargs)
+
+        dataset = self.dataset
+        if isinstance(dataset, xu.UgridDataset):
+            if mdal_compliant:
+                dataset = dataset.ugrid.to_dataset()
+                mdal_dataset = imod.util.spatial.mdal_compliant_ugrid2d(
+                    dataset, crs=crs
+                )
+                mdal_dataset.to_netcdf(*args, **kwargs)
+            else:
+                dataset.ugrid.to_netcdf(*args, **kwargs)
+        else:
+            if is_spatial_grid(dataset):
+                dataset = imod.util.spatial.gdal_compliant_grid(dataset, crs=crs)
+            dataset.to_netcdf(*args, **kwargs)
 
     def _netcdf_encoding(self):
         """
@@ -123,11 +173,14 @@ def from_file(cls, path, **kwargs):
         if dataset.ugrid_roles.topology:
             dataset = xu.UgridDataset(dataset)
             dataset = imod.util.spatial.from_mdal_compliant_ugrid2d(dataset)
+            # Drop node dim, as we don't need in the UgridDataset (it will be
+            # preserved in the ``grid`` property, so topology stays intact!)
+            node_dim = dataset.grid.node_dimension
+            dataset = dataset.drop_dims(node_dim, errors="ignore")
 
         # Replace NaNs by None
         for key, value in dataset.items():
-            stripped_value = value.values[()]
-            if isinstance(stripped_value, numbers.Real) and np.isnan(stripped_value):  # type: ignore[call-overload]
+            if _is_scalar_nan(value):
                 dataset[key] = None
 
         return cls._from_dataset(dataset)

imod/tests/conftest.py

@@ -81,6 +81,7 @@
     transient_unconfined_twri_model,
     transient_unconfined_twri_result,
     twri_model,
+    twri_model_hfb,
     twri_result,
 )
 from .fixtures.mf6_welltest_fixture import (

imod/tests/fixtures/mf6_twri_fixture.py

@@ -1,7 +1,9 @@
+import geopandas as gpd
 import numpy as np
 import pandas as pd
 import pytest
 import xarray as xr
+from shapely import linestrings
 
 import imod
 from imod.typing.grid import zeros_like
@@ -232,6 +234,28 @@ def transient_twri_model():
     return simulation
 
 
+@pytest.fixture(scope="function")
+def twri_model_hfb():
+    """Returns steady-state confined model."""
+    simulation = make_twri_model()
+    gwf_model = simulation["GWF_1"]
+
+    barrier_y = [150_000.0, 50_000.0, 0.0]
+    barrier_x = [50_000.0, 50_000.0, 50_000.0]
+
+    geometry = gpd.GeoDataFrame(
+        geometry=[linestrings(barrier_x, barrier_y)],
+        data={
+            "resistance": [1200.0],
+            "layer": [1],
+        },
+    )
+
+    gwf_model["hfb"] = imod.mf6.SingleLayerHorizontalFlowBarrierResistance(geometry)
+
+    return simulation
+
+
 @pytest.fixture(scope="function")
 def transient_unconfined_twri_model():
     """Returns transient unconfined model, also saves specific discharges."""

imod/tests/test_mf6/test_mf6_simulation.py

@@ -52,6 +52,11 @@ def test_twri_roundtrip(twri_model, tmpdir_factory):
     roundtrip(twri_model, tmpdir_factory, "twri")
 
 
+@pytest.mark.usefixtures("twri_model_hfb")
+def test_twri_hfb_roundtrip(twri_model_hfb, tmpdir_factory):
+    roundtrip(twri_model_hfb, tmpdir_factory, "twri")
+
+
 @pytest.mark.usefixtures("transient_twri_model")
 def test_twri_transient_roundtrip(transient_twri_model, tmpdir_factory):
     roundtrip(transient_twri_model, tmpdir_factory, "twri_transient")