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

1""" 

2This module contains all kinds of utilities to prepare rivers 

3""" 

4 

5from enum import Enum 

6 

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 

14 

15 

16class ALLOCATION_OPTION(Enum): 

17 """ 

18 Enumerator for allocation settings. Numbers match the IDEFLAYER options in 

19 iMOD 5.6. 

20 """ 

21 

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 

27 

28 

29PLANAR_GRID = ( 

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

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

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

33 | DimsSchema("{face_dim}") 

34) 

35 

36 

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) 

62 

63 

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 ) 

83 

84 

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 ) 

104 

105 

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 ) 

119 

120 

121def _is_layered(grid: GridDataArray): 

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

123 

124 

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. 

134 

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

136 dimension. 

137 

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 

148 

149 Returns 

150 ------- 

151 GridDataArray 

152 River cells 

153 """ 

154 PLANAR_GRID.validate(stage) 

155 PLANAR_GRID.validate(bottom_elevation) 

156 

157 if _is_layered(top): 

158 top_layered = top 

159 else: 

160 top_layered = create_layered_top(bottom, top) 

161 

162 return (stage <= top_layered) & (bottom_elevation >= bottom) 

163 

164 

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. 

172 

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

174 

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 

183 

184 Returns 

185 ------- 

186 GridDataArray 

187 River cells 

188 """ 

189 PLANAR_GRID.validate(bottom_elevation) 

190 

191 upper_active_layer = get_upper_active_layer_number(active) 

192 layer = active.coords["layer"] 

193 

194 return (layer >= upper_active_layer) & (bottom_elevation >= bottom) 

195 

196 

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. 

208 

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

210 dimension. 

211 

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 

224 

225 Returns 

226 ------- 

227 riv_cells: GridDataArray 

228 River cells (below stage) 

229 drn_cells: GridDataArray 

230 Drainage cells (above stage) 

231 """ 

232 

233 PLANAR_GRID.validate(stage) 

234 PLANAR_GRID.validate(bottom_elevation) 

235 

236 if _is_layered(top): 

237 top_layered = top 

238 else: 

239 top_layered = create_layered_top(bottom, top) 

240 

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 

247 

248 return riv_cells, drn_cells 

249 

250 

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. 

257 

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

259 

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. 

268 

269 Returns 

270 ------- 

271 GridDataArray 

272 River cells 

273 """ 

274 

275 PLANAR_GRID.validate(elevation) 

276 

277 if _is_layered(top): 

278 top_layered = top 

279 else: 

280 top_layered = create_layered_top(bottom, top) 

281 

282 return (elevation < top_layered) & (elevation >= bottom) 

283 

284 

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. 

290 

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

292 

293 Parameters 

294 ---------- 

295 active: GridDataArray 

296 active model cells 

297 

298 Returns 

299 ------- 

300 GridDataArray 

301 River cells 

302 """ 

303 

304 return get_upper_active_grid_cells(active)