Coverage for C:\src\imod-python\imod\prepare\topsystem\allocation.py: 27%
78 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 10:26 +0200
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-08 10:26 +0200
1"""
2This module contains all kinds of utilities to prepare rivers
3"""
5from enum import Enum
7from imod.prepare.layer import (
8 create_layered_top,
9 get_upper_active_grid_cells,
10 get_upper_active_layer_number,
11)
12from imod.schemata import DimsSchema
13from imod.typing import GridDataArray
16class ALLOCATION_OPTION(Enum):
17 """
18 Enumerator for allocation settings. Numbers match the IDEFLAYER options in
19 iMOD 5.6.
20 """
22 stage_to_riv_bot = 0
23 first_active_to_riv_bot = -1
24 first_active_to_riv_bot__drn = 1
25 at_elevation = 2
26 at_first_active = 9 # Not an iMOD 5.6 option
29PLANAR_GRID = (
30 DimsSchema("time", "y", "x")
31 | DimsSchema("y", "x")
32 | DimsSchema("time", "{face_dim}")
33 | DimsSchema("{face_dim}")
34)
37def allocate_river_cells(
38 allocation_option: ALLOCATION_OPTION,
39 active: GridDataArray,
40 top: GridDataArray,
41 bottom: GridDataArray,
42 stage: GridDataArray,
43 bottom_elevation: GridDataArray,
44):
45 match allocation_option:
46 case ALLOCATION_OPTION.stage_to_riv_bot:
47 return _allocate_cells__stage_to_riv_bot(
48 top, bottom, stage, bottom_elevation
49 )
50 case ALLOCATION_OPTION.first_active_to_riv_bot:
51 return _allocate_cells__first_active_to_riv_bot(
52 active, bottom, bottom_elevation
53 )
54 case ALLOCATION_OPTION.first_active_to_riv_bot__drn:
55 return _allocate_cells__first_active_to_riv_bot__drn(
56 active, top, bottom, stage, bottom_elevation
57 )
58 case ALLOCATION_OPTION.at_elevation:
59 return _allocate_cells__at_elevation(top, bottom, bottom_elevation)
60 case ALLOCATION_OPTION.at_first_active:
61 return _allocate_cells__at_first_active(active)
64def allocate_drain_cells(
65 allocation_option: ALLOCATION_OPTION,
66 active: GridDataArray,
67 top: GridDataArray,
68 bottom: GridDataArray,
69 elevation: GridDataArray,
70):
71 match allocation_option:
72 case ALLOCATION_OPTION.at_elevation:
73 return _allocate_cells__at_elevation(top, bottom, elevation)
74 case ALLOCATION_OPTION.at_first_active:
75 return _allocate_cells__at_first_active(active)
76 case _:
77 raise ValueError(
78 "Received incompatible setting for drains, only"
79 f"'{ALLOCATION_OPTION.at_elevation.name}' and"
80 f"'{ALLOCATION_OPTION.at_first_active.name}' supported."
81 f"got: '{allocation_option.name}'"
82 )
85def allocate_ghb_cells(
86 allocation_option: ALLOCATION_OPTION,
87 active: GridDataArray,
88 top: GridDataArray,
89 bottom: GridDataArray,
90 head: GridDataArray,
91):
92 match allocation_option:
93 case ALLOCATION_OPTION.at_elevation:
94 return _allocate_cells__at_elevation(top, bottom, head)
95 case ALLOCATION_OPTION.at_first_active:
96 return _allocate_cells__at_first_active(active)
97 case _:
98 raise ValueError(
99 "Received incompatible setting for drains, only"
100 f"'{ALLOCATION_OPTION.at_elevation.name}' and"
101 f"'{ALLOCATION_OPTION.at_first_active.name}' supported."
102 f"got: '{allocation_option.name}'"
103 )
106def allocate_rch_cells(
107 allocation_option: ALLOCATION_OPTION,
108 active: GridDataArray,
109):
110 match allocation_option:
111 case ALLOCATION_OPTION.at_first_active:
112 return _allocate_cells__at_first_active(active)
113 case _:
114 raise ValueError(
115 "Received incompatible setting for drains, only"
116 f"'{ALLOCATION_OPTION.at_first_active.name}' supported."
117 f"got: '{allocation_option.name}'"
118 )
121def _is_layered(grid: GridDataArray):
122 return "layer" in grid.sizes and grid.sizes["layer"] > 1
125def _allocate_cells__stage_to_riv_bot(
126 top: GridDataArray,
127 bottom: GridDataArray,
128 stage: GridDataArray,
129 bottom_elevation: GridDataArray,
130):
131 """
132 Allocate cells inbetween river stage and river bottom_elevation. Compared to
133 iMOD5.6, this is similar to setting IDEFFLAYER=0 in the RUNFILE function.
135 Note that ``stage`` and ``bottom_elevation`` are not allowed to have a layer
136 dimension.
138 Parameters
139 ----------
140 top: GridDataArray
141 top of model layers
142 bottom: GridDataArray
143 bottom of model layers
144 stage: GridDataArray
145 river stage
146 bottom_elevation: GridDataArray
147 river bottom elevation
149 Returns
150 -------
151 GridDataArray
152 River cells
153 """
154 PLANAR_GRID.validate(stage)
155 PLANAR_GRID.validate(bottom_elevation)
157 if _is_layered(top):
158 top_layered = top
159 else:
160 top_layered = create_layered_top(bottom, top)
162 return (stage <= top_layered) & (bottom_elevation >= bottom)
165def _allocate_cells__first_active_to_riv_bot(
166 active: GridDataArray, bottom: GridDataArray, bottom_elevation: GridDataArray
167):
168 """
169 Allocate cells inbetween first active layer and river bottom elevation.
170 Compared to iMOD5.6, this is similar to setting IDEFFLAYER=-1 in the RUNFILE
171 function.
173 Note that ``bottom_elevation`` is not allowed to have a layer dimension.
175 Parameters
176 ----------
177 active: GridDataArray
178 active model cells
179 bottom: GridDataArray
180 bottom of model layers
181 bottom_elevation: GridDataArray
182 river bottom elevation
184 Returns
185 -------
186 GridDataArray
187 River cells
188 """
189 PLANAR_GRID.validate(bottom_elevation)
191 upper_active_layer = get_upper_active_layer_number(active)
192 layer = active.coords["layer"]
194 return (layer >= upper_active_layer) & (bottom_elevation >= bottom)
197def _allocate_cells__first_active_to_riv_bot__drn(
198 active: GridDataArray,
199 top: GridDataArray,
200 bottom: GridDataArray,
201 stage: GridDataArray,
202 bottom_elevation: GridDataArray,
203):
204 """
205 Allocate cells inbetween first active layer and river bottom elevation.
206 Cells above river stage are deactivated and returned as drn cells. Compared
207 to iMOD5.6, this is similar to setting IDEFFLAYER=1 in the RUNFILE function.
209 Note that ``stage`` and ``bottom_elevation`` are not allowed to have a layer
210 dimension.
212 Parameters
213 ----------
214 active: GridDataArray
215 active grid cells
216 top: GridDataArray
217 top of model layers
218 bottom: GridDataArray
219 bottom of model layers
220 stage: GridDataArray
221 river stage
222 bottom_elevation: GridDataArray
223 river bottom elevation
225 Returns
226 -------
227 riv_cells: GridDataArray
228 River cells (below stage)
229 drn_cells: GridDataArray
230 Drainage cells (above stage)
231 """
233 PLANAR_GRID.validate(stage)
234 PLANAR_GRID.validate(bottom_elevation)
236 if _is_layered(top):
237 top_layered = top
238 else:
239 top_layered = create_layered_top(bottom, top)
241 upper_active_layer = get_upper_active_layer_number(active)
242 layer = active.coords["layer"]
243 drn_cells = (layer >= upper_active_layer) & (stage <= top_layered)
244 riv_cells = (layer >= upper_active_layer) & (
245 bottom_elevation >= bottom
246 ) != drn_cells
248 return riv_cells, drn_cells
251def _allocate_cells__at_elevation(
252 top: GridDataArray, bottom: GridDataArray, elevation: GridDataArray
253):
254 """
255 Allocate cells in river bottom elevation layer. Compared to iMOD5.6, this is
256 similar to setting IDEFFLAYER=2 in the RUNFILE function.
258 Note that ``bottom_elevation`` is not allowed to have a layer dimension.
260 Parameters
261 ----------
262 top: GridDataArray
263 top of model layers
264 bottom: GridDataArray
265 bottom of model layers
266 elevation: GridDataArray
267 elevation. Can be river bottom, drain elevation or head of GHB.
269 Returns
270 -------
271 GridDataArray
272 River cells
273 """
275 PLANAR_GRID.validate(elevation)
277 if _is_layered(top):
278 top_layered = top
279 else:
280 top_layered = create_layered_top(bottom, top)
282 return (elevation < top_layered) & (elevation >= bottom)
285def _allocate_cells__at_first_active(active: GridDataArray):
286 """
287 Allocate cells inbetween first active layer and river bottom elevation.
288 Compared to iMOD5.6, this is similar to setting IDEFFLAYER=-1 in the RUNFILE
289 function.
291 Note that ``bottom_elevation`` is not allowed to have a layer dimension.
293 Parameters
294 ----------
295 active: GridDataArray
296 active model cells
298 Returns
299 -------
300 GridDataArray
301 River cells
302 """
304 return get_upper_active_grid_cells(active)