diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst index a5b40a49a..03684c90f 100644 --- a/docs/api/changelog.rst +++ b/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 ~~~~~ diff --git a/imod/prepare/layerregrid.py b/imod/prepare/layerregrid.py index c239e3240..00e4b5c42 100644 --- a/imod/prepare/layerregrid.py +++ b/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 diff --git a/imod/prepare/topsystem/allocation.py b/imod/prepare/topsystem/allocation.py index b74bec2f9..2bd12cdd2 100644 --- a/imod/prepare/topsystem/allocation.py +++ b/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,6 +150,13 @@ 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, @@ -144,6 +164,7 @@ def allocate_drn_cells( 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,7 +334,7 @@ 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" @@ -302,6 +342,8 @@ def allocate_rch_cells( 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) diff --git a/imod/tests/test_prepare/test_topsystem_layer_preservation.py b/imod/tests/test_prepare/test_topsystem_layer_preservation.py new file mode 100644 index 000000000..d6de6873e --- /dev/null +++ b/imod/tests/test_prepare/test_topsystem_layer_preservation.py @@ -0,0 +1,353 @@ +import numpy as np +import pytest +import xarray as xr + +from imod.prepare import LayerRegridder +from imod.prepare.topsystem import ( + ALLOCATION_OPTION, + allocate_drn_cells, + allocate_rch_cells, + allocate_riv_cells, +) +from imod.typing import GridDataArray + + +def make_model_grid(n_layers, nrow=10, ncol=10, dx=100.0): + x = np.arange(ncol) * dx + y = np.arange(nrow) * -dx + layer = np.arange(1, n_layers + 1) + + top = xr.DataArray( + np.stack([np.full((nrow, ncol), -float(k)) for k in range(n_layers)]), + {"layer": layer, "y": y, "x": x}, + ("layer", "y", "x"), + ) + bottom = top - 1.0 + active = xr.full_like(top, True, dtype=bool) + return active, top, bottom + + +def make_sparse_riv(nrow=10, ncol=10, dx=100.0, active_layer=1): + """Planar river stage/bottom_elevation - genuinely intersects only + ~1 layer of a deep model.""" + x = np.arange(ncol) * dx + y = np.arange(nrow) * -dx + stage = xr.DataArray(np.full((nrow, ncol), -0.2), {"y": y, "x": x}, ("y", "x")) + bottom_elevation = xr.DataArray( + np.full((nrow, ncol), -0.8), {"y": y, "x": x}, ("y", "x") + ) + return stage, bottom_elevation + + +@pytest.fixture(params=[2, 10, 30]) +def n_layers(request): + return request.param + + +def n_nonempty_layers(da: GridDataArray) -> int: + """Number of layers that contain at least one meaningfully "present" value. + + Handles the fact that boolean arrays get upcast to float by xarray's + .where() (NaN has no bool representation) - after such a coercion, + False becomes 0.0, which must still be treated as "no data", not as + a valid float value. + """ + reduce_dims = [d for d in da.dims if d != "layer"] + + if da.dtype == bool: + has_data_per_layer = da.any(dim=reduce_dims) + else: + # Treat both NaN and 0.0 as "no data" - this covers arrays that + # started boolean and were upcast to float by .where()/masking. + is_present = (~da.isnull()) & (da != 0) + has_data_per_layer = is_present.any(dim=reduce_dims) + + return int(has_data_per_layer.sum()) + + +def reindex_to_full_layers( + da: xr.DataArray, full_layer: xr.DataArray, dtype +) -> xr.DataArray: + """ + Re-expand a layer-trimmed result back onto the full model layer + coordinate, so it can be compared against expectations written for + the untrimmed (drop_empty_layers=False) behaviour. + """ + fill_value = False if dtype is bool else np.nan + return da.reindex(layer=full_layer, fill_value=fill_value) + + +def take_nth_layer_column(grid, n): + if "time" in grid.dims: + grid = grid.isel(time=-1) + return grid.values[:, n, n] + + +@pytest.fixture +def basic_riv_inputs(): + nlayer, nrow, ncol = 4, 3, 3 + layer = np.array([1, 2, 3, 4]) + y = np.arange(nrow) * -10.0 + x = np.arange(ncol) * 10.0 + + top = xr.DataArray( + np.stack([np.full((nrow, ncol), -float(k)) for k in range(nlayer)]), + {"layer": layer, "y": y, "x": x}, + ("layer", "y", "x"), + ) + bottom = top - 1.0 + active = xr.full_like(top, True, dtype=bool) + + # Stage/bottom_elevation only intersect layer 1: planar, no layer dim. + stage = xr.DataArray(np.full((nrow, ncol), -0.2), {"y": y, "x": x}, ("y", "x")) + bottom_elevation = xr.DataArray( + np.full((nrow, ncol), -0.8), {"y": y, "x": x}, ("y", "x") + ) + return active, top, bottom, stage, bottom_elevation, layer + + +def test_allocate_riv_cells_drop_empty_layers_matches_full(basic_riv_inputs): + active, top, bottom, stage, bottom_elevation, layer = basic_riv_inputs + + full, _ = allocate_riv_cells( + ALLOCATION_OPTION.stage_to_riv_bot, + active, + top, + bottom, + stage, + bottom_elevation, + drop_empty_layers=False, + ) + trimmed, _ = allocate_riv_cells( + ALLOCATION_OPTION.stage_to_riv_bot, + active, + top, + bottom, + stage, + bottom_elevation, + drop_empty_layers=True, + ) + + # Trimmed result should have fewer (or equal) layers than the full one. + assert trimmed.sizes["layer"] <= full.sizes["layer"] + + # Once re-expanded, trimmed result must be identical to the full one. + re_expanded = reindex_to_full_layers(trimmed, full["layer"], dtype=bool) + xr.testing.assert_equal(re_expanded, full) + + +def test_allocate_drn_cells_drop_empty_layers_matches_full(basic_riv_inputs): + active, top, bottom, _, elevation, layer = basic_riv_inputs + + full = allocate_drn_cells( + ALLOCATION_OPTION.at_elevation, + active, + top, + bottom, + elevation, + drop_empty_layers=False, + ) + trimmed = allocate_drn_cells( + ALLOCATION_OPTION.at_elevation, + active, + top, + bottom, + elevation, + drop_empty_layers=True, + ) + + assert trimmed.sizes["layer"] <= full.sizes["layer"] + re_expanded = reindex_to_full_layers(trimmed, full["layer"], dtype=bool) + xr.testing.assert_equal(re_expanded, full) + + +def test_allocate_rch_cells_drop_empty_layers_matches_full(basic_riv_inputs): + active, _, _, _, _, layer = basic_riv_inputs + nrow, ncol = active.sizes["y"], active.sizes["x"] + rate = xr.DataArray( + np.full((nrow, ncol), 0.001), {"y": active.y, "x": active.x}, ("y", "x") + ) + active2d_active = ( + active # active already has layer dim, at_first_active uses it directly + ) + + full = allocate_rch_cells( + ALLOCATION_OPTION.at_first_active, + active2d_active, + rate, + drop_empty_layers=False, + ) + trimmed = allocate_rch_cells( + ALLOCATION_OPTION.at_first_active, active2d_active, rate, drop_empty_layers=True + ) + + assert trimmed.sizes["layer"] <= full.sizes["layer"] + re_expanded = reindex_to_full_layers(trimmed, full["layer"], dtype=bool) + xr.testing.assert_equal(re_expanded, full) + + +class TestAllocationLayerCount: + def test_allocate_riv_cells_does_not_grow_beyond_real_extent(self, n_layers): + """ + The allocated result may keep the full model `layer` coordinate + (that part is unavoidable, see investigation doc), but the number + of layers that actually contain True values should not depend on + total model layer count - it should stay pinned to how many + layers the stage/bottom_elevation genuinely intersect (here: 1). + """ + active, top, bottom = make_model_grid(n_layers) + stage, bottom_elevation = make_sparse_riv() + + riv_cells, _ = allocate_riv_cells( + ALLOCATION_OPTION.stage_to_riv_bot, + active, + top, + bottom, + stage, + bottom_elevation, + ) + + assert n_nonempty_layers(riv_cells) == 1, ( + "Number of allocated layers should not scale with total model " + f"layers (got layers with True values for n_layers={n_layers})" + ) + + +class TestRegridLayerCount: + def test_regrid_preserves_sparse_layer_count(self, n_layers): + """ + A source array pre-trimmed to 2 real layers should not become + denser after regridding onto a destination grid with n_layers + model layers - only 2 destination layers should end up non-nan. + """ + real_layers = 2 + _, src_top, src_bot = make_model_grid(n_layers) + _, dst_top, dst_bot = make_model_grid(n_layers) # same discretization here + + # Sparse source: only first `real_layers` layers have data, rest all-nan. + source = xr.full_like(src_top, np.nan) + source.values[:real_layers] = 1.0 + + regridder = LayerRegridder(method="mean") + result = regridder.regrid(source, src_top, src_bot, dst_top, dst_bot) + + assert n_nonempty_layers(result) == real_layers, ( + "Regridding introduced extra non-nan layers beyond the " + f"source's real extent (n_layers={n_layers})" + ) + + def test_regrid_sparse_input_matches_dense_input_result(self, n_layers): + """ + Regression guard: regridding a package pre-trimmed to its real + layers should give the same numerical result as regridding the + same package padded out to the full model layer range with nan. + This is the property that allows allocation to safely trim layers + before regridding without changing behaviour. + """ + real_layers = 2 + _, src_top, src_bot = make_model_grid(n_layers) + _, dst_top, dst_bot = make_model_grid(n_layers) + + dense_source = xr.full_like(src_top, np.nan) + dense_source.values[:real_layers] = 1.0 + + sparse_source = dense_source.isel(layer=slice(0, real_layers)) + sparse_top = src_top.isel(layer=slice(0, real_layers)) + sparse_bot = src_bot.isel(layer=slice(0, real_layers)) + + regridder = LayerRegridder(method="mean") + dense_result = regridder.regrid( + dense_source, src_top, src_bot, dst_top, dst_bot + ) + sparse_result = regridder.regrid( + sparse_source, sparse_top, sparse_bot, dst_top, dst_bot + ) + + xr.testing.assert_allclose(dense_result, sparse_result) + + +class TestClipLayerCount: + def test_clip_by_grid_preserves_sparse_layers(self, n_layers): + """ + Clipping a package to a smaller planar extent should not + reintroduce layers that had no data before clipping. + """ + active, top, bottom = make_model_grid(n_layers) + stage, bottom_elevation = make_sparse_riv() + + riv_cells, _ = allocate_riv_cells( + ALLOCATION_OPTION.stage_to_riv_bot, + active, + top, + bottom, + stage, + bottom_elevation, + ) + + # Clip to a smaller planar window. + x_slice = slice(0, 500.0) + y_slice = slice(0.0, -500.0) + clipped = riv_cells.sel(x=x_slice, y=y_slice) + + assert n_nonempty_layers(clipped) <= n_nonempty_layers(riv_cells), ( + "Clipping should never increase the number of non-empty layers" + ) + + +class TestMaskLayerCount: + def test_mask_does_not_densify_layers(self, n_layers): + """ + Masking with idomain (full n_layers) should not turn a + sparse-layer package dense via alignment/broadcasting. + """ + active, top, bottom = make_model_grid(n_layers) + stage, bottom_elevation = make_sparse_riv() + + riv_cells, _ = allocate_riv_cells( + ALLOCATION_OPTION.stage_to_riv_bot, + active, + top, + bottom, + stage, + bottom_elevation, + ) + + idomain = active.astype(int) # full n_layers, all active + masked = riv_cells.where(idomain > 0) + + assert n_nonempty_layers(masked) == n_nonempty_layers(riv_cells), ( + "Masking against a full-layer idomain changed the number of " + "non-empty layers - likely due to alignment/broadcasting" + ) + + +class TestSplitLayerCount: + def test_split_preserves_sparse_layers_per_partition(self, n_layers): + """ + Partitioning by a planar label array should not force a + sparse-layer package to become dense in any partition. + """ + active, top, bottom = make_model_grid(n_layers) + stage, bottom_elevation = make_sparse_riv() + + riv_cells, _ = allocate_riv_cells( + ALLOCATION_OPTION.stage_to_riv_bot, + active, + top, + bottom, + stage, + bottom_elevation, + ) + + # Simple 2-partition planar label: left half / right half. + label = xr.zeros_like(riv_cells.isel(layer=0, drop=True), dtype=int) + ncol = label.sizes["x"] + label[:, ncol // 2 :] = 1 + + for part in [0, 1]: + part_mask = label == part + partitioned = riv_cells.where(part_mask) + assert n_nonempty_layers(partitioned) <= n_nonempty_layers(riv_cells), ( + f"Partition {part} has more non-empty layers than the " + "original unpartitioned array" + )