Issue #1132 performance hfb (#1157)

Fixes #1132 

# Description
This PR adds the following:

- Increase performance for writing the HFB files, by using new xugrid
create_snap_to_grid_dataframe.
- Fixes error when multiple barriers in a single dataset would be
snapped to same edge. These lines are now aggregated.
- Add some basic infrastructure for user acceptance test for writing
LHM, the path to the LHM projectfile should be set in a .env as LHM_PRJ.
This can later be extended to TeamCity.

This reduces the time spent writing the HFB packages from 34 minutes to
12.5 minutes on my local machine. I've profiled the causes of the
remaining bottlenecks to writing. I'll post a follow-up issue for that.

# Checklist

- [x] Links to correct issue
- [x] Update changelog, if changes affect users
- [x] PR title starts with ``Issue #nr``, e.g. ``Issue #737``
- [x] Unit tests were added
- [ ] **If feature added**: Added/extended example

Author: JoerivanEngelen

docs/api/changelog.rst

@@ -15,7 +15,8 @@ Fixed
 - Multiple ``HorizontalFlowBarrier`` objects attached to
   :class:`imod.mf6.GroundwaterFlowModel` are merged into a single horizontal
   flow barrier for MODFLOW 6
-
+- Bug where error would be thrown when barriers in a ``HorizontalFlowBarrier``
+  would be snapped to the same cell edge. These are now summed.
 
 Changed
 ~~~~~~~

imod/mf6/hfb.py

@@ -8,6 +8,7 @@
 
 import cftime
 import numpy as np
+import numpy.typing as npt
 import pandas as pd
 import xarray as xr
 import xugrid as xu
@@ -27,6 +28,7 @@
 )
 from imod.schemata import EmptyIndexesSchema, MaxNUniqueValuesSchema
 from imod.typing import GeoDataFrameType, GridDataArray, LineStringType
+from imod.typing.grid import as_ugrid_dataarray
 from imod.util.imports import MissingOptionalModule
 
 if TYPE_CHECKING:
@@ -374,6 +376,53 @@ def _prepare_barrier_dataset_for_mf6_adapter(dataset: xr.Dataset) -> xr.Dataset:
     return dataset
 
 
+def _snap_to_grid_and_aggregate(
+    barrier_dataframe: GeoDataFrameType, grid2d: xu.Ugrid2d, vardict_agg: dict[str, str]
+) -> tuple[xu.UgridDataset, npt.NDArray]:
+    """
+    Snap barrier dataframe to grid and aggregate multiple lines with a list of
+    methods per variable.
+
+    Parameters
+    ----------
+    barrier_dataframe: geopandas.GeoDataFrame
+        GeoDataFrame with barriers, should have variable "line_index".
+    grid2d: xugrid.Ugrid2d
+        Grid to snap lines to
+    vardict_agg: dict
+        Mapping of variable name to aggregation method
+
+    Returns
+    -------
+    snapping_dataset: xugrid.UgridDataset
+        Dataset with all variables snapped and aggregated to cell edges
+    edge_index: numpy.array
+        1D array with indices of cell edges that lines were snapped to
+    """
+    snapping_df = xu.create_snap_to_grid_dataframe(
+        barrier_dataframe, grid2d, max_snap_distance=0.5
+    )
+    # Map other variables to snapping_df with line indices
+    line_index = snapping_df["line_index"]
+    vars_to_snap = list(vardict_agg.keys())
+    if "line_index" in vars_to_snap:
+        vars_to_snap.remove("line_index")  # snapping_df already has line_index
+    for varname in vars_to_snap:
+        snapping_df[varname] = barrier_dataframe[varname].iloc[line_index].to_numpy()
+    # Aggregate to grid edges
+    gb_edge = snapping_df.groupby("edge_index")
+    # Initialize dataset and dataarray with the right shape and dims
+    snapped_dataset = xu.UgridDataset(grids=[grid2d])
+    new = xr.DataArray(np.full(grid2d.n_edge, np.nan), dims=[grid2d.edge_dimension])
+    edge_index = np.array(list(gb_edge.indices.keys()))
+    # Aggregate with different methods per variable
+    for varname, method in vardict_agg.items():
+        snapped_dataset[varname] = new.copy()
+        snapped_dataset[varname].data[edge_index] = gb_edge[varname].aggregate(method)
+
+    return snapped_dataset, edge_index
+
+
 class BarrierType(Enum):
     HydraulicCharacteristic = 0
     Multiplier = 1
