diff --git a/docs/api/changelog.rst b/docs/api/changelog.rst index 83417922e..f4f798c35 100644 --- a/docs/api/changelog.rst +++ b/docs/api/changelog.rst @@ -9,6 +9,15 @@ The format is based on `Keep a Changelog`_, and this project adheres to [Unreleased] ------------ +Added +~~~~~ +- Added functions to allocate planar grids over layers for the topsystem in + :func:`imod.prepare.allocate_drn_cells`, + :func:`imod.prepare.allocate_ghb_cells`, + :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`. + Fixed ~~~~~ - No ``ValidationError`` thrown anymore in :class:`imod.mf6.River` when diff --git a/docs/api/prepare.rst b/docs/api/prepare.rst index 236c1ee4b..28cbb2c5f 100644 --- a/docs/api/prepare.rst +++ b/docs/api/prepare.rst @@ -31,3 +31,12 @@ Prepare model input get_lower_active_layer_number get_upper_active_grid_cells get_upper_active_layer_number + create_layered_top + + ALLOCATION_OPTION + allocate_drn_cells + allocate_ghb_cells + allocate_rch_cells + allocate_riv_cells + c_leakage + c_radial \ No newline at end of file diff --git a/imod/mf6/utilities/grid.py b/imod/mf6/utilities/grid.py index 0f6718ac9..d5d5ac8be 100644 --- a/imod/mf6/utilities/grid.py +++ b/imod/mf6/utilities/grid.py @@ -7,6 +7,7 @@ import xugrid as xu import imod +from imod.prepare.layer import create_layered_top from imod.typing import GridDataArray from imod.typing.grid import zeros_like from imod.util.spatial import spatial_reference @@ -48,17 +49,6 @@ def broadcast_to_full_domain( return top, bottom -def create_layered_top(bottom: GridDataArray, top: GridDataArray) -> GridDataArray: - """ - Create a top array with layers from a single top array and a full bottom array - """ - new_top = zeros_like(bottom) - new_top[0] = top - new_top[1:] = bottom[0:-1].values - - return new_top - - def to_cell_idx(idomain: xr.DataArray) -> xr.DataArray: """ Assigns an unique index to each cell in the domain diff --git a/imod/mf6/wel.py b/imod/mf6/wel.py index 655f44069..5c3d9fc43 100644 --- a/imod/mf6/wel.py +++ b/imod/mf6/wel.py @@ -21,11 +21,11 @@ from imod.mf6.mf6_wel_adapter import Mf6Wel from imod.mf6.package import Package from imod.mf6.utilities.dataset import remove_inactive -from imod.mf6.utilities.grid import create_layered_top from imod.mf6.utilities.regrid import RegridderType from imod.mf6.validation import validation_pkg_error_message from imod.mf6.write_context import WriteContext from imod.prepare import assign_wells +from imod.prepare.layer import create_layered_top from imod.schemata import ( AnyNoDataSchema, DTypeSchema, diff --git a/imod/prepare/__init__.py b/imod/prepare/__init__.py index e648576e8..004612578 100644 --- a/imod/prepare/__init__.py +++ b/imod/prepare/__init__.py @@ -15,6 +15,7 @@ from imod.prepare import spatial, subsoil, surface_water from imod.prepare.layer import ( + create_layered_top, get_lower_active_grid_cells, get_lower_active_layer_number, get_upper_active_grid_cells, @@ -34,5 +35,14 @@ zonal_aggregate_polygons, zonal_aggregate_raster, ) +from imod.prepare.topsystem import ( + ALLOCATION_OPTION, + allocate_drn_cells, + allocate_ghb_cells, + allocate_rch_cells, + allocate_riv_cells, + c_leakage, + c_radial, +) from imod.prepare.voxelize import Voxelizer from imod.prepare.wells import assign_wells diff --git a/imod/prepare/layer.py b/imod/prepare/layer.py index a1586c5ce..35088d54e 100644 --- a/imod/prepare/layer.py +++ b/imod/prepare/layer.py @@ -3,8 +3,10 @@ """ from imod.typing import GridDataArray +from imod.typing.grid import preserve_gridtype, zeros_like +@preserve_gridtype def get_upper_active_layer_number(active: GridDataArray) -> GridDataArray: """ Returns planar grid of integers with the layer number of the upper most @@ -22,6 +24,7 @@ def get_upper_active_layer_number(active: GridDataArray) -> GridDataArray: return layer.where(active, nodata).min("layer") +@preserve_gridtype def get_upper_active_grid_cells(active: GridDataArray) -> GridDataArray: """ Returns grid of booleans designating location of the uppermost active grid @@ -37,6 +40,7 @@ def get_upper_active_grid_cells(active: GridDataArray) -> GridDataArray: return layer == upper_active_layer +@preserve_gridtype def get_lower_active_layer_number(active: GridDataArray) -> GridDataArray: """ Returns two-dimensional grid of integers with the layer number of the lower @@ -54,6 +58,7 @@ def get_lower_active_layer_number(active: GridDataArray) -> GridDataArray: return layer.where(active, nodata).max("layer") +@preserve_gridtype def get_lower_active_grid_cells(active: GridDataArray) -> GridDataArray: """ Returns grid of booleans designating location of the lowermost active grid @@ -67,3 +72,28 @@ def get_lower_active_grid_cells(active: GridDataArray) -> GridDataArray: layer = active.coords["layer"] lower_active_layer = get_lower_active_layer_number(active) return layer == lower_active_layer + + +def create_layered_top(bottom: GridDataArray, top: GridDataArray) -> GridDataArray: + """ + Create a top array with a layer dimension, from a top array with no layer + dimension and a bottom array with a layer dimension. The (output) top of + layer n is assigned the bottom of layer n-1. + + Parameters + ---------- + bottom: {DataArray, UgridDataArray} + Bottoms with layer dimension + top: {DataArray, UgridDataArray} + Top, without layer dimension + + Returns + ------- + new_top: {DataArray, UgridDataArray} + Top with layer dimension. + """ + new_top = zeros_like(bottom) + new_top[0] = top + new_top[1:] = bottom[0:-1].values + + return new_top diff --git a/imod/prepare/spatial.py b/imod/prepare/spatial.py index 1ea005747..74387c361 100644 --- a/imod/prepare/spatial.py +++ b/imod/prepare/spatial.py @@ -1,5 +1,5 @@ import pathlib -from typing import Callable, Union +from typing import TYPE_CHECKING, Callable, Union import dask import numba @@ -11,6 +11,7 @@ import imod from imod.prepare import common, pcg from imod.util.imports import MissingOptionalModule +from imod.util.spatial import _polygonize # since rasterio is a big dependency that is sometimes hard to install # and not always required, we made this an optional dependency @@ -21,6 +22,9 @@ except ImportError: rasterio = MissingOptionalModule("rasterio") +if TYPE_CHECKING: + import geopandas as gpd + def round_extent(extent, cellsize): """Increases the extent until all sides lie on a coordinate @@ -252,7 +256,7 @@ def rasterize(geodataframe, like, column=None, fill=np.nan, **kwargs): return xr.DataArray(raster, like.coords, like.dims) -def polygonize(da): +def polygonize(da: xr.DataArray) -> "gpd.GeoDataFrame": """ Polygonize a 2D-DataArray into a GeoDataFrame of polygons. @@ -264,28 +268,7 @@ def polygonize(da): ------- polygonized : geopandas.GeoDataFrame """ - import geopandas as gpd - import shapely.geometry as sg - - if da.dims != ("y", "x"): - raise ValueError('Dimensions must be ("y", "x")') - - values = da.values - if values.dtype == np.float64: - values = values.astype(np.float32) - - transform = imod.util.spatial.transform(da) - shapes = rasterio.features.shapes(values, transform=transform) - - geometries = [] - colvalues = [] - for geom, colval in shapes: - geometries.append(sg.Polygon(geom["coordinates"][0])) - colvalues.append(colval) - - gdf = gpd.GeoDataFrame({"value": colvalues, "geometry": geometries}) - gdf.crs = da.attrs.get("crs") - return gdf + return _polygonize(da) def _handle_dtype(dtype, nodata): diff --git a/imod/prepare/surface_water.py b/imod/prepare/surface_water.py index 5e3ff4298..d4f3bc3a1 100644 --- a/imod/prepare/surface_water.py +++ b/imod/prepare/surface_water.py @@ -1,109 +1,21 @@ -import numpy as np +# Import for backwards compatibility +import warnings +from imod.prepare.topsystem import c_leakage as _c_leakage +from imod.prepare.topsystem import c_radial as _c_radial -def c_radial(L, kh, kv, B, D): - """ - Ernst's radial resistance term to a drain. +WARNING_MESSAGE = ( + "function has been moved from imod.prepare.surface_water to" + "imod.prepare.topsystem, please update your scripts." + "imod.prepare.surface_water is going to be removed in version 1.0" +) - Parameters - ---------- - L : xr.DataArray of floats - distance between water features - kh : xr.DataArray of floats - horizontal conductivity - kv : xr.DataArray of floats - vertical conductivity - B : xr.DataArray of floats - water feature wetted perimeter - D : xr.DataArray of floats - saturated thickness of the top system - Returns - ------- - radial_c : xr.DataArray - Ernst's radial resistance for a drain - """ - # Take anisotropy into account fully - c = ( - L - / (np.pi * np.sqrt(kh * kv)) - * np.log((4.0 * D) / (np.pi * B) * np.sqrt(kh / kv)) - ) - c = c.where(~(c < 0.0), other=0.0) - return c +def c_radial(L, kh, kv, B, D): + warnings.warn(WARNING_MESSAGE, DeprecationWarning) + return _c_radial(L, kh, kv, B, D) def c_leakage(kh, kv, D, c0, c1, B, length, dx, dy): - """ - Computes the phreatic leakage resistance. - - Parameters - ---------- - kh : xr.DataArray of floats - horizontal conductivity of phreatic aquifer - kv : xr.DataArray of floats - vertical conductivity of phreatic aquifer - c0 : xr.DataArray of floats - hydraulic bed resistance of water feature - c1 : xr.DataArray of floats - hydraulic resistance of the first aquitard - D : xr.DataArray of floats - saturated thickness of the top system - B : xr.DataArray of floats - water feature wetted perimeter - length : xr.DataArray of floats - water feature length per cell - dx : xr.DataArray of floats - cellsize in x - dy : xr.DataArray of floats - cellsize in y - - Returns - ------- - c_leakage: xr.DataArray of floats - Hydraulic resistance of water features corrected for intra-cell - hydraulic resistance and surface water interaction. - """ - - def f(x): - """ - two x times cotangens hyperbolicus of x - """ - # Avoid overflow for large x values - # 20 is safe for 32 bit floats - x = x.where(~(x > 20.0), other=20.0) - exp2x = np.exp(2.0 * x) - return x * (exp2x + 1) / (exp2x - 1.0) - - # TODO: raise ValueError when cells are far from square? - # Since method of average ditch distance will not work well - # Compute geometry of ditches within cells - cell_area = abs(dy * dx) - wetted_area = length * B - fraction_wetted = wetted_area / cell_area - # Compute average distance between ditches - L = (cell_area - wetted_area) / length - - # Compute total resistance to aquifer - c1_prime = c1 + (D / kv) - # Compute influence for the part below the ditch - labda_B = np.sqrt((kh * D * c1_prime * c0) / (c1_prime + c0)) - # ... and the field - labda_L = np.sqrt(c1_prime * kh * D) - - x_B = B / (2.0 * labda_B) - x_L = L / (2.0 * labda_L) - - # Compute feeding resistance - c_rad = c_radial(L, kv, kh, B, D) - c_L = (c0 + c1_prime) * f(x_L) + (c0 * L / B) * f(x_B) - c_B = (c1_prime + c0) / (c_L - c0 * L / B) * c_L - # total feeding resistance equals the harmonic mean of resistances below - # ditch (B) and field (L) plus the radical resistance - # Can also be regarded as area weighted sum of conductances of B and L - c_total = 1.0 / (fraction_wetted / c_B + (1.0 - fraction_wetted) / c_L) + c_rad - # Subtract aquifer and aquitard resistance from feeding resistance - c = c_total - c1_prime - # Replace areas where cells are fully covered by water - c = c.where(~(L == 0.0), other=c0) - return c + warnings.warn(WARNING_MESSAGE, DeprecationWarning) + return _c_leakage(kh, kv, D, c0, c1, B, length, dx, dy) diff --git a/imod/prepare/topsystem/__init__.py b/imod/prepare/topsystem/__init__.py new file mode 100644 index 000000000..a63b10718 --- /dev/null +++ b/imod/prepare/topsystem/__init__.py @@ -0,0 +1,8 @@ +from imod.prepare.topsystem.allocation import ( + ALLOCATION_OPTION, + allocate_drn_cells, + allocate_ghb_cells, + allocate_rch_cells, + allocate_riv_cells, +) +from imod.prepare.topsystem.resistance import c_leakage, c_radial diff --git a/imod/prepare/topsystem/allocation.py b/imod/prepare/topsystem/allocation.py new file mode 100644 index 000000000..792bf3547 --- /dev/null +++ b/imod/prepare/topsystem/allocation.py @@ -0,0 +1,339 @@ +""" +This module contains all kinds of utilities to prepare rivers +""" + +from enum import Enum +from typing import Optional + +from imod.prepare.layer import ( + create_layered_top, + get_upper_active_grid_cells, + get_upper_active_layer_number, +) +from imod.schemata import DimsSchema +from imod.typing import GridDataArray +from imod.util.dims import enforced_dim_order + + +class ALLOCATION_OPTION(Enum): + """ + Enumerator for allocation settings. Numbers match the IDEFLAYER options in + iMOD 5.6. + """ + + stage_to_riv_bot = 0 + first_active_to_riv_bot = -1 + first_active_to_riv_bot__drn = 1 + at_elevation = 2 + at_first_active = 9 # Not an iMOD 5.6 option + + +PLANAR_GRID = ( + DimsSchema("time", "y", "x") + | DimsSchema("y", "x") + | DimsSchema("time", "{face_dim}") + | DimsSchema("{face_dim}") +) + + +def allocate_riv_cells( + allocation_option: ALLOCATION_OPTION, + active: GridDataArray, + top: GridDataArray, + bottom: GridDataArray, + stage: GridDataArray, + bottom_elevation: GridDataArray, +) -> tuple[GridDataArray, Optional[GridDataArray]]: + match allocation_option: + case ALLOCATION_OPTION.stage_to_riv_bot: + return _allocate_cells__stage_to_riv_bot( + top, bottom, stage, bottom_elevation + ) + case ALLOCATION_OPTION.first_active_to_riv_bot: + return _allocate_cells__first_active_to_riv_bot( + active, top, bottom, bottom_elevation + ) + case ALLOCATION_OPTION.first_active_to_riv_bot__drn: + return _allocate_cells__first_active_to_riv_bot__drn( + active, top, bottom, stage, bottom_elevation + ) + case ALLOCATION_OPTION.at_elevation: + return _allocate_cells__at_elevation(top, bottom, bottom_elevation) + case ALLOCATION_OPTION.at_first_active: + return _allocate_cells__at_first_active(active) + case _: + raise ValueError( + "Received incompatible setting for rivers, only" + f"'{ALLOCATION_OPTION.stage_to_riv_bot.name}' and" + f"'{ALLOCATION_OPTION.first_active_to_riv_bot.name}' and" + f"'{ALLOCATION_OPTION.first_active_to_riv_bot__drn.name}' and" + f"'{ALLOCATION_OPTION.at_elevation.name}' and" + f"'{ALLOCATION_OPTION.at_first_active.name}' supported." + f"got: '{allocation_option.name}'" + ) + + +def allocate_drn_cells( + allocation_option: ALLOCATION_OPTION, + active: GridDataArray, + top: GridDataArray, + bottom: GridDataArray, + elevation: GridDataArray, +) -> GridDataArray: + match allocation_option: + case ALLOCATION_OPTION.at_elevation: + return _allocate_cells__at_elevation(top, bottom, elevation)[0] + case ALLOCATION_OPTION.at_first_active: + return _allocate_cells__at_first_active(active)[0] + case _: + raise ValueError( + "Received incompatible setting for drains, only" + f"'{ALLOCATION_OPTION.at_elevation.name}' and" + f"'{ALLOCATION_OPTION.at_first_active.name}' supported." + f"got: '{allocation_option.name}'" + ) + + +def allocate_ghb_cells( + allocation_option: ALLOCATION_OPTION, + active: GridDataArray, + top: GridDataArray, + bottom: GridDataArray, + head: GridDataArray, +) -> GridDataArray: + match allocation_option: + case ALLOCATION_OPTION.at_elevation: + return _allocate_cells__at_elevation(top, bottom, head)[0] + case ALLOCATION_OPTION.at_first_active: + return _allocate_cells__at_first_active(active)[0] + case _: + raise ValueError( + "Received incompatible setting for general head boundary, only" + f"'{ALLOCATION_OPTION.at_elevation.name}' and" + f"'{ALLOCATION_OPTION.at_first_active.name}' supported." + f"got: '{allocation_option.name}'" + ) + + +def allocate_rch_cells( + allocation_option: ALLOCATION_OPTION, + active: GridDataArray, +) -> GridDataArray: + match allocation_option: + case ALLOCATION_OPTION.at_first_active: + return _allocate_cells__at_first_active(active)[0] + case _: + raise ValueError( + "Received incompatible setting for recharge, only" + f"'{ALLOCATION_OPTION.at_first_active.name}' supported." + f"got: '{allocation_option.name}'" + ) + + +def _is_layered(grid: GridDataArray): + return "layer" in grid.sizes and grid.sizes["layer"] > 1 + + +@enforced_dim_order +def _allocate_cells__stage_to_riv_bot( + top: GridDataArray, + bottom: GridDataArray, + stage: GridDataArray, + bottom_elevation: GridDataArray, +) -> tuple[GridDataArray, None]: + """ + Allocate cells inbetween river stage and river bottom_elevation. Compared to + iMOD5.6, this is similar to setting IDEFFLAYER=0 in the RUNFILE function. + + Note that ``stage`` and ``bottom_elevation`` are not allowed to have a layer + dimension. + + Parameters + ---------- + top: GridDataArray + top of model layers + bottom: GridDataArray + bottom of model layers + stage: GridDataArray + river stage + bottom_elevation: GridDataArray + river bottom elevation + + Returns + ------- + GridDataArray + River cells + """ + PLANAR_GRID.validate(stage) + PLANAR_GRID.validate(bottom_elevation) + + if _is_layered(top): + top_layered = top + else: + top_layered = create_layered_top(bottom, top) + + riv_cells = (stage > bottom) & (bottom_elevation < top_layered) + + return riv_cells, None + + +@enforced_dim_order +def _allocate_cells__first_active_to_riv_bot( + active: GridDataArray, + top: GridDataArray, + bottom: GridDataArray, + bottom_elevation: GridDataArray, +) -> tuple[GridDataArray, None]: + """ + Allocate cells inbetween first active layer and river bottom elevation. + Compared to iMOD5.6, this is similar to setting IDEFFLAYER=-1 in the RUNFILE + function. + + Note that ``bottom_elevation`` is not allowed to have a layer dimension. + + Parameters + ---------- + active: GridDataArray + active model cells + top: GridDataArray + top of model layers + bottom: GridDataArray + bottom of model layers + bottom_elevation: GridDataArray + river bottom elevation + + Returns + ------- + GridDataArray + River cells + """ + PLANAR_GRID.validate(bottom_elevation) + + 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) + + riv_cells = (upper_active_layer <= layer) & (bottom_elevation < top_layered) + + return riv_cells, None + + +@enforced_dim_order +def _allocate_cells__first_active_to_riv_bot__drn( + active: GridDataArray, + top: GridDataArray, + bottom: GridDataArray, + stage: GridDataArray, + bottom_elevation: GridDataArray, +) -> tuple[GridDataArray, GridDataArray]: + """ + Allocate cells inbetween first active layer and river bottom elevation. + Cells above river stage are deactivated and returned as drn cells. Compared + to iMOD5.6, this is similar to setting IDEFFLAYER=1 in the RUNFILE function. + + Note that ``stage`` and ``bottom_elevation`` are not allowed to have a layer + dimension. + + Parameters + ---------- + active: GridDataArray + active grid cells + top: GridDataArray + top of model layers + bottom: GridDataArray + bottom of model layers + stage: GridDataArray + river stage + bottom_elevation: GridDataArray + river bottom elevation + + Returns + ------- + riv_cells: GridDataArray + River cells (below stage) + drn_cells: GridDataArray + Drainage cells (above stage) + """ + + PLANAR_GRID.validate(stage) + PLANAR_GRID.validate(bottom_elevation) + + if _is_layered(top): + top_layered = top + else: + top_layered = create_layered_top(bottom, top) + + upper_active_layer = get_upper_active_layer_number(active) + layer = active.coords["layer"] + drn_cells = (upper_active_layer <= layer) & (bottom >= stage) + riv_cells = ( + (upper_active_layer <= layer) & (bottom_elevation < top_layered) + ) != drn_cells + + return riv_cells, drn_cells + + +@enforced_dim_order +def _allocate_cells__at_elevation( + top: GridDataArray, bottom: GridDataArray, elevation: GridDataArray +) -> tuple[GridDataArray, None]: + """ + Allocate cells in river bottom elevation layer. Compared to iMOD5.6, this is + similar to setting IDEFFLAYER=2 in the RUNFILE function. + + Note that ``bottom_elevation`` is not allowed to have a layer dimension. + + Parameters + ---------- + top: GridDataArray + top of model layers + bottom: GridDataArray + bottom of model layers + elevation: GridDataArray + elevation. Can be river bottom, drain elevation or head of GHB. + + Returns + ------- + GridDataArray + River cells + """ + + PLANAR_GRID.validate(elevation) + + if _is_layered(top): + top_layered = top + else: + top_layered = create_layered_top(bottom, top) + + riv_cells = (elevation < top_layered) & (elevation >= bottom) + + return riv_cells, None + + +@enforced_dim_order +def _allocate_cells__at_first_active( + active: GridDataArray, +) -> tuple[GridDataArray, None]: + """ + Allocate cells inbetween first active layer and river bottom elevation. + Compared to iMOD5.6, this is similar to setting IDEFFLAYER=-1 in the RUNFILE + function. + + Note that ``bottom_elevation`` is not allowed to have a layer dimension. + + Parameters + ---------- + active: GridDataArray + active model cells + + Returns + ------- + GridDataArray + River cells + """ + + return get_upper_active_grid_cells(active), None diff --git a/imod/prepare/topsystem/resistance.py b/imod/prepare/topsystem/resistance.py new file mode 100644 index 000000000..c340baa51 --- /dev/null +++ b/imod/prepare/topsystem/resistance.py @@ -0,0 +1,127 @@ +import numpy as np + +from imod.typing import GridDataArray + + +def c_radial( + L: GridDataArray, + kh: GridDataArray, + kv: GridDataArray, + B: GridDataArray, + D: GridDataArray, +) -> GridDataArray: + """ + Ernst's radial resistance term to a drain. + + Parameters + ---------- + L : xr.DataArray of floats + distance between water features + kh : xr.DataArray of floats + horizontal conductivity + kv : xr.DataArray of floats + vertical conductivity + B : xr.DataArray of floats + water feature wetted perimeter + D : xr.DataArray of floats + saturated thickness of the top system + + Returns + ------- + radial_c : xr.DataArray + Ernst's radial resistance for a drain + """ + # Take anisotropy into account fully + c = ( + L + / (np.pi * np.sqrt(kh * kv)) + * np.log((4.0 * D) / (np.pi * B) * np.sqrt(kh / kv)) + ) + c = c.where(~(c < 0.0), other=0.0) + return c + + +def c_leakage( + kh: GridDataArray, + kv: GridDataArray, + D: GridDataArray, + c0: GridDataArray, + c1: GridDataArray, + B: GridDataArray, + length: GridDataArray, + dx: GridDataArray, + dy: GridDataArray, +) -> GridDataArray: + """ + Computes the phreatic leakage resistance. + + Parameters + ---------- + kh : xr.DataArray of floats + horizontal conductivity of phreatic aquifer + kv : xr.DataArray of floats + vertical conductivity of phreatic aquifer + c0 : xr.DataArray of floats + hydraulic bed resistance of water feature + c1 : xr.DataArray of floats + hydraulic resistance of the first aquitard + D : xr.DataArray of floats + saturated thickness of the top system + B : xr.DataArray of floats + water feature wetted perimeter + length : xr.DataArray of floats + water feature length per cell + dx : xr.DataArray of floats + cellsize in x + dy : xr.DataArray of floats + cellsize in y + + Returns + ------- + c_leakage: xr.DataArray of floats + Hydraulic resistance of water features corrected for intra-cell + hydraulic resistance and surface water interaction. + """ + + def f(x): + """ + two x times cotangens hyperbolicus of x + """ + # Avoid overflow for large x values + # 20 is safe for 32 bit floats + x = x.where(~(x > 20.0), other=20.0) + exp2x = np.exp(2.0 * x) + return x * (exp2x + 1) / (exp2x - 1.0) + + # TODO: raise ValueError when cells are far from square? + # Since method of average ditch distance will not work well + # Compute geometry of ditches within cells + cell_area = abs(dy * dx) + wetted_area = length * B + fraction_wetted = wetted_area / cell_area + # Compute average distance between ditches + L = (cell_area - wetted_area) / length + + # Compute total resistance to aquifer + c1_prime = c1 + (D / kv) + # Compute influence for the part below the ditch + labda_B = np.sqrt((kh * D * c1_prime * c0) / (c1_prime + c0)) + # ... and the field + labda_L = np.sqrt(c1_prime * kh * D) + + x_B = B / (2.0 * labda_B) + x_L = L / (2.0 * labda_L) + + # Compute feeding resistance + c_rad = c_radial(L, kv, kh, B, D) + c_L = (c0 + c1_prime) * f(x_L) + (c0 * L / B) * f(x_B) + c_B = (c1_prime + c0) / (c_L - c0 * L / B) * c_L + # total feeding resistance equals the harmonic mean of resistances below + # ditch (B) and field (L) plus the radical resistance + # Can also be regarded as area weighted sum of conductances of B and L + c_total = 1.0 / (fraction_wetted / c_B + (1.0 - fraction_wetted) / c_L) + c_rad + # Subtract aquifer and aquitard resistance from feeding resistance + c = c_total - c1_prime + # Replace areas where cells are fully covered by water + c = c.where(~(L == 0.0), other=c0) + return c diff --git a/imod/tests/test_topsystem/test_allocation.py b/imod/tests/test_topsystem/test_allocation.py new file mode 100644 index 000000000..7713ce15c --- /dev/null +++ b/imod/tests/test_topsystem/test_allocation.py @@ -0,0 +1,219 @@ +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) diff --git a/imod/tests/test_typing/test_typing_grid.py b/imod/tests/test_typing/test_typing_grid.py new file mode 100644 index 000000000..10a0bb953 --- /dev/null +++ b/imod/tests/test_typing/test_typing_grid.py @@ -0,0 +1,66 @@ +import xarray as xr +import xugrid as xu + +from imod.typing.grid import enforce_dim_order, preserve_gridtype + + +def test_preserve_gridtype__single_output(basic_unstructured_dis): + uda, _, da = basic_unstructured_dis + + @preserve_gridtype + def to_be_decorated(a, b): + return a * b + + result1 = to_be_decorated(uda, da) + result2 = to_be_decorated(da, uda) + + # Verify fixture provides expected type + assert isinstance(da, xr.DataArray) + assert isinstance(uda, xu.UgridDataArray) + + assert isinstance(result1, xu.UgridDataArray) + assert isinstance(result2, xu.UgridDataArray) + + +def test_preserve_gridtype__multiple_outputs(basic_unstructured_dis): + uda, _, da = basic_unstructured_dis + + @preserve_gridtype + def to_be_decorated(a, b): + return a * b, a + b + + result1a, result1b = to_be_decorated(uda, da) + result2a, result2b = to_be_decorated(da, uda) + + # Verify fixture provides expected type + assert isinstance(da, xr.DataArray) + assert isinstance(uda, xu.UgridDataArray) + + assert isinstance(result1a, xu.UgridDataArray) + assert isinstance(result1b, xu.UgridDataArray) + assert isinstance(result2a, xu.UgridDataArray) + assert isinstance(result2b, xu.UgridDataArray) + + +def test_enforce_dim_order__structured(basic_dis): + ibound, _, _ = basic_dis + + ibound_wrong_order = ibound.transpose("x", "y", "layer") + + actual = enforce_dim_order(ibound_wrong_order) + + assert actual.dims == ibound.dims + assert isinstance(actual, type(ibound)) + + +def test_enforce_dim_order__unstructured(basic_unstructured_dis): + ibound, _, _ = basic_unstructured_dis + + face_dim = ibound.ugrid.grid.face_dimension + + ibound_wrong_order = ibound.transpose(face_dim, "layer") + + actual = enforce_dim_order(ibound_wrong_order) + + assert actual.dims == ibound.dims + assert isinstance(actual, type(ibound)) diff --git a/imod/tests/test_util/test_util_dims.py b/imod/tests/test_util/test_util_dims.py new file mode 100644 index 000000000..74e91a808 --- /dev/null +++ b/imod/tests/test_util/test_util_dims.py @@ -0,0 +1,33 @@ +from imod.util.dims import enforced_dim_order + + +def test_enforced_dim_order_structured(basic_dis): + ibound, _, _ = basic_dis + + @enforced_dim_order + def to_be_decorated(da): + return da + + ibound_wrong_order = ibound.transpose("x", "y", "layer") + + actual = to_be_decorated(ibound_wrong_order) + + assert actual.dims == ibound.dims + assert isinstance(actual, type(ibound)) + + +def test_enforced_dim_order_unstructured(basic_unstructured_dis): + ibound, _, _ = basic_unstructured_dis + + @enforced_dim_order + def to_be_decorated(da): + return da + + face_dim = ibound.ugrid.grid.face_dimension + + ibound_wrong_order = ibound.transpose(face_dim, "layer") + + actual = to_be_decorated(ibound_wrong_order) + + assert actual.dims == ibound.dims + assert isinstance(actual, type(ibound)) diff --git a/imod/typing/grid.py b/imod/typing/grid.py index a59b0d4d1..2a74d7b25 100644 --- a/imod/typing/grid.py +++ b/imod/typing/grid.py @@ -7,8 +7,8 @@ import xugrid as xu from fastcore.dispatch import typedispatch -from imod.prepare import polygonize from imod.typing import GridDataArray, GridDataset, structured +from imod.util.spatial import _polygonize @typedispatch @@ -228,7 +228,7 @@ def merge_with_dictionary( def bounding_polygon(active: xr.DataArray): """Return bounding polygon of active cells""" to_polygonize = active.where(active, other=np.nan) - polygons_gdf = polygonize(to_polygonize) + polygons_gdf = _polygonize(to_polygonize) # Filter polygons with inactive values (NaN) is_active_polygon = polygons_gdf["value"] == 1.0 return polygons_gdf.loc[is_active_polygon] @@ -336,3 +336,57 @@ def get_grid_geometry_hash(grid: xu.UgridDataArray) -> int: @typedispatch def get_grid_geometry_hash(grid: object) -> int: raise ValueError("get_grid_geometry_hash not supported for this object.") + + +@typedispatch +def enforce_dim_order(grid: xr.DataArray) -> xr.DataArray: + """Enforce dimension order to iMOD Python standard""" + return grid.transpose("species", "time", "layer", "y", "x", missing_dims="ignore") + + +@typedispatch +def enforce_dim_order(grid: xu.UgridDataArray) -> xu.UgridDataArray: + """Enforce dimension order to iMOD Python standard""" + face_dimension = grid.ugrid.grid.face_dimension + return grid.transpose( + "species", "time", "layer", face_dimension, missing_dims="ignore" + ) + + +def _enforce_unstructured(obj: GridDataArray, ugrid2d=xu.Ugrid2d) -> xu.UgridDataArray: + """Force obj to unstructured""" + return xu.UgridDataArray(xr.DataArray(obj), ugrid2d) + + +def preserve_gridtype(func): + """ + Decorator to preserve gridtype, this is to work around the following xugrid + behavior: + + >>> UgridDataArray() * DataArray() -> UgridDataArray + >>> DataArray() * UgridDataArray() -> DataArray + + with this decorator: + + >>> UgridDataArray() * DataArray() -> UgridDataArray + >>> DataArray() * UgridDataArray() -> UgridDataArray + """ + + def decorator(*args, **kwargs): + unstructured = False + grid = None + for arg in args: + if is_unstructured(arg): + unstructured = True + grid = arg.ugrid.grid + + x = func(*args, **kwargs) + + if unstructured: + # Multiple grids returned + if isinstance(x, tuple): + return tuple(_enforce_unstructured(i, grid) for i in x) + return _enforce_unstructured(x, grid) + return x + + return decorator diff --git a/imod/util/dims.py b/imod/util/dims.py new file mode 100644 index 000000000..7db0be77e --- /dev/null +++ b/imod/util/dims.py @@ -0,0 +1,14 @@ +from imod.typing.grid import enforce_dim_order + + +def enforced_dim_order(func): + """Decorator to enforce dimension order after function call""" + + def decorator(*args, **kwargs): + x = func(*args, **kwargs) + # Multiple grids returned + if isinstance(x, tuple): + return tuple(enforce_dim_order(i) for i in x) + return enforce_dim_order(x) + + return decorator diff --git a/imod/util/spatial.py b/imod/util/spatial.py index 6a00564bd..391e1ebf3 100644 --- a/imod/util/spatial.py +++ b/imod/util/spatial.py @@ -7,7 +7,7 @@ import collections import re -from typing import Any, Dict, Sequence, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, Sequence, Tuple, Union import affine import numpy as np @@ -16,6 +16,28 @@ import xugrid as xu from imod.typing import FloatArray, GridDataset, IntArray +from imod.util.imports import MissingOptionalModule + +# since rasterio, shapely, and geopandas are a big dependencies that are +# sometimes hard to install and not always required, we made this an optional +# dependency +try: + import rasterio +except ImportError: + rasterio = MissingOptionalModule("rasterio") + +try: + import shapely +except ImportError: + shapely = MissingOptionalModule("shapely") + +if TYPE_CHECKING: + import geopandas as gpd +else: + try: + import geopandas as gpd + except ImportError: + gpd = MissingOptionalModule("geopandas") def _xycoords(bounds, cellsizes) -> Dict[str, Any]: @@ -612,3 +634,31 @@ def is_divisor(numerator: FloatArray, denominator: float) -> bool: denominator = abs(denominator) remainder = np.abs(numerator) % denominator return (np.isclose(remainder, 0.0) | np.isclose(remainder, denominator)).all() + + +def _polygonize(da: xr.DataArray) -> "gpd.GeoDataFrame": + """ + Polygonize a 2D-DataArray into a GeoDataFrame of polygons. + + Private method located in util.spatial to work around circular imports. + """ + + if da.dims != ("y", "x"): + raise ValueError('Dimensions must be ("y", "x")') + + values = da.values + if values.dtype == np.float64: + values = values.astype(np.float32) + + affine_transform = transform(da) + shapes = rasterio.features.shapes(values, transform=affine_transform) + + geometries = [] + colvalues = [] + for geom, colval in shapes: + geometries.append(shapely.geometry.Polygon(geom["coordinates"][0])) + colvalues.append(colval) + + gdf = gpd.GeoDataFrame({"value": colvalues, "geometry": geometries}) + gdf.crs = da.attrs.get("crs") + return gdf