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

docs/api/changelog.rst

@@ -18,6 +18,13 @@ Added
   :meth:`imod.msw.SprinklingPoints.from_imod5_data`.
 - :class:`imod.mf6.LayeredWell.from_imod5_cap_data` now also supports loading
   wells from IPF files in an iMOD5 CAP dataset.
+- Added ``drop_empty_layers: bool = False`` to various cell allocation functions 
+  in :mod:`imod.prepare.topsystem.allocation` to remove fully empty layers from the grids.
+  Setting this to True, strips the empty layers before they are passed along to 
+  reprojection/regridding operations. Which can save considerable time for models with 
+  many empty layers. :meth:`imod.prepare.topsystem.allocation.allocate_riv_cells`, 
+  :meth:`imod.prepare.topsystem.allocation.allocate_drn_cells`,
+  :meth:`imod.prepare.topsystem.allocation.allocate_rch_cells`
 
 Fixed
 ~~~~~

imod/prepare/layerregrid.py

@@ -12,14 +12,32 @@
 
 
 @numba.njit(cache=True)
-def _regrid_layers(src, dst, src_top, dst_top, src_bot, dst_bot, method):
+def _valid_layer_indices(top_col, bot_col, out):
     """
-    Maps one set of layers unto the other.
+    Fill `out` (int64 array, same length as top_col) with the indices of
+    layers that have non-nan top AND bottom, at a single (row, col) column.
+    Returns the count of valid entries. Avoids allocating a new array per
+    column call.
     """
+    count = 0
+    n = top_col.shape[0]
+    for k in range(n):
+        if not (np.isnan(top_col[k]) or np.isnan(bot_col[k])):
+            out[count] = k
+            count += 1
+    return count
+
+
+@numba.njit(cache=True)
+def _regrid_layers(src, dst, src_top, dst_top, src_bot, dst_bot, method):
+    """Maps one set of layers onto the other."""
     nlayer_src, nrow, ncol = src.shape
     nlayer_dst = dst.shape[0]
+
     values = np.zeros(nlayer_src)
     weights = np.zeros(nlayer_src)
+    src_valid_idx = np.empty(nlayer_src, dtype=np.int64)
+    dst_valid_idx = np.empty(nlayer_dst, dtype=np.int64)
 
     for i in range(nrow):
         for j in range(ncol):
@@ -28,23 +46,27 @@ def _regrid_layers(src, dst, src_top, dst_top, src_bot, dst_bot, method):
             src_b = src_bot[:, i, j]
             dst_b = dst_bot[:, i, j]
 
-            # ii is index of dst
-            for ii in range(nlayer_dst):
+            # Precompute valid layer indices ONCE per column, instead of
+            # re-checking isnan for every (ii, jj) pair.
+            n_src_valid = _valid_layer_indices(src_t, src_b, src_valid_idx)
+            if n_src_valid == 0:
+                continue
+            n_dst_valid = _valid_layer_indices(dst_t, dst_b, dst_valid_idx)
+            if n_dst_valid == 0:
+                continue
+
+            for di in range(n_dst_valid):
+                ii = dst_valid_idx[di]
                 dt = dst_t[ii]
                 db = dst_b[ii]
-                if np.isnan(dt) or np.isnan(db):
-                    continue
 
                 count = 0
                 has_value = False
-                # jj is index of src
-                for jj in range(nlayer_src):
+                for sj in range(n_src_valid):
+                    jj = src_valid_idx[sj]
                     st = src_t[jj]
                     sb = src_b[jj]
 
-                    if np.isnan(st) or np.isnan(sb):
-                        continue
-
                     overlap = common._overlap((db, dt), (sb, st))
                     if overlap == 0:
                         continue
@@ -53,12 +75,11 @@ def _regrid_layers(src, dst, src_top, dst_top, src_bot, dst_bot, method):
                     values[count] = src[jj, i, j]
                     weights[count] = overlap
                     count += 1
-                else:
-                    if has_value:
-                        dst[ii, i, j] = method(values, weights)
-                        # Reset
-                        values[:count] = 0
-                        weights[:count] = 0
+
+                if has_value:
+                    dst[ii, i, j] = method(values, weights)
+                    values[:count] = 0
+                    weights[:count] = 0
 
     return dst
 

