Repository: Deltares/imod-python
Issue category: Performance Issue

Bottlenecks writing HFBs LHM
https://github.com/Deltares/imod-python/pull/1157 considerably improves the speed with which the LHM is written: From 34 minutes to 12.5 minutes. This is still slower than expected however. Based on profiling, I've identified 3 main bottlenecks:

- 4.5 minutes: ``xu.DataArray.from_structured`` is called 70 times. This is quite a computationally intensive operation. It can probably be reworked to be called only once.
- 3 minutes are spent in ``DTypeSchema.validate``. This make as call to the ``scalar_None`` function, which is supposed to shortcut if len(da.shape)  != 0, but because I used the ``&`` operator accidentily instead of ``and``, this is not happening.
- 3 minutes: Spent in ``mask_all_models``. This is called because the automatic regridding might introduce inconsistencies in where idomain is located and where data is specified.  I think we can remove this function, because:
     1. in this case no regridding takes place and the masking is unnecessary, besides I see that ``_regrid_like(model: IModel, *args)`` also does masking?
     2. automatic masking can hide issues already present in the original iMOD5 model. 
     3. Regridding might introduce other issues as well: For example, model tops and bottoms might be regridded in such a manner that the bottom exceeds top. Automatically cleaning up all these potentially introduced artifacts is excessive. Probably better is to create cleanup utilities and guide the user to these in the ValidationError messages. 

So tasks:

- [x] Ensure  ``xu.DataArray.from_structured`` is only called once
- [x] Make sure ``scalar_None`` shortcuts appropriately
- [x] Remove call to mask_all_models


Functions to edit:
imod/mf6/simulation.py:Modflow6Simulation.from_imod5_data
imod/schemata.py:scalar_None
imod/typing/grid.py:as_ugrid_dataarray
Functions added: imod/typing/grid.py:GridCache
imod/typing/grid.py:GridCache.__init__
imod/typing/grid.py:GridCache.get_grid
imod/typing/grid.py:GridCache.remove_first
imod/typing/grid.py:GridCache.clear

Patch:
diff --git a/.gitignore b/.gitignore
index fd95f7d8d..28f9b76a9 100644
--- a/.gitignore
+++ b/.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
diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst
index 95985f29f..525f330ef 100644
--- a/docs/api/changelog.rst
+++ b/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
 ~~~~~~~
diff --git a/imod/mf6/simulation.py b/imod/mf6/simulation.py
index 810a42a1a..54613eb3c 100644
--- a/imod/mf6/simulation.py
+++ b/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
diff --git a/imod/schemata.py b/imod/schemata.py
index 95b5d6a3a..fd9d8d737 100644
--- a/imod/schemata.py
+++ b/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(
diff --git a/imod/typing/grid.py b/imod/typing/grid.py
index 3381b32ae..509da6aef 100644
--- a/imod/typing/grid.py
+++ b/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]
diff --git a/pixi.toml b/pixi.toml
index eac3c0f4d..e7c867b60 100644
--- a/pixi.toml
+++ b/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",
diff --git a/pyproject.toml b/pyproject.toml
index dcaaaa437..c265cabe8 100644
--- a/pyproject.toml
+++ b/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]

Test patch: diff --git a/imod/tests/test_mf6/test_mf6_chd.py b/imod/tests/test_mf6/test_mf6_chd.py
index ef0a04108..c3b96b982 100644
--- a/imod/tests/test_mf6/test_mf6_chd.py
+++ b/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):
diff --git a/imod/tests/test_mf6/test_mf6_simulation.py b/imod/tests/test_mf6/test_mf6_simulation.py
index b60643793..b9b78fdcb 100644
--- a/imod/tests/test_mf6/test_mf6_simulation.py
+++ b/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)
diff --git a/imod/tests/test_typing/test_typing_grid.py b/imod/tests/test_typing/test_typing_grid.py
index a1129940b..ff7f398b0 100644
--- a/imod/tests/test_typing/test_typing_grid.py
+++ b/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