Issue #949 allocate river layers (#952)

Fixes #949 

# Description

- Adds functionality to allocate river cells. 
- These can also be used to allocate DRN, GHB, and RCH cells
- I had to move the ``polygonize`` util function to work around a nasty
circular import. The solution also prevents changes to the public API.
- A decorator was added to enforce an iMOD Python compatible dimension
order (Related to: #239 )
- A decorator was added to preserve grid type (xu.UgridDataArray or
xr.DataArray).


# Checklist

- [x] Links to correct issue
- [x] Update changelog, if changes affect users
- [x] PR title starts with ``Issue #nr``, e.g. ``Issue #737``
- [x] Unit tests were added
- [ ] **If feature added**: Added/extended example

Author: JoerivanEngelen

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

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

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

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,

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

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

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):

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)

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