Issue #1821: Improved performance for allocation & sparse/empty layer handling.

.dvc/config

@@ -1,5 +1,6 @@
-[core]
-    remote = minio
-['remote "minio"']
-    url = s3://imod-python-test-data
-    endpointurl = https://s3.deltares.nl
+[core]
+    remote = minio
+['remote "minio"']
+    url = s3://imod-python-test-data
+    endpointurl = https://s3.deltares.nl
+    allow_anonymous_login = true

docs/api/changelog.rst

@@ -22,6 +22,23 @@ Added
   :meth:`imod.mf6.GroundwaterTransportModel.mask_packages` to mask specific
   packages of a groundwater flow model and a groundwater transport model
   respectively.
+- Added ``drop_empty_layers: bool = True`` to various cell allocation functions 
+  in :mod:`imod.prepare.topsystem.allocation` to remove fully empty layers from the grids.
+  Strips the empty layers before they are passed along to 
+  reprojection/regridding operations, which can save considerable time for models with 
+  many empty layers. Set to False to keep the previous full-layer-coordinate
+  behaviour. :meth:`imod.prepare.topsystem.allocation.allocate_riv_cells`, 
+  :meth:`imod.prepare.topsystem.allocation.allocate_drn_cells`,
+  :meth:`imod.prepare.topsystem.allocation.allocate_ghb_cells`,
+  :meth:`imod.prepare.topsystem.allocation.allocate_rch_cells`
+- Added ``drop_empty_layers: bool = True`` to
+  :meth:`imod.mf6.River.reallocate`, :meth:`imod.mf6.Drainage.reallocate`,
+  :meth:`imod.mf6.GeneralHeadBoundary.reallocate`, and
+  :meth:`imod.mf6.Recharge.reallocate`. Allocation and conductance
+  distribution are always computed over the full layer range first; only
+  the final package has fully empty layers trimmed off afterwards, so this
+  does not affect computed values. Set to False to keep the previous
+  full-layer-coordinate behaviour.
 
 Fixed
 ~~~~~
@@ -35,6 +52,11 @@ Fixed
   :class:`imod.mf6.Recharge`) where IBOUND is less than 0.
 - :meth:`imod.msw.MetaSwapModel.from_imod5_data` now masks cells where IBOUND is
   less than 0.
+- Fixed :func:`imod.prepare.cleanup.align_interface_levels` (used by
+  ``cleanup_riv``, and therefore :meth:`imod.mf6.River.cleanup`) raising an
+  alignment error when a package's own layer coordinate is a subset of the
+  model's full layer range, e.g. after :meth:`imod.mf6.River.reallocate` with
+  ``drop_empty_layers=True``.
 - :class:`imod.msw.FileCopier` and :class:`imod.msw.MeteoGridCopy` now force
   paths to be stored as strings in the dataset. ``pathlib.Path`` objects could
   cause errors when calling :meth:`imod.msw.MetaSwapModel.dump`.

imod/mf6/drn.py

@@ -17,7 +17,11 @@
 from imod.mf6.utilities.package import set_repeat_stress_if_available
 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA
 from imod.prepare.cleanup import cleanup_drn
