diff --git a/.dvc/config b/.dvc/config index 3fd415324..1d1b3075a 100644 --- a/.dvc/config +++ b/.dvc/config @@ -1,5 +1,6 @@ -[core] - remote = minio -['remote "minio"'] - url = s3://imod-python-test-data - endpointurl = https://s3.deltares.nl +[core] + remote = minio +['remote "minio"'] + url = s3://imod-python-test-data + endpointurl = https://s3.deltares.nl + allow_anonymous_login = true diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst index aac18f99b..18f834f50 100644 --- a/docs/api/changelog.rst +++ b/docs/api/changelog.rst @@ -22,6 +22,23 @@ Added :meth:`imod.mf6.GroundwaterTransportModel.mask_packages` to mask specific packages of a groundwater flow model and a groundwater transport model respectively. +- Added ``drop_empty_layers: bool = True`` to various cell allocation functions + in :mod:`imod.prepare.topsystem.allocation` to remove fully empty layers from the grids. + Strips the empty layers before they are passed along to + reprojection/regridding operations, which can save considerable time for models with + many empty layers. Set to False to keep the previous full-layer-coordinate + behaviour. :meth:`imod.prepare.topsystem.allocation.allocate_riv_cells`, + :meth:`imod.prepare.topsystem.allocation.allocate_drn_cells`, + :meth:`imod.prepare.topsystem.allocation.allocate_ghb_cells`, + :meth:`imod.prepare.topsystem.allocation.allocate_rch_cells` +- Added ``drop_empty_layers: bool = True`` to + :meth:`imod.mf6.River.reallocate`, :meth:`imod.mf6.Drainage.reallocate`, + :meth:`imod.mf6.GeneralHeadBoundary.reallocate`, and + :meth:`imod.mf6.Recharge.reallocate`. Allocation and conductance + distribution are always computed over the full layer range first; only + the final package has fully empty layers trimmed off afterwards, so this + does not affect computed values. Set to False to keep the previous + full-layer-coordinate behaviour. Fixed ~~~~~ @@ -35,6 +52,11 @@ Fixed :class:`imod.mf6.Recharge`) where IBOUND is less than 0. - :meth:`imod.msw.MetaSwapModel.from_imod5_data` now masks cells where IBOUND is less than 0. +- Fixed :func:`imod.prepare.cleanup.align_interface_levels` (used by + ``cleanup_riv``, and therefore :meth:`imod.mf6.River.cleanup`) raising an + alignment error when a package's own layer coordinate is a subset of the + model's full layer range, e.g. after :meth:`imod.mf6.River.reallocate` with + ``drop_empty_layers=True``. - :class:`imod.msw.FileCopier` and :class:`imod.msw.MeteoGridCopy` now force paths to be stored as strings in the dataset. ``pathlib.Path`` objects could cause errors when calling :meth:`imod.msw.MetaSwapModel.dump`. diff --git a/imod/mf6/drn.py b/imod/mf6/drn.py index b58ac0dd9..ca2a58ae8 100644 --- a/imod/mf6/drn.py +++ b/imod/mf6/drn.py @@ -17,7 +17,11 @@ from imod.mf6.utilities.package import set_repeat_stress_if_available from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA from imod.prepare.cleanup import cleanup_drn -from imod.prepare.topsystem.allocation import ALLOCATION_OPTION, allocate_drn_cells +from imod.prepare.topsystem.allocation import ( + ALLOCATION_OPTION, + allocate_drn_cells, + drop_empty_layers_from_dict, +) from imod.prepare.topsystem.conductance import ( DISTRIBUTING_OPTION, distribute_drn_conductance, @@ -193,6 +197,7 @@ def _allocate_and_distribute_planar_data( npf: NodePropertyFlow, allocation_option: ALLOCATION_OPTION, distributing_option: DISTRIBUTING_OPTION, + drop_empty_layers: bool = True, ) -> dict[str, GridDataArray]: """ Allocate and distribute planar data for given discretization and npf @@ -214,6 +219,13 @@ def _allocate_and_distribute_planar_data( ALLOCATION_OPTION.at_first_active. distributing_option: DISTRIBUTING_OPTION distributing option. + drop_empty_layers: bool, default True + If True, drop layers that contain no allocated cells anywhere in the + domain (or at any time), so the returned grids only span the layers that + are actually used. Reduces memory use and speeds up later operations such + as regridding, clipping and splitting. If no cells are allocated in any + layer, nothing is dropped and the full layer range is returned, as a layer + dimension of size 0 is not valid. Returns ------- @@ -240,6 +252,7 @@ def _allocate_and_distribute_planar_data( top, bottom, planar_data["elevation"], + drop_empty_layers=False, # Keep full layer range, drop empty layers below ) layered_data = {} layered_data["conductance"] = distribute_drn_conductance( @@ -253,6 +266,10 @@ def _allocate_and_distribute_planar_data( ) layered_data["elevation"] = planar_data["elevation"].where(drn_allocation) layered_data["elevation"] = enforce_dim_order(layered_data["elevation"]) + + if drop_empty_layers: + layered_data = drop_empty_layers_from_dict(layered_data, drn_allocation) + return layered_data @classmethod diff --git a/imod/mf6/ghb.py b/imod/mf6/ghb.py index 955a9e89b..9b554e955 100644 --- a/imod/mf6/ghb.py +++ b/imod/mf6/ghb.py @@ -19,7 +19,11 @@ from imod.mf6.utilities.package import set_repeat_stress_if_available from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA from imod.prepare.cleanup import cleanup_ghb -from imod.prepare.topsystem.allocation import ALLOCATION_OPTION, allocate_ghb_cells +from imod.prepare.topsystem.allocation import ( + ALLOCATION_OPTION, + allocate_ghb_cells, + drop_empty_layers_from_dict, +) from imod.prepare.topsystem.conductance import ( DISTRIBUTING_OPTION, distribute_ghb_conductance, @@ -198,6 +202,7 @@ def _allocate_and_distribute_planar_data( npf: NodePropertyFlow, allocation_option: ALLOCATION_OPTION, distributing_option: DISTRIBUTING_OPTION, + drop_empty_layers: bool = True, ) -> dict[str, GridDataArray]: """ Allocate and distribute planar data for given discretization and npf @@ -219,6 +224,13 @@ def _allocate_and_distribute_planar_data( ALLOCATION_OPTION.at_first_active. distributing_option: DISTRIBUTING_OPTION distributing option. + drop_empty_layers: bool, default True + If True, drop layers that contain no allocated cells anywhere in the + domain (or at any time), so the returned grids only span the layers that + are actually used. Reduces memory use and speeds up later operations such + as regridding, clipping and splitting. If no cells are allocated in any + layer, nothing is dropped and the full layer range is returned, as a layer + dimension of size 0 is not valid. Returns ------- @@ -245,6 +257,7 @@ def _allocate_and_distribute_planar_data( top, bottom, planar_data["head"], + drop_empty_layers=False, # Keep full layer range, drop empty layers below ) layered_data = {} @@ -259,6 +272,8 @@ def _allocate_and_distribute_planar_data( bottom, npf.dataset["k"], ) + if drop_empty_layers: + layered_data = drop_empty_layers_from_dict(layered_data, ghb_allocation) return layered_data @classmethod diff --git a/imod/mf6/rch.py b/imod/mf6/rch.py index 731676776..a32e6c28a 100644 --- a/imod/mf6/rch.py +++ b/imod/mf6/rch.py @@ -25,7 +25,11 @@ get_cell_area_from_imod5_data, is_msw_active_cell, ) -from imod.prepare.topsystem.allocation import ALLOCATION_OPTION, allocate_rch_cells +from imod.prepare.topsystem.allocation import ( + ALLOCATION_OPTION, + allocate_rch_cells, + drop_empty_layers_from_dict, +) from imod.schemata import ( AllCoordsValueSchema, AllInsideNoDataSchema, @@ -182,6 +186,7 @@ def _allocate_planar_data( planar_data: dict[str, GridDataArray], dis: StructuredDiscretization | VerticesDiscretization, allocation_option: ALLOCATION_OPTION, + drop_empty_layers: bool = True, ) -> dict[str, GridDataArray]: """ Allocate and distribute planar data for given discretization and npf @@ -196,6 +201,13 @@ def _allocate_planar_data( Model discretization package. allocation_option: ALLOCATION_OPTION The allocation option to use for the reallocation. + drop_empty_layers: bool, default True + If True, drop layers that contain no allocated cells anywhere in the + domain (or at any time), so the returned grids only span the layers that + are actually used. Reduces memory use and speeds up later operations such + as regridding, clipping and splitting. If no cells are allocated in any + layer, nothing is dropped and the full layer range is returned, as a layer + dimension of size 0 is not valid. Returns ------- @@ -210,11 +222,14 @@ def _allocate_planar_data( allocation_option, idomain > 0, planar_data["rate"], + drop_empty_layers=False, # Keep full here, drop empty layers below ) # remove rch from cells where it is not allocated and broadcast over layers. layered_data = {} layered_data["rate"] = planar_data["rate"].where(is_rch_cell) layered_data["rate"] = enforce_dim_order(layered_data["rate"]) + if drop_empty_layers: + layered_data = drop_empty_layers_from_dict(layered_data, is_rch_cell) return layered_data @classmethod diff --git a/imod/mf6/riv.py b/imod/mf6/riv.py index 0957ebbad..c80e89be2 100644 --- a/imod/mf6/riv.py +++ b/imod/mf6/riv.py @@ -19,7 +19,11 @@ from imod.mf6.utilities.package import set_repeat_stress_if_available from imod.mf6.validation import BOUNDARY_DIMS_SCHEMA, CONC_DIMS_SCHEMA from imod.prepare.cleanup import AlignLevelsMode, align_interface_levels, cleanup_riv -from imod.prepare.topsystem.allocation import ALLOCATION_OPTION, allocate_riv_cells +from imod.prepare.topsystem.allocation import ( + ALLOCATION_OPTION, + allocate_riv_cells, + drop_empty_layers_from_dict, +) from imod.prepare.topsystem.conductance import ( DISTRIBUTING_OPTION, distribute_drn_conductance, @@ -317,6 +321,7 @@ def _allocate_and_distribute_planar_data( npf: NodePropertyFlow, allocation_option: ALLOCATION_OPTION, distributing_option: DISTRIBUTING_OPTION, + drop_empty_layers: bool = True, ) -> tuple[GridDataDict, GridDataDict]: """ Allocate and distribute planar data for given discretization and npf @@ -338,6 +343,13 @@ def _allocate_and_distribute_planar_data( ALLOCATION_OPTION.at_first_active. distributing_option: DISTRIBUTING_OPTION distributing option. + drop_empty_layers: bool, default True + If True, drop layers that contain no allocated cells anywhere in the + domain (or at any time), so the returned grids only span the layers that + are actually used. Reduces memory use and speeds up later operations such + as regridding, clipping and splitting. If no cells are allocated in any + layer, nothing is dropped and the full layer range is returned, as a layer + dimension of size 0 is not valid. Returns ------- @@ -364,6 +376,7 @@ def _allocate_and_distribute_planar_data( bottom, planar_data["stage"], planar_data["bottom_elevation"], + drop_empty_layers=False, # Keep full layer range, drop empty layers below ) drn_is_allocated = drn_allocated is not None # Distribution of conductances @@ -409,6 +422,15 @@ def _allocate_and_distribute_planar_data( layered_data_riv["bottom_elevation"], bottom ) + if drop_empty_layers: + layered_data_riv = drop_empty_layers_from_dict( + layered_data_riv, riv_allocated + ) + if drn_allocated is not None: + layered_data_drn = drop_empty_layers_from_dict( + layered_data_drn, drn_allocated + ) + return layered_data_riv, layered_data_drn @classmethod diff --git a/imod/mf6/topsystem.py b/imod/mf6/topsystem.py index 692fa1d7a..a6ff44c42 100644 --- a/imod/mf6/topsystem.py +++ b/imod/mf6/topsystem.py @@ -58,6 +58,7 @@ def reallocate( npf: Optional[NodePropertyFlow] = None, allocation_option: Optional[ALLOCATION_OPTION] = None, distributing_option: Optional[DISTRIBUTING_OPTION] = None, + drop_empty_layers: bool = True, ) -> Self: """ Reallocates topsystem data across layers and create new package with it. @@ -83,6 +84,13 @@ def reallocate( The distributing option to use for the reallocation. Required for packages with a conductance variable. If None, the default is taken from :class:`imod.prepare.SimulationDistributingOptions`. + drop_empty_layers: bool, default True + If True, drop layers that contain no allocated cells anywhere in the + domain (or at any time), so the returned grids only span the layers that + are actually used. Reduces memory use and speeds up later operations such + as regridding, clipping and splitting. If no cells are allocated in any + layer, nothing is dropped and the full layer range is returned, as a layer + dimension of size 0 is not valid. Returns ------- @@ -102,10 +110,17 @@ def reallocate( npf = cast(NodePropertyFlow, npf) distributing_option = cast(DISTRIBUTING_OPTION, distributing_option) grid_dict = self._allocate_and_distribute_planar_data( - planar_data, dis, npf, allocation_option, distributing_option + planar_data, + dis, + npf, + allocation_option, + distributing_option, + drop_empty_layers, ) else: - grid_dict = self._allocate_planar_data(planar_data, dis, allocation_option) + grid_dict = self._allocate_planar_data( + planar_data, dis, allocation_option, drop_empty_layers + ) # River package returns a tuple (second argument can also be Drainage # package) if isinstance(grid_dict, tuple): @@ -122,6 +137,7 @@ def _allocate_and_distribute_planar_data( npf: NodePropertyFlow, allocation_option: ALLOCATION_OPTION, distributing_option: DISTRIBUTING_OPTION, + drop_empty_layers: bool = True, ) -> tuple[GridDataDict, GridDataDict] | GridDataDict: raise NotImplementedError( "This method should be implemented in the specific boundary condition " @@ -134,6 +150,7 @@ def _allocate_planar_data( planar_data: GridDataDict, dis: StructuredDiscretization | VerticesDiscretization, allocation_option: ALLOCATION_OPTION, + drop_empty_layers: bool = True, ) -> tuple[GridDataDict, GridDataDict] | GridDataDict: raise NotImplementedError( "This method should be implemented in the specific boundary condition " diff --git a/imod/prepare/cleanup.py b/imod/prepare/cleanup.py index 4200eba59..f49a02c21 100644 --- a/imod/prepare/cleanup.py +++ b/imod/prepare/cleanup.py @@ -27,6 +27,16 @@ def align_interface_levels( bottom: GridDataArray, method: AlignLevelsMode = AlignLevelsMode.TOPDOWN, ) -> tuple[GridDataArray, GridDataArray]: + # `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` + # down to `top`'s own layers first, so the comparison below doesn't fail + # with an alignment error; `top`'s layers are always the ones we want to + # keep, matching the ``join="left"`` pattern used in + # ``imod.common.utilities.mask.mask_da``. + if "layer" in top.dims and "layer" in bottom.dims: + bottom = bottom.sel(layer=top["layer"]) + to_align = top < bottom match method: 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..ea40858a9 100644 --- a/imod/prepare/topsystem/allocation.py +++ b/imod/prepare/topsystem/allocation.py @@ -8,12 +8,13 @@ import numpy as np from imod.common.utilities.layer import create_layered_top +from imod.logging import logger from imod.schemata import DimsSchema from imod.select.layers import ( get_upper_active_grid_cells, get_upper_active_layer_number, ) -from imod.typing import GridDataArray +from imod.typing import GridDataArray, GridDataDict from imod.util.dims import enforced_dim_order @@ -56,6 +57,19 @@ class ALLOCATION_OPTION(Enum): at_first_active = 9 # Not an iMOD 5.6 option +class LAYERS_USED(Enum): + """ + What this result represents. It can be either: + - ALL: nothing to trim (all layers used) + - SOME: some layers used, some empty + - NONE: nothing to keep (no layer has any allocated cell) + """ + + ALL = 0 + SOME = 1 + NONE = 2 + + PLANAR_GRID = ( DimsSchema("time", "y", "x") | DimsSchema("y", "x") @@ -71,6 +85,7 @@ def allocate_riv_cells( bottom: GridDataArray, stage: GridDataArray, bottom_elevation: GridDataArray, + drop_empty_layers: bool = True, ) -> tuple[GridDataArray, Optional[GridDataArray]]: """ Allocate river cells from a planar grid across the vertical dimension. @@ -96,6 +111,15 @@ 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 True + If True, drop layers that contain no allocated cells anywhere in the + domain (or at any time), so the returned grids only span the layers that + are actually used. Reduces memory use and speeds up later operations such + as regridding, clipping and splitting. If no cells are allocated in any + layer, nothing is dropped and the full layer range is returned, as a layer + dimension of size 0 is not valid. Note that allocation and conductance + distribution are always computed over the full layer range first, so + dropping layers does not change the computed values. Returns ------- @@ -111,21 +135,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 +165,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 +179,7 @@ def allocate_drn_cells( top: GridDataArray, bottom: GridDataArray, elevation: GridDataArray, + drop_empty_layers: bool = True, ) -> GridDataArray: """ Allocate drain cells from a planar grid across the vertical dimension. @@ -166,6 +202,15 @@ 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 True + If True, drop layers that contain no allocated cells anywhere in the + domain (or at any time), so the returned grids only span the layers that + are actually used. Reduces memory use and speeds up later operations such + as regridding, clipping and splitting. If no cells are allocated in any + layer, nothing is dropped and the full layer range is returned, as a layer + dimension of size 0 is not valid. Note that allocation and conductance + distribution are always computed over the full layer range first, so + dropping layers does not change the computed values. Returns ------- @@ -181,13 +226,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 +242,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, @@ -204,6 +251,7 @@ def allocate_ghb_cells( top: GridDataArray, bottom: GridDataArray, head: GridDataArray, + drop_empty_layers: bool = True, ) -> GridDataArray: """ Allocate general head boundary (GHB) cells from a planar grid across the @@ -227,6 +275,15 @@ def allocate_ghb_cells( head: DataArray | UgridDatarray Planar grid containing general head boundary's head. Is not allowed to have a layer dimension. + drop_empty_layers: bool, default True + If True, drop layers that contain no allocated cells anywhere in the + domain (or at any time), so the returned grids only span the layers that + are actually used. Reduces memory use and speeds up later operations such + as regridding, clipping and splitting. If no cells are allocated in any + layer, nothing is dropped and the full layer range is returned, as a layer + dimension of size 0 is not valid. Note that allocation and conductance + distribution are always computed over the full layer range first, so + dropping layers does not change the computed values. Returns ------- @@ -242,13 +299,13 @@ def allocate_ghb_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, head )[0] case ALLOCATION_OPTION.at_elevation: - return _allocate_cells__at_elevation(top, bottom, head)[0] + result = _allocate_cells__at_elevation(top, bottom, head)[0] case ALLOCATION_OPTION.at_first_active: - return _allocate_cells__at_first_active(active, head)[0] + result = _allocate_cells__at_first_active(active, head)[0] case _: raise ValueError( "Received incompatible setting for general head boundary, only" @@ -258,11 +315,14 @@ def allocate_ghb_cells( f"got: '{allocation_option.name}'" ) + return _drop_empty_layers(result) if drop_empty_layers else result + def allocate_rch_cells( allocation_option: ALLOCATION_OPTION, active: GridDataArray, rate: GridDataArray, + drop_empty_layers: bool = True, ) -> GridDataArray: """ Allocate recharge cells from a planar grid across the vertical dimension. @@ -279,6 +339,15 @@ 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 True + If True, drop layers that contain no allocated cells anywhere in the + domain (or at any time), so the returned grids only span the layers that + are actually used. Reduces memory use and speeds up later operations such + as regridding, clipping and splitting. If no cells are allocated in any + layer, nothing is dropped and the full layer range is returned, as a layer + dimension of size 0 is not valid. Note that allocation and conductance + distribution are always computed over the full layer range first, so + dropping layers does not change the computed values. Returns ------- @@ -294,7 +363,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 +371,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 +608,146 @@ def _allocate_cells__at_first_active( topsystem_upper_active = upper_active & ~np.isnan(planar_topsystem_grid) return topsystem_upper_active, None + + +def _used_layers(mask: GridDataArray) -> tuple[Optional[GridDataArray], LAYERS_USED]: + """ + Return the layer coordinate values of ``mask`` that contain at least one + True value anywhere in the domain (and, if present, at any timestep), or + None if there is nothing to trim (no layer dimension, or every layer has + data). + + Parameters + ---------- + mask: GridDataArray + Boolean array with a "layer" dimension, typically one of the + ``allocated`` grids returned by an ``_allocate_cells__*`` function. + + Returns + ------- + GridDataArray + Layer coordinate values with data. If all or none of the layers have + data, this will be None. In the latter case, do not return an empty + layer coordinate: a layer dimension of size 0 fails package validation + and breaks downstream operations. Keep the full layer range instead, so + empty packages can be removed by callers, e.g. + ``mask_package__drop_if_empty``. + LAYERS_USED + Enumerator indicating whether all, none, or some layers are used. + """ + if "layer" not in mask.dims: + return None, LAYERS_USED.ALL + + if mask.dtype != bool: + raise ValueError( + f"Expected a boolean grid to drop empty layers from, got: {mask.dtype}" + ) + + # Reduce the mask to the first time step if a time dimension exists for + # performance reasons + mask = mask.isel(time=0, drop=True, missing_dims="ignore") + reduce_dims = [d for d in mask.dims if d != "layer"] + has_data_per_layer = mask.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() + + # Never return an empty layer coordinate: keep the full range and + # let callers (e.g. # mask_package__drop_if_empty) remove empty packages. + if bool(has_data_per_layer.all()): + return None, LAYERS_USED.ALL + elif not bool(has_data_per_layer.any()): + return None, LAYERS_USED.NONE + + return mask["layer"].where(has_data_per_layer, drop=True), LAYERS_USED.SOME + + +def _drop_empty_layers(grid: GridDataArray) -> GridDataArray: + """ + Drop layers that contain no True 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 + Boolean array with a "layer" dimension, typically the output of one + of the ``_allocate_cells__*`` functions. + + Returns + ------- + GridDataArray + Same array, subset to layers with data. + """ + used_layers, layers_used_option = _used_layers(grid) + name = grid.name if hasattr(grid, "name") else "" + match layers_used_option: + case LAYERS_USED.NONE: + logger.warning( + f"No layers have data in grid '{name}', the package should be removed by the caller." + ) + return grid + case LAYERS_USED.ALL: + return grid + case LAYERS_USED.SOME: + logger.debug( + f"Dropping empty layers from grid '{name}', keeping only used layers." + ) + return grid.sel(layer=used_layers) + case _: + raise ValueError( + f"Unexpected value for layers_used_option: {layers_used_option}" + ) + + +def drop_empty_layers_from_dict( + data: GridDataDict, mask: GridDataArray +) -> GridDataDict: + """ + Trim every layered grid in ``data`` down to the layers that contain data + in ``mask``. Unlike :func:`_drop_empty_layers`, this can be applied to + grids (such as distributed conductances) that are not themselves + boolean, as long as a boolean ``mask`` with the same (full) layer range + is available to decide which layers to keep. + + Parameters + ---------- + data: GridDataDict + Dictionary of grids, e.g. the layered package data returned by + ``_allocate_and_distribute_planar_data``. Grids without a "layer" + dimension are returned unchanged. + mask: GridDataArray + Boolean array with a "layer" dimension, e.g. the ``allocated`` grid + that was used to build the grids in ``data``. + + Returns + ------- + GridDataDict + Same dictionary, with every layered grid subset to layers with data. + If mask has no True values, data is returned unchanged. + To fully remove empty empty packages, additional logic outside this function is required. + """ + used_layers, layers_used = _used_layers(mask) + names = list(data.keys()) + match layers_used: + case LAYERS_USED.NONE: + logger.warning( + f"No layers have data in grid '{names}', the package should be removed by the caller." + ) + return data + case LAYERS_USED.ALL: + return data # return as-is, let caller handle what to do with LAYERS_USED.NONE or LAYERS_USED.ALL + case LAYERS_USED.SOME: + logger.debug( + f"Dropping empty layers from grid '{names}', keeping only used layers." + ) + return { + key: grid.sel(layer=used_layers) if "layer" in grid.dims else grid + for key, grid in data.items() + } + case _: + raise ValueError(f"Unexpected value for layers_used: {layers_used}") diff --git a/imod/tests/test_mf6/test_mf6_drn.py b/imod/tests/test_mf6/test_mf6_drn.py index 62d2138aa..355496824 100644 --- a/imod/tests/test_mf6/test_mf6_drn.py +++ b/imod/tests/test_mf6/test_mf6_drn.py @@ -9,7 +9,6 @@ from pytest_cases import parametrize_with_cases import imod -import imod.mf6.drn from imod.common.utilities.version import get_version from imod.logging import LoggerType, LogLevel from imod.mf6.dis import StructuredDiscretization @@ -486,6 +485,39 @@ def test_reallocate(drainage): ) +def test_reallocate_drop_empty_layers(drainage): + """ + drop_empty_layers=True should trim layers off the final package without + changing the values of the layers that remain. + """ + drn = imod.mf6.Drainage(**drainage) + idomain = drainage["elevation"].astype(np.int16) + top = 1.0 + bottom = top - idomain.coords["layer"] + + dis = imod.mf6.StructuredDiscretization(top=top, bottom=bottom, idomain=idomain) + npf = imod.mf6.NodePropertyFlow(icelltype=0, k=1.0) + allocation_option = ALLOCATION_OPTION.first_active_to_elevation + distributing_option = DISTRIBUTING_OPTION.by_corrected_transmissivity + + full = drn.reallocate( + dis, npf, allocation_option, distributing_option, drop_empty_layers=False + ) + trimmed = drn.reallocate( + dis, npf, allocation_option, distributing_option, drop_empty_layers=True + ) + + full_layers = full.dataset["layer"].values + trimmed_layers = trimmed.dataset["layer"].values + assert set(trimmed_layers) <= set(full_layers) + assert len(trimmed_layers) < len(full_layers) + np.testing.assert_allclose( + trimmed["conductance"].sel(layer=trimmed_layers).values, + full["conductance"].sel(layer=trimmed_layers).values, + equal_nan=True, + ) + + def test_repr(drainage): repr_string = imod.mf6.Drainage(**drainage).__repr__() assert isinstance(repr_string, str) diff --git a/imod/tests/test_mf6/test_mf6_ghb.py b/imod/tests/test_mf6/test_mf6_ghb.py index 1005113f6..3143b3b3f 100644 --- a/imod/tests/test_mf6/test_mf6_ghb.py +++ b/imod/tests/test_mf6/test_mf6_ghb.py @@ -13,6 +13,49 @@ from imod.prepare.topsystem.conductance import DISTRIBUTING_OPTION +def test_reallocate_drop_empty_layers(): + """ + drop_empty_layers=True should trim layers off the final package without + changing the values of the layers that remain. + """ + layer = [1, 2, 3] + y = [25.0, 15.0, 5.0] + x = [5.0, 15.0, 25.0] + dx, dy = 10.0, -10.0 + head = xr.DataArray( + np.full((3, 3, 3), 1.0), + coords={"layer": layer, "y": y, "x": x, "dx": dx, "dy": dy}, + dims=("layer", "y", "x"), + ) + conductance = head.copy() + ghb = imod.mf6.GeneralHeadBoundary(head=head, conductance=conductance) + + idomain = head.astype(np.int16) + top = 1.0 + bottom = top - idomain.coords["layer"] + dis = imod.mf6.StructuredDiscretization(top=top, bottom=bottom, idomain=idomain) + npf = imod.mf6.NodePropertyFlow(icelltype=0, k=1.0) + allocation_option = ALLOCATION_OPTION.at_first_active + distributing_option = DISTRIBUTING_OPTION.by_layer_thickness + + full = ghb.reallocate( + dis, npf, allocation_option, distributing_option, drop_empty_layers=False + ) + trimmed = ghb.reallocate( + dis, npf, allocation_option, distributing_option, drop_empty_layers=True + ) + + full_layers = full.dataset["layer"].values + trimmed_layers = trimmed.dataset["layer"].values + assert set(trimmed_layers) <= set(full_layers) + assert len(trimmed_layers) < len(full_layers) + np.testing.assert_allclose( + trimmed["conductance"].sel(layer=trimmed_layers).values, + full["conductance"].sel(layer=trimmed_layers).values, + equal_nan=True, + ) + + @pytest.mark.unittest_jit def test_from_imod5_non_planar(imod5_dataset_periods, tmp_path): period_data = imod5_dataset_periods[1] diff --git a/imod/tests/test_mf6/test_mf6_rch.py b/imod/tests/test_mf6/test_mf6_rch.py index 5867234f0..393bba46a 100644 --- a/imod/tests/test_mf6/test_mf6_rch.py +++ b/imod/tests/test_mf6/test_mf6_rch.py @@ -402,6 +402,54 @@ def test_reallocate(rch_dict, allocation_option): assert rch_reallocated.dataset.equals(rch.dataset) +def test_reallocate_drop_empty_layers(): + """ + drop_empty_layers=True should trim layers off the final package without + changing the values of the layers that remain. + """ + x = [5.0, 15.0, 25.0] + y = [25.0, 15.0, 5.0] + layer = [1, 2, 3] + dx, dy = 10.0, -10.0 + + idomain = xr.DataArray( + np.ones((3, 3, 3), dtype=np.int16), + coords={"layer": layer, "y": y, "x": x, "dx": dx, "dy": dy}, + dims=("layer", "y", "x"), + ) + top = 1.0 + bottom = top - idomain.coords["layer"] + dis = imod.mf6.StructuredDiscretization(top=top, bottom=bottom, idomain=idomain) + + rate = xr.DataArray( + 1.0, + coords={"y": y, "x": x, "dx": dx, "dy": dy}, + dims=("y", "x"), + ).expand_dims(layer=[1]) + rch = imod.mf6.Recharge(rate=rate) + + full = rch.reallocate( + dis, + allocation_option=ALLOCATION_OPTION.at_first_active, + drop_empty_layers=False, + ) + trimmed = rch.reallocate( + dis, + allocation_option=ALLOCATION_OPTION.at_first_active, + drop_empty_layers=True, + ) + + full_layers = full.dataset["layer"].values + trimmed_layers = trimmed.dataset["layer"].values + assert set(trimmed_layers) <= set(full_layers) + assert len(trimmed_layers) < len(full_layers) + np.testing.assert_allclose( + trimmed["rate"].sel(layer=trimmed_layers).values, + full["rate"].sel(layer=trimmed_layers).values, + equal_nan=True, + ) + + @pytest.mark.unittest_jit def test_planar_rch_from_imod5_constant(imod5_dataset, tmp_path): data = deepcopy(imod5_dataset[0]) diff --git a/imod/tests/test_mf6/test_mf6_riv.py b/imod/tests/test_mf6/test_mf6_riv.py index 828d76e02..65352cdc8 100644 --- a/imod/tests/test_mf6/test_mf6_riv.py +++ b/imod/tests/test_mf6/test_mf6_riv.py @@ -486,6 +486,72 @@ def test_reallocate__wrong_allocation_option(riv_data, dis_data): river.reallocate(dis, npf, allocation_option, distributing_option) +def test_reallocate_drop_empty_layers(): + """ + drop_empty_layers=True should trim layers off the final package without + changing the values of the layers that remain. + """ + x = [5.0, 15.0, 25.0] + y = [25.0, 15.0, 5.0] + dx, dy = 10.0, -10.0 + layer = [1, 2, 3, 4] + + top = xr.DataArray( + 0.0, coords={"y": y, "x": x, "dx": dx, "dy": dy}, dims=("y", "x") + ) + bottom = xr.DataArray( + np.array([-1.0, -2.0, -3.0, -4.0])[:, None, None] * np.ones((4, 3, 3)), + coords={"layer": layer, "y": y, "x": x, "dx": dx, "dy": dy}, + dims=("layer", "y", "x"), + ) + idomain = xr.DataArray( + np.ones((4, 3, 3), dtype=int), + coords={"layer": layer, "y": y, "x": x, "dx": dx, "dy": dy}, + dims=("layer", "y", "x"), + ) + dis = imod.mf6.StructuredDiscretization(top=top, bottom=bottom, idomain=idomain) + npf = imod.mf6.NodePropertyFlow(icelltype=0, k=1.0) + + # Stage and bottom entirely confined to layer 2 (-1.0 to -2.0). + planar_coords = {"y": y, "x": x, "dx": dx, "dy": dy} + river = imod.mf6.River( + stage=xr.DataArray(-1.1, coords=planar_coords, dims=("y", "x")).expand_dims( + layer=[1] + ), + conductance=xr.DataArray( + 10.0, coords=planar_coords, dims=("y", "x") + ).expand_dims(layer=[1]), + bottom_elevation=xr.DataArray( + -1.9, coords=planar_coords, dims=("y", "x") + ).expand_dims(layer=[1]), + ) + + allocation_option = ALLOCATION_OPTION.stage_to_riv_bot + distributing_option = DISTRIBUTING_OPTION.by_corrected_transmissivity + + full = river.reallocate( + dis, npf, allocation_option, distributing_option, drop_empty_layers=False + ) + trimmed = river.reallocate( + dis, npf, allocation_option, distributing_option, drop_empty_layers=True + ) + + full_layers = full.dataset["layer"].values + trimmed_layers = trimmed.dataset["layer"].values + assert set(trimmed_layers) <= set(full_layers) + assert len(trimmed_layers) < len(full_layers) + np.testing.assert_allclose( + trimmed["conductance"].sel(layer=trimmed_layers).values, + full["conductance"].sel(layer=trimmed_layers).values, + equal_nan=True, + ) + np.testing.assert_allclose( + trimmed["stage"].sel(layer=trimmed_layers).values, + full["stage"].sel(layer=trimmed_layers).values, + equal_nan=True, + ) + + def test_check_dim_monotonicity(): """ Test if dimensions are monotonically increasing or, in case of the y coord, diff --git a/imod/tests/test_prepare/test_cleanup.py b/imod/tests/test_prepare/test_cleanup.py index 7e2de1929..5346a9edd 100644 --- a/imod/tests/test_prepare/test_cleanup.py +++ b/imod/tests/test_prepare/test_cleanup.py @@ -196,6 +196,38 @@ def test_cleanup_riv__stage_equals_bottom_elevation(riv_data: dict, dis_data: di ) +@parametrize_with_cases("riv_data, dis_data", cases=RivDisCases) +def test_cleanup_riv__trimmed_layers(riv_data: dict, dis_data: dict): + """ + A river package's own grids may have fewer layers than the model's full + ``bottom`` (e.g. produced via ``drop_empty_layers=True``, see + ``imod.prepare.topsystem``). ``cleanup_riv`` should still work in that + case, keeping the package's own (trimmed) layer coordinate rather than + raising an alignment error. + """ + dis_dict = _prepare_dis_dict(dis_data, cleanup_riv) + # Trim the river package down to a single layer, model bottom stays full range. + trimmed_layer = riv_data["stage"]["layer"].isel(layer=[0]) + riv_data = { + key: (value.sel(layer=trimmed_layer) if "layer" in value.dims else value) + for key, value in riv_data.items() + } + # Force a bottom_elevation/bottom mismatch, to also exercise align_interface_levels. + riv_data["bottom_elevation"] -= 3.0 + + riv_data_cleaned = cleanup_riv(**dis_dict, **riv_data) + + # Returned grids keep the package's own (trimmed) layers. + for value in riv_data_cleaned.values(): + if value is not None and "layer" in value.dims: + np.testing.assert_equal(value["layer"].values, trimmed_layer.values) + riv_active = riv_data_cleaned["stage"].notnull() + expected = dis_dict["bottom"].sel(layer=trimmed_layer).where(riv_active) + np.testing.assert_equal( + riv_data_cleaned["bottom_elevation"].values, expected.values + ) + + @parametrize_with_cases("riv_data, dis_data", cases=RivDisCases) def test_cleanup_riv__raise_error(riv_data: dict, dis_data: dict): """ diff --git a/imod/tests/test_prepare/test_topsystem.py b/imod/tests/test_prepare/test_topsystem.py index 65295c7ff..9b0ab8a0e 100644 --- a/imod/tests/test_prepare/test_topsystem.py +++ b/imod/tests/test_prepare/test_topsystem.py @@ -52,7 +52,7 @@ 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 + option, active, top, bottom, stage, bottom_elevation, drop_empty_layers=False ) actual_riv = take_nth_layer_column(actual_riv_da, 0) @@ -71,6 +71,20 @@ def test_riv_allocation( if empty_drn is not None: assert np.all(~empty_drn) + # drop_empty_layers=True should keep only the layers with allocated cells + actual_riv_da, actual_drn_da = allocate_riv_cells( + option, active, top, bottom, stage, bottom_elevation, drop_empty_layers=True + ) + expected_riv_layers = np.nonzero(expected_riv)[0] + 1 + np.testing.assert_array_equal( + actual_riv_da.coords["layer"].values, expected_riv_layers + ) + if actual_drn_da is not None: + expected_drn_layers = np.nonzero(expected_drn)[0] + 1 + np.testing.assert_array_equal( + actual_drn_da.coords["layer"].values, expected_drn_layers + ) + @parametrize_with_cases( argnames="active,top,bottom,drn_elevation", @@ -80,7 +94,9 @@ def test_riv_allocation( argnames="option,expected,_", prefix="allocation_", has_tag="drn" ) def test_drn_allocation(active, top, bottom, drn_elevation, option, expected, _): - actual_da = allocate_drn_cells(option, active, top, bottom, drn_elevation) + actual_da = allocate_drn_cells( + option, active, top, bottom, drn_elevation, drop_empty_layers=False + ) actual = take_nth_layer_column(actual_da, 0) empty = take_nth_layer_column(actual_da, 1) @@ -88,6 +104,13 @@ def test_drn_allocation(active, top, bottom, drn_elevation, option, expected, _) np.testing.assert_equal(actual, expected) assert np.all(~empty) + # drop_empty_layers=True should keep only the layers with allocated cells + actual_da = allocate_drn_cells( + option, active, top, bottom, drn_elevation, drop_empty_layers=True + ) + expected_layers = np.nonzero(expected)[0] + 1 + np.testing.assert_array_equal(actual_da.coords["layer"].values, expected_layers) + @parametrize_with_cases( argnames="active,top,bottom,head", @@ -97,7 +120,9 @@ def test_drn_allocation(active, top, bottom, drn_elevation, option, expected, _) argnames="option,expected,_", prefix="allocation_", has_tag="ghb" ) def test_ghb_allocation(active, top, bottom, head, option, expected, _): - actual_da = allocate_ghb_cells(option, active, top, bottom, head) + actual_da = allocate_ghb_cells( + option, active, top, bottom, head, drop_empty_layers=False + ) actual = take_nth_layer_column(actual_da, 0) empty = take_nth_layer_column(actual_da, 1) @@ -105,6 +130,13 @@ def test_ghb_allocation(active, top, bottom, head, option, expected, _): np.testing.assert_equal(actual, expected) assert np.all(~empty) + # drop_empty_layers=True should keep only the layers with allocated cells + actual_da = allocate_ghb_cells( + option, active, top, bottom, head, drop_empty_layers=True + ) + expected_layers = np.nonzero(expected)[0] + 1 + np.testing.assert_array_equal(actual_da.coords["layer"].values, expected_layers) + @parametrize_with_cases( argnames="active,rate", @@ -114,7 +146,7 @@ def test_ghb_allocation(active, top, bottom, head, option, expected, _): argnames="option,expected,_", prefix="allocation_", has_tag="rch" ) def test_rch_allocation(active, rate, option, expected, _): - actual_da = allocate_rch_cells(option, active, rate) + actual_da = allocate_rch_cells(option, active, rate, drop_empty_layers=False) actual = take_nth_layer_column(actual_da, 0) empty = take_nth_layer_column(actual_da, 1) @@ -122,6 +154,11 @@ def test_rch_allocation(active, rate, option, expected, _): np.testing.assert_equal(actual, expected) assert np.all(~empty) + # drop_empty_layers=True should keep only the layers with allocated cells + actual_da = allocate_rch_cells(option, active, rate, drop_empty_layers=True) + expected_layers = np.nonzero(expected)[0] + 1 + np.testing.assert_array_equal(actual_da.coords["layer"].values, expected_layers) + @parametrize_with_cases( argnames="active,top,bottom,stage,bottom_elevation", @@ -211,7 +248,13 @@ def test_riv_allocation__elevation_above_surface_level( # Put elevations a lot above surface level. Need to be allocated to first # layer. actual_riv_da, actual_drn_da = allocate_riv_cells( - option, active, top, bottom, stage + 100.0, bottom_elevation + 100.0 + option, + active, + top, + bottom, + stage + 100.0, + bottom_elevation + 100.0, + drop_empty_layers=False, ) # Override expected values @@ -235,6 +278,38 @@ def test_riv_allocation__elevation_above_surface_level( if empty_drn is not None: assert np.all(~empty_drn) + # drop_empty_layers=True should keep only the layers with allocated cells + actual_riv_da, actual_drn_da = allocate_riv_cells( + option, + active, + top, + bottom, + stage + 100.0, + bottom_elevation + 100.0, + drop_empty_layers=True, + ) + + expected_riv_layers = np.nonzero(expected_riv)[0] + 1 + expected_riv_layers = ( + expected_riv_layers + if expected_riv_layers.size > 0 + else np.arange(1, len(expected_riv) + 1) + ) + np.testing.assert_array_equal( + actual_riv_da.coords["layer"].values, expected_riv_layers + ) + if actual_drn_da is not None: + expected_drn_layers = np.nonzero(expected_drn)[0] + 1 + expected_drn_layers = ( + expected_drn_layers + if expected_drn_layers.size > 0 + else np.arange(1, len(expected_drn) + 1) + ) + np.testing.assert_array_equal( + actual_drn_da.coords["layer"].values, + expected_drn_layers, + ) + @parametrize_with_cases( argnames="active,top,bottom,stage,bottom_elevation", @@ -248,7 +323,7 @@ def test_riv_allocation__stage_equals_bottom_elevation( ): # Bottom elevation equals stage here. actual_riv_da, actual_drn_da = allocate_riv_cells( - option, active, top, bottom, stage, stage + option, active, top, bottom, stage, stage, drop_empty_layers=False ) # Override expected values @@ -275,6 +350,20 @@ def test_riv_allocation__stage_equals_bottom_elevation( if empty_drn is not None: assert np.all(~empty_drn) + # drop_empty_layers=True should keep only the layers with allocated cells + actual_riv_da, actual_drn_da = allocate_riv_cells( + option, active, top, bottom, stage, stage, drop_empty_layers=True + ) + expected_riv_layers = np.nonzero(expected_riv)[0] + 1 + np.testing.assert_array_equal( + actual_riv_da.coords["layer"].values, expected_riv_layers + ) + if actual_drn_da is not None: + expected_drn_layers = np.nonzero(expected_drn)[0] + 1 + np.testing.assert_array_equal( + actual_drn_da.coords["layer"].values, expected_drn_layers + ) + @parametrize_with_cases( argnames="active,top,bottom,stage,bottom_elevation", @@ -291,7 +380,7 @@ def test_riv_allocation__stage_equals_bottom_elevation_equals_bottom( # Bottom elevation equals stage here. actual_riv_da, actual_drn_da = allocate_riv_cells( - option, active, top, bottom, stage, stage + option, active, top, bottom, stage, stage, drop_empty_layers=False ) # Override expected values @@ -318,6 +407,20 @@ def test_riv_allocation__stage_equals_bottom_elevation_equals_bottom( if empty_drn is not None: assert np.all(~empty_drn) + # drop_empty_layers=True should keep only the layers with allocated cells + actual_riv_da, actual_drn_da = allocate_riv_cells( + option, active, top, bottom, stage, stage, drop_empty_layers=True + ) + expected_riv_layers = np.nonzero(expected_riv)[0] + 1 + np.testing.assert_array_equal( + actual_riv_da.coords["layer"].values, expected_riv_layers + ) + if actual_drn_da is not None: + expected_drn_layers = np.nonzero(expected_drn)[0] + 1 + np.testing.assert_array_equal( + actual_drn_da.coords["layer"].values, expected_drn_layers + ) + @parametrize_with_cases( argnames="active,top,bottom,elevation", @@ -337,6 +440,7 @@ def test_drn_allocation__elevation_above_surface_level( top, bottom, elevation + 100.0, + drop_empty_layers=False, ) # Override expected @@ -350,6 +454,18 @@ def test_drn_allocation__elevation_above_surface_level( if empty is not None: assert np.all(~empty) + # drop_empty_layers=True should keep only the layers with allocated cells + actual_da = allocate_drn_cells( + option, + active, + top, + bottom, + elevation + 100.0, + drop_empty_layers=True, + ) + expected_layers = np.nonzero(expected)[0] + 1 + np.testing.assert_array_equal(actual_da.coords["layer"].values, expected_layers) + @parametrize_with_cases( argnames="active,top,bottom,drn_elevation", @@ -365,7 +481,9 @@ def test_drn_allocation__elevation_equal_to_bottom( # elevation is equal everywhere.) bottom.loc[bottom.coords["layer"] == 3] = drn_elevation.values.ravel()[0] - actual_da = allocate_drn_cells(option, active, top, bottom, drn_elevation) + actual_da = allocate_drn_cells( + option, active, top, bottom, drn_elevation, drop_empty_layers=False + ) actual = take_nth_layer_column(actual_da, 0) empty = take_nth_layer_column(actual_da, 1) @@ -373,6 +491,13 @@ def test_drn_allocation__elevation_equal_to_bottom( np.testing.assert_equal(actual, expected) assert np.all(~empty) + # drop_empty_layers=True should keep only the layers with allocated cells + actual_da = allocate_drn_cells( + option, active, top, bottom, drn_elevation, drop_empty_layers=True + ) + expected_layers = np.nonzero(expected)[0] + 1 + np.testing.assert_array_equal(actual_da.coords["layer"].values, expected_layers) + @parametrize_with_cases( argnames="active,top,bottom,head", @@ -392,6 +517,7 @@ def test_ghb_allocation__elevation_above_surface_level( top, bottom, head + 100.0, + drop_empty_layers=False, ) # Override expected @@ -405,6 +531,18 @@ def test_ghb_allocation__elevation_above_surface_level( if empty is not None: assert np.all(~empty) + # drop_empty_layers=True should keep only the layers with allocated cells + actual_da = allocate_ghb_cells( + option, + active, + top, + bottom, + head + 100.0, + drop_empty_layers=True, + ) + expected_layers = np.nonzero(expected)[0] + 1 + np.testing.assert_array_equal(actual_da.coords["layer"].values, expected_layers) + @parametrize_with_cases( argnames="active,top,bottom,elevation", 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" + )