Coverage for C:\src\imod-python\imod\prepare\topsystem\conductance.py: 95%

142 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-08 14:15 +0200

1from enum import Enum 

2from typing import Optional 

3 

4import numpy as np 

5 

6from imod.prepare.topsystem.allocation import _enforce_layered_top 

7from imod.schemata import DimsSchema 

8from imod.typing import GridDataArray 

9from imod.typing.grid import ones_like, preserve_gridtype, zeros_like 

10from imod.util.dims import enforced_dim_order 

11 

12 

13class DISTRIBUTING_OPTION(Enum): 

14 """ 

15 Enumerator containing settings to distribute 2D conductance grids over 

16 vertical layers for the RIV, DRN or GHB package. 

17 

18 * ``by_corrected_transmissivity``: RIV. Distribute the conductance by 

19 corrected transmissivities. Crosscut thicknesses are used to compute 

20 transmissivities. The crosscut thicknesses is computed based on the 

21 overlap of bottom_elevation over the bottom allocated layer. Same holds 

22 for the stage and top allocated layer. Furthermore the method corrects 

23 distribution weights for the mismatch between the midpoints of crosscut 

24 areas and model layer midpoints. This is the default method in iMOD 5.6, 

25 thus DISTRCOND = 0. 

26 * ``equally``: RIV, DRN, GHB. Distribute conductances equally over layers. 

27 This matches iMOD 5.6 DISTRCOND = 1 option. 

28 * ``by_crosscut_thickness``: RIV. Distribute the conductance by crosscut 

29 thicknesses. The crosscut thicknesses is computed based on the overlap of 

30 bottom_elevation over the bottom allocated layer. Same holds for the stage 

31 and top allocated layer. This matches iMOD 5.6 DISTRCOND = 2 option. 

32 * ``by_layer_thickness``: RIV, DRN, GHB. Distribute the conductance by model 

33 layer thickness. This matches iMOD 5.6 DISTRCOND = 3 option. 

34 * ``by_crosscut_transmissivity``: RIV. Distribute the conductance by 

35 crosscut transmissivity. Crosscut thicknesses are used to compute 

36 transmissivities. The crosscut thicknesses is computed based on the 

37 overlap of bottom_elevation over the bottom allocated layer. Same holds 

38 for the stage and top allocated layer. This matches iMOD 5.6 DISTRCOND = 4 

39 option. 

40 * ``by_conductivity``: RIV, DRN, GHB. Distribute the conductance weighted by 

41 model layer hydraulic conductivities. This matches iMOD 5.6 DISTRCOND = 5 

42 option. 

43 * ``by_layer_transmissivity``: RIV, DRN, GHB. Distribute the conductance by 

44 model layer transmissivity. This has no equivalent in iMOD 5.6. 

45 """ 

46 

47 by_corrected_transmissivity = 0 

48 equally = 1 

49 by_crosscut_thickness = 2 

50 by_layer_thickness = 3 

51 by_crosscut_transmissivity = 4 

52 by_conductivity = 5 

53 by_layer_transmissivity = 9 # Not an iMOD 5.6 option 

54 

55 

56PLANAR_GRID = ( 

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

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

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

60 | DimsSchema("{face_dim}") 

61) 

62 

63 

64@enforced_dim_order 

65def distribute_riv_conductance( 

66 distributing_option: DISTRIBUTING_OPTION, 

67 allocated: GridDataArray, 

68 conductance: GridDataArray, 

69 top: GridDataArray, 

70 bottom: GridDataArray, 

71 k: GridDataArray, 

72 stage: GridDataArray, 

73 bottom_elevation: GridDataArray, 

74) -> GridDataArray: 

