Issue #1158 remaining hfb bottlenecks (#1159)

Fixes #1158 

# Description
Fix remaining HFB bottlenecks. This reduces writing the HFB package from
the LHM from 12.5 minutes to 2 minutes.

- Cache results ``xu.Ugrid2d.from_structured``, as this is a costly
operation
- Use ``and`` instead of ``&`` operator in the ``scalar_None`` function,
to enable shortcutting.
- Remove ``mask_all_packages`` function in ``from_imod5_data``, call in
tests.
- Create separate pixi task for slow unittests, where just in time
compilation is enabled. ``pixi run unittests`` still runs all unittests,
by starting two pixi tasks ``unittests_njit`` & ``unittests_jit``.

# 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

.gitignore

@@ -141,5 +141,4 @@ examples/data
 .pixi
 
 /imod/tests/mydask.png
-/imod/tests/unittest_report.xml
-/imod/tests/examples_report.xml
+/imod/tests/*_report.xml

docs/api/changelog.rst

@@ -17,6 +17,8 @@ Fixed
   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.
+- Improve performance validation upon Package initialization
+- Improve performance writing ``HorizontalFlowBarrier`` objects
 
 Changed
 ~~~~~~~

imod/mf6/simulation.py

@@ -1378,8 +1378,5 @@ def from_imod5_data(
         )
         simulation["ims"] = solution
 
-        # cleanup packages for validation
-        idomain = groundwaterFlowModel.domain
-        simulation.mask_all_models(idomain)
         simulation.create_time_discretization(additional_times=times)
         return simulation

imod/schemata.py

@@ -74,7 +74,7 @@ def scalar_None(obj):
     if not isinstance(obj, (xr.DataArray, xu.UgridDataArray)):
         return False
     else:
-        return (len(obj.shape) == 0) & (~obj.notnull()).all()
+        return (len(obj.shape) == 0) and (obj.isnull()).all()
 
 
 def align_other_obj_with_coords(

imod/tests/test_mf6/test_mf6_chd.py

@@ -238,6 +238,7 @@ def test_from_imod5_shd(imod5_dataset, tmp_path):
     chd_shd.write("chd_shd", [1], write_context)
 
 
+@pytest.mark.unittest_jit
 @pytest.mark.parametrize("remove_merged_packages", [True, False])
 @pytest.mark.usefixtures("imod5_dataset")
 def test_concatenate_chd(imod5_dataset, tmp_path, remove_merged_packages):

imod/tests/test_mf6/test_mf6_simulation.py

@@ -476,6 +476,7 @@ def compare_submodel_partition_info(first: PartitionInfo, second: PartitionInfo)
     )
 
 
+@pytest.mark.unittest_jit
 @pytest.mark.usefixtures("imod5_dataset")
 def test_import_from_imod5(imod5_dataset, tmp_path):
     imod5_data = imod5_dataset[0]
@@ -495,18 +496,20 @@ def test_import_from_imod5(imod5_dataset, tmp_path):
     simulation["imported_model"]["oc"] = OutputControl(
         save_head="last", save_budget="last"
     )
-
     simulation.create_time_discretization(["01-01-2003", "02-01-2003"])
-
+    # Cleanup
     # Remove HFB packages outside domain
     # TODO: Build in support for hfb packages outside domain
     for hfb_outside in ["hfb-24", "hfb-26"]:
         simulation["imported_model"].pop(hfb_outside)
-
+    # Align NoData to domain
+    idomain = simulation["imported_model"].domain
+    simulation.mask_all_models(idomain)
     # write and validate the simulation.
     simulation.write(tmp_path, binary=False, validate=True)
 
 
+@pytest.mark.unittest_jit
 @pytest.mark.usefixtures("imod5_dataset")
 def test_import_from_imod5__correct_well_type(imod5_dataset):
     # Unpack
@@ -537,6 +540,7 @@ def test_import_from_imod5__correct_well_type(imod5_dataset):
     assert isinstance(simulation["imported_model"]["wel-WELLS_L5"], LayeredWell)
 
 
+@pytest.mark.unittest_jit
 @pytest.mark.usefixtures("imod5_dataset")
 def test_import_from_imod5__nonstandard_regridding(imod5_dataset, tmp_path):
     imod5_data = imod5_dataset[0]
@@ -558,22 +562,23 @@ def test_import_from_imod5__nonstandard_regridding(imod5_dataset, tmp_path):
         times,
         regridding_option,
     )
-
     simulation["imported_model"]["oc"] = OutputControl(
         save_head="last", save_budget="last"
     )
-
     simulation.create_time_discretization(["01-01-2003", "02-01-2003"])
-
+    # Cleanup
     # Remove HFB packages outside domain
     # TODO: Build in support for hfb packages outside domain
     for hfb_outside in ["hfb-24", "hfb-26"]:
         simulation["imported_model"].pop(hfb_outside)
-
+    # Align NoData to domain
+    idomain = simulation["imported_model"].domain
+    simulation.mask_all_models(idomain)
     # write and validate the simulation.
     simulation.write(tmp_path, binary=False, validate=True)
 
 
+@pytest.mark.unittest_jit
 @pytest.mark.usefixtures("imod5_dataset")
 def test_import_from_imod5_no_storage_no_recharge(imod5_dataset, tmp_path):
     # this test imports an imod5 simulation, but it has no recharge and no storage package.
@@ -594,23 +599,22 @@ def test_import_from_imod5_no_storage_no_recharge(imod5_dataset, tmp_path):
         default_simulation_distributing_options,
         times,
     )
-
     simulation["imported_model"]["oc"] = OutputControl(
         save_head="last", save_budget="last"
     )
-
     simulation.create_time_discretization(["01-01-2003", "02-01-2003"])
-
+    # Cleanup
     # Remove HFB packages outside domain
     # TODO: Build in support for hfb packages outside domain
     for hfb_outside in ["hfb-24", "hfb-26"]:
         simulation["imported_model"].pop(hfb_outside)
-
     # check storage is present and rch is absent
     assert not simulation["imported_model"]["sto"].dataset["transient"].values[()]
     package_keys = simulation["imported_model"].keys()
     for key in package_keys:
         assert key[0:3] != "rch"
-
+    # Align NoData to domain
+    idomain = simulation["imported_model"].domain
+    simulation.mask_all_models(idomain)
     # write and validate the simulation.
     simulation.write(tmp_path, binary=False, validate=True)

imod/tests/test_typing/test_typing_grid.py

@@ -1,7 +1,11 @@
+import numpy as np
 import xarray as xr
 import xugrid as xu
 
 from imod.typing.grid import (
+    UGRID2D_FROM_STRUCTURED_CACHE,
+    GridCache,
+    as_ugrid_dataarray,
     enforce_dim_order,
     is_planar_grid,
     is_spatial_grid,
@@ -145,3 +149,73 @@ def test_merge_dictionary__unstructured(basic_unstructured_dis):
     assert isinstance(uds["bottom"], xr.DataArray)
     assert uds["ibound"].dims == ("layer", "mesh2d_nFaces")
     assert uds["bottom"].dims == ("layer",)
+
+
+def test_as_ugrid_dataarray__structured(basic_dis):
+    # Arrange
+    ibound, top, bottom = basic_dis
+    top_3d = top * ibound
+    bottom_3d = bottom * ibound
+    # Clear cache
+    UGRID2D_FROM_STRUCTURED_CACHE.clear()
+    # Act
+    ibound_disv = as_ugrid_dataarray(ibound)
+    top_disv = as_ugrid_dataarray(top_3d)
+    bottom_disv = as_ugrid_dataarray(bottom_3d)
+    # Assert
+    # Test types
+    assert isinstance(ibound_disv, xu.UgridDataArray)
+    assert isinstance(top_disv, xu.UgridDataArray)
+    assert isinstance(bottom_disv, xu.UgridDataArray)
+    # Test cache proper size
+    assert len(UGRID2D_FROM_STRUCTURED_CACHE.grid_cache) == 1
+    # Test that data is different
+    assert np.all(ibound_disv != top_disv)
+    assert np.all(top_disv != bottom_disv)
+    # Test that grid is equal
+    assert np.all(ibound_disv.grid == top_disv.grid)
+    assert np.all(top_disv.grid == bottom_disv.grid)
+
+
+def test_as_ugrid_dataarray__unstructured(basic_unstructured_dis):
+    # Arrange
+    ibound, top, bottom = basic_unstructured_dis
+    top_3d = enforce_dim_order(ibound * top)
+    bottom_3d = enforce_dim_order(ibound * bottom)
+    # Clear cache
+    UGRID2D_FROM_STRUCTURED_CACHE.clear()
+    # Act
+    ibound_disv = as_ugrid_dataarray(ibound)
+    top_disv = as_ugrid_dataarray(top_3d)
+    bottom_disv = as_ugrid_dataarray(bottom_3d)
+    # Assert
+    # Test types
+    assert isinstance(ibound_disv, xu.UgridDataArray)
+    assert isinstance(top_disv, xu.UgridDataArray)
+    assert isinstance(bottom_disv, xu.UgridDataArray)
+    assert len(UGRID2D_FROM_STRUCTURED_CACHE.grid_cache) == 0
+
+
+def test_ugrid2d_cache(basic_dis):
+    # Arrange
+    ibound, _, _ = basic_dis
+    # Act
+    cache = GridCache(xu.Ugrid2d.from_structured, max_cache_size=3)
+    for i in range(5):
+        ugrid2d = cache.get_grid(ibound[:, i:, :])
+    # Assert
+    # Test types
+    assert isinstance(ugrid2d, xu.Ugrid2d)
+    # Test cache proper size
+    assert cache.max_cache_size == 3
+    assert len(cache.grid_cache) == 3
+    # Check if smallest grid in last cache list by checking if amount of faces
+    # correct
+    expected_size = ibound[0, i:, :].size
+    keys = list(cache.grid_cache.keys())
+    last_ugrid = cache.grid_cache[keys[-1]]
+    actual_size = last_ugrid.n_face
+    assert expected_size == actual_size
+    # Test clear cache
+    cache.clear()
+    assert len(cache.grid_cache) == 0

imod/typing/grid.py

@@ -435,10 +435,58 @@ def is_transient_data_grid(
     return False
 
 
+class GridCache:
+    """
+    Cache grids in this object for a specific function, lookup grids based on
+    unique geometry hash.
+    """
+
+    def __init__(self, func: Callable, max_cache_size=5):
+        self.max_cache_size = max_cache_size
+        self.grid_cache: dict[int, GridDataArray] = {}
+        self.func = func
+
+    def get_grid(self, grid: GridDataArray):
+        geom_hash = get_grid_geometry_hash(grid)
+        if geom_hash not in self.grid_cache.keys():
+            if len(self.grid_cache.keys()) >= self.max_cache_size:
+                self.remove_first()
+            self.grid_cache[geom_hash] = self.func(grid)
+        return self.grid_cache[geom_hash]
+
+    def remove_first(self):
+        keys = list(self.grid_cache.keys())
+        self.grid_cache.pop(keys[0])
+
+    def clear(self):
+        self.grid_cache = {}
+
+
+UGRID2D_FROM_STRUCTURED_CACHE = GridCache(xu.Ugrid2d.from_structured)
+
+
 @typedispatch
 def as_ugrid_dataarray(grid: xr.DataArray) -> xu.UgridDataArray:
-    """Enforce GridDataArray to UgridDataArray"""
-    return xu.UgridDataArray.from_structured(grid)
+    """
+    Enforce GridDataArray to UgridDataArray, calls
+    xu.UgridDataArray.from_structured, which is a costly operation. Therefore
+    cache results.
+    """
+
+    topology = UGRID2D_FROM_STRUCTURED_CACHE.get_grid(grid)
+
+    # Copied from:
+    # https://github.com/Deltares/xugrid/blob/3dee693763da1c4c0859a4f53ac38d4b99613a33/xugrid/core/wrap.py#L236
+    # Note that "da" is renamed to "grid" and "grid" to "topology"
+    dims = grid.dims[:-2]
+    coords = {k: grid.coords[k] for k in dims}
+    face_da = xr.DataArray(
+        grid.data.reshape(*grid.shape[:-2], -1),
+        coords=coords,
+        dims=[*dims, topology.face_dimension],
+        name=grid.name,
+    )
+    return xu.UgridDataArray(face_da, topology)
 
 
 @typedispatch  # type: ignore[no-redef]

pixi.toml

@@ -19,11 +19,12 @@ install_with_deps = "python -m pip install --editable ."
 format = "ruff check --fix .; ruff format ."
 lint = "ruff check . ; ruff format --check ."
 tests = { depends_on = ["unittests", "examples"] }
-unittests = { cmd = [
+unittests = { depends_on = ["unittests_njit", "unittests_jit"] }
+unittests_njit = { cmd = [
     "NUMBA_DISABLE_JIT=1",
     "pytest",
     "-n", "auto",
-    "-m", "not example and not user_acceptance",
+    "-m", "not example and not user_acceptance and not unittest_jit",
     "--cache-clear",
     "--verbose",
     "--junitxml=unittest_report.xml",
@@ -32,6 +33,14 @@ unittests = { cmd = [
     "--cov-report=html:coverage",
     "--cov-config=.coveragerc"
 ], depends_on = ["install"], cwd = "imod/tests" }
+unittests_jit = { cmd = [
+    "pytest",
+    "-n", "auto",
+    "-m", "unittest_jit",
+    "--cache-clear",
+    "--verbose",
+    "--junitxml=unittest_jit_report.xml",
+], depends_on = ["install"], cwd = "imod/tests" }
 examples = { cmd = [
     "NUMBA_DISABLE_JIT=1",
     "pytest",

pyproject.toml

@@ -158,6 +158,7 @@ ignore_missing_imports = true
 markers = [
     "example: marks test as example (deselect with '-m \"not example\"')",
     "user_acceptance: marks user acceptance tests (deselect with '-m \"not user_acceptance\"')",
+    "unittest_jit: marks unit tests that should be jitted (deselect with '-m \"not unittest_jit\"')"
 ]
 
 [tool.hatch.version]