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

1""" 

2This module contains all kinds of utilities to prepare rivers 

3""" 

4 

5from enum import Enum 

6from typing import Optional 

7 

8import numpy as np 

9 

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 

18 

19 

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. 

25 

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. 

43 

44 Examples 

45 -------- 

46 

47 >>> from imod.prepare.topsystem import ALLOCATION_OPTION 

48 >>> setting = ALLOCATION_OPTION.at_first_active 

49 >>> allocated = allocate_rch_cells(setting, active, rate) 

50 """ 

51 

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 

57 

58 

59PLANAR_GRID = ( 

60 DimsSchema("time", "y", "x") 

61 | DimsSchema("y", "x") 

62 | DimsSchema("time", "{face_dim}") 

63 | DimsSchema("{face_dim}") 

64) 

65 

66 

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. 

78 

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. 

99 

100 Returns 

101 ------- 

102 tuple(DataArray | UgridDatarray) 

103 Allocated river cells 

104 

105 Examples 

106 -------- 

107 

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 ) 

139 

140 

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. 

151 

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. 

169 

170 Returns 

171 ------- 

172 DataArray | UgridDatarray 

173 Allocated drain cells 

174 

175 Examples 

176 -------- 

177 

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 ) 

199 

200 

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. 

212 

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. 

230 

231 Returns 

232 ------- 

233 DataArray | UgridDatarray 

234 Allocated general head boundary cells 

235 

236 Examples 

237 -------- 

238 

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 ) 

260 

261 

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. 

270 

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. 

282 

283 Returns 

284 ------- 

285 DataArray | UgridDataArray 

286 Allocated recharge cells 

287 

288 Examples 

289 -------- 

290 

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 ) 

304 

305 

306def _is_layered(grid: GridDataArray): 

307 return "layer" in grid.sizes and grid.sizes["layer"] > 1 

308 

309 

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) 

315 

316 

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. 

327 

328 Note that ``stage`` and ``bottom_elevation`` are not allowed to have a layer 

329 dimension. 

330 

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. 

343 

344 Returns 

345 ------- 

346 GridDataArray 

347 River cells 

348 """ 

349 PLANAR_GRID.validate(stage) 

350 PLANAR_GRID.validate(bottom_elevation) 

351 

352 top_layered = _enforce_layered_top(top, bottom) 

353 

354 riv_cells = (stage > bottom) & (bottom_elevation < top_layered) 

355 

356 return riv_cells, None 

357 

358 

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. 

370 

371 Note that ``bottom_elevation`` is not allowed to have a layer dimension. 

372 

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. 

384 

385 Returns 

386 ------- 

387 GridDataArray 

388 River cells 

389 """ 

390 PLANAR_GRID.validate(bottom_elevation) 

391 

392 upper_active_layer = get_upper_active_layer_number(active) 

393 layer = active.coords["layer"] 

394 

395 top_layered = _enforce_layered_top(top, bottom) 

396 

397 riv_cells = (upper_active_layer <= layer) & (bottom_elevation < top_layered) 

398 

399 return riv_cells, None 

400 

401 

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. 

414 

415 Note that ``stage`` and ``bottom_elevation`` are not allowed to have a layer 

416 dimension. 

417 

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. 

432 

433 Returns 

434 ------- 

435 riv_cells: GridDataArray 

436 River cells (below stage) 

437 drn_cells: GridDataArray 

438 Drainage cells (above stage) 

439 """ 

440 

441 PLANAR_GRID.validate(stage) 

442 PLANAR_GRID.validate(bottom_elevation) 

443 

444 top_layered = _enforce_layered_top(top, bottom) 

445 

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 

452 

453 return riv_cells, drn_cells 

454 

455 

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. 

463 

464 Note that ``bottom_elevation`` is not allowed to have a layer dimension. 

465 

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. 

475 

476 Returns 

477 ------- 

478 GridDataArray 

479 Topsystem cells 

480 """ 

481 

482 PLANAR_GRID.validate(elevation) 

483 

484 top_layered = _enforce_layered_top(top, bottom) 

485 

486 riv_cells = (elevation < top_layered) & (elevation >= bottom) 

487 

488 return riv_cells, None 

489 

490 

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. 

500 

501 Note that ``bottom_elevation`` is not allowed to have a layer dimension. 

502 

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. 

509 

510 Returns 

511 ------- 

512 GridDataArray 

513 Topsystem cells 

514 """ 

515 PLANAR_GRID.validate(planar_topsystem_grid) 

516 

517 upper_active = get_upper_active_grid_cells(active) 

518 topsystem_upper_active = upper_active & ~np.isnan(planar_topsystem_grid) 

519 

520 return topsystem_upper_active, None