From 0920768e1d020a1662ce7ef287deeee4c8f5b357 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Fri, 19 Apr 2024 17:45:52 +0200 Subject: [PATCH 1/2] Add docstrings for public API --- imod/prepare/topsystem/allocation.py | 161 ++++++++++++++++++++++++++- 1 file changed, 159 insertions(+), 2 deletions(-) diff --git a/imod/prepare/topsystem/allocation.py b/imod/prepare/topsystem/allocation.py index 792bf3547..a21001131 100644 --- a/imod/prepare/topsystem/allocation.py +++ b/imod/prepare/topsystem/allocation.py @@ -17,8 +17,34 @@ class ALLOCATION_OPTION(Enum): """ - Enumerator for allocation settings. Numbers match the IDEFLAYER options in - iMOD 5.6. + Enumerator for settings to allocate planar grid with RIV, DRN, GHB, or RCH + cells over the vertical layer dimensions. Numbers match the IDEFLAYER + options in iMOD 5.6. + + * ``stage_to_riv_bot``: RIV. Allocate cells spanning from the river stage up + to the river bottom elevation. This matches the iMOD 5.6 IDEFLAYER = 0 + option. + * ``first_active_to_riv_bot``: RIV. Allocate cells spanning from first upper + active cell up to the river bottom elevation. This matches the iMOD 5.6 + IDEFLAYER = -1 option. + * ``first_active_to_riv_bot__drn``: RIV. Allocate cells spanning from first + upper active cell up to the river bottom elevation. Method returns both + allocated cells for a river package as well as a drain package. Cells + above river stage are allocated as drain cells, cells below are as river + cells. This matches the iMOD 5.6 IDEFLAYER = 1 option. + * ``at_elevation``: RIV, DRN, GHB. Allocate cells containing the river + bottom elevation, drain elevation, or head respectively for river, drain + and general head boundary. This matches the iMOD 5.6 IDEFLAYER = 2 + option. + * ``at_first_active``: RIV, DRN, GHB, RCH. Allocate cells at the upper + active cells. This has no equivalent option in iMOD 5.6. + + Examples + -------- + + >>> from imod.prepare.topsystem import ALLOCATION_OPTION + >>> setting = ALLOCATION_OPTION.at_first_active + >>> allocated = allocate_rch_cells(setting, active) """ stage_to_riv_bot = 0 @@ -44,6 +70,43 @@ def allocate_riv_cells( stage: GridDataArray, bottom_elevation: GridDataArray, ) -> tuple[GridDataArray, Optional[GridDataArray]]: + """ + Allocate river cells from a planar grid across the vertical dimension. + Multiple options are available, which can be selected in ALLOCATION_OPTION. + + Parameters + ---------- + allocation_option: ALLOCATION_OPTION + Chosen allocation option, can be selected using the ALLOCATION_OPTION + enumerator. + active: DataArray | UgridDatarray + Boolean array containing active model cells. For Modflow 6, this is the + equivalent of ``idomain == 1``. + top: DataArray | UgridDatarray + Grid containing tops of model layers. If has no layer dimension, is + assumed as top of upper layer and the other layers are padded with + bottom values of the overlying model layer. + bottom: DataArray | UgridDatarray + Grid containing bottoms of model layers. + stage: DataArray | UgridDatarray + Planar grid containing river stages. Is not allowed to have a layer + dimension. + bottom_elevation: DataArray | UgridDatarray + Planar grid containing river bottom elevations. Is not allowed to have a + layer dimension. + + Returns + ------- + DataArray | UgridDatarray + Allocated river cells + + Examples + -------- + + >>> from imod.prepare.topsystem import ALLOCATION_OPTION, allocate_riv_cells + >>> setting = ALLOCATION_OPTION.stage_to_riv_bot + >>> allocated = allocate_riv_cells(setting, active, top, bottom, stage, bottom_elevation) + """ match allocation_option: case ALLOCATION_OPTION.stage_to_riv_bot: return _allocate_cells__stage_to_riv_bot( @@ -80,6 +143,40 @@ def allocate_drn_cells( bottom: GridDataArray, elevation: GridDataArray, ) -> GridDataArray: + """ + Allocate drain cells from a planar grid across the vertical dimension. + Multiple options are available, which can be selected in ALLOCATION_OPTION. + + Parameters + ---------- + allocation_option: ALLOCATION_OPTION + Chosen allocation option, can be selected using the ALLOCATION_OPTION + enumerator. + active: DataArray | UgridDatarray + Boolean array containing active model cells. For Modflow 6, this is the + equivalent of ``idomain == 1``. + top: DataArray | UgridDatarray + Grid containing tops of model layers. If has no layer dimension, is + assumed as top of upper layer and the other layers are padded with + bottom values of the overlying model layer. + bottom: DataArray | UgridDatarray + Grid containing bottoms of model layers. + elevation: DataArray | UgridDatarray + Planar grid containing drain elevation. Is not allowed to have a layer + dimension. + + Returns + ------- + DataArray | UgridDatarray + Allocated drain cells + + Examples + -------- + + >>> from imod.prepare.topsystem import ALLOCATION_OPTION, allocate_drn_cells + >>> setting = ALLOCATION_OPTION.at_elevation + >>> allocated = allocate_drn_cells(setting, active, top, bottom, stage, drain_elevation) + """ match allocation_option: case ALLOCATION_OPTION.at_elevation: return _allocate_cells__at_elevation(top, bottom, elevation)[0] @@ -101,6 +198,41 @@ def allocate_ghb_cells( bottom: GridDataArray, head: GridDataArray, ) -> GridDataArray: + """ + Allocate general head boundary (GHB) cells from a planar grid across the + vertical dimension. Multiple options are available, which can be selected in + ALLOCATION_OPTION. + + Parameters + ---------- + allocation_option: ALLOCATION_OPTION + Chosen allocation option, can be selected using the ALLOCATION_OPTION + enumerator. + active: DataArray | UgridDatarray + Boolean array containing active model cells. For Modflow 6, this is the + equivalent of ``idomain == 1``. + top: DataArray | UgridDatarray + Grid containing tops of model layers. If has no layer dimension, is + assumed as top of upper layer and the other layers are padded with + bottom values of the overlying model layer. + bottom: DataArray | UgridDatarray + Grid containing bottoms of model layers. + head: DataArray | UgridDatarray + Planar grid containing general head boundary's head. Is not allowed to + have a layer dimension. + + Returns + ------- + DataArray | UgridDatarray + Allocated general head boundary cells + + Examples + -------- + + >>> from imod.prepare.topsystem import ALLOCATION_OPTION, allocate_ghb_cells + >>> setting = ALLOCATION_OPTION.at_elevation + >>> allocated = allocate_ghb_cells(setting, active, top, bottom, head) + """ match allocation_option: case ALLOCATION_OPTION.at_elevation: return _allocate_cells__at_elevation(top, bottom, head)[0] @@ -119,6 +251,31 @@ def allocate_rch_cells( allocation_option: ALLOCATION_OPTION, active: GridDataArray, ) -> GridDataArray: + """ + Allocate recharge cells from a planar grid across the vertical dimension. + Multiple options are available, which can be selected in ALLOCATION_OPTION. + + Parameters + ---------- + allocation_option: ALLOCATION_OPTION + Chosen allocation option, can be selected using the ALLOCATION_OPTION + enumerator. + active: DataArray | UgridDatarray + Boolean array containing active model cells. For Modflow 6, this is the + equivalent of ``idomain == 1``. + + Returns + ------- + DataArray | UgridDatarray + Allocated recharge cells + + Examples + -------- + + >>> from imod.prepare.topsystem import ALLOCATION_OPTION, allocate_rch_cells + >>> setting = ALLOCATION_OPTION.at_first_active + >>> allocated = allocate_rch_cells(setting, active) + """ match allocation_option: case ALLOCATION_OPTION.at_first_active: return _allocate_cells__at_first_active(active)[0] From cee51598662a5c5929e4abac95bd7cf64e0660a2 Mon Sep 17 00:00:00 2001 From: Joeri van Engelen Date: Mon, 22 Apr 2024 09:45:57 +0200 Subject: [PATCH 2/2] format --- imod/prepare/topsystem/allocation.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/imod/prepare/topsystem/allocation.py b/imod/prepare/topsystem/allocation.py index a21001131..a8a9dc00a 100644 --- a/imod/prepare/topsystem/allocation.py +++ b/imod/prepare/topsystem/allocation.py @@ -38,10 +38,10 @@ class ALLOCATION_OPTION(Enum): option. * ``at_first_active``: RIV, DRN, GHB, RCH. Allocate cells at the upper active cells. This has no equivalent option in iMOD 5.6. - + Examples -------- - + >>> from imod.prepare.topsystem import ALLOCATION_OPTION >>> setting = ALLOCATION_OPTION.at_first_active >>> allocated = allocate_rch_cells(setting, active) @@ -72,7 +72,7 @@ def allocate_riv_cells( ) -> tuple[GridDataArray, Optional[GridDataArray]]: """ Allocate river cells from a planar grid across the vertical dimension. - Multiple options are available, which can be selected in ALLOCATION_OPTION. + Multiple options are available, which can be selected in ALLOCATION_OPTION. Parameters ---------- @@ -94,12 +94,12 @@ def allocate_riv_cells( bottom_elevation: DataArray | UgridDatarray Planar grid containing river bottom elevations. Is not allowed to have a layer dimension. - + Returns ------- DataArray | UgridDatarray Allocated river cells - + Examples -------- @@ -145,7 +145,7 @@ def allocate_drn_cells( ) -> GridDataArray: """ Allocate drain cells from a planar grid across the vertical dimension. - Multiple options are available, which can be selected in ALLOCATION_OPTION. + Multiple options are available, which can be selected in ALLOCATION_OPTION. Parameters ---------- @@ -164,12 +164,12 @@ def allocate_drn_cells( elevation: DataArray | UgridDatarray Planar grid containing drain elevation. Is not allowed to have a layer dimension. - + Returns ------- DataArray | UgridDatarray Allocated drain cells - + Examples -------- @@ -201,7 +201,7 @@ def allocate_ghb_cells( """ Allocate general head boundary (GHB) cells from a planar grid across the vertical dimension. Multiple options are available, which can be selected in - ALLOCATION_OPTION. + ALLOCATION_OPTION. Parameters ---------- @@ -220,12 +220,12 @@ def allocate_ghb_cells( head: DataArray | UgridDatarray Planar grid containing general head boundary's head. Is not allowed to have a layer dimension. - + Returns ------- DataArray | UgridDatarray Allocated general head boundary cells - + Examples -------- @@ -253,7 +253,7 @@ def allocate_rch_cells( ) -> GridDataArray: """ Allocate recharge cells from a planar grid across the vertical dimension. - Multiple options are available, which can be selected in ALLOCATION_OPTION. + Multiple options are available, which can be selected in ALLOCATION_OPTION. Parameters ---------- @@ -263,12 +263,12 @@ def allocate_rch_cells( active: DataArray | UgridDatarray Boolean array containing active model cells. For Modflow 6, this is the equivalent of ``idomain == 1``. - + Returns ------- DataArray | UgridDatarray Allocated recharge cells - + Examples --------