-from imod.prepare.topsystem.allocation import ALLOCATION_OPTION, allocate_drn_cells
+from imod.prepare.topsystem.allocation import (
+    ALLOCATION_OPTION,
+    allocate_drn_cells,
+    drop_empty_layers_from_dict,
+)
 from imod.prepare.topsystem.conductance import (
     DISTRIBUTING_OPTION,
     distribute_drn_conductance,
@@ -193,6 +197,7 @@ def _allocate_and_distribute_planar_data(
         npf: NodePropertyFlow,
         allocation_option: ALLOCATION_OPTION,
         distributing_option: DISTRIBUTING_OPTION,
+        drop_empty_layers: bool = True,
     ) -> dict[str, GridDataArray]:
         """
         Allocate and distribute planar data for given discretization and npf
@@ -214,6 +219,13 @@ def _allocate_and_distribute_planar_data(
             ALLOCATION_OPTION.at_first_active.
         distributing_option: DISTRIBUTING_OPTION
             distributing option.
+        drop_empty_layers: bool, default True
+            If True, drop layers that contain no allocated cells anywhere in the
+            domain (or at any time), so the returned grids only span the layers that
+            are actually used. Reduces memory use and speeds up later operations such
+            as regridding, clipping and splitting. If no cells are allocated in any
+            layer, nothing is dropped and the full layer range is returned, as a layer
+            dimension of size 0 is not valid.
 
         Returns
         -------
@@ -240,6 +252,7 @@ def _allocate_and_distribute_planar_data(
             top,
             bottom,
             planar_data["elevation"],
+            drop_empty_layers=False,  # Keep full layer range, drop empty layers below
         )
         layered_data = {}
         layered_data["conductance"] = distribute_drn_conductance(
@@ -253,6 +266,10 @@ def _allocate_and_distribute_planar_data(
         )
         layered_data["elevation"] = planar_data["elevation"].where(drn_allocation)
         layered_data["elevation"] = enforce_dim_order(layered_data["elevation"])
+
+        if drop_empty_layers:
+            layered_data = drop_empty_layers_from_dict(layered_data, drn_allocation)
+
         return layered_data
 
     @classmethod

imod/mf6/ghb.py

@@ -19,7 +19,11 @@
 from imod.mf6.utilities.package import set_repeat_stress_if_available
 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA
 from imod.prepare.cleanup import cleanup_ghb
-from imod.prepare.topsystem.allocation import ALLOCATION_OPTION, allocate_ghb_cells
+from imod.prepare.topsystem.allocation import (
+    ALLOCATION_OPTION,
+    allocate_ghb_cells,
+    drop_empty_layers_from_dict,
+)
 from imod.prepare.topsystem.conductance import (
     DISTRIBUTING_OPTION,
     distribute_ghb_conductance,
@@ -198,6 +202,7 @@ def _allocate_and_distribute_planar_data(
         npf: NodePropertyFlow,
         allocation_option: ALLOCATION_OPTION,
         distributing_option: DISTRIBUTING_OPTION,
+        drop_empty_layers: bool = True,
     ) -> dict[str, GridDataArray]:
         """
         Allocate and distribute planar data for given discretization and npf
@@ -219,6 +224,13 @@ def _allocate_and_distribute_planar_data(
             ALLOCATION_OPTION.at_first_active.
         distributing_option: DISTRIBUTING_OPTION
             distributing option.
+        drop_empty_layers: bool, default True
+            If True, drop layers that contain no allocated cells anywhere in the
+            domain (or at any time), so the returned grids only span the layers that
+            are actually used. Reduces memory use and speeds up later operations such
+            as regridding, clipping and splitting. If no cells are allocated in any
+            layer, nothing is dropped and the full layer range is returned, as a layer
+            dimension of size 0 is not valid.
 
         Returns
         -------
@@ -245,6 +257,7 @@ def _allocate_and_distribute_planar_data(
             top,
             bottom,
             planar_data["head"],
+            drop_empty_layers=False,  # Keep full layer range, drop empty layers below
         )
 
         layered_data = {}
@@ -259,6 +272,8 @@ def _allocate_and_distribute_planar_data(
             bottom,
             npf.dataset["k"],
         )
+        if drop_empty_layers:
+            layered_data = drop_empty_layers_from_dict(layered_data, ghb_allocation)
         return layered_data
 
     @classmethod

imod/mf6/rch.py

@@ -25,7 +25,11 @@
     get_cell_area_from_imod5_data,
     is_msw_active_cell,
 )
-from imod.prepare.topsystem.allocation import ALLOCATION_OPTION, allocate_rch_cells
+from imod.prepare.topsystem.allocation import (
+    ALLOCATION_OPTION,
+    allocate_rch_cells,
+    drop_empty_layers_from_dict,
+)
 from imod.schemata import (
     AllCoordsValueSchema,
     AllInsideNoDataSchema,
@@ -182,6 +186,7 @@ def _allocate_planar_data(
         planar_data: dict[str, GridDataArray],
         dis: StructuredDiscretization | VerticesDiscretization,
         allocation_option: ALLOCATION_OPTION,
+        drop_empty_layers: bool = True,
     ) -> dict[str, GridDataArray]:
         """
         Allocate and distribute planar data for given discretization and npf
@@ -196,6 +201,13 @@ def _allocate_planar_data(
             Model discretization package.
         allocation_option: ALLOCATION_OPTION
             The allocation option to use for the reallocation.
+        drop_empty_layers: bool, default True
+            If True, drop layers that contain no allocated cells anywhere in the
+            domain (or at any time), so the returned grids only span the layers that
+            are actually used. Reduces memory use and speeds up later operations such
+            as regridding, clipping and splitting. If no cells are allocated in any
+            layer, nothing is dropped and the full layer range is returned, as a layer
+            dimension of size 0 is not valid.
 
         Returns
         -------
@@ -210,11 +222,14 @@ def _allocate_planar_data(
             allocation_option,
             idomain > 0,
             planar_data["rate"],
+            drop_empty_layers=False,  # Keep full here, drop empty layers below
         )
         # remove rch from cells where it is not allocated and broadcast over layers.
         layered_data = {}
         layered_data["rate"] = planar_data["rate"].where(is_rch_cell)
         layered_data["rate"] = enforce_dim_order(layered_data["rate"])
+        if drop_empty_layers:
+            layered_data = drop_empty_layers_from_dict(layered_data, is_rch_cell)
         return layered_data
 
     @classmethod

imod/mf6/riv.py

@@ -19,7 +19,11 @@
 from imod.mf6.utilities.package import set_repeat_stress_if_available
 from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA
 from imod.prepare.cleanup import AlignLevelsMode, align_interface_levels, cleanup_riv
-from imod.prepare.topsystem.allocation import ALLOCATION_OPTION, allocate_riv_cells
+from imod.prepare.topsystem.allocation import (
+    ALLOCATION_OPTION,
+    allocate_riv_cells,
+    drop_empty_layers_from_dict,
+)
 from imod.prepare.topsystem.conductance import (
     DISTRIBUTING_OPTION,
     distribute_drn_conductance,
@@ -317,6 +321,7 @@ def _allocate_and_distribute_planar_data(
         npf: NodePropertyFlow,
         allocation_option: ALLOCATION_OPTION,
         distributing_option: DISTRIBUTING_OPTION,
+        drop_empty_layers: bool = True,
     ) -> tuple[GridDataDict, GridDataDict]:
         """
         Allocate and distribute planar data for given discretization and npf
@@ -338,6 +343,13 @@ def _allocate_and_distribute_planar_data(
             ALLOCATION_OPTION.at_first_active.
         distributing_option: DISTRIBUTING_OPTION
             distributing option.
+        drop_empty_layers: bool, default True
+            If True, drop layers that contain no allocated cells anywhere in the
+            domain (or at any time), so the returned grids only span the layers that
+            are actually used. Reduces memory use and speeds up later operations such
+            as regridding, clipping and splitting. If no cells are allocated in any
+            layer, nothing is dropped and the full layer range is returned, as a layer
+            dimension of size 0 is not valid.
 
         Returns
         -------
@@ -364,6 +376,7 @@ def _allocate_and_distribute_planar_data(
             bottom,
             planar_data["stage"],
             planar_data["bottom_elevation"],
+            drop_empty_layers=False,  # Keep full layer range, drop empty layers below
         )
         drn_is_allocated = drn_allocated is not None
         # Distribution of conductances
@@ -409,6 +422,15 @@ def _allocate_and_distribute_planar_data(
             layered_data_riv["bottom_elevation"], bottom
         )
 
+        if drop_empty_layers:
+            layered_data_riv = drop_empty_layers_from_dict(
+                layered_data_riv, riv_allocated
+            )
+            if drn_allocated is not None:
+                layered_data_drn = drop_empty_layers_from_dict(
+                    layered_data_drn, drn_allocated
+                )
+
         return layered_data_riv, layered_data_drn
 
     @classmethod

imod/mf6/topsystem.py

@@ -58,6 +58,7 @@ def reallocate(
         npf: Optional[NodePropertyFlow] = None,
         allocation_option: Optional[ALLOCATION_OPTION] = None,
         distributing_option: Optional[DISTRIBUTING_OPTION] = None,
+        drop_empty_layers: bool = True,
     ) -> Self:
         """
         Reallocates topsystem data across layers and create new package with it.
@@ -83,6 +84,13 @@ def reallocate(
             The distributing option to use for the reallocation. Required for
             packages with a conductance variable. If None, the default is taken
             from :class:`imod.prepare.SimulationDistributingOptions`.
+        drop_empty_layers: bool, default True
+            If True, drop layers that contain no allocated cells anywhere in the
+            domain (or at any time), so the returned grids only span the layers that
+            are actually used. Reduces memory use and speeds up later operations such
+            as regridding, clipping and splitting. If no cells are allocated in any
+            layer, nothing is dropped and the full layer range is returned, as a layer
+            dimension of size 0 is not valid.
 
         Returns
         -------
@@ -102,10 +110,17 @@ def reallocate(
             npf = cast(NodePropertyFlow, npf)
             distributing_option = cast(DISTRIBUTING_OPTION, distributing_option)
             grid_dict = self._allocate_and_distribute_planar_data(
-                planar_data, dis, npf, allocation_option, distributing_option
+                planar_data,
+                dis,
+                npf,
+                allocation_option,
+                distributing_option,
+                drop_empty_layers,
             )
         else:
-            grid_dict = self._allocate_planar_data(planar_data, dis, allocation_option)
+            grid_dict = self._allocate_planar_data(
+                planar_data, dis, allocation_option, drop_empty_layers
+            )
         # River package returns a tuple (second argument can also be Drainage
         # package)
         if isinstance(grid_dict, tuple):
@@ -122,6 +137,7 @@ def _allocate_and_distribute_planar_data(
         npf: NodePropertyFlow,
         allocation_option: ALLOCATION_OPTION,
         distributing_option: DISTRIBUTING_OPTION,
+        drop_empty_layers: bool = True,
     ) -> tuple[GridDataDict, GridDataDict] | GridDataDict:
         raise NotImplementedError(
             "This method should be implemented in the specific boundary condition "
@@ -134,6 +150,7 @@ def _allocate_planar_data(
         planar_data: GridDataDict,
         dis: StructuredDiscretization | VerticesDiscretization,
         allocation_option: ALLOCATION_OPTION,
+        drop_empty_layers: bool = True,
     ) -> tuple[GridDataDict, GridDataDict] | GridDataDict:
         raise NotImplementedError(
             "This method should be implemented in the specific boundary condition "

imod/prepare/cleanup.py

@@ -27,6 +27,16 @@ def align_interface_levels(
     bottom: GridDataArray,
     method: AlignLevelsMode = AlignLevelsMode.TOPDOWN,
 ) -> tuple[GridDataArray, GridDataArray]:
+    # `bottom` (e.g. a model's full layer range) may have more layers than
+    # `top` (e.g. a package trimmed to only its allocated layers, see
+    # ``drop_empty_layers`` in ``imod.prepare.topsystem``). Reindex `bottom`
+    # down to `top`'s own layers first, so the comparison below doesn't fail
+    # with an alignment error; `top`'s layers are always the ones we want to
+    # keep, matching the ``join="left"`` pattern used in
+    # ``imod.common.utilities.mask.mask_da``.
+    if "layer" in top.dims and "layer" in bottom.dims:
+        bottom = bottom.sel(layer=top["layer"])
+
     to_align = top < bottom
 
     match method: