Coverage for C:\src\imod-python\imod\prepare\topsystem\allocation.py: 92%
97 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +0200
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +0200
1"""
2This module contains all kinds of utilities to prepare rivers
3"""
5from enum import Enum
6from typing import Optional
8import numpy as np
10from imod.prepare.layer import (
11 create_layered_top,
12 get_upper_active_grid_cells,
13 get_upper_active_layer_number,
14)
15from imod.schemata import DimsSchema
16from imod.typing import GridDataArray
17from imod.util.dims import enforced_dim_order
20class ALLOCATION_OPTION(Enum):
21 """
22 Enumerator for settings to allocate planar grid with RIV, DRN, GHB, or RCH
23 cells over the vertical layer dimensions. Numbers match the IDEFLAYER
24 options in iMOD 5.6.
26 * ``stage_to_riv_bot``: RIV. Allocate cells spanning from the river stage up
27 to the river bottom elevation. This matches the iMOD 5.6 IDEFLAYER = 0
28 option.
29 * ``first_active_to_elevation``: RIV, DRN, GHB. Allocate cells spanning from
30 first upper active cell up to the river bottom elevation. This matches the
31 iMOD 5.6 IDEFLAYER = -1 option.
32 * ``stage_to_riv_bot_drn_above``: RIV. Allocate cells spanning from first
33 upper active cell up to the river bottom elevation. Method returns both
34 allocated cells for a river package as well as a drain package. Cells
35 above river stage are allocated as drain cells, cells below are as river
36 cells. This matches the iMOD 5.6 IDEFLAYER = 1 option.
37 * ``at_elevation``: RIV, DRN, GHB. Allocate cells containing the river
38 bottom elevation, drain elevation, or head respectively for river, drain
39 and general head boundary. This matches the iMOD 5.6 IDEFLAYER = 2
40 option.
41 * ``at_first_active``: RIV, DRN, GHB, RCH. Allocate cells at the upper
42 active cells. This has no equivalent option in iMOD 5.6.
44 Examples
45 --------
47 >>> from imod.prepare.topsystem import ALLOCATION_OPTION
48 >>> setting = ALLOCATION_OPTION.at_first_active
49 >>> allocated = allocate_rch_cells(setting, active, rate)
50 """
52 stage_to_riv_bot = 0
53 first_active_to_elevation = -1
54 stage_to_riv_bot_drn_above = 1
55 at_elevation = 2
56 at_first_active = 9 # Not an iMOD 5.6 option
59PLANAR_GRID = (
60 DimsSchema("time", "y", "x")
61 | DimsSchema("y", "x")
62 | DimsSchema("time", "{face_dim}")
63 | DimsSchema("{face_dim}")
64)
67def allocate_riv_cells(
68 allocation_option: ALLOCATION_OPTION,
69 active: GridDataArray,
70 top: GridDataArray,
71 bottom: GridDataArray,
72 stage: GridDataArray,
73 bottom_elevation: GridDataArray,
74) -> tuple[GridDataArray, Optional[GridDataArray]]:
75 """
76 Allocate river cells from a planar grid across the vertical dimension.
77 Multiple options are available, which can be selected in ALLOCATION_OPTION.
79 Parameters
80 ----------
81 allocation_option: ALLOCATION_OPTION
82 Chosen allocation option, can be selected using the ALLOCATION_OPTION
83 enumerator.
84 active: DataArray | UgridDatarray
85 Boolean array containing active model cells. For Modflow 6, this is the
86 equivalent of ``idomain == 1``.
87 top: DataArray | UgridDatarray
88 Grid containing tops of model layers. If has no layer dimension, is
89 assumed as top of upper layer and the other layers are padded with
90 bottom values of the overlying model layer.
91 bottom: DataArray | UgridDatarray
92 Grid containing bottoms of model layers.
93 stage: DataArray | UgridDatarray
94 Planar grid containing river stages. Is not allowed to have a layer
95 dimension.
96 bottom_elevation: DataArray | UgridDatarray
97 Planar grid containing river bottom elevations. Is not allowed to have a
98 layer dimension.
100 Returns
101 -------
102 tuple(DataArray | UgridDatarray)
103 Allocated river cells
105 Examples
106 --------
108 >>> from imod.prepare.topsystem import ALLOCATION_OPTION, allocate_riv_cells
109 >>> setting = ALLOCATION_OPTION.stage_to_riv_bot
110 >>> allocated = allocate_riv_cells(setting, active, top, bottom, stage, bottom_elevation)
111 """
112 match allocation_option:
113 case ALLOCATION_OPTION.stage_to_riv_bot:
114 return _allocate_cells__stage_to_riv_bot(
115 top, bottom, stage, bottom_elevation
116 )
117 case ALLOCATION_OPTION.first_active_to_elevation:
118 return _allocate_cells__first_active_to_elevation(
119 active, top, bottom, bottom_elevation
120 )
121 case ALLOCATION_OPTION.stage_to_riv_bot_drn_above:
122 return _allocate_cells__stage_to_riv_bot_drn_above(
123 active, top, bottom, stage, bottom_elevation
124 )
125 case ALLOCATION_OPTION.at_elevation:
126 return _allocate_cells__at_elevation(top, bottom, bottom_elevation)
127 case ALLOCATION_OPTION.at_first_active:
128 return _allocate_cells__at_first_active(active, bottom_elevation)
129 case _:
130 raise ValueError(
131 "Received incompatible setting for rivers, only"
132 f"'{ALLOCATION_OPTION.stage_to_riv_bot.name}' and"
133 f"'{ALLOCATION_OPTION.first_active_to_elevation.name}' and"
134 f"'{ALLOCATION_OPTION.stage_to_riv_bot_drn_above.name}' and"
135 f"'{ALLOCATION_OPTION.at_elevation.name}' and"
136 f"'{ALLOCATION_OPTION.at_first_active.name}' supported."
137 f"got: '{allocation_option.name}'"
138 )
141def allocate_drn_cells(
142 allocation_option: ALLOCATION_OPTION,
143 active: GridDataArray,
144 top: GridDataArray,
145 bottom: GridDataArray,
146 elevation: GridDataArray,
147) -> GridDataArray:
148 """
149 Allocate drain cells from a planar grid across the vertical dimension.
150 Multiple options are available, which can be selected in ALLOCATION_OPTION.
152 Parameters
153 ----------
154 allocation_option: ALLOCATION_OPTION
155 Chosen allocation option, can be selected using the ALLOCATION_OPTION
156 enumerator.
157 active: DataArray | UgridDatarray
158 Boolean array containing active model cells. For Modflow 6, this is the
159 equivalent of ``idomain == 1``.
160 top: DataArray | UgridDatarray
161 Grid containing tops of model layers. If has no layer dimension, is
162 assumed as top of upper layer and the other layers are padded with
163 bottom values of the overlying model layer.
164 bottom: DataArray | UgridDatarray
165 Grid containing bottoms of model layers.
166 elevation: DataArray | UgridDatarray
167 Planar grid containing drain elevation. Is not allowed to have a layer
168 dimension.
170 Returns
171 -------
172 DataArray | UgridDatarray
173 Allocated drain cells
175 Examples
176 --------
178 >>> from imod.prepare.topsystem import ALLOCATION_OPTION, allocate_drn_cells
179 >>> setting = ALLOCATION_OPTION.at_elevation
180 >>> allocated = allocate_drn_cells(setting, active, top, bottom, stage, drain_elevation)
181 """
182 match allocation_option:
183 case ALLOCATION_OPTION.first_active_to_elevation:
184 return _allocate_cells__first_active_to_elevation(
185 active, top, bottom, elevation
186 )[0]
187 case ALLOCATION_OPTION.at_elevation:
188 return _allocate_cells__at_elevation(top, bottom, elevation)[0]
189 case ALLOCATION_OPTION.at_first_active:
190 return _allocate_cells__at_first_active(active, elevation)[0]
191 case _:
192 raise ValueError(
193 "Received incompatible setting for drains, only"
194 f"'{ALLOCATION_OPTION.first_active_to_elevation.name}', "
195 f"'{ALLOCATION_OPTION.at_elevation.name}' and"
196 f"'{ALLOCATION_OPTION.at_first_active.name}' supported."
197 f"got: '{allocation_option.name}'"
198 )
201def allocate_ghb_cells(
202 allocation_option: ALLOCATION_OPTION,
203 active: GridDataArray,
204 top: GridDataArray,
205 bottom: GridDataArray,
206 head: GridDataArray,
207) -> GridDataArray:
208 """
209 Allocate general head boundary (GHB) cells from a planar grid across the
210 vertical dimension. Multiple options are available, which can be selected in
211 ALLOCATION_OPTION.
213 Parameters
214 ----------
215 allocation_option: ALLOCATION_OPTION
216 Chosen allocation option, can be selected using the ALLOCATION_OPTION
217 enumerator.
218 active: DataArray | UgridDatarray
219 Boolean array containing active model cells. For Modflow 6, this is the
220 equivalent of ``idomain == 1``.
221 top: DataArray | UgridDatarray
222 Grid containing tops of model layers. If has no layer dimension, is
223 assumed as top of upper layer and the other layers are padded with
224 bottom values of the overlying model layer.
225 bottom: DataArray | UgridDatarray
226 Grid containing bottoms of model layers.
227 head: DataArray | UgridDatarray
228 Planar grid containing general head boundary's head. Is not allowed to
229 have a layer dimension.
231 Returns
232 -------
233 DataArray | UgridDatarray
234 Allocated general head boundary cells
236 Examples
237 --------
239 >>> from imod.prepare.topsystem import ALLOCATION_OPTION, allocate_ghb_cells
240 >>> setting = ALLOCATION_OPTION.at_elevation
241 >>> allocated = allocate_ghb_cells(setting, active, top, bottom, head)
242 """
243 match allocation_option:
244 case ALLOCATION_OPTION.first_active_to_elevation:
245 return _allocate_cells__first_active_to_elevation(
246 active, top, bottom, head
247 )[0]
248 case ALLOCATION_OPTION.at_elevation:
249 return _allocate_cells__at_elevation(top, bottom, head)[0]
250 case ALLOCATION_OPTION.at_first_active:
251 return _allocate_cells__at_first_active(active, head)[0]
252 case _:
253 raise ValueError(
254 "Received incompatible setting for general head boundary, only"
255 f"'{ALLOCATION_OPTION.first_active_to_elevation.name}', "
256 f"'{ALLOCATION_OPTION.at_elevation.name}' and"
257 f"'{ALLOCATION_OPTION.at_first_active.name}' supported."
258 f"got: '{allocation_option.name}'"
259 )
262def allocate_rch_cells(
263 allocation_option: ALLOCATION_OPTION,
264 active: GridDataArray,
265 rate: GridDataArray,
266) -> GridDataArray:
267 """
268 Allocate recharge cells from a planar grid across the vertical dimension.
269 Multiple options are available, which can be selected in ALLOCATION_OPTION.
271 Parameters
272 ----------
273 allocation_option: ALLOCATION_OPTION
274 Chosen allocation option, can be selected using the ALLOCATION_OPTION
275 enumerator.
276 active: DataArray | UgridDataArray
277 Boolean array containing active model cells. For Modflow 6, this is the
278 equivalent of ``idomain == 1``.
279 rate: DataArray | UgridDataArray
280 Array with recharge rates. This will only be used to infer where
281 recharge cells are defined.
283 Returns
284 -------
285 DataArray | UgridDataArray
286 Allocated recharge cells
288 Examples
289 --------
291 >>> from imod.prepare.topsystem import ALLOCATION_OPTION, allocate_rch_cells
292 >>> setting = ALLOCATION_OPTION.at_first_active
293 >>> allocated = allocate_rch_cells(setting, active, rate)
294 """
295 match allocation_option:
296 case ALLOCATION_OPTION.at_first_active:
297 return _allocate_cells__at_first_active(active, rate)[0]
298 case _:
299 raise ValueError(
300 "Received incompatible setting for recharge, only"
301 f"'{ALLOCATION_OPTION.at_first_active.name}' supported."
302 f"got: '{allocation_option.name}'"
303 )
306def _is_layered(grid: GridDataArray):
307 return "layer" in grid.sizes and grid.sizes["layer"] > 1
310def _enforce_layered_top(top: GridDataArray, bottom: GridDataArray):
311 if _is_layered(top):
312 return top
313 else:
314 return create_layered_top(bottom, top)
317@enforced_dim_order
318def _allocate_cells__stage_to_riv_bot(
319 top: GridDataArray,
320 bottom: GridDataArray,
321 stage: GridDataArray,
322 bottom_elevation: GridDataArray,
323) -> tuple[GridDataArray, None]:
324 """
325 Allocate cells inbetween river stage and river bottom_elevation. Compared to
326 iMOD5.6, this is similar to setting IDEFFLAYER=0 in the RUNFILE function.
328 Note that ``stage`` and ``bottom_elevation`` are not allowed to have a layer
329 dimension.
331 Parameters
332 ----------
333 top: GridDataArray
334 top of model layers
335 bottom: GridDataArray
336 bottom of model layers
337 stage: GridDataArray
338 river stage, cannot contain a layer dimension. Can contain a time
339 dimension.
340 bottom_elevation: GridDataArray
341 river bottom elevation, cannot contain a layer dimension. Can contain a
342 time dimension.
344 Returns
345 -------
346 GridDataArray
347 River cells
348 """
349 PLANAR_GRID.validate(stage)
350 PLANAR_GRID.validate(bottom_elevation)
352 top_layered = _enforce_layered_top(top, bottom)
354 riv_cells = (stage > bottom) & (bottom_elevation < top_layered)
356 return riv_cells, None
359@enforced_dim_order
360def _allocate_cells__first_active_to_elevation(
361 active: GridDataArray,
362 top: GridDataArray,
363 bottom: GridDataArray,
364 bottom_elevation: GridDataArray,
365) -> tuple[GridDataArray, None]:
366 """
367 Allocate cells inbetween first active layer and river bottom elevation.
368 Compared to iMOD5.6, this is similar to setting IDEFFLAYER=-1 in the RUNFILE
369 function.
371 Note that ``bottom_elevation`` is not allowed to have a layer dimension.
373 Parameters
374 ----------
375 active: GridDataArray
376 active model cells
377 top: GridDataArray
378 top of model layers
379 bottom: GridDataArray
380 bottom of model layers
381 bottom_elevation: GridDataArray
382 river bottom elevation, cannot contain a layer dimension. Can contain a
383 time dimension.
385 Returns
386 -------
387 GridDataArray
388 River cells
389 """
390 PLANAR_GRID.validate(bottom_elevation)
392 upper_active_layer = get_upper_active_layer_number(active)
393 layer = active.coords["layer"]
395 top_layered = _enforce_layered_top(top, bottom)
397 riv_cells = (upper_active_layer <= layer) & (bottom_elevation < top_layered)
399 return riv_cells, None
402@enforced_dim_order
403def _allocate_cells__stage_to_riv_bot_drn_above(
404 active: GridDataArray,
405 top: GridDataArray,
406 bottom: GridDataArray,
407 stage: GridDataArray,
408 bottom_elevation: GridDataArray,
409) -> tuple[GridDataArray, GridDataArray]:
410 """
411 Allocate cells inbetween first active layer and river bottom elevation.
412 Cells above river stage are deactivated and returned as drn cells. Compared
413 to iMOD5.6, this is similar to setting IDEFFLAYER=1 in the RUNFILE function.
415 Note that ``stage`` and ``bottom_elevation`` are not allowed to have a layer
416 dimension.
418 Parameters
419 ----------
420 active: GridDataArray
421 active grid cells
422 top: GridDataArray
423 top of model layers
424 bottom: GridDataArray
425 bottom of model layers
426 stage: GridDataArray
427 river stage, cannot contain a layer dimension. Can contain a time
428 dimension.
429 bottom_elevation: GridDataArray
430 river bottom elevation, cannot contain a layer dimension. Can contain a
431 time dimension.
433 Returns
434 -------
435 riv_cells: GridDataArray
436 River cells (below stage)
437 drn_cells: GridDataArray
438 Drainage cells (above stage)
439 """
441 PLANAR_GRID.validate(stage)
442 PLANAR_GRID.validate(bottom_elevation)
444 top_layered = _enforce_layered_top(top, bottom)
446 upper_active_layer = get_upper_active_layer_number(active)
447 layer = active.coords["layer"]
448 drn_cells = (upper_active_layer <= layer) & (bottom >= stage)
449 riv_cells = (
450 (upper_active_layer <= layer) & (bottom_elevation < top_layered)
451 ) != drn_cells
453 return riv_cells, drn_cells
456@enforced_dim_order
457def _allocate_cells__at_elevation(
458 top: GridDataArray, bottom: GridDataArray, elevation: GridDataArray
459) -> tuple[GridDataArray, None]:
460 """
461 Allocate cells in river bottom elevation layer. Compared to iMOD5.6, this is
462 similar to setting IDEFFLAYER=2 in the RUNFILE function.
464 Note that ``bottom_elevation`` is not allowed to have a layer dimension.
466 Parameters
467 ----------
468 top: GridDataArray
469 top of model layers
470 bottom: GridDataArray
471 bottom of model layers
472 elevation: GridDataArray
473 elevation. Can be river bottom, drain elevation or head of GHB. Cannot
474 contain a layer dimension. Can contain a time dimension.
476 Returns
477 -------
478 GridDataArray
479 Topsystem cells
480 """
482 PLANAR_GRID.validate(elevation)
484 top_layered = _enforce_layered_top(top, bottom)
486 riv_cells = (elevation < top_layered) & (elevation >= bottom)
488 return riv_cells, None
491@enforced_dim_order
492def _allocate_cells__at_first_active(
493 active: GridDataArray,
494 planar_topsystem_grid: GridDataArray,
495) -> tuple[GridDataArray, None]:
496 """
497 Allocate cells inbetween first active layer and river bottom elevation.
498 Compared to iMOD5.6, this is similar to setting IDEFFLAYER=-1 in the RUNFILE
499 function.
501 Note that ``bottom_elevation`` is not allowed to have a layer dimension.
503 Parameters
504 ----------
505 active: GridDataArray
506 active model cells
507 planar_topsystem_grid: GridDataArray
508 Grid with planar topsystem cells, assumed active where not nan.
510 Returns
511 -------
512 GridDataArray
513 Topsystem cells
514 """
515 PLANAR_GRID.validate(planar_topsystem_grid)
517 upper_active = get_upper_active_grid_cells(active)
518 topsystem_upper_active = upper_active & ~np.isnan(planar_topsystem_grid)
520 return topsystem_upper_active, None