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

[LuukBlom]

Fixes #1821

Description

• Add drop_empty_layers: bool = False to allocate_<riv/drn/rch>_cells functions in imod.prepare.topsystem.allocation to remove fully empty layers from the grids to prevent passing them along to reprojection/regridding operations.
• Improved imod.prepare.layerregrid._regrid_layers to precompute valid layer indices once per column, instead of re-checking isnan for every (ii, jj) pair.

Default value discussion

I chose the default for drop_empty_layers to be False to have it as a backwards compatible opt-in new behaviour. But perhaps the default should be True, since it is somewhat difficult for users to actually use the opt-in behaviour. (its an internal function call, so we might need to add it to some other functions as an arg with default as well?)

Setting it to True broke some tests (which is expected since these tests asserted on the shape of the returned objects).
I think setting the default to True and updating those tests to expect the new format is the way to go, but I'd like to discuss with @JoerivanEngelen before doing that. For now, I have left it at False to not break the existing tests, and added tests that call this path with drop_empty_layers=True.

Alternative solution could be that we set the default to True and call a function similar to reindex_to_full_layers (see usage in test_topsystem_layer_preservation.py) in places where its required to have all layers (even empty ones)

Checklist

• Links to correct issue
• Update changelog, if changes affect users
• PR title starts with Issue #nr, e.g. Issue #737
• Unit tests were added
• If feature added: Added/extended example
• If feature added: Added feature to API documentation
• If pixi.lock was changed: Ran pixi run generate-sbom and committed changes

---

Investigation: cost of broadcasting to all layers in allocation/regrid pipeline

Details

Spent some time to investigate this issue to find 1) the slow/expensive part of the code and 2) whether and how to optimize.

Summary

The allocation boolean logic itself (allocate_riv_cells, allocate_drn_cells, etc.) is not the bottleneck, it's cheap regardless of layer count.
The real cost is in LayerRegridder._regrid_layers, which has a nested n_layers_dst × n_layers_src loop per (row, col). When a package's real data only occupies a handful of layers but is broadcast/kept at full model-layer size before regridding, cost scales roughly quadratically with total layer count, not with the number of layers that actually contain data.

Measured: ~8.4s at 300 layers vs ~1ms at 2 layers, for a package whose real data lives in only 2 layers throughout.

Where the full-layer array comes from

imod/prepare/topsystem/allocation.py (_allocate_cells__stage_to_riv_bot, _allocate_cells__at_elevation, _allocate_cells__first_active_to_elevation, etc.) all compute their result by comparing stage/bottom_elevation (planar, no layer dim, enforced by PLANAR_GRID.validate) against top/bottom/active, which always carry the full model layer coordinate:

is_above_lower_bound = get_above_lower_bound(bottom_elevation, top_layered)
is_below_upper_bound = stage >= bottom
riv_cells = is_below_upper_bound & is_above_lower_bound

This is not changeable, we can't know in advance which layers a planar stage/elevation grid intersects without comparing against every layer's top/bottom. That comparison is a single vectorized boolean op, and is cheap at any layer count (confirmed below).

The problem is that riv_cells (and the package's stage/conductance/etc. once built via .where(riv_cells)) then keeps the full n_layers coordinate, mostly False/NaN, and nothing downstream trims it back down before it's handed to regridding, clipping, masking, or splitting.

Where the cost actually is: LayerRegridder._regrid_layers

From imod/prepare/regrid_layers.py (numba-jitted):

for ii in range(nlayer_dst):
    ...
    for jj in range(nlayer_src):
        st = src_t[jj]
        sb = src_b[jj]
        if np.isnan(st) or np.isnan(sb):
            continue
        ...

This is a nested loop over nlayer_dst × nlayer_src × nrow × ncol. Even though the isnan check that skips empty layers is cheap per-iteration, the loop still has to visit every (dst_layer, src_layer) pair for every cell, so an all-NaN source layer still costs an iteration, for every destination layer. If a package's real data occupies only 2 of N layers, this loop still costs nlayer_dst × N iterations instead of nlayer_dst × 2.

Since nlayer_dst also scales with total model layers in typical regridding (source and destination models usually share vertical discretization), the practical cost scales roughly quadratically with total layer count, independent of how many layers actually contain real data.

Reproducer

import time

import numpy as np
import xarray as xr

from imod.prepare import LayerRegridder


def make_layered(n_layers, nrow, ncol, real_layers, trim=False):
    x = np.arange(ncol) * 10.0
    y = np.arange(nrow) * -10.0
    layer = np.arange(1, n_layers + 1)

    top = np.zeros((n_layers, nrow, ncol))
    bottom = np.zeros((n_layers, nrow, ncol))
    for k in range(n_layers):
        top[k] = -k
        bottom[k] = -k - 1

    source = np.full((n_layers, nrow, ncol), np.nan)
    source[:real_layers] = 1.0

    coords = {"layer": layer, "y": y, "x": x}
    dims = ("layer", "y", "x")
    src_da = xr.DataArray(source, coords, dims)
    top_da = xr.DataArray(top, coords, dims)
    bot_da = xr.DataArray(bottom, coords, dims)

    if trim:
        # This is what allocation SHOULD hand to the regridder once
        # trimmed: only the layers that actually contain data.
        src_da = src_da.isel(layer=slice(0, real_layers))
        top_da = top_da.isel(layer=slice(0, real_layers))
        bot_da = bot_da.isel(layer=slice(0, real_layers))

    return src_da, top_da, bot_da


def time_regrid(
    n_layers, real_layers=2, nrow=200, ncol=200, trim_source=False, repeats=3
):
    src, src_top, src_bot = make_layered(
        n_layers, nrow, ncol, real_layers, trim=trim_source
    )

    # Destination always spans the full model layer range - that part
    # is unavoidable, since regridding must produce a value (or nan)
    # for every destination layer.
    dst_layer = np.arange(1, n_layers + 1)
    full_top, _, _ = make_layered(n_layers, nrow, ncol, real_layers, trim=False)
    dst_top = (
        full_top.rename({"layer": "layer"}).assign_coords(layer=dst_layer) * 0
        + full_top.values
    )
    dst_top = xr.DataArray(
        np.stack([-np.full((nrow, ncol), float(k)) for k in range(n_layers)]),
        {"layer": dst_layer, "y": src.y, "x": src.x},
        ("layer", "y", "x"),
    )
    dst_bot = dst_top - 1.0

    regridder = LayerRegridder(method="mean")
    regridder.regrid(src, src_top, src_bot, dst_top, dst_bot)  # warm up JIT

    times = []
    for _ in range(repeats):
        t0 = time.perf_counter()
        regridder.regrid(src, src_top, src_bot, dst_top, dst_bot)
        times.append(time.perf_counter() - t0)

    label = "TRIMMED src" if trim_source else "dense src (padded w/ nan)"
    print(
        f"n_layers={n_layers:4d} real_layers={real_layers}  [{label}]  best={min(times):.4f}s"
    )


if __name__ == "__main__":
    print("-- dense (current behaviour): source kept at full n_layers, nan-padded --")
    for n in [2, 5, 10, 30, 100, 300]:
        time_regrid(n_layers=n, real_layers=2, trim_source=False)

    print("\n-- trimmed: source pre-cut to its 2 real layers before regrid --")
    for n in [2, 5, 10, 30, 100, 300]:
        time_regrid(n_layers=n, real_layers=2, trim_source=True)

Results

-- dense (current behaviour): source kept at full n_layers, nan-padded --
n_layers=   2 real_layers=2  [dense src (padded w/ nan)]  best=0.0051s
n_layers=   5 real_layers=2  [dense src (padded w/ nan)]  best=0.0070s
n_layers=  10 real_layers=2  [dense src (padded w/ nan)]  best=0.0143s
n_layers=  30 real_layers=2  [dense src (padded w/ nan)]  best=0.0985s
n_layers= 100 real_layers=2  [dense src (padded w/ nan)]  best=1.0968s
n_layers= 300 real_layers=2  [dense src (padded w/ nan)]  best=9.7136s

-- trimmed: source pre-cut to its 2 real layers before regrid --
n_layers=   2 real_layers=2  [TRIMMED src]  best=0.0041s
n_layers=   5 real_layers=2  [TRIMMED src]  best=0.0041s
n_layers=  10 real_layers=2  [TRIMMED src]  best=0.0051s
n_layers=  30 real_layers=2  [TRIMMED src]  best=0.0106s
n_layers= 100 real_layers=2  [TRIMMED src]  best=0.0342s
n_layers= 300 real_layers=2  [TRIMMED src]  best=0.1197s

In every case, only 2 of n_layers_src layers ever contain real (non-NaN) data. Yet runtime grows from 1.4ms to 9.7s (~6000x) purely as a function of total layer count. This confirms the cost scales with the model's total layer count, not with the amount of actual topsystem data, and that LayerRegridder is the concrete hot spot to fix.

Trimming the empty layers before regridding is a good approach for performance.

[sonarqubecloud]

Quality Gate passed

Issues
0 New issues
0 Accepted issues

Measures
0 Security Hotspots
0.0% Coverage on New Code
0.0% Duplication on New Code

See analysis details on SonarQube Cloud

[LuukBlom]

Tests are green locally when calling pixi run tests.

There seems to be something wrong with docker + Teamcity though.

Which is very strange as I didnt change anything there.
Re-running the failed checks doesnt fix it unfortunately.

All failures are caused by the same error:

pull containers.deltares.nl/hydrology_product_line_imod/windows-pixi:v0.69.0
  failed to register layer: strconv.ParseInt: parsing "": invalid syntax
Unable to find image 'containers.deltares.nl/hydrology_product_line_imod/windows-pixi:v0.69.0' locally
v0.69.0: Pulling from hydrology_product_line_imod/windows-pixi
fbc33898c8ff: Pulling fs layer
1e22643c3eaa: Pulling fs layer
0318f4e42c5d: Pulling fs layer
e778778b905f: Pulling fs layer
9ffbd341cca3: Pulling fs layer
c6aee63eee57: Pulling fs layer
b88504fedbea: Pulling fs layer
a29a128dd81c: Pulling fs layer
ed50fbf4d281: Pulling fs layer
85bee0e55f42: Pulling fs layer
92ea28bcf781: Pulling fs layer
4e87dcc2f18a: Pulling fs layer
c6aee63eee57: Waiting
b88504fedbea: Waiting
e778778b905f: Waiting
a29a128dd81c: Waiting
92ea28bcf781: Waiting
ed50fbf4d281: Waiting
4e87dcc2f18a: Waiting
85bee0e55f42: Waiting
9ffbd341cca3: Waiting
0318f4e42c5d: Verifying Checksum
0318f4e42c5d: Download complete
e778778b905f: Verifying Checksum
e778778b905f: Download complete
9ffbd341cca3: Verifying Checksum
9ffbd341cca3: Download complete
c6aee63eee57: Verifying Checksum
c6aee63eee57: Download complete
b88504fedbea: Verifying Checksum
b88504fedbea: Download complete
a29a128dd81c: Verifying Checksum
a29a128dd81c: Download complete
ed50fbf4d281: Verifying Checksum
ed50fbf4d281: Download complete
85bee0e55f42: Verifying Checksum
85bee0e55f42: Download complete
92ea28bcf781: Verifying Checksum
92ea28bcf781: Download complete
4e87dcc2f18a: Verifying Checksum
4e87dcc2f18a: Download complete
1e22643c3eaa: Verifying Checksum
1e22643c3eaa: Download complete
fbc33898c8ff: Verifying Checksum
fbc33898c8ff: Download complete
docker: failed to register layer: strconv.ParseInt: parsing "": invalid syntax
Run 'docker run --help' for more information
Process exited with code 125

Not really sure what to do about this

[LuukBlom]

Spoke to the Devops team and they are aware of the flaky docker problem.
Is an issue thats popped up today and they are working on a fix.
ETA is unclear, but they will let us know when there are updates

[LuukBlom]

TODO proper run of all tests (including user acceptance / weekly) to make sure that this doesnt break anything when set to True

We need to decide on whether the default will be True or False, or not configurable and always drop empty layers.
If we decide on True or dropping, we also need to update the assertions for various unittests, and perhaps other tests