Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
187 changes: 110 additions & 77 deletions imod/prepare/topsystem/conductance.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
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
from imod.typing.grid import ones_like, preserve_gridtype, zeros_like
from imod.util.dims import enforced_dim_order


Expand Down Expand Up @@ -67,9 +68,9 @@ def distribute_riv_conductance(
conductance: GridDataArray,
top: GridDataArray,
bottom: GridDataArray,
k: GridDataArray,
stage: GridDataArray,
bottom_elevation: GridDataArray,
k: GridDataArray,
) -> GridDataArray:
"""
Function to distribute 2D conductance over vertical layers for the RIV
Expand All @@ -90,14 +91,14 @@ def distribute_riv_conductance(
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.
k: DataArray | UgridDataArray
Hydraulic conductivities

Returns
-------
Expand All @@ -117,39 +118,39 @@ def distribute_riv_conductance(
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 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.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}'"
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)

Expand All @@ -162,6 +163,7 @@ def distribute_drn_conductance(
top: GridDataArray,
bottom: GridDataArray,
k: GridDataArray,
elevation: GridDataArray,
) -> GridDataArray:
"""
Function to distribute 2D conductance over vertical layers for the DRN
Expand All @@ -184,6 +186,8 @@ def distribute_drn_conductance(
Model layer bottoms
k: DataArray | UgridDataArray
Hydraulic conductivities
elevation: DataArray | UgridDataArray
Drain elevation

Returns
-------
Expand All @@ -197,7 +201,7 @@ def distribute_drn_conductance(
)
>>> conductances_distributed = distribute_drn_conductance(
DISTRIBUTING_OPTION.by_layer_transmissivity, allocated,
conductance, top, bottom, k
conductance, top, bottom, k, drain_elevation
)
"""
PLANAR_GRID.validate(conductance)
Expand All @@ -213,14 +217,29 @@ def distribute_drn_conductance(
)
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}' 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}'"
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)

Expand Down Expand Up @@ -286,17 +305,19 @@ def distribute_ghb_conductance(
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}'"
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, top, bottom):
def _compute_layer_thickness(
allocated: GridDataArray, top: GridDataArray, bottom: GridDataArray
):
"""
Compute 3D grid of thicknesses in allocated cells
"""
Expand All @@ -307,28 +328,35 @@ def _compute_layer_thickness(allocated, top, bottom):


@preserve_gridtype
def _compute_crosscut_thickness(allocated, top, bottom, stage, bottom_elevation):
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 river in allocated cells. So the
upper allocated layer thickness is stage - bottom, the lower allocated layer
is top - river_bottom_elevation.
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.
"""
top_layered = _enforce_layered_top(top, 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)

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)
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(
~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
Expand All @@ -338,9 +366,9 @@ def _distribute_weights__by_corrected_transmissivity(
allocated: GridDataArray,
top: GridDataArray,
bottom: GridDataArray,
stage: GridDataArray,
bottom_elevation: GridDataArray,
k: GridDataArray,
bc_top: Optional[GridDataArray] = None,
bc_bottom: Optional[GridDataArray] = None,
):
"""
Distribute conductances according to default method in iMOD 5.6, as
Expand All @@ -349,27 +377,28 @@ def _distribute_weights__by_corrected_transmissivity(
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
allocated, top, bottom, bc_top=bc_top, bc_bottom=bc_bottom
)
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
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)

# 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)
Expand Down Expand Up @@ -398,15 +427,17 @@ def _distribute_weights__by_crosscut_thickness(
allocated: GridDataArray,
top: GridDataArray,
bottom: GridDataArray,
stage: GridDataArray,
bottom_elevation: GridDataArray,
bc_top: Optional[GridDataArray] = None,
bc_bottom: Optional[GridDataArray] = None,
):
"""Compute weights based on crosscut thickness"""
PLANAR_GRID.validate(stage)
PLANAR_GRID.validate(bottom_elevation)
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, stage, bottom_elevation
allocated, top, bottom, bc_top, bc_bottom
).where(allocated)

return crosscut_thickness / crosscut_thickness.sum(dim="layer")
Expand All @@ -429,16 +460,18 @@ def _distribute_weights__by_crosscut_transmissivity(
allocated: GridDataArray,
top: GridDataArray,
bottom: GridDataArray,
stage: GridDataArray,
bottom_elevation: GridDataArray,
k: GridDataArray,
bc_top: Optional[GridDataArray] = None,
bc_bottom: Optional[GridDataArray] = None,
):
"""Compute weights based on crosscut transmissivity"""
PLANAR_GRID.validate(stage)
PLANAR_GRID.validate(bottom_elevation)
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, stage, bottom_elevation
allocated, top, bottom, bc_top=bc_top, bc_bottom=bc_bottom
)
transmissivity = crosscut_thickness * k

Expand Down
4 changes: 2 additions & 2 deletions imod/tests/test_prepare/test_topsystem.py
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ def test_distribute_riv_conductance(
conductance = zeros_like(bottom_elevation) + 1.0

actual_da = distribute_riv_conductance(
option, allocated, conductance, top, bottom, stage, bottom_elevation, k
option, allocated, conductance, top, bottom, k, stage, bottom_elevation
)
actual = take_first_cell_in_xy_plane(actual_da)

Expand All @@ -136,7 +136,7 @@ def test_distribute_drn_conductance(
conductance = zeros_like(elevation) + 1.0

actual_da = distribute_drn_conductance(
option, allocated, conductance, top, bottom, k
option, allocated, conductance, top, bottom, k, elevation
)
actual = take_first_cell_in_xy_plane(actual_da)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

you only compare "actual" and "expected" in the first cell. So it would make the tests more strict if arrays like "bottom" and "elevation" have a different value in the first cell than in the rest of the array. Otherwise the tests would still pass even if something went wrong in the indexing.

@JoerivanEngelen JoerivanEngelen Apr 22, 2024

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

By doing this, we would mainly test if the testing utility function take_first_cell_in_xy_plane does what it's supposed to do. I think this adds quite some complexity to the test cases for relatively little gains. I think it is better to allocate our constrained resources to fixing other issues.


Expand Down
Loading