diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst index cfd932663..483d30ba9 100644 --- a/docs/api/changelog.rst +++ b/docs/api/changelog.rst @@ -17,6 +17,12 @@ Added :func:`imod.prepare.allocate_rch_cells`, :func:`imod.prepare.allocate_riv_cells`, for this multiple options can be selected, available in :func:`imod.prepare.ALLOCATION_OPTION`. +- Added functions to distribute conductances of planar grids over layers for the + topsystem in :func:`imod.prepare.distribute_riv_conductance`, + :func:`imod.prepare.distribute_drn_conductance`, + :func:`imod.prepare.distribute_ghb_conductance`, for this multiple options can + be selected, available in :func:`imod.prepare.DISTRIBUTION_OPTION`. + Fixed ~~~~~ diff --git a/docs/api/prepare.rst b/docs/api/prepare.rst index 28cbb2c5f..efc863159 100644 --- a/docs/api/prepare.rst +++ b/docs/api/prepare.rst @@ -34,9 +34,13 @@ Prepare model input create_layered_top ALLOCATION_OPTION + DISTRIBUTION_OPTION allocate_drn_cells allocate_ghb_cells allocate_rch_cells allocate_riv_cells c_leakage - c_radial \ No newline at end of file + c_radial + distribute_drn_conductance + distribute_ghb_conductance + distribute_riv_conductance diff --git a/imod/prepare/__init__.py b/imod/prepare/__init__.py index 004612578..13032246d 100644 --- a/imod/prepare/__init__.py +++ b/imod/prepare/__init__.py @@ -37,12 +37,16 @@ ) from imod.prepare.topsystem import ( ALLOCATION_OPTION, + DISTRIBUTING_OPTION, allocate_drn_cells, allocate_ghb_cells, allocate_rch_cells, allocate_riv_cells, c_leakage, c_radial, + distribute_drn_conductance, + distribute_ghb_conductance, + distribute_riv_conductance, ) from imod.prepare.voxelize import Voxelizer from imod.prepare.wells import assign_wells diff --git a/imod/prepare/topsystem/__init__.py b/imod/prepare/topsystem/__init__.py index a63b10718..e3445e6eb 100644 --- a/imod/prepare/topsystem/__init__.py +++ b/imod/prepare/topsystem/__init__.py @@ -5,4 +5,10 @@ allocate_rch_cells, allocate_riv_cells, ) +from imod.prepare.topsystem.conductance import ( + DISTRIBUTING_OPTION, + distribute_drn_conductance, + distribute_ghb_conductance, + distribute_riv_conductance, +) from imod.prepare.topsystem.resistance import c_leakage, c_radial diff --git a/imod/prepare/topsystem/allocation.py b/imod/prepare/topsystem/allocation.py index 792bf3547..35c8f59b0 100644 --- a/imod/prepare/topsystem/allocation.py +++ b/imod/prepare/topsystem/allocation.py @@ -134,6 +134,13 @@ def _is_layered(grid: GridDataArray): return "layer" in grid.sizes and grid.sizes["layer"] > 1 +def _enforce_layered_top(top: GridDataArray, bottom: GridDataArray): + if _is_layered(top): + return top + else: + return create_layered_top(bottom, top) + + @enforced_dim_order def _allocate_cells__stage_to_riv_bot( top: GridDataArray, @@ -155,9 +162,11 @@ def _allocate_cells__stage_to_riv_bot( bottom: GridDataArray bottom of model layers stage: GridDataArray - river stage + river stage, cannot contain a layer dimension. Can contain a time + dimension. bottom_elevation: GridDataArray - river bottom elevation + river bottom elevation, cannot contain a layer dimension. Can contain a + time dimension. Returns ------- @@ -167,10 +176,7 @@ def _allocate_cells__stage_to_riv_bot( PLANAR_GRID.validate(stage) PLANAR_GRID.validate(bottom_elevation) - if _is_layered(top): - top_layered = top - else: - top_layered = create_layered_top(bottom, top) + top_layered = _enforce_layered_top(top, bottom) riv_cells = (stage > bottom) & (bottom_elevation < top_layered) @@ -200,7 +206,8 @@ def _allocate_cells__first_active_to_riv_bot( bottom: GridDataArray bottom of model layers bottom_elevation: GridDataArray - river bottom elevation + river bottom elevation, cannot contain a layer dimension. Can contain a + time dimension. Returns ------- @@ -212,10 +219,7 @@ def _allocate_cells__first_active_to_riv_bot( upper_active_layer = get_upper_active_layer_number(active) layer = active.coords["layer"] - if _is_layered(top): - top_layered = top - else: - top_layered = create_layered_top(bottom, top) + top_layered = _enforce_layered_top(top, bottom) riv_cells = (upper_active_layer <= layer) & (bottom_elevation < top_layered) @@ -247,9 +251,11 @@ def _allocate_cells__first_active_to_riv_bot__drn( bottom: GridDataArray bottom of model layers stage: GridDataArray - river stage + river stage, cannot contain a layer dimension. Can contain a time + dimension. bottom_elevation: GridDataArray - river bottom elevation + river bottom elevation, cannot contain a layer dimension. Can contain a + time dimension. Returns ------- @@ -262,10 +268,7 @@ def _allocate_cells__first_active_to_riv_bot__drn( PLANAR_GRID.validate(stage) PLANAR_GRID.validate(bottom_elevation) - if _is_layered(top): - top_layered = top - else: - top_layered = create_layered_top(bottom, top) + top_layered = _enforce_layered_top(top, bottom) upper_active_layer = get_upper_active_layer_number(active) layer = active.coords["layer"] @@ -294,7 +297,8 @@ def _allocate_cells__at_elevation( bottom: GridDataArray bottom of model layers elevation: GridDataArray - elevation. Can be river bottom, drain elevation or head of GHB. + elevation. Can be river bottom, drain elevation or head of GHB. Cannot + contain a layer dimension. Can contain a time dimension. Returns ------- @@ -304,10 +308,7 @@ def _allocate_cells__at_elevation( PLANAR_GRID.validate(elevation) - if _is_layered(top): - top_layered = top - else: - top_layered = create_layered_top(bottom, top) + top_layered = _enforce_layered_top(top, bottom) riv_cells = (elevation < top_layered) & (elevation >= bottom) diff --git a/imod/prepare/topsystem/conductance.py b/imod/prepare/topsystem/conductance.py new file mode 100644 index 000000000..1b3a408bc --- /dev/null +++ b/imod/prepare/topsystem/conductance.py @@ -0,0 +1,452 @@ +from enum import Enum + +import numpy as np + +from imod.prepare.topsystem.allocation import _enforce_layered_top +from imod.schemata import DimsSchema +from imod.typing import GridDataArray +from imod.typing.grid import ones_like, preserve_gridtype +from imod.util.dims import enforced_dim_order + + +class DISTRIBUTING_OPTION(Enum): + """ + Enumerator containing settings to distribute 2D conductance grids over + vertical layers for the RIV, DRN or GHB package. + + * ``by_corrected_transmissivity``: RIV. Distribute the conductance by + corrected transmissivities. Crosscut thicknesses are used to compute + transmissivities. The crosscut thicknesses is computed based on the + overlap of bottom_elevation over the bottom allocated layer. Same holds + for the stage and top allocated layer. Furthermore the method corrects + distribution weights for the mismatch between the midpoints of crosscut + areas and model layer midpoints. This is the default method in iMOD 5.6, + thus DISTRCOND = 0. + * ``equally``: RIV, DRN, GHB. Distribute conductances equally over layers. + This matches iMOD 5.6 DISTRCOND = 1 option. + * ``by_crosscut_thickness``: RIV. Distribute the conductance by crosscut + thicknesses. The crosscut thicknesses is computed based on the overlap of + bottom_elevation over the bottom allocated layer. Same holds for the stage + and top allocated layer. This matches iMOD 5.6 DISTRCOND = 2 option. + * ``by_layer_thickness``: RIV, DRN, GHB. Distribute the conductance by model + layer thickness. This matches iMOD 5.6 DISTRCOND = 3 option. + * ``by_crosscut_transmissivity``: RIV. Distribute the conductance by + crosscut transmissivity. Crosscut thicknesses are used to compute + transmissivities. The crosscut thicknesses is computed based on the + overlap of bottom_elevation over the bottom allocated layer. Same holds + for the stage and top allocated layer. This matches iMOD 5.6 DISTRCOND = 4 + option. + * ``by_conductivity``: RIV, DRN, GHB. Distribute the conductance weighted by + model layer hydraulic conductivities. This matches iMOD 5.6 DISTRCOND = 5 + option. + * ``by_layer_transmissivity``: RIV, DRN, GHB. Distribute the conductance by + model layer transmissivity. This has no equivalent in iMOD 5.6. + """ + + by_corrected_transmissivity = 0 + equally = 1 + by_crosscut_thickness = 2 + by_layer_thickness = 3 + by_crosscut_transmissivity = 4 + by_conductivity = 5 + by_layer_transmissivity = 9 # Not an iMOD 5.6 option + + +PLANAR_GRID = ( + DimsSchema("time", "y", "x") + | DimsSchema("y", "x") + | DimsSchema("time", "{face_dim}") + | DimsSchema("{face_dim}") +) + + +@enforced_dim_order +def distribute_riv_conductance( + distributing_option: DISTRIBUTING_OPTION, + allocated: GridDataArray, + conductance: GridDataArray, + top: GridDataArray, + bottom: GridDataArray, + stage: GridDataArray, + bottom_elevation: GridDataArray, + k: GridDataArray, +) -> GridDataArray: + """ + Function to distribute 2D conductance over vertical layers for the RIV + package. Multiple options are available, which need to be selected in the + DISTRIBUTING_OPTION enumerator. + + Parameters + ---------- + distributing_option : DISTRIBUTING_OPTION + Distributing option available in the DISTRIBUTING_OPTION enumerator. + allocated: DataArray | UgridDataArray + 3D boolean array with river cell locations. This can be made with the + :func:`imod.prepare.allocate_riv_cells` function. + conductance: DataArray | UgridDataArray + Planar grid with conductances that need to be distributed over layers, + so grid cannot contain a layer dimension. Can contain a time dimension. + top: DataArray | UgridDataArray + Model top + bottom: DataArray | UgridDataArray + Model layer bottoms + stage: DataArray | UgridDataArray + Planar grid with river stages, cannot contain a layer dimension. Can + contain a time dimension. + bottom_elevation: DataArray | UgridDataArray + Planar grid with river bottom elevations, cannot contain a layer + dimension. Can contain a time dimension. + k: DataArray | UgridDataArray + Hydraulic conductivities + + Returns + ------- + Conductances distributed over depth. + + Examples + -------- + >>> from imod.prepare import allocate_riv_cells, distribute_riv_conductance, ALLOCATION_OPTION, DISTRIBUTING_OPTION + >>> allocated = allocate_riv_cells( + ALLOCATION_OPTION.stage_to_riv_bot, active, top, bottom, stage, bottom_elevation + ) + >>> conductances_distributed = distribute_riv_conductance( + DISTRIBUTING_OPTION.by_corrected_transmissivity, allocated, + conductance, top, bottom, stage, bottom_elevation, k + ) + """ + PLANAR_GRID.validate(conductance) + + match distributing_option: + case DISTRIBUTING_OPTION.by_corrected_transmissivity: + weights = _distribute_weights__by_corrected_transmissivity( + allocated, top, bottom, stage, bottom_elevation, k + ) + case DISTRIBUTING_OPTION.equally: + weights = _distribute_weights__equally(allocated) + case DISTRIBUTING_OPTION.by_crosscut_thickness: + weights = _distribute_weights__by_crosscut_thickness( + allocated, top, bottom, stage, bottom_elevation + ) + case DISTRIBUTING_OPTION.by_layer_thickness: + weights = _distribute_weights__by_layer_thickness(allocated, top, bottom) + case DISTRIBUTING_OPTION.by_crosscut_transmissivity: + weights = _distribute_weights__by_crosscut_transmissivity( + allocated, top, bottom, stage, bottom_elevation, k + ) + case DISTRIBUTING_OPTION.by_layer_transmissivity: + weights = _distribute_weights__by_layer_transmissivity( + allocated, top, bottom, k + ) + case DISTRIBUTING_OPTION.by_conductivity: + weights = _distribute_weights__by_conductivity(allocated, k) + case _: + raise ValueError( + "Received incompatible setting for rivers, only" + f"'{DISTRIBUTING_OPTION.by_corrected_transmissivity.name}' and" + f"'{DISTRIBUTING_OPTION.equally.name}' and" + f"'{DISTRIBUTING_OPTION.by_crosscut_thickness.name}' and" + f"'{DISTRIBUTING_OPTION.by_layer_thickness.name}' and" + f"'{DISTRIBUTING_OPTION.by_crosscut_transmissivity.name}' and" + f"'{DISTRIBUTING_OPTION.by_layer_transmissivity.name}' and" + f"'{DISTRIBUTING_OPTION.by_conductivity.name}' supported." + f"got: '{distributing_option.name}'" + ) + return (weights * conductance).where(allocated) + + +@enforced_dim_order +def distribute_drn_conductance( + distributing_option: DISTRIBUTING_OPTION, + allocated: GridDataArray, + conductance: GridDataArray, + top: GridDataArray, + bottom: GridDataArray, + k: GridDataArray, +) -> GridDataArray: + """ + Function to distribute 2D conductance over vertical layers for the DRN + package. Multiple options are available, which need to be selected in the + DISTRIBUTING_OPTION enumerator. + + Parameters + ---------- + distributing_option : DISTRIBUTING_OPTION + Distributing option available in the DISTRIBUTING_OPTION enumerator. + allocated: DataArray | UgridDataArray + 3D boolean array with drain cell locations. This can be made with the + :func:`imod.prepare.allocate_drn_cells` function. + conductance: DataArray | UgridDataArray + Planar grid with conductances that need to be distributed over layers, + so grid cannot contain a layer dimension. Can contain a time dimension. + top: DataArray | UgridDataArray + Model top + bottom: DataArray | UgridDataArray + Model layer bottoms + k: DataArray | UgridDataArray + Hydraulic conductivities + + Returns + ------- + Conductances distributed over depth. + + Examples + -------- + >>> from imod.prepare import allocate_drn_cells, distribute_drn_conductance, ALLOCATION_OPTION, DISTRIBUTING_OPTION + >>> allocated = allocate_drn_cells( + ALLOCATION_OPTION.at_elevation, active, top, bottom, drain_elevation + ) + >>> conductances_distributed = distribute_drn_conductance( + DISTRIBUTING_OPTION.by_layer_transmissivity, allocated, + conductance, top, bottom, k + ) + """ + PLANAR_GRID.validate(conductance) + + match distributing_option: + case DISTRIBUTING_OPTION.equally: + weights = _distribute_weights__equally(allocated) + case DISTRIBUTING_OPTION.by_layer_thickness: + weights = _distribute_weights__by_layer_thickness(allocated, top, bottom) + case DISTRIBUTING_OPTION.by_layer_transmissivity: + weights = _distribute_weights__by_layer_transmissivity( + allocated, top, bottom, k + ) + case DISTRIBUTING_OPTION.by_conductivity: + weights = _distribute_weights__by_conductivity(allocated, k) + case _: + raise ValueError( + "Received incompatible setting for drains, only" + f"'{DISTRIBUTING_OPTION.equally.name}' and" + f"'{DISTRIBUTING_OPTION.by_layer_thickness.name}' and" + f"'{DISTRIBUTING_OPTION.by_layer_transmissivity.name}' and" + f"'{DISTRIBUTING_OPTION.by_conductivity.name}' supported." + f"got: '{distributing_option.name}'" + ) + return (weights * conductance).where(allocated) + + +@enforced_dim_order +def distribute_ghb_conductance( + distributing_option: DISTRIBUTING_OPTION, + allocated: GridDataArray, + conductance: GridDataArray, + top: GridDataArray, + bottom: GridDataArray, + k: GridDataArray, +) -> GridDataArray: + PLANAR_GRID.validate(conductance) + """ + Function to distribute 2D conductance over vertical layers for the GHB + package. Multiple options are available, which need to be selected in the + DISTRIBUTING_OPTION enumerator. + + Parameters + ---------- + distributing_option : DISTRIBUTING_OPTION + Distributing option available in the DISTRIBUTING_OPTION enumerator. + allocated: DataArray | UgridDataArray + 3D boolean array with GHB cell locations. This can be made with the + :func:`imod.prepare.allocate_ghb_cells` function. + conductance: DataArray | UgridDataArray + Planar grid with conductances that need to be distributed over layers, + so grid cannot contain a layer dimension. Can contain a time dimension. + top: DataArray | UgridDataArray + Model top + bottom: DataArray | UgridDataArray + Model layer bottoms + k: DataArray | UgridDataArray + Hydraulic conductivities + + Returns + ------- + Conductances distributed over depth. + + Examples + -------- + >>> from imod.prepare import allocate_ghb_cells, distribute_drn_conductance, ALLOCATION_OPTION, DISTRIBUTING_OPTION + >>> allocated = allocate_ghb_cells( + ALLOCATION_OPTION.at_elevation, active, top, bottom, ghb_head + ) + >>> conductances_distributed = distribute_ghb_conductance( + DISTRIBUTING_OPTION.by_layer_transmissivity, allocated, + conductance, top, bottom, k + ) + """ + match distributing_option: + case DISTRIBUTING_OPTION.equally: + weights = _distribute_weights__equally(allocated) + case DISTRIBUTING_OPTION.by_layer_thickness: + weights = _distribute_weights__by_layer_thickness(allocated, top, bottom) + case DISTRIBUTING_OPTION.by_layer_transmissivity: + weights = _distribute_weights__by_layer_transmissivity( + allocated, top, bottom, k + ) + case DISTRIBUTING_OPTION.by_conductivity: + weights = _distribute_weights__by_conductivity(allocated, k) + case _: + raise ValueError( + "Received incompatible setting for general head boundary, only" + f"'{DISTRIBUTING_OPTION.equally.name}' and" + f"'{DISTRIBUTING_OPTION.by_layer_thickness.name}' and" + f"'{DISTRIBUTING_OPTION.by_layer_transmissivity.name}' and" + f"'{DISTRIBUTING_OPTION.by_conductivity.name}' supported." + f"got: '{distributing_option.name}'" + ) + return (weights * conductance).where(allocated) + + +@preserve_gridtype +def _compute_layer_thickness(allocated, top, bottom): + """ + Compute 3D grid of thicknesses in allocated cells + """ + top_layered = _enforce_layered_top(top, bottom) + + thickness = top_layered - bottom + return thickness.where(allocated) + + +@preserve_gridtype +def _compute_crosscut_thickness(allocated, top, bottom, stage, bottom_elevation): + """ + Compute 3D grid of thicknesses crosscut by river in allocated cells. So the + upper allocated layer thickness is stage - bottom, the lower allocated layer + is top - river_bottom_elevation. + """ + top_layered = _enforce_layered_top(top, bottom) + + thickness = _compute_layer_thickness(allocated, top, bottom) + + upper_layer_allocated = (stage < top_layered) & (stage > bottom) + lower_layer_allocated = (bottom_elevation < top_layered) & ( + bottom_elevation > bottom + ) + outside = (stage < bottom) | (bottom_elevation > top_layered) + + thickness = thickness.where( + ~upper_layer_allocated, thickness - (top_layered - stage) + ) + thickness = thickness.where( + ~lower_layer_allocated, thickness - (bottom_elevation - bottom) + ) + thickness = thickness.where(~outside, 0.0) + + return thickness + + +def _distribute_weights__by_corrected_transmissivity( + allocated: GridDataArray, + top: GridDataArray, + bottom: GridDataArray, + stage: GridDataArray, + bottom_elevation: GridDataArray, + k: GridDataArray, +): + """ + Distribute conductances according to default method in iMOD 5.6, as + described page 690 of the iMOD 5.6 manual (but then to distribute WEL + rates). The method uses crosscut thicknesses to compute transmissivities. + Furthermore it corrects distribution weights for the mismatch between the + midpoints of crosscut areas and layer midpoints. + """ + PLANAR_GRID.validate(stage) + PLANAR_GRID.validate(bottom_elevation) + + crosscut_thickness = _compute_crosscut_thickness( + allocated, top, bottom, stage, bottom_elevation + ) + transmissivity = crosscut_thickness * k + + top_layered = _enforce_layered_top(top, bottom) + + upper_layer_allocated = (stage < top_layered) & (stage > bottom) + lower_layer_allocated = (bottom_elevation < top_layered) & ( + bottom_elevation > bottom + ) + + layer_thickness = _compute_layer_thickness(allocated, top, bottom) + midpoints = (top_layered + bottom) / 2 + + # Computing vertical midpoint of river crosscutting layers. + Fc = midpoints.where(~upper_layer_allocated, (bottom + stage) / 2) + Fc = Fc.where(~lower_layer_allocated, (top_layered + bottom_elevation) / 2) + # Correction factor for mismatch between midpoints of crosscut layers and + # layer midpoints. + F = 1.0 - np.abs(midpoints - Fc) / (layer_thickness * 0.5) + + transmissivity_corrected = transmissivity * F + return transmissivity_corrected / transmissivity_corrected.sum(dim="layer") + + +def _distribute_weights__equally(allocated: GridDataArray): + """Compute weights over layers equally.""" + return ones_like(allocated) / allocated.sum(dim="layer") + + +def _distribute_weights__by_layer_thickness( + allocated: GridDataArray, + top: GridDataArray, + bottom: GridDataArray, +): + """Compute weights based on layer thickness""" + layer_thickness = _compute_layer_thickness(allocated, top, bottom) + + return layer_thickness / layer_thickness.sum(dim="layer") + + +def _distribute_weights__by_crosscut_thickness( + allocated: GridDataArray, + top: GridDataArray, + bottom: GridDataArray, + stage: GridDataArray, + bottom_elevation: GridDataArray, +): + """Compute weights based on crosscut thickness""" + PLANAR_GRID.validate(stage) + PLANAR_GRID.validate(bottom_elevation) + + crosscut_thickness = _compute_crosscut_thickness( + allocated, top, bottom, stage, bottom_elevation + ).where(allocated) + + return crosscut_thickness / crosscut_thickness.sum(dim="layer") + + +def _distribute_weights__by_layer_transmissivity( + allocated: GridDataArray, + top: GridDataArray, + bottom: GridDataArray, + k: GridDataArray, +): + """Compute weights based on layer transmissivity""" + layer_thickness = _compute_layer_thickness(allocated, top, bottom) + transmissivity = layer_thickness * k + + return transmissivity / transmissivity.sum(dim="layer") + + +def _distribute_weights__by_crosscut_transmissivity( + allocated: GridDataArray, + top: GridDataArray, + bottom: GridDataArray, + stage: GridDataArray, + bottom_elevation: GridDataArray, + k: GridDataArray, +): + """Compute weights based on crosscut transmissivity""" + PLANAR_GRID.validate(stage) + PLANAR_GRID.validate(bottom_elevation) + + crosscut_thickness = _compute_crosscut_thickness( + allocated, top, bottom, stage, bottom_elevation + ) + transmissivity = crosscut_thickness * k + + return transmissivity / transmissivity.sum(dim="layer") + + +def _distribute_weights__by_conductivity(allocated: GridDataArray, k: GridDataArray): + """Compute weights based on hydraulic conductivity""" + k_allocated = allocated * k + + return k_allocated / k_allocated.sum(dim="layer") diff --git a/imod/tests/conftest.py b/imod/tests/conftest.py index 7f0e781d8..80abb3889 100644 --- a/imod/tests/conftest.py +++ b/imod/tests/conftest.py @@ -2,6 +2,7 @@ from .fixtures.flow_basic_fixture import ( basic_dis, + basic_dis__topsystem, get_render_dict, horizontal_flow_barrier_gdf, metaswap_dict, @@ -10,7 +11,11 @@ two_days, well_df, ) -from .fixtures.flow_basic_unstructured_fixture import basic_unstructured_dis, circle_dis +from .fixtures.flow_basic_unstructured_fixture import ( + basic_disv__topsystem, + basic_unstructured_dis, + circle_dis, +) from .fixtures.flow_example_fixture import imodflow_model from .fixtures.flow_transport_simulation_fixture import flow_transport_simulation from .fixtures.mf6_circle_fixture import ( diff --git a/imod/tests/fixtures/flow_basic_fixture.py b/imod/tests/fixtures/flow_basic_fixture.py index 24a0307ad..3c7aa8b30 100644 --- a/imod/tests/fixtures/flow_basic_fixture.py +++ b/imod/tests/fixtures/flow_basic_fixture.py @@ -6,15 +6,15 @@ import xarray as xr -@pytest.fixture(scope="module") -def basic_dis(): +def make_basic_dis(dz, nrow, ncol): """Basic model discretization""" - - shape = nlay, nrow, ncol = 3, 9, 9 - dx = 10.0 dy = -10.0 - dz = np.array([5, 30, 100]) + + nlay = len(dz) + + shape = nlay, nrow, ncol + xmin = 0.0 xmax = dx * ncol ymin = 0.0 @@ -37,6 +37,16 @@ def basic_dis(): return ibound, top, bottom +@pytest.fixture(scope="module") +def basic_dis(): + return make_basic_dis(dz=[5, 30, 100], nrow=9, ncol=9) + + +@pytest.fixture(scope="function") +def basic_dis__topsystem(): + return make_basic_dis(dz=[1.0, 2.0, 4.0, 10.0], nrow=9, ncol=9) + + class BasicDisSettings(NamedTuple): nlay: int = 1 nrow: int = 1 diff --git a/imod/tests/fixtures/flow_basic_unstructured_fixture.py b/imod/tests/fixtures/flow_basic_unstructured_fixture.py index 83ecc98be..44e7fae07 100644 --- a/imod/tests/fixtures/flow_basic_unstructured_fixture.py +++ b/imod/tests/fixtures/flow_basic_unstructured_fixture.py @@ -15,6 +15,14 @@ def basic_unstructured_dis(basic_dis): return idomain_ugrid, top_mf6, bottom +@pytest.fixture(scope="function") +def basic_disv__topsystem(basic_dis__topsystem): + ibound, top, bottom = basic_dis__topsystem + idomain_ugrid = xu.UgridDataArray.from_structured(ibound) + + return idomain_ugrid, top, bottom + + @pytest.fixture(scope="module") def circle_dis(): """Basic model discretization""" diff --git a/imod/tests/test_prepare/test_topsystem.py b/imod/tests/test_prepare/test_topsystem.py new file mode 100644 index 000000000..ed31e5ae6 --- /dev/null +++ b/imod/tests/test_prepare/test_topsystem.py @@ -0,0 +1,168 @@ +import numpy as np +import xarray as xr +from pytest_cases import parametrize_with_cases + +from imod.prepare.topsystem import ( + allocate_drn_cells, + allocate_ghb_cells, + allocate_rch_cells, + allocate_riv_cells, + distribute_drn_conductance, + distribute_ghb_conductance, + distribute_riv_conductance, +) +from imod.typing import GridDataArray +from imod.typing.grid import is_unstructured, zeros_like +from imod.util.dims import enforce_dim_order + + +def take_first_cell_in_xy_plane(grid: GridDataArray) -> GridDataArray: + if is_unstructured(grid): + return grid.values[:, 0] + else: + return grid.values[:, 0, 0] + + +@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_first_cell_in_xy_plane(actual_riv_da) + + if actual_drn_da is None: + actual_drn = actual_drn_da + else: + actual_drn = take_first_cell_in_xy_plane(actual_drn_da) + + np.testing.assert_equal(actual_riv, expected_riv) + np.testing.assert_equal(actual_drn, expected_drn) + + +@parametrize_with_cases( + argnames="active,top,bottom,drn_elevation", + prefix="drn_", +) +@parametrize_with_cases( + 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 = take_first_cell_in_xy_plane(actual_da) + + np.testing.assert_equal(actual, expected) + + +@parametrize_with_cases( + argnames="active,top,bottom,head", + prefix="ghb_", +) +@parametrize_with_cases( + 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 = take_first_cell_in_xy_plane(actual_da) + + np.testing.assert_equal(actual, expected) + + +@parametrize_with_cases( + argnames="active", + prefix="rch_", +) +@parametrize_with_cases( + argnames="option,expected,_", prefix="allocation_", has_tag="rch" +) +def test_rch_allocation(active, option, expected, _): + actual_da = allocate_rch_cells(option, active) + + actual = take_first_cell_in_xy_plane(actual_da) + + np.testing.assert_equal(actual, expected) + + +@parametrize_with_cases( + argnames="active,top,bottom,stage,bottom_elevation", + prefix="riv_", +) +@parametrize_with_cases( + argnames="option,allocated_layer,expected", prefix="distribution_", has_tag="riv" +) +def test_distribute_riv_conductance( + active, top, bottom, stage, bottom_elevation, option, allocated_layer, expected +): + allocated = enforce_dim_order(active & allocated_layer) + k = xr.DataArray( + [2.0, 2.0, 1.0, 1.0], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + + conductance = zeros_like(bottom_elevation) + 1.0 + + actual_da = distribute_riv_conductance( + option, allocated, conductance, top, bottom, stage, bottom_elevation, k + ) + actual = take_first_cell_in_xy_plane(actual_da) + + np.testing.assert_equal(actual, expected) + + +@parametrize_with_cases( + argnames="active,top,bottom,elevation", + prefix="drn_", +) +@parametrize_with_cases( + argnames="option,allocated_layer,expected", prefix="distribution_", has_tag="drn" +) +def test_distribute_drn_conductance( + active, top, bottom, elevation, option, allocated_layer, expected +): + allocated = enforce_dim_order(active & allocated_layer) + k = xr.DataArray( + [2.0, 2.0, 1.0, 1.0], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + + conductance = zeros_like(elevation) + 1.0 + + actual_da = distribute_drn_conductance( + option, allocated, conductance, top, bottom, k + ) + actual = take_first_cell_in_xy_plane(actual_da) + + np.testing.assert_equal(actual, expected) + + +@parametrize_with_cases( + argnames="active,top,bottom,elevation", + prefix="ghb_", +) +@parametrize_with_cases( + argnames="option,allocated_layer,expected", prefix="distribution_", has_tag="ghb" +) +def test_distribute_ghb_conductance( + active, top, bottom, elevation, option, allocated_layer, expected +): + allocated = enforce_dim_order(active & allocated_layer) + k = xr.DataArray( + [2.0, 2.0, 1.0, 1.0], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + + conductance = zeros_like(elevation) + 1.0 + + actual_da = distribute_ghb_conductance( + option, allocated, conductance, top, bottom, k + ) + actual = take_first_cell_in_xy_plane(actual_da) + + np.testing.assert_equal(actual, expected) diff --git a/imod/tests/test_prepare/test_topsystem_cases.py b/imod/tests/test_prepare/test_topsystem_cases.py new file mode 100644 index 000000000..ed8db2a2a --- /dev/null +++ b/imod/tests/test_prepare/test_topsystem_cases.py @@ -0,0 +1,411 @@ +import numpy as np +import xarray as xr +from pytest_cases import case + +from imod.prepare.topsystem import ( + ALLOCATION_OPTION, + DISTRIBUTING_OPTION, +) +from imod.typing.grid import zeros_like + + +def riv_structured(basic_dis__topsystem): + ibound, top, bottom = basic_dis__topsystem + top = top.sel(layer=1) + elevation = zeros_like(ibound.sel(layer=1)) + stage = elevation - 2.0 + bottom_elevation = elevation - 4.0 + active = ibound == 1 + return active, top, bottom, stage, bottom_elevation + + +def riv_unstructured(basic_disv__topsystem): + ibound, top, bottom = basic_disv__topsystem + elevation = zeros_like(ibound.sel(layer=1)) + stage = elevation - 2.0 + bottom_elevation = elevation - 4.0 + active = ibound == 1 + return active, top, bottom, stage, bottom_elevation + + +def drn_structured(basic_dis__topsystem): + ibound, top, bottom = basic_dis__topsystem + top = top.sel(layer=1) + elevation = zeros_like(ibound.sel(layer=1)) + drain_elevation = elevation - 2.0 + active = ibound == 1 + return active, top, bottom, drain_elevation + + +def drn_unstructured(basic_disv__topsystem): + ibound, top, bottom = basic_disv__topsystem + elevation = zeros_like(ibound.sel(layer=1)) + drain_elevation = elevation - 2.0 + active = ibound == 1 + return active, top, bottom, drain_elevation + + +def ghb_structured(basic_dis__topsystem): + ibound, top, bottom = basic_dis__topsystem + top = top.sel(layer=1) + elevation = zeros_like(ibound.sel(layer=1)) + head = elevation - 2.0 + active = ibound == 1 + return active, top, bottom, head + + +def ghb_unstructured(basic_disv__topsystem): + ibound, top, bottom = basic_disv__topsystem + elevation = zeros_like(ibound.sel(layer=1)) + head = elevation - 2.0 + active = ibound == 1 + return active, top, bottom, head + + +def rch_structured(basic_dis__topsystem): + ibound, _, _ = basic_dis__topsystem + active = ibound == 1 + return active + + +def rch_unstructured(basic_disv__topsystem): + ibound, _, _ = basic_disv__topsystem + active = ibound == 1 + return active + + +@case(tags=["riv"]) +def allocation_stage_to_riv_bot(): + option = ALLOCATION_OPTION.stage_to_riv_bot + expected = [False, True, True, False] + + return option, expected, None + + +@case(tags=["riv"]) +def allocation_first_active_to_riv_bot(): + option = ALLOCATION_OPTION.first_active_to_riv_bot + expected = [True, True, True, False] + + return option, expected, None + + +@case(tags=["riv"]) +def allocation_first_active_to_riv_bot__drn(): + option = ALLOCATION_OPTION.first_active_to_riv_bot__drn + expected = [False, True, True, False] + expected__drn = [True, False, False, False] + + return option, expected, expected__drn + + +@case(tags=["drn", "ghb"]) +def allocation_at_elevation(): + option = ALLOCATION_OPTION.at_elevation + expected = [False, True, False, False] + + return option, expected, None + + +@case(tags=["riv"]) +def allocation_at_riv_bottom_elevation(): + option = ALLOCATION_OPTION.at_elevation + expected = [False, False, True, False] + + return option, expected, None + + +@case(tags=["riv", "drn", "ghb", "rch"]) +def allocation_at_first_active(): + option = ALLOCATION_OPTION.at_first_active + expected = [True, False, False, False] + + return option, expected, None + + +@case(tags=["riv"]) +def distribution_by_corrected_transmissivity(): + option = DISTRIBUTING_OPTION.by_corrected_transmissivity + allocated_layer = xr.DataArray( + [False, True, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [np.nan, (4 / 5), (1 / 5), np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv"]) +def distribution_by_corrected_transmissivity__first_active(): + """First active cell allocated, while stage in second layer""" + option = DISTRIBUTING_OPTION.by_corrected_transmissivity + allocated_layer = xr.DataArray( + [True, True, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [0.0, (4 / 5), (1 / 5), np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv"]) +def distribution_by_corrected_transmissivity__third_only(): + """Third layer active only, while stage in second layer""" + option = DISTRIBUTING_OPTION.by_corrected_transmissivity + allocated_layer = xr.DataArray( + [False, False, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [np.nan, np.nan, 1.0, np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv"]) +def distribution_by_corrected_transmissivity__TFTF(): + option = DISTRIBUTING_OPTION.by_corrected_transmissivity + allocated_layer = xr.DataArray( + [True, False, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [0.0, np.nan, 1.0, np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv"]) +def distribution_by_crosscut_transmissivity(): + option = DISTRIBUTING_OPTION.by_crosscut_transmissivity + allocated_layer = xr.DataArray( + [False, True, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [np.nan, (2 / 3), (1 / 3), np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv"]) +def distribution_by_crosscut_transmissivity__first_active(): + """First active cell allocated, while stage in second layer""" + option = DISTRIBUTING_OPTION.by_crosscut_transmissivity + allocated_layer = xr.DataArray( + [True, True, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [0.0, (2 / 3), (1 / 3), np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv"]) +def distribution_by_crosscut_transmissivity__third_only(): + """Third layer active only, while stage in second layer""" + option = DISTRIBUTING_OPTION.by_crosscut_transmissivity + allocated_layer = xr.DataArray( + [False, False, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [np.nan, np.nan, 1.0, np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv"]) +def distribution_by_crosscut_transmissivity__TFTF(): + """Third layer active only, while stage in second layer""" + option = DISTRIBUTING_OPTION.by_crosscut_transmissivity + allocated_layer = xr.DataArray( + [True, False, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [0.0, np.nan, 1.0, np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv"]) +def distribution_by_crosscut_thickness(): + option = DISTRIBUTING_OPTION.by_crosscut_thickness + allocated_layer = xr.DataArray( + [False, True, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [np.nan, 0.5, 0.5, np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv"]) +def distribution_by_crosscut_thickness__first_active(): + """First active cell allocated, while stage in second layer""" + option = DISTRIBUTING_OPTION.by_crosscut_thickness + allocated_layer = xr.DataArray( + [True, True, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [0.0, 0.5, 0.5, np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv"]) +def distribution_by_crosscut_thickness__third_only(): + """Third layer active only, while stage in second layer""" + option = DISTRIBUTING_OPTION.by_crosscut_thickness + allocated_layer = xr.DataArray( + [False, False, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [np.nan, np.nan, 1.0, np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv"]) +def distribution_by_crosscut_thickness__TFTF(): + """Third layer active only, while stage in second layer""" + option = DISTRIBUTING_OPTION.by_crosscut_thickness + allocated_layer = xr.DataArray( + [True, False, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [0.0, np.nan, 1.0, np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv", "drn", "ghb"]) +def distribution_by_conductivity(): + option = DISTRIBUTING_OPTION.by_conductivity + allocated_layer = xr.DataArray( + [False, True, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [np.nan, (2 / 3), (1 / 3), np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv", "drn", "ghb"]) +def distribution_by_conductivity__first_active(): + option = DISTRIBUTING_OPTION.by_conductivity + allocated_layer = xr.DataArray( + [True, True, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [(2 / 5), (2 / 5), (1 / 5), np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv", "drn", "ghb"]) +def distribution_by_conductivity__third_only(): + option = DISTRIBUTING_OPTION.by_conductivity + allocated_layer = xr.DataArray( + [False, False, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [np.nan, np.nan, 1.0, np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv", "drn", "ghb"]) +def distribution_by_conductivity__TFTF(): + option = DISTRIBUTING_OPTION.by_conductivity + allocated_layer = xr.DataArray( + [True, False, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [(2 / 3), np.nan, (1 / 3), np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv", "drn", "ghb"]) +def distribution_by_layer_thickness(): + option = DISTRIBUTING_OPTION.by_layer_thickness + allocated_layer = xr.DataArray( + [False, True, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [np.nan, (1 / 3), (2 / 3), np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv", "drn", "ghb"]) +def distribution_by_layer_thickness__first_active(): + option = DISTRIBUTING_OPTION.by_layer_thickness + allocated_layer = xr.DataArray( + [True, True, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [(1 / 7), (2 / 7), (4 / 7), np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv", "drn", "ghb"]) +def distribution_by_layer_thickness__third_only(): + option = DISTRIBUTING_OPTION.by_layer_thickness + allocated_layer = xr.DataArray( + [False, False, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [np.nan, np.nan, 1.0, np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv", "drn", "ghb"]) +def distribution_by_layer_thickness__TFTF(): + option = DISTRIBUTING_OPTION.by_layer_thickness + allocated_layer = xr.DataArray( + [True, False, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [(1 / 5), np.nan, (4 / 5), np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv", "drn", "ghb"]) +def distribution_by_layer_transmissivity(): + option = DISTRIBUTING_OPTION.by_layer_transmissivity + allocated_layer = xr.DataArray( + [False, True, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [np.nan, 0.5, 0.5, np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv", "drn", "ghb"]) +def distribution_by_layer_transmissivity__first_active(): + option = DISTRIBUTING_OPTION.by_layer_transmissivity + allocated_layer = xr.DataArray( + [True, True, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [0.2, 0.4, 0.4, np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv", "drn", "ghb"]) +def distribution_by_layer_transmissivity__third_only(): + option = DISTRIBUTING_OPTION.by_layer_transmissivity + allocated_layer = xr.DataArray( + [False, False, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [np.nan, np.nan, 1.0, np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv", "drn", "ghb"]) +def distribution_by_layer_transmissivity__TFTF(): + option = DISTRIBUTING_OPTION.by_layer_transmissivity + allocated_layer = xr.DataArray( + [True, False, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [(1 / 3), np.nan, (2 / 3), np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv", "drn", "ghb"]) +def distribution_equally(): + option = DISTRIBUTING_OPTION.equally + allocated_layer = xr.DataArray( + [False, True, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [np.nan, 0.5, 0.5, np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv", "drn", "ghb"]) +def distribution_equally__first_active(): + option = DISTRIBUTING_OPTION.equally + allocated_layer = xr.DataArray( + [True, True, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [(1 / 3), (1 / 3), (1 / 3), np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv", "drn", "ghb"]) +def distribution_equally__third_only(): + option = DISTRIBUTING_OPTION.equally + allocated_layer = xr.DataArray( + [False, False, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [np.nan, np.nan, 1.0, np.nan] + return option, allocated_layer, expected + + +@case(tags=["riv", "drn", "ghb"]) +def distribution_equally__TFTF(): + option = DISTRIBUTING_OPTION.equally + allocated_layer = xr.DataArray( + [True, False, True, False], coords={"layer": [1, 2, 3, 4]}, dims=("layer",) + ) + expected = [0.5, np.nan, 0.5, np.nan] + return option, allocated_layer, expected diff --git a/imod/tests/test_topsystem/test_allocation.py b/imod/tests/test_topsystem/test_allocation.py deleted file mode 100644 index 7713ce15c..000000000 --- a/imod/tests/test_topsystem/test_allocation.py +++ /dev/null @@ -1,219 +0,0 @@ -import numpy as np -from pytest_cases import parametrize_with_cases - -from imod.prepare.topsystem import ( - ALLOCATION_OPTION, - allocate_drn_cells, - allocate_ghb_cells, - allocate_rch_cells, - allocate_riv_cells, -) -from imod.typing import GridDataArray -from imod.typing.grid import is_unstructured, zeros_like - - -def take_first_planar_cell(grid: GridDataArray): - if is_unstructured(grid): - return grid.values[:, 0] - else: - return grid.values[:, 0, 0] - - -class RiverCases: - def case_structured(self, basic_dis): - ibound, top, bottom = basic_dis - top = top.sel(layer=1) - elevation = zeros_like(ibound.sel(layer=1)) - stage = elevation - 7.5 - bottom_elevation = elevation - 10.0 - active = ibound == 1 - return active, top, bottom, stage, bottom_elevation - - def case_unstructured(self, basic_unstructured_dis): - ibound, top, bottom = basic_unstructured_dis - elevation = zeros_like(ibound.sel(layer=1)) - stage = elevation - 7.5 - bottom_elevation = elevation - 10.0 - active = ibound == 1 - return active, top, bottom, stage, bottom_elevation - - -class DrainCases: - def case_structured(self, basic_dis): - ibound, top, bottom = basic_dis - top = top.sel(layer=1) - elevation = zeros_like(ibound.sel(layer=1)) - drain_elevation = elevation - 7.5 - active = ibound == 1 - return active, top, bottom, drain_elevation - - def case_unstructured(self, basic_unstructured_dis): - ibound, top, bottom = basic_unstructured_dis - elevation = zeros_like(ibound.sel(layer=1)) - drain_elevation = elevation - 7.5 - active = ibound == 1 - return active, top, bottom, drain_elevation - - -class GeneralHeadBoundaryCases: - def case_structured(self, basic_dis): - ibound, top, bottom = basic_dis - top = top.sel(layer=1) - elevation = zeros_like(ibound.sel(layer=1)) - head = elevation - 7.5 - active = ibound == 1 - return active, top, bottom, head - - def case_unstructured(self, basic_unstructured_dis): - ibound, top, bottom = basic_unstructured_dis - elevation = zeros_like(ibound.sel(layer=1)) - head = elevation - 7.5 - active = ibound == 1 - return active, top, bottom, head - - -class RechargeCases: - def case_structured(self, basic_dis): - ibound, _, _ = basic_dis - active = ibound == 1 - return active - - def case_unstructured(self, basic_unstructured_dis): - ibound, _, _ = basic_unstructured_dis - active = ibound == 1 - return active - - -class AllocationOptionRiverCases: - def case_stage_to_riv_bot(self): - option = ALLOCATION_OPTION.stage_to_riv_bot - expected = [False, True, False] - - return option, expected, None - - def case_first_active_to_riv_bot(self): - option = ALLOCATION_OPTION.first_active_to_riv_bot - expected = [True, True, False] - - return option, expected, None - - def case_first_active_to_riv_bot__drn(self): - option = ALLOCATION_OPTION.first_active_to_riv_bot__drn - expected = [False, True, False] - expected__drn = [True, False, False] - - return option, expected, expected__drn - - def case_at_elevation(self): - option = ALLOCATION_OPTION.at_elevation - expected = [False, True, False] - - return option, expected, None - - def case_at_first_active(self): - option = ALLOCATION_OPTION.at_first_active - expected = [True, False, False] - - return option, expected, None - - -class AllocationOptionDrainCases: - def case_at_elevation(self): - option = ALLOCATION_OPTION.at_elevation - expected = [False, True, False] - - return option, expected - - def case_at_first_active(self): - option = ALLOCATION_OPTION.at_first_active - expected = [True, False, False] - - return option, expected - - -class AllocationOptionGeneralHeadCases: - def case_at_elevation(self): - option = ALLOCATION_OPTION.at_elevation - expected = [False, True, False] - - return option, expected - - def case_at_first_active(self): - option = ALLOCATION_OPTION.at_first_active - expected = [True, False, False] - - return option, expected - - -class AllocationOptionRechargeCases: - def case_at_first_active(self): - option = ALLOCATION_OPTION.at_first_active - expected = [True, False, False] - - return option, expected - - -@parametrize_with_cases( - argnames="active,top,bottom,stage,bottom_elevation", - cases=RiverCases, -) -@parametrize_with_cases( - argnames="option,expected_riv,expected_drn", cases=AllocationOptionRiverCases -) -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_first_planar_cell(actual_riv_da) - - if actual_drn_da is None: - actual_drn = actual_drn_da - else: - actual_drn = take_first_planar_cell(actual_drn_da) - - np.testing.assert_equal(actual_riv, expected_riv) - np.testing.assert_equal(actual_drn, expected_drn) - - -@parametrize_with_cases( - argnames="active,top,bottom,drn_elevation", - cases=DrainCases, -) -@parametrize_with_cases(argnames="option,expected", cases=AllocationOptionDrainCases) -def test_drn_allocation(active, top, bottom, drn_elevation, option, expected): - actual_da = allocate_drn_cells(option, active, top, bottom, drn_elevation) - - actual = take_first_planar_cell(actual_da) - - np.testing.assert_equal(actual, expected) - - -@parametrize_with_cases( - argnames="active,top,bottom,head", - cases=GeneralHeadBoundaryCases, -) -@parametrize_with_cases( - argnames="option,expected", cases=AllocationOptionGeneralHeadCases -) -def test_ghb_allocation(active, top, bottom, head, option, expected): - actual_da = allocate_ghb_cells(option, active, top, bottom, head) - - actual = take_first_planar_cell(actual_da) - - np.testing.assert_equal(actual, expected) - - -@parametrize_with_cases( - argnames="active", - cases=RechargeCases, -) -@parametrize_with_cases(argnames="option,expected", cases=AllocationOptionRechargeCases) -def test_rch_allocation(active, option, expected): - actual_da = allocate_rch_cells(option, active) - - actual = take_first_planar_cell(actual_da) - - np.testing.assert_equal(actual, expected)