Conversation
- Add drop_empty_layers: bool = False to 锟絣locate_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
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.
| # 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 | ||
| ) |
There was a problem hiding this comment.
Probably good to call a compute here, but the check for hasattr compute is not necessary, xr.DataArray.compute() is always available and always returns the same thing.
| # 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 | |
| ) | |
| # 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 grid.dtype == bool: | ||
| has_data_per_layer = grid.any(dim=reduce_dims) | ||
| else: | ||
| has_data_per_layer = (~grid.isnull()).any(dim=reduce_dims) |
There was a problem hiding this comment.
I think grid.dtype should always be bool here, so the if else statement should not be necessary here, maybe throw an error if not the case? I might be missing something though.
There was a problem hiding this comment.
Thanks for optimizing and improving this. Unfortunately, this is not the regridder we are using commonly. We are relying on the stuff under imod/common/utilities/regrid. This is a specific older regridder for which xugrid doesn't have a stand-in replacement (yet). I thing these optimizations are not harmful though
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.