from enum import Enum
from typing import Optional
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, zeros_like
from imod.util.dims import enforced_dim_order
[docs]
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}")
)
[docs]
@enforced_dim_order
def distribute_riv_conductance(
distributing_option: DISTRIBUTING_OPTION,
allocated: GridDataArray,
conductance: GridDataArray,
top: GridDataArray,
bottom: GridDataArray,
k: GridDataArray,
stage: GridDataArray,
bottom_elevation: 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
k: DataArray | UgridDataArray
Hydraulic conductivities
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.
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.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 DISTRIBUTING_OPTION.by_crosscut_thickness:
weights = _distribute_weights__by_crosscut_thickness(
allocated, top, bottom, stage, bottom_elevation
)
case DISTRIBUTING_OPTION.by_crosscut_transmissivity:
weights = _distribute_weights__by_crosscut_transmissivity(
allocated, top, bottom, k, stage, bottom_elevation
)
case DISTRIBUTING_OPTION.by_corrected_transmissivity:
weights = _distribute_weights__by_corrected_transmissivity(
allocated, top, bottom, k, stage, bottom_elevation
)
case _:
raise ValueError(
"Received incompatible setting for rivers, only"
f"'{DISTRIBUTING_OPTION.equally.name}', "
f"'{DISTRIBUTING_OPTION.by_layer_thickness.name}', "
f"'{DISTRIBUTING_OPTION.by_layer_transmissivity.name}', "
f"'{DISTRIBUTING_OPTION.by_conductivity.name}', "
f"'{DISTRIBUTING_OPTION.by_crosscut_thickness.name}', "
f"'{DISTRIBUTING_OPTION.by_crosscut_transmissivity.name}', and "
f"'{DISTRIBUTING_OPTION.by_corrected_transmissivity.name}' supported. "
f"Got: '{distributing_option.name}'"
)
return (weights * conductance).where(allocated)
[docs]
@enforced_dim_order
def distribute_drn_conductance(
distributing_option: DISTRIBUTING_OPTION,
allocated: GridDataArray,
conductance: GridDataArray,
top: GridDataArray,
bottom: GridDataArray,
k: GridDataArray,
elevation: 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
elevation: DataArray | UgridDataArray
Drain elevation
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, drain_elevation
)
"""
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 DISTRIBUTING_OPTION.by_crosscut_thickness:
weights = _distribute_weights__by_crosscut_thickness(
allocated, top, bottom, bc_bottom=elevation
)
case DISTRIBUTING_OPTION.by_crosscut_transmissivity:
weights = _distribute_weights__by_crosscut_transmissivity(
allocated, top, bottom, k, bc_bottom=elevation
)
case DISTRIBUTING_OPTION.by_corrected_transmissivity:
weights = _distribute_weights__by_corrected_transmissivity(
allocated, top, bottom, k, bc_bottom=elevation
)
case _:
raise ValueError(
"Received incompatible setting for drains, only"
f"'{DISTRIBUTING_OPTION.equally.name}', "
f"'{DISTRIBUTING_OPTION.by_layer_thickness.name}', "
f"'{DISTRIBUTING_OPTION.by_layer_transmissivity.name}', "
f"'{DISTRIBUTING_OPTION.by_conductivity.name}', "
f"'{DISTRIBUTING_OPTION.by_crosscut_thickness.name}', "
f"'{DISTRIBUTING_OPTION.by_crosscut_transmissivity.name}', and "
f"'{DISTRIBUTING_OPTION.by_corrected_transmissivity.name}' supported. "
f"Got: '{distributing_option.name}'"
)
return (weights * conductance).where(allocated)
[docs]
@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}', "
f"'{DISTRIBUTING_OPTION.by_layer_thickness.name}', "
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: GridDataArray, top: GridDataArray, bottom: GridDataArray
):
"""
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: GridDataArray,
top: GridDataArray,
bottom: GridDataArray,
bc_top: Optional[GridDataArray] = None,
bc_bottom: Optional[GridDataArray] = None,
):
"""
Compute 3D grid of thicknesses crosscut by boundary condition (river/drain)
in allocated cells. So the thickness in the upper allocated layer is bc_top
- bottom and the lower allocated layer is top - bc_bottom.
"""
if (bc_top is None) & (bc_bottom is None):
raise ValueError("`bc_top` and `bc_bottom` cannot both be None.")
top_layered = _enforce_layered_top(top, bottom)
thickness = _compute_layer_thickness(allocated, top, bottom)
outside = zeros_like(allocated).astype(bool)
if bc_top is not None:
upper_layer_bc = (bc_top < top_layered) & (bc_top > bottom)
outside = outside | (bc_top < bottom)
thickness = thickness.where(~upper_layer_bc, thickness - (top_layered - bc_top))
if bc_bottom is not None:
lower_layer_bc = (bc_bottom < top_layered) & (bc_bottom > bottom)
outside = outside | (bc_bottom > top_layered)
thickness = thickness.where(~lower_layer_bc, thickness - (bc_bottom - bottom))
thickness = thickness.where(~outside, 0.0)
return thickness
def _distribute_weights__by_corrected_transmissivity(
allocated: GridDataArray,
top: GridDataArray,
bottom: GridDataArray,
k: GridDataArray,
bc_top: Optional[GridDataArray] = None,
bc_bottom: Optional[GridDataArray] = None,
):
"""
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.
"""
crosscut_thickness = _compute_crosscut_thickness(
allocated, top, bottom, bc_top=bc_top, bc_bottom=bc_bottom
)
transmissivity = crosscut_thickness * k
top_layered = _enforce_layered_top(top, bottom)
layer_thickness = _compute_layer_thickness(allocated, top, bottom)
midpoints = (top_layered + bottom) / 2
Fc = midpoints.copy()
if bc_top is not None:
PLANAR_GRID.validate(bc_top)
upper_layer_bc = (bc_top < top_layered) & (bc_top > bottom)
# Computing vertical midpoint of river crosscutting layers.
Fc = Fc.where(~upper_layer_bc, (bottom + bc_top) / 2)
if bc_bottom is not None:
PLANAR_GRID.validate(bc_bottom)
lower_layer_bc = (bc_bottom < top_layered) & (bc_bottom > bottom)
# Computing vertical midpoint of river crosscutting layers.
Fc = Fc.where(~lower_layer_bc, (top_layered + bc_bottom) / 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,
bc_top: Optional[GridDataArray] = None,
bc_bottom: Optional[GridDataArray] = None,
):
"""Compute weights based on crosscut thickness"""
if bc_top is not None:
PLANAR_GRID.validate(bc_top)
if bc_bottom is not None:
PLANAR_GRID.validate(bc_bottom)
crosscut_thickness = _compute_crosscut_thickness(
allocated, top, bottom, bc_top, bc_bottom
).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,
k: GridDataArray,
bc_top: Optional[GridDataArray] = None,
bc_bottom: Optional[GridDataArray] = None,
):
"""Compute weights based on crosscut transmissivity"""
if bc_top is not None:
PLANAR_GRID.validate(bc_top)
if bc_bottom is not None:
PLANAR_GRID.validate(bc_bottom)
crosscut_thickness = _compute_crosscut_thickness(
allocated, top, bottom, bc_top=bc_top, bc_bottom=bc_bottom
)
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")