imod/prepare/topsystem/allocation.py

@@ -71,6 +71,7 @@ def allocate_riv_cells(
     bottom: GridDataArray,
     stage: GridDataArray,
     bottom_elevation: GridDataArray,
+    drop_empty_layers: bool = False,
 ) -> tuple[GridDataArray, Optional[GridDataArray]]:
     """
     Allocate river cells from a planar grid across the vertical dimension.
@@ -96,6 +97,14 @@ def allocate_riv_cells(
     bottom_elevation: DataArray | UgridDatarray
         Planar grid containing river bottom elevations. Is not allowed to have a
         layer dimension.
+    drop_empty_layers: bool, default False
+        If True, drop layers from the result that contain no allocated
+        cells anywhere in the domain. This avoids carrying the package's
+        arrays at full model-layer size through downstream regridding,
+        clipping, masking, and splitting, which can otherwise become
+        expensive for models with many layers relative to how many
+        layers the topsystem package actually occupies. Set to False to
+        keep the previous full-layer-coordinate behaviour.
 
     Returns
     -------
@@ -111,21 +120,25 @@ def allocate_riv_cells(
     """
     match allocation_option:
         case ALLOCATION_OPTION.stage_to_riv_bot:
-            return _allocate_cells__stage_to_riv_bot(
+            riv_cells, drn_cells = _allocate_cells__stage_to_riv_bot(
                 top, bottom, stage, bottom_elevation
             )
         case ALLOCATION_OPTION.first_active_to_elevation:
-            return _allocate_cells__first_active_to_elevation(
+            riv_cells, drn_cells = _allocate_cells__first_active_to_elevation(
                 active, top, bottom, bottom_elevation
             )
         case ALLOCATION_OPTION.stage_to_riv_bot_drn_above:
-            return _allocate_cells__stage_to_riv_bot_drn_above(
+            riv_cells, drn_cells = _allocate_cells__stage_to_riv_bot_drn_above(
                 active, top, bottom, stage, bottom_elevation
             )
         case ALLOCATION_OPTION.at_elevation:
-            return _allocate_cells__at_elevation(top, bottom, bottom_elevation)
+            riv_cells, drn_cells = _allocate_cells__at_elevation(
+                top, bottom, bottom_elevation
+            )
         case ALLOCATION_OPTION.at_first_active:
-            return _allocate_cells__at_first_active(active, bottom_elevation)
+            riv_cells, drn_cells = _allocate_cells__at_first_active(
+                active, bottom_elevation
+            )
         case _:
             raise ValueError(
                 "Received incompatible setting for rivers, only"
@@ -137,13 +150,21 @@ def allocate_riv_cells(
                 f"got: '{allocation_option.name}'"
             )
 
+    if drop_empty_layers:
+        riv_cells = _drop_empty_layers(riv_cells)
+        if drn_cells is not None:
+            drn_cells = _drop_empty_layers(drn_cells)
+
+    return riv_cells, drn_cells
+
 
 def allocate_drn_cells(
     allocation_option: ALLOCATION_OPTION,
     active: GridDataArray,
     top: GridDataArray,
     bottom: GridDataArray,
     elevation: GridDataArray,
+    drop_empty_layers: bool = False,
 ) -> GridDataArray:
     """
     Allocate drain cells from a planar grid across the vertical dimension.
@@ -166,6 +187,14 @@ def allocate_drn_cells(
     elevation: DataArray | UgridDatarray
         Planar grid containing drain elevation. Is not allowed to have a layer
         dimension.
+    drop_empty_layers: bool, default False
+        If True, drop layers from the result that contain no allocated
+        cells anywhere in the domain. This avoids carrying the package's
+        arrays at full model-layer size through downstream regridding,
+        clipping, masking, and splitting, which can otherwise become
+        expensive for models with many layers relative to how many
+        layers the topsystem package actually occupies. Set to False to
+        keep the previous full-layer-coordinate behaviour.
 
     Returns
     -------
@@ -181,13 +210,13 @@ def allocate_drn_cells(
     """
     match allocation_option:
         case ALLOCATION_OPTION.first_active_to_elevation:
-            return _allocate_cells__first_active_to_elevation(
+            result = _allocate_cells__first_active_to_elevation(
                 active, top, bottom, elevation
             )[0]
         case ALLOCATION_OPTION.at_elevation:
-            return _allocate_cells__at_elevation(top, bottom, elevation)[0]
+            result = _allocate_cells__at_elevation(top, bottom, elevation)[0]
         case ALLOCATION_OPTION.at_first_active:
-            return _allocate_cells__at_first_active(active, elevation)[0]
+            result = _allocate_cells__at_first_active(active, elevation)[0]
         case _:
             raise ValueError(
                 "Received incompatible setting for drains, only"
@@ -197,6 +226,8 @@ def allocate_drn_cells(
                 f"got: '{allocation_option.name}'"
             )
 
+    return _drop_empty_layers(result) if drop_empty_layers else result
+
 
 def allocate_ghb_cells(
     allocation_option: ALLOCATION_OPTION,
@@ -263,6 +294,7 @@ def allocate_rch_cells(
     allocation_option: ALLOCATION_OPTION,
     active: GridDataArray,
     rate: GridDataArray,
+    drop_empty_layers: bool = False,
 ) -> GridDataArray:
     """
     Allocate recharge cells from a planar grid across the vertical dimension.
@@ -279,6 +311,14 @@ def allocate_rch_cells(
     rate: DataArray | UgridDataArray
         Array with recharge rates. This will only be used to infer where
         recharge cells are defined.
+    drop_empty_layers: bool, default False
+        If True, drop layers from the result that contain no allocated
+        cells anywhere in the domain. This avoids carrying the package's
+        arrays at full model-layer size through downstream regridding,
+        clipping, masking, and splitting, which can otherwise become
+        expensive for models with many layers relative to how many
+        layers the topsystem package actually occupies. Set to False to
+        keep the previous full-layer-coordinate behaviour.
 
     Returns
     -------
@@ -294,14 +334,16 @@ def allocate_rch_cells(
     """
     match allocation_option:
         case ALLOCATION_OPTION.at_first_active:
-            return _allocate_cells__at_first_active(active, rate)[0]
+            result = _allocate_cells__at_first_active(active, rate)[0]
         case _:
             raise ValueError(
                 "Received incompatible setting for recharge, only"
                 f"'{ALLOCATION_OPTION.at_first_active.name}' supported."
                 f"got: '{allocation_option.name}'"
             )
 
+    return _drop_empty_layers(result) if drop_empty_layers else result
+
 
 def _is_layered(grid: GridDataArray):
     return "layer" in grid.sizes and grid.sizes["layer"] > 1
@@ -537,3 +579,48 @@ def _allocate_cells__at_first_active(
     topsystem_upper_active = upper_active & ~np.isnan(planar_topsystem_grid)
 
     return topsystem_upper_active, None
+
+
+def _drop_empty_layers(grid: GridDataArray) -> GridDataArray:
+    """
+    Drop layers that contain no True/non-nan values in any spatial cell
+    (and, if present, at any timestep). Keeps the `layer` coordinate but
+    only for layers that actually contain data - this is what lets
+    downstream regridding/clipping/masking/splitting operate over a much
+    smaller layer range when the topsystem package only spans a handful
+    of the model's total layers.
+
+    Parameters
+    ----------
+    grid: GridDataArray
+        Array with a "layer" dimension, typically the output of one of the
+        ``_allocate_cells__*`` functions.
+
+    Returns
+    -------
+    GridDataArray
+        Same array, subset to layers with data.
+    """
+    if "layer" not in grid.dims:
+        return grid
+
+    reduce_dims = [d for d in grid.dims if d != "layer"]
+
+    if grid.dtype == bool:
+        has_data_per_layer = grid.any(dim=reduce_dims)
+    else:
+        has_data_per_layer = (~grid.isnull()).any(dim=reduce_dims)
+
+    # Force to plain numpy/bool to avoid triggering a dask compute deep
+    # inside indexing logic more than once.
+    has_data_per_layer = (
+        has_data_per_layer.compute()
+        if hasattr(has_data_per_layer, "compute")
+        else has_data_per_layer
+    )
+
+    if bool(has_data_per_layer.all()):
+        return grid  # nothing to trim, skip the extra indexing op
+
+    used_layers = grid["layer"].where(has_data_per_layer, drop=True)
+    return grid.sel(layer=used_layers)