Issue #950 Distribute conductances

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
 ~~~~~

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
+    c_radial
+    distribute_drn_conductance
+    distribute_ghb_conductance
+    distribute_riv_conductance

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

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

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)