@@ -464,10 +513,9 @@ def _to_connected_cells_dataset(
         """
         top, bottom = broadcast_to_full_domain(idomain, top, bottom)
         k = idomain * k
+        # Enforce unstructured
         unstructured_grid, top, bottom, k = (
-            self.__to_unstructured(idomain, top, bottom, k)
-            if isinstance(idomain, xr.DataArray)
-            else [idomain, top, bottom, k]
+            as_ugrid_dataarray(grid) for grid in [idomain, top, bottom, k]
         )
         snapped_dataset, edge_index = self._snap_to_grid(idomain)
         edge_index = self.__remove_invalid_edges(unstructured_grid, edge_index)
@@ -736,61 +784,36 @@ def mask(self, _) -> Package:
         """
         return deepcopy(self)
 
-    @staticmethod
-    def __to_unstructured(
-        idomain: xr.DataArray, top: xr.DataArray, bottom: xr.DataArray, k: xr.DataArray
-    ) -> Tuple[
-        xu.UgridDataArray, xu.UgridDataArray, xu.UgridDataArray, xu.UgridDataArray
-    ]:
-        unstruct = xu.UgridDataArray.from_structured(idomain)
-        top = xu.UgridDataArray.from_structured(top)
-        bottom = xu.UgridDataArray.from_structured(bottom)
-        k = xu.UgridDataArray.from_structured(k)
-
-        return unstruct, top, bottom, k
-
     def _snap_to_grid(
         self, idomain: GridDataArray
     ) -> Tuple[xu.UgridDataset, np.ndarray]:
-        if "layer" in self._get_vertical_variables():
-            variable_names = [self._get_variable_name(), "geometry", "layer"]
+        variable_name = self._get_variable_name()
+        has_layer = "layer" in self._get_vertical_variables()
+        # Create geodataframe with barriers
+        if has_layer:
+            varnames_for_df = [variable_name, "geometry", "layer"]
         else:
-            variable_names = [self._get_variable_name(), "geometry"]
+            varnames_for_df = [variable_name, "geometry"]
         barrier_dataframe = gpd.GeoDataFrame(
-            self.dataset[variable_names].to_dataframe()
+            self.dataset[varnames_for_df].to_dataframe()
         )
-
-        if "layer" not in self._get_vertical_variables():
+        # Convert vertical polygon to linestring
+        if not has_layer:
             lower, _ = _extract_hfb_bounds_from_zpolygons(barrier_dataframe)
             linestring = _create_zlinestring_from_bound_df(lower)
             barrier_dataframe["geometry"] = linestring["geometry"]
-
+        # Clip barriers outside domain
         barrier_dataframe = clip_line_gdf_by_grid(
             barrier_dataframe, idomain.sel(layer=1)
         )
-
-        # Work around issue where xu.snap_to_grid cannot handle snapping a
-        # dataset with multiple lines appropriately. This can be later replaced
-        # to a single call to xu.snap_to_grid if this bug is fixed:
-        # <link_to_github_issue_here>
-        snapped_dataset_rows: List[xr.Dataset] = []
-        for index in barrier_dataframe.index:
-            # Index with list to return pd.DataFrame instead of pd.Series for
-            # xu.snap_to_grid.
-            row_df = barrier_dataframe.loc[[index]]
-            snapped_dataset_row, _ = typing.cast(
-                xu.UgridDataset,
-                xu.snap_to_grid(row_df, grid=idomain, max_snap_distance=0.5),
-            )
-            snapped_dataset_row["line_index"] += index[0]  # Set to original line index.
-            snapped_dataset_rows.append(snapped_dataset_row)
-        snapped_dataset = xu.merge(snapped_dataset_rows)
-
-        edge_index = np.argwhere(
-            snapped_dataset[self._get_variable_name()].notnull().values
-        ).ravel()
-
-        return snapped_dataset, edge_index
+        # Prepare variable names and methods for aggregation
+        vardict_agg = {"line_index": "first", variable_name: "sum"}
+        if has_layer:
+            vardict_agg["layer"] = "first"
+        # Create grid from structured
+        grid2d = as_ugrid_dataarray(idomain.sel(layer=1)).grid
+
+        return _snap_to_grid_and_aggregate(barrier_dataframe, grid2d, vardict_agg)
 
     @staticmethod
     def __remove_invalid_edges(

imod/tests/test_mf6/test_mf6_LHM.py

@@ -0,0 +1,94 @@
+"""
+LHM tests, these are pytest-marked with 'user_acceptance'.
+
+These require the LHM model to be available on the local drive. The tests looks
+for the path to the projectfile needs to be included in a .env file, with the
+environmental variable "LHM_PRJ" with the path to the projectfile.
+"""
+
+import os
+import sys
+
+import pandas as pd
+import pytest
+
+import imod
+from imod.formats.prj.prj import open_projectfile_data
+from imod.logging.config import LoggerType
+from imod.logging.loglevel import LogLevel
+from imod.mf6.oc import OutputControl
+from imod.mf6.regrid.regrid_schemes import (
+    DiscretizationRegridMethod,
+    NodePropertyFlowRegridMethod,
+    StorageCoefficientRegridMethod,
+)
+from imod.mf6.simulation import Modflow6Simulation
+from imod.mf6.utilities.mf6hfb import merge_hfb_packages
+from imod.mf6.write_context import WriteContext
+from imod.prepare.topsystem.default_allocation_methods import (
+    SimulationAllocationOptions,
+    SimulationDistributingOptions,
+)
+
+
+# In function, not a fixture, to allow logging of the import.
+def LHM_imod5_data():
+    lhm_prjfile = os.environ["LHM_PRJ"]
+    data = open_projectfile_data(lhm_prjfile)
+
+    imod5_data = data[0]
+    period_data = data[1]
+    default_simulation_allocation_options = SimulationAllocationOptions
+    default_simulation_distributing_options = SimulationDistributingOptions
+
+    regridding_option = {}
+    regridding_option["npf"] = NodePropertyFlowRegridMethod()
+    regridding_option["dis"] = DiscretizationRegridMethod()
+    regridding_option["sto"] = StorageCoefficientRegridMethod()
+    times = pd.date_range(start="1/1/2018", end="12/1/2018", freq="ME")
+
+    simulation = Modflow6Simulation.from_imod5_data(
+        imod5_data,
+        period_data,
+        default_simulation_allocation_options,
+        default_simulation_distributing_options,
+        times,
+        regridding_option,
+    )
+    simulation["imported_model"]["oc"] = OutputControl(
+        save_head="last", save_budget="last"
+    )
+    return simulation
+
+
+@pytest.mark.user_acceptance
+def test_mf6_LHM_write_HFB(tmp_path):
+    logfile_path = tmp_path / "logfile.txt"
+    with open(logfile_path, "w") as sys.stdout:
+        imod.logging.configure(
+            LoggerType.PYTHON,
+            log_level=LogLevel.DEBUG,
+            add_default_file_handler=False,
+            add_default_stream_handler=True,
+        )
+        simulation = LHM_imod5_data()
+        model = simulation["imported_model"]
+
+        mf6_hfb_ls = []
+        for key, pkg in model.items():
+            if issubclass(type(pkg), imod.mf6.HorizontalFlowBarrierBase):
+                mf6_hfb_ls.append(pkg)
+            pkg.dataset.load()
+
+        top, bottom, idomain = model._Modflow6Model__get_domain_geometry()
+        k = model._Modflow6Model__get_k()
+
+        mf6_hfb = merge_hfb_packages(mf6_hfb_ls, idomain, top, bottom, k)
+
+        times = pd.date_range(start="1/1/2018", end="12/1/2018", freq="ME")
+
+        out_dir = tmp_path / "LHM"
+        out_dir.mkdir(parents=True, exist_ok=True)
+        write_context = WriteContext(out_dir, use_binary=True, use_absolute_paths=False)
+
+        mf6_hfb.write("hfb", times, write_context)

imod/tests/test_mf6/test_mf6_hfb.py

@@ -22,6 +22,7 @@
     _extract_mean_hfb_bounds_from_dataframe,
     _make_linestring_from_polygon,
     _prepare_barrier_dataset_for_mf6_adapter,
+    _snap_to_grid_and_aggregate,
     to_connected_cells_dataset,
 )
 from imod.mf6.ims import SolutionPresetSimple
@@ -697,6 +698,77 @@ def test_hfb_from_imod5(imod5_dataset, tmp_path):
     assert list(np.unique(hfb_package["layer"].values)) == [7]
 
 
+@pytest.mark.usefixtures("structured_flow_model")
+def test_snap_to_grid_and_aggregate(structured_flow_model):
+    idomain = structured_flow_model["dis"]["idomain"]
+    grid2d = xu.Ugrid2d.from_structured(idomain)
+
+    barrier_y = [11.0, 5.0, -1.0]
+    barrier_x = [5.0, 5.0, 5.0]
+    line = linestrings(barrier_x, barrier_y)
+    layer = [1, 1, 1]
+
+    geometry_triple = gpd.GeoDataFrame(
+        geometry=[line, line, line],
+        data={
+            "resistance": [400.0, 400.0, 400.0],
+            "layer": layer,
+            "line_index": [0, 1, 2],
+        },
+    )
+    geometry_triple = geometry_triple.set_index("line_index")
+
+    vardict_agg = {"resistance": "sum", "layer": "first"}
+
+    snapped_dataset, edge_index = _snap_to_grid_and_aggregate(
+        geometry_triple, grid2d, vardict_agg
+    )
+
+    argwhere_summed_expected = np.array([7, 20, 33, 46, 59, 72], dtype=np.int64)
+    argwhere_summed_actual = np.nonzero((snapped_dataset["resistance"] == 1200).values)[
+        0
+    ]
+
+    np.testing.assert_array_equal(argwhere_summed_actual, argwhere_summed_expected)
+    np.testing.assert_array_equal(edge_index, argwhere_summed_expected)
+
+
+@pytest.mark.usefixtures("structured_flow_model")
+def test_combine_linestrings(structured_flow_model):
+    dis = structured_flow_model["dis"]
+    top, bottom, idomain = (
+        dis["top"],
+        dis["bottom"],
+        dis["idomain"],
+    )
+    k = xr.ones_like(idomain)
+
+    barrier_y = [11.0, 5.0, -1.0]
+    barrier_x = [5.0, 5.0, 5.0]
+    line = linestrings(barrier_x, barrier_y)
+
+    geometry_single = gpd.GeoDataFrame(
+        geometry=[line],
+        data={
+            "resistance": [1200.0],
+            "layer": [1],
+        },
+    )
+    geometry_triple = gpd.GeoDataFrame(
+        geometry=[line, line, line],
+        data={
+            "resistance": [400.0, 400.0, 400.0],
+            "layer": [1, 1, 1],
+        },
+    )
+    hfb_single = SingleLayerHorizontalFlowBarrierResistance(geometry_single)
+    hfb_triple = SingleLayerHorizontalFlowBarrierResistance(geometry_triple)
+    mf6_hfb_single = hfb_single.to_mf6_pkg(idomain, top, bottom, k)
+    mf6_hfb_triple = hfb_triple.to_mf6_pkg(idomain, top, bottom, k)
+
+    xr.testing.assert_equal(mf6_hfb_single.dataset, mf6_hfb_triple.dataset)
+
+
 @pytest.mark.usefixtures("structured_flow_model")
 def test_run_multiple_hfbs(tmp_path, structured_flow_model):
     # Single layered model