Issue #1821: Improved performance for allocation & sparse/empty layer handling. - #1912
Conversation
- Add drop_empty_layers: bool = False to �llocate_x_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.
|
Tests are green locally when calling There seems to be something wrong with docker + Teamcity though. Which is very strange as I didnt change anything there. All failures are caused by the same error: Not really sure what to do about this |
|
Spoke to the Devops team and they are aware of the flaky docker problem. |
|
TODO proper run of all tests (including user acceptance / weekly) to make sure that this doesnt break anything when set to We need to decide on whether the default will be True or False, or not configurable and always drop empty layers. |
JoerivanEngelen
left a comment
JoerivanEngelen
left a comment
There was a problem hiding this comment.
The changeset looks good, my plan was indeed to make the default of dropping layers True, so I'm fine with updating the test fixtures/test cases, except for the tests under tests/test_prepare/test_topsystem.py (or change the test and its cases with to test for a list of layer numbers instead of a list with True/False). Testing against the user acceptance tests is a must.
The performance investigation went a bit on the wrong track: LayerRegridder is not the regridder that is called when regridding MODFLOW 6 simulations. I was interested if feeding this into imod.mf6.Modflow6Simulation.regrid_like lead to any problems. I guess we'll run soon enough into this if you change the default of drop_empty_layers to True.
|
Pseudocode to extend tests for allocation: @parametrize_with_cases(
argnames="active,top,bottom,stage,bottom_elevation",
prefix="riv_",
)
@parametrize_with_cases(
argnames="option,expected_riv,expected_drn", prefix="allocation_", has_tag="riv"
)
def test_riv_allocation(
active, top, bottom, stage, bottom_elevation, option, expected_riv, expected_drn
):
actual_riv_da, actual_drn_da = allocate_riv_cells(
option, active, top, bottom, stage, bottom_elevation
)
actual_riv = take_nth_layer_column(actual_riv_da, 0)
empty_riv = take_nth_layer_column(actual_riv_da, 1)
if actual_drn_da is None:
actual_drn = None
empty_drn = None
else:
actual_drn = take_nth_layer_column(actual_drn_da, 0)
empty_drn = take_nth_layer_column(actual_drn_da, 1)
np.testing.assert_equal(actual_riv, expected_riv)
np.testing.assert_equal(actual_drn, expected_drn)
assert np.all(~empty_riv)
if empty_drn is not None:
assert np.all(~empty_drn)
expected_layers = np.nonzero(expected_riv) + 1
actual_riv_da, actual_drn_da = allocate_riv_cells(
option, active, top, bottom, stage, bottom_elevation, drop_empty_layers=True
)
assert np.array_equal(actual_riv_da.coords["layer"].values, expected_layers) |
…ug in 'imod.prepare.cleanup.align_interface_levels' to prevent raising alignment error during cleanup.
Summary
@JoerivanEngelen What do you think about this? Im not sure about it:
|
JoerivanEngelen
left a comment
JoerivanEngelen
left a comment
There was a problem hiding this comment.
Looks good! I think this improves matters without introducing too much new code.
I think it's fine to not have drop_empty_layers configurable when calling from_imod5_data.
Two questions:
- Why is the extra selection in the cleanup utility necessary? How can "bottom" get the full layer index again? I guess because of a comparison we do somewhere with model bottom? I'd like to know, because it could mean we missed something and it might be worth seeing if we can fix this upstream.
- The edge case when the package is entirely empty: would be nice if we could drop that for performance reasons, as it would save us some costly checks later. Would that introduce a lot of extra code to deal with this?
| # `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` |
There was a problem hiding this comment.
I didn't expect this to happen, which tests were failing because of this? Just wondering what the cause of this is as it might mean we overlooked something.
There was a problem hiding this comment.
While developing I got errors, but I cannot remember exactly which ones to be honest.
I have run the tests again with this fix commented out, and got 2 errors:
test_prepare\test_cleanup.py::test_cleanup_riv__trimmed_layers[structured & unstructured]
This is however the test I added to produce this bug and show it does not error anymore. (Which is good news: all callers pass matching layers to this function.)
The bug (introduced by trimming layers):
bottom has the full model layers, while top/to_align can be a trimmed package with a subset of layers.
Calling xr.where(~to_align, bottom) will then raise an alignment error since xarray uses join="exact" by default.
This is why other places in the code use xr.align(obj1, obj2, join="left"), which fills missing layers with NaN (mask.py and schemata.py).
Since the inputs to xr.where need to be the same shape, and we are interested in the top layers, we can throw away the irrelevant layers from bottom before doing this.
In this case,
If by trimming layers, no layers would remain, there would be errors during package validation. See this Teamcity run. |
|
Approved in advance, still things to do:
|
…dle it automatically.
…ay to LAYERS_USED
|
I've made some small fixes to improve matters and fix mypy issues:
|
…hon into feat/drop-unused-layers
|
2 participants
Fixes #1821
Description
drop_empty_layers: bool = Falsetoallocate_<riv/drn/rch>_cellsfunctions inimod.prepare.topsystem.allocationto remove fully empty layers from the grids to prevent passing them along to reprojection/regridding operations.Default value discussion
I chose the default for
drop_empty_layersto beFalseto 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
Truebroke some tests (which is expected since these tests asserted on the shape of the returned objects).I think setting the default to
Trueand 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 withdrop_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 intest_topsystem_layer_preservation.py) in places where its required to have all layers (even empty ones)Checklist
Issue #nr, e.g.Issue #737pixi run generate-sbomand committed changesInvestigation: 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 nestedn_layers_dst × n_layers_srcloop 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 comparingstage/bottom_elevation(planar, no layer dim, enforced byPLANAR_GRID.validate) againsttop/bottom/active, which always carry the full model layer coordinate: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'sstage/conductance/etc. once built via.where(riv_cells)) then keeps the fulln_layerscoordinate, mostlyFalse/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_layersFrom
imod/prepare/regrid_layers.py(numba-jitted):This is a nested loop over
nlayer_dst × nlayer_src × nrow × ncol. Even though theisnancheck 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 costsnlayer_dst × Niterations instead ofnlayer_dst × 2.Since
nlayer_dstalso 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
Results
In every case, only 2 of
n_layers_srclayers 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 thatLayerRegridderis the concrete hot spot to fix.Trimming the empty layers before regridding is a good approach for performance.