75 """ 

76 Function to distribute 2D conductance over vertical layers for the RIV 

77 package. Multiple options are available, which need to be selected in the 

78 DISTRIBUTING_OPTION enumerator. 

79 

80 Parameters 

81 ---------- 

82 distributing_option : DISTRIBUTING_OPTION 

83 Distributing option available in the DISTRIBUTING_OPTION enumerator. 

84 allocated: DataArray | UgridDataArray 

85 3D boolean array with river cell locations. This can be made with the 

86 :func:`imod.prepare.allocate_riv_cells` function. 

87 conductance: DataArray | UgridDataArray 

88 Planar grid with conductances that need to be distributed over layers, 

89 so grid cannot contain a layer dimension. Can contain a time dimension. 

90 top: DataArray | UgridDataArray 

91 Model top 

92 bottom: DataArray | UgridDataArray 

93 Model layer bottoms 

94 k: DataArray | UgridDataArray 

95 Hydraulic conductivities 

96 stage: DataArray | UgridDataArray 

97 Planar grid with river stages, cannot contain a layer dimension. Can 

98 contain a time dimension. 

99 bottom_elevation: DataArray | UgridDataArray 

100 Planar grid with river bottom elevations, cannot contain a layer 

101 dimension. Can contain a time dimension. 

102 

103 Returns 

104 ------- 

105 Conductances distributed over depth. 

106 

107 Examples 

108 -------- 

109 >>> from imod.prepare import allocate_riv_cells, distribute_riv_conductance, ALLOCATION_OPTION, DISTRIBUTING_OPTION 

110 >>> allocated = allocate_riv_cells( 

111 ALLOCATION_OPTION.stage_to_riv_bot, active, top, bottom, stage, bottom_elevation 

112 ) 

113 >>> conductances_distributed = distribute_riv_conductance( 

114 DISTRIBUTING_OPTION.by_corrected_transmissivity, allocated, 

115 conductance, top, bottom, stage, bottom_elevation, k 

116 ) 

117 """ 

118 PLANAR_GRID.validate(conductance) 

119 

120 match distributing_option: 

121 case DISTRIBUTING_OPTION.equally: 

122 weights = _distribute_weights__equally(allocated) 

123 case DISTRIBUTING_OPTION.by_layer_thickness: 

124 weights = _distribute_weights__by_layer_thickness(allocated, top, bottom) 

125 case DISTRIBUTING_OPTION.by_layer_transmissivity: 

126 weights = _distribute_weights__by_layer_transmissivity( 

127 allocated, top, bottom, k 

128 ) 

129 case DISTRIBUTING_OPTION.by_conductivity: 

130 weights = _distribute_weights__by_conductivity(allocated, k) 

131 case DISTRIBUTING_OPTION.by_crosscut_thickness: 

132 weights = _distribute_weights__by_crosscut_thickness( 

133 allocated, top, bottom, stage, bottom_elevation 

134 ) 

135 case DISTRIBUTING_OPTION.by_crosscut_transmissivity: 

136 weights = _distribute_weights__by_crosscut_transmissivity( 

137 allocated, top, bottom, k, stage, bottom_elevation 

138 ) 

139 case DISTRIBUTING_OPTION.by_corrected_transmissivity: 

140 weights = _distribute_weights__by_corrected_transmissivity( 

141 allocated, top, bottom, k, stage, bottom_elevation 

142 ) 

143 case _: 

144 raise ValueError( 

145 "Received incompatible setting for rivers, only" 

146 f"'{DISTRIBUTING_OPTION.equally.name}', " 

147 f"'{DISTRIBUTING_OPTION.by_layer_thickness.name}', " 

148 f"'{DISTRIBUTING_OPTION.by_layer_transmissivity.name}', " 

149 f"'{DISTRIBUTING_OPTION.by_conductivity.name}', " 

150 f"'{DISTRIBUTING_OPTION.by_crosscut_thickness.name}', " 

151 f"'{DISTRIBUTING_OPTION.by_crosscut_transmissivity.name}', and " 

152 f"'{DISTRIBUTING_OPTION.by_corrected_transmissivity.name}' supported. " 

153 f"Got: '{distributing_option.name}'" 

154 ) 

155 return (weights * conductance).where(allocated) 

156 

157 

158@enforced_dim_order 

159def distribute_drn_conductance( 

160 distributing_option: DISTRIBUTING_OPTION, 

161 allocated: GridDataArray, 

162 conductance: GridDataArray, 

163 top: GridDataArray, 

164 bottom: GridDataArray, 

165 k: GridDataArray, 

166 elevation: GridDataArray, 

167) -> GridDataArray: 

168 """ 

169 Function to distribute 2D conductance over vertical layers for the DRN 

170 package. Multiple options are available, which need to be selected in the 

171 DISTRIBUTING_OPTION enumerator. 

172 

173 Parameters 

174 ---------- 

175 distributing_option : DISTRIBUTING_OPTION 

176 Distributing option available in the DISTRIBUTING_OPTION enumerator. 

177 allocated: DataArray | UgridDataArray 

178 3D boolean array with drain cell locations. This can be made with the 

179 :func:`imod.prepare.allocate_drn_cells` function. 

180 conductance: DataArray | UgridDataArray 

181 Planar grid with conductances that need to be distributed over layers, 

182 so grid cannot contain a layer dimension. Can contain a time dimension. 

183 top: DataArray | UgridDataArray 

184 Model top 

185 bottom: DataArray | UgridDataArray 

186 Model layer bottoms 

187 k: DataArray | UgridDataArray 

188 Hydraulic conductivities 

189 elevation: DataArray | UgridDataArray 

190 Drain elevation 

191 

192 Returns 

193 ------- 

194 Conductances distributed over depth. 

195 

196 Examples 

197 -------- 

198 >>> from imod.prepare import allocate_drn_cells, distribute_drn_conductance, ALLOCATION_OPTION, DISTRIBUTING_OPTION 

199 >>> allocated = allocate_drn_cells( 

200 ALLOCATION_OPTION.at_elevation, active, top, bottom, drain_elevation 

201 ) 

202 >>> conductances_distributed = distribute_drn_conductance( 

203 DISTRIBUTING_OPTION.by_layer_transmissivity, allocated, 

204 conductance, top, bottom, k, drain_elevation 

205 ) 

206 """ 

207 PLANAR_GRID.validate(conductance) 

208 

209 match distributing_option: 

210 case DISTRIBUTING_OPTION.equally: 

211 weights = _distribute_weights__equally(allocated) 

212 case DISTRIBUTING_OPTION.by_layer_thickness: 

213 weights = _distribute_weights__by_layer_thickness(allocated, top, bottom) 

214 case DISTRIBUTING_OPTION.by_layer_transmissivity: 

215 weights = _distribute_weights__by_layer_transmissivity( 

216 allocated, top, bottom, k 

217 ) 

218 case DISTRIBUTING_OPTION.by_conductivity: 

219 weights = _distribute_weights__by_conductivity(allocated, k) 

220 case DISTRIBUTING_OPTION.by_crosscut_thickness: 

221 weights = _distribute_weights__by_crosscut_thickness( 

222 allocated, top, bottom, bc_bottom=elevation 

223 ) 

224 case DISTRIBUTING_OPTION.by_crosscut_transmissivity: 

225 weights = _distribute_weights__by_crosscut_transmissivity( 

226 allocated, top, bottom, k, bc_bottom=elevation 

227 ) 

228 case DISTRIBUTING_OPTION.by_corrected_transmissivity: 

229 weights = _distribute_weights__by_corrected_transmissivity( 

230 allocated, top, bottom, k, bc_bottom=elevation 

231 ) 

232 case _: 

233 raise ValueError( 

234 "Received incompatible setting for drains, only" 

235 f"'{DISTRIBUTING_OPTION.equally.name}', " 

236 f"'{DISTRIBUTING_OPTION.by_layer_thickness.name}', " 

237 f"'{DISTRIBUTING_OPTION.by_layer_transmissivity.name}', " 

238 f"'{DISTRIBUTING_OPTION.by_conductivity.name}', " 

239 f"'{DISTRIBUTING_OPTION.by_crosscut_thickness.name}', " 

240 f"'{DISTRIBUTING_OPTION.by_crosscut_transmissivity.name}', and " 

241 f"'{DISTRIBUTING_OPTION.by_corrected_transmissivity.name}' supported. " 

242 f"Got: '{distributing_option.name}'" 

243 ) 

244 return (weights * conductance).where(allocated) 

245 

246 

247@enforced_dim_order 

248def distribute_ghb_conductance( 

249 distributing_option: DISTRIBUTING_OPTION, 

250 allocated: GridDataArray, 

251 conductance: GridDataArray, 

252 top: GridDataArray, 

253 bottom: GridDataArray, 

254 k: GridDataArray, 

255) -> GridDataArray: 

256 PLANAR_GRID.validate(conductance) 

257 """ 

258 Function to distribute 2D conductance over vertical layers for the GHB 

259 package. Multiple options are available, which need to be selected in the 

260 DISTRIBUTING_OPTION enumerator. 

261 

262 Parameters 

263 ---------- 

264 distributing_option : DISTRIBUTING_OPTION 

265 Distributing option available in the DISTRIBUTING_OPTION enumerator. 

266 allocated: DataArray | UgridDataArray 

267 3D boolean array with GHB cell locations. This can be made with the 

268 :func:`imod.prepare.allocate_ghb_cells` function. 

269 conductance: DataArray | UgridDataArray 

270 Planar grid with conductances that need to be distributed over layers, 

271 so grid cannot contain a layer dimension. Can contain a time dimension. 

272 top: DataArray | UgridDataArray 

273 Model top 

274 bottom: DataArray | UgridDataArray 

275 Model layer bottoms 

276 k: DataArray | UgridDataArray 

277 Hydraulic conductivities 

278  

279 Returns 

280 ------- 

281 Conductances distributed over depth. 

282 

283 Examples 

284 -------- 

285 >>> from imod.prepare import allocate_ghb_cells, distribute_drn_conductance, ALLOCATION_OPTION, DISTRIBUTING_OPTION 

286 >>> allocated = allocate_ghb_cells( 

287 ALLOCATION_OPTION.at_elevation, active, top, bottom, ghb_head 

288 ) 

289 >>> conductances_distributed = distribute_ghb_conductance( 

290 DISTRIBUTING_OPTION.by_layer_transmissivity, allocated,  

291 conductance, top, bottom, k 

292 ) 

293 """ 

294 match distributing_option: 

295 case DISTRIBUTING_OPTION.equally: 

296 weights = _distribute_weights__equally(allocated) 

297 case DISTRIBUTING_OPTION.by_layer_thickness: 

298 weights = _distribute_weights__by_layer_thickness(allocated, top, bottom) 

299 case DISTRIBUTING_OPTION.by_layer_transmissivity: 

300 weights = _distribute_weights__by_layer_transmissivity( 

301 allocated, top, bottom, k 

302 ) 

303 case DISTRIBUTING_OPTION.by_conductivity: 

304 weights = _distribute_weights__by_conductivity(allocated, k) 

305 case _: 

306 raise ValueError( 

307 "Received incompatible setting for general head boundary, only" 

308 f"'{DISTRIBUTING_OPTION.equally.name}', " 

309 f"'{DISTRIBUTING_OPTION.by_layer_thickness.name}', " 

310 f"'{DISTRIBUTING_OPTION.by_layer_transmissivity.name}', and " 

311 f"'{DISTRIBUTING_OPTION.by_conductivity.name}' supported. " 

312 f"Got: '{distributing_option.name}'" 

313 ) 

314 return (weights * conductance).where(allocated) 

315 

316 

317@preserve_gridtype 

318def _compute_layer_thickness( 

319 allocated: GridDataArray, top: GridDataArray, bottom: GridDataArray 

320): 

321 """ 

322 Compute 3D grid of thicknesses in allocated cells 

323 """ 

324 top_layered = _enforce_layered_top(top, bottom) 

325 

326 thickness = top_layered - bottom 

327 return thickness.where(allocated) 

328 

329 

330@preserve_gridtype 

331def _compute_crosscut_thickness( 

332 allocated: GridDataArray, 

333 top: GridDataArray, 

334 bottom: GridDataArray, 

335 bc_top: Optional[GridDataArray] = None, 

336 bc_bottom: Optional[GridDataArray] = None, 

337): 

338 """ 

339 Compute 3D grid of thicknesses crosscut by boundary condition (river/drain) 

340 in allocated cells. So the thickness in the upper allocated layer is bc_top 

341 - bottom and the lower allocated layer is top - bc_bottom. 

342 """ 

343 if (bc_top is None) & (bc_bottom is None): 

344 raise ValueError("`bc_top` and `bc_bottom` cannot both be None.") 

345 

346 top_layered = _enforce_layered_top(top, bottom) 

347 thickness = _compute_layer_thickness(allocated, top, bottom) 

348 outside = zeros_like(allocated).astype(bool) 

349 

350 if bc_top is not None: 

351 upper_layer_bc = (bc_top < top_layered) & (bc_top > bottom) 

352 outside = outside | (bc_top < bottom) 

353 thickness = thickness.where(~upper_layer_bc, thickness - (top_layered - bc_top)) 

354 

355 if bc_bottom is not None: 

356 lower_layer_bc = (bc_bottom < top_layered) & (bc_bottom > bottom) 

357 outside = outside | (bc_bottom > top_layered) 

358 thickness = thickness.where(~lower_layer_bc, thickness - (bc_bottom - bottom)) 

359 

360 thickness = thickness.where(~outside, 0.0) 

361 

362 return thickness 

363 

364 

365def _distribute_weights__by_corrected_transmissivity( 

366 allocated: GridDataArray, 

367 top: GridDataArray, 

368 bottom: GridDataArray, 

369 k: GridDataArray, 

370 bc_top: Optional[GridDataArray] = None, 

371 bc_bottom: Optional[GridDataArray] = None, 

372): 

373 """ 

374 Distribute conductances according to default method in iMOD 5.6, as 

375 described page 690 of the iMOD 5.6 manual (but then to distribute WEL 

376 rates). The method uses crosscut thicknesses to compute transmissivities. 

377 Furthermore it corrects distribution weights for the mismatch between the 

378 midpoints of crosscut areas and layer midpoints. 

379 """ 

380 crosscut_thickness = _compute_crosscut_thickness( 

381 allocated, top, bottom, bc_top=bc_top, bc_bottom=bc_bottom 

382 ) 

383 transmissivity = crosscut_thickness * k 

384 

385 top_layered = _enforce_layered_top(top, bottom) 

386 layer_thickness = _compute_layer_thickness(allocated, top, bottom) 

387 midpoints = (top_layered + bottom) / 2 

388 Fc = midpoints.copy() 

389 

390 if bc_top is not None: 

391 PLANAR_GRID.validate(bc_top) 

392 upper_layer_bc = (bc_top < top_layered) & (bc_top > bottom) 

393 # Computing vertical midpoint of river crosscutting layers. 

394 Fc = Fc.where(~upper_layer_bc, (bottom + bc_top) / 2) 

395 

396 if bc_bottom is not None: 

397 PLANAR_GRID.validate(bc_bottom) 

398 lower_layer_bc = (bc_bottom < top_layered) & (bc_bottom > bottom) 

399 # Computing vertical midpoint of river crosscutting layers. 

400 Fc = Fc.where(~lower_layer_bc, (top_layered + bc_bottom) / 2) 

401 

402 # Correction factor for mismatch between midpoints of crosscut layers and 

403 # layer midpoints. 

404 F = 1.0 - np.abs(midpoints - Fc) / (layer_thickness * 0.5) 

405 

406 transmissivity_corrected = transmissivity * F 

407 return transmissivity_corrected / transmissivity_corrected.sum(dim="layer") 

408 

409 

410def _distribute_weights__equally(allocated: GridDataArray): 

411 """Compute weights over layers equally.""" 

412 return ones_like(allocated) / allocated.sum(dim="layer") 

413 

414 

415def _distribute_weights__by_layer_thickness( 

416 allocated: GridDataArray, 

417 top: GridDataArray, 

418 bottom: GridDataArray, 

419): 

420 """Compute weights based on layer thickness""" 

421 layer_thickness = _compute_layer_thickness(allocated, top, bottom) 

422 

423 return layer_thickness / layer_thickness.sum(dim="layer") 

424 

425 

426def _distribute_weights__by_crosscut_thickness( 

427 allocated: GridDataArray, 

428 top: GridDataArray, 

429 bottom: GridDataArray, 

430 bc_top: Optional[GridDataArray] = None, 

431 bc_bottom: Optional[GridDataArray] = None, 

432): 

433 """Compute weights based on crosscut thickness""" 

434 if bc_top is not None: 

435 PLANAR_GRID.validate(bc_top) 

436 if bc_bottom is not None: 

437 PLANAR_GRID.validate(bc_bottom) 

438 

439 crosscut_thickness = _compute_crosscut_thickness( 

440 allocated, top, bottom, bc_top, bc_bottom 

441 ).where(allocated) 

442 

443 return crosscut_thickness / crosscut_thickness.sum(dim="layer") 

444 

445 

446def _distribute_weights__by_layer_transmissivity( 

447 allocated: GridDataArray, 

448 top: GridDataArray, 

449 bottom: GridDataArray, 

450 k: GridDataArray, 

451): 

452 """Compute weights based on layer transmissivity""" 

453 layer_thickness = _compute_layer_thickness(allocated, top, bottom) 

454 transmissivity = layer_thickness * k 

455 

456 return transmissivity / transmissivity.sum(dim="layer") 

457 

458 

459def _distribute_weights__by_crosscut_transmissivity( 

460 allocated: GridDataArray, 

461 top: GridDataArray, 

462 bottom: GridDataArray, 

463 k: GridDataArray, 

464 bc_top: Optional[GridDataArray] = None, 

465 bc_bottom: Optional[GridDataArray] = None, 

466): 

467 """Compute weights based on crosscut transmissivity""" 

468 if bc_top is not None: 

469 PLANAR_GRID.validate(bc_top) 

470 if bc_bottom is not None: 

471 PLANAR_GRID.validate(bc_bottom) 

472 

473 crosscut_thickness = _compute_crosscut_thickness( 

474 allocated, top, bottom, bc_top=bc_top, bc_bottom=bc_bottom 

475 ) 

476 transmissivity = crosscut_thickness * k 

477 

478 return transmissivity / transmissivity.sum(dim="layer") 

479 

480 

481def _distribute_weights__by_conductivity(allocated: GridDataArray, k: GridDataArray): 

482 """Compute weights based on hydraulic conductivity""" 

483 k_allocated = allocated * k 

484 

485 return k_allocated / k_allocated.sum(dim="layer")