Coverage for C:\src\imod-python\imod\evaluate\budget.py: 89%

181 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-16 11:41 +0200

1from typing import Tuple 

2 

3import dask 

4import dask.array 

5import numba 

6import numpy as np 

7import scipy.ndimage 

8import xarray as xr 

9 

10DIM_Z = 0 

11DIM_Y = 1 

12DIM_X = 2 

13LOWER = 3 

14UPPER = 4 

15FRONT = 5 

16BACK = 6 

17RIGHT = 7 

18LEFT = 8 

19 

20 

21def _outer_edge(da): 

22 faces = np.full(da.shape, np.nan) 

23 unique_ids = np.unique(da.values) 

24 unique_ids = unique_ids[~np.isnan(unique_ids)] 

25 # Number the faces by their id 

26 for unique_id in unique_ids: 

27 data = (da == unique_id).values 

28 from_edge = ~scipy.ndimage.binary_dilation(~data) 

29 is_edge = (data == 1) & (from_edge == 0) 

30 faces[is_edge] = unique_id 

31 return xr.DataArray(faces, da.coords, da.dims) 

32 

33 

34@numba.njit 

35def _face_indices(face, budgetzone): 

36 nface = int(np.isfinite(face).sum()) 

37 shape = (nface, 9) 

38 indices = np.zeros(shape, dtype=np.int32) 

39 if nface == 0: 

40 return indices 

41 # Loop over cells 

42 nlay, nrow, ncol = budgetzone.shape 

43 count = 0 

44 for k in range(nlay): 

45 for i in range(nrow): 

46 for j in range(ncol): 

47 face_value = face[k, i, j] 

48 if ~np.isnan(face_value): 

49 # Store indices 

50 indices[count, DIM_Z] = k 

51 indices[count, DIM_Y] = i 

52 indices[count, DIM_X] = j 

53 # Default value: part of domain for edges 

54 lower = front = right = upper = back = left = face_value 

55 if k > 0: 

56 upper = budgetzone[k - 1, i, j] 

57 if k < (nlay - 1): 

58 lower = budgetzone[k + 1, i, j] 

59 if i > 0: 

60 back = budgetzone[k, i - 1, j] 

61 if i < (nrow - 1): 

62 front = budgetzone[k, i + 1, j] 

63 if j > 0: 

64 left = budgetzone[k, i, j - 1] 

65 if j < (ncol - 1): 

66 right = budgetzone[k, i, j + 1] 

67 

68 # Test if cell is a control surface cell for the direction 

69 if lower != face_value: 

70 indices[count, LOWER] = 1 

71 if upper != face_value: 

72 indices[count, UPPER] = 1 

73 if front != face_value: 

74 indices[count, FRONT] = 1 

75 if back != face_value: 

76 indices[count, BACK] = 1 

77 if right != face_value: 

78 indices[count, RIGHT] = 1 

79 if left != face_value: 

80 indices[count, LEFT] = 1 

81 

82 # Incrementer counter 

83 count += 1 

84 return indices 

85 

86 

87@numba.njit 

88def _collect_flowfront(indices, flow): 

89 result = np.zeros(flow.shape) 

90 nface = indices.shape[0] 

91 for count in range(nface): 

92 k = indices[count, DIM_Z] 

93 i = indices[count, DIM_Y] 

94 j = indices[count, DIM_X] 

95 if indices[count, FRONT]: 

96 result[k, i, j] += flow[k, i, j] 

97 if indices[count, BACK]: 

98 result[k, i, j] -= flow[k, i - 1, j] 

99 return result 

100 

101 

102@numba.njit 

103def _collect_flowlower(indices, flow): 

104 result = np.zeros(flow.shape) 

105 nface = indices.shape[0] 

106 for count in range(nface): 

107 k = indices[count, DIM_Z] 

108 i = indices[count, DIM_Y] 

109 j = indices[count, DIM_X] 

110 if indices[count, LOWER]: 

111 result[k, i, j] += flow[k, i, j] 

112 if indices[count, UPPER]: 

113 result[k, i, j] -= flow[k - 1, i, j] 

114 return result 

115 

116 

117@numba.njit 

118def _collect_flowright(indices, flow): 

119 result = np.zeros(flow.shape) 

120 nface = indices.shape[0] 

121 for count in range(nface): 

122 k = indices[count, DIM_Z] 

123 i = indices[count, DIM_Y] 

124 j = indices[count, DIM_X] 

125 if indices[count, RIGHT]: 

126 result[k, i, j] += flow[k, i, j] 

127 if indices[count, LEFT]: 

128 result[k, i, j] -= flow[k, i, j - 1] 

129 return result 

130 

131 

132def delayed_collect(indices, front, lower, right): 

133 indices = indices.compute() 

134 result_front = dask.delayed(_collect_flowfront, nout=1)(indices, front.values) 

135 result_lower = dask.delayed(_collect_flowlower, nout=1)(indices, lower.values) 

136 result_right = dask.delayed(_collect_flowright, nout=1)(indices, right.values) 

137 dask_front = dask.array.from_delayed(result_front, front.shape, dtype=np.float64) 

138 dask_lower = dask.array.from_delayed(result_lower, lower.shape, dtype=np.float64) 

139 dask_right = dask.array.from_delayed(result_right, right.shape, dtype=np.float64) 

140 return dask_front, dask_lower, dask_right 

141 

142 

143def facebudget(budgetzone, front=None, lower=None, right=None, netflow=True): 

144 """ 

145 Computes net face flow into a control volume, as defined by ``budgetzone``. 

146 

147 Returns a three dimensional DataArray with in- and outgoing flows for every 

148 cell that is located on the edge of the control volume, thereby calculating 

149 the flow through the control surface of the control volume. 

150 

151 Front, lower, and right arguments refer to iMOD face flow budgets, in cubic 

152 meters per day. In terms of flow direction these are defined as: 

153 

154 * ``front``: positive with ``y`` (negative with row index) 

155 * ``lower``: positive with ``layer`` (positive with layer index) 

156 * ``right``: negative with ``x`` (negative with column index) 

157 

158 Only a single face flow has to be defined, for example, if only vertical 

159 fluxes (``lower``) are to be considered. 

160 

161 Note that you generally should not define a budgetzone that is only one cell 

162 wide. In that case, flow will both enter and leave the cell, resulting in a 

163 net flow of zero (given there are no boundary conditions present). 

164 

165 The output of this function defines ingoing flow as positive, and outgoing 

166 flow as negative. The output is a 3D array with net flows for every control 

167 surface cell. You can select specific parts of the control surface, for 

168 example only the east-facing side of the control volume. Please refer to the 

169 examples. 

170 

171 Parameters 

172 ---------- 

173 budgetzone: xr.DataAray of floats 

174 Array defining zones. Non-zones should be with a ``np.nan`` value. 

175 Dimensions must be exactly ``("layer", "y", "x")``. 

176 front: xr.DataArray of floats, optional 

177 Dimensions must be exactly ``("layer", "y", "x")`` or 

178 ``("time", "layer", "y", "x")``. 

179 lower: xr.DataArray of floats, optional 

180 Dimensions must be exactly ``("layer", "y", "x")`` or 

181 ``("time", "layer", "y", "x")``. 

182 right: xr.DataArray of floats, optional 

183 Dimensions must be exactly ``("layer", "y", "x")`` or 

184 ``("time", "layer", "y", "x")``. 

185 netflow : bool, optional 

186 Whether to split flows by direction (front, lower, right). 

187 True: sum all flows. False: return individual directions. 

188 

189 Returns 

190 ------- 

191 facebudget_front, facebudget_lower, face_budget_right : xr.DataArray of floats 

192 Only returned if `netflow` is False. 

193 facebudget_net : xr.DataArray of floats 

194 Only returned if `netflow` is True. 

195 

196 Examples 

197 -------- 

198 Load the face flows, and select the last time using index selection: 

199 

200 >>> import imod 

201 >>> lower = imod.idf.open("bdgflf*.idf").isel(time=-1) 

202 >>> right = imod.idf.open("bdgfrf*.idf").isel(time=-1) 

203 >>> front = imod.idf.open("bdgfff*.idf").isel(time=-1) 

204 

205 Define the zone of interest, e.g. via rasterizing a shapefile: 

206 

207 >>> import geopandas as gpd 

208 >>> gdf = gpd.read_file("zone_of_interest.shp") 

209 >>> zone2D = imod.prepare.rasterize(gdf, like=lower.isel(layer=0)) 

210 

211 Broadcast it to three dimensions: 

212 

213 >>> zone = xr.ones_like(flow) * zone2D 

214 

215 Compute net flow through the (control) surface of the budget zone: 

216 

217 >>> flow = imod.evaluate.facebudget( 

218 >>> budgetzone=zone, front=front, lower=lower, right=right 

219 >>> ) 

220 

221 Or evaluate only horizontal fluxes: 

222 

223 >>> flow = imod.evaluate.facebudget( 

224 >>> budgetzone=zone, front=front, right=right 

225 >>> ) 

226 

227 Extract the net flow, only on the right side of the zone, for example as 

228 defined by x > 10000: 

229 

230 >>> netflow_right = flow.where(flow["x"] > 10_000.0).sum() 

231 

232 Extract the net flow for only the first four layers: 

233 

234 >>> netflow_layers = netflow_right.sel(layer=slice(1, 4)).sum() 

235 

236 Extract the net flow to the right of an arbitrary diagonal in ``x`` and 

237 ``y`` is simple using the equation for a straight line: 

238 

239 >>> netflow_right_of_diagonal = flow.where( 

240 >>> flow["y"] < (line_slope * flow["x"] + line_intercept) 

241 >>> ) 

242 

243 There are many ways to extract flows for a certain part of the zone of 

244 interest. The most flexible one with regards to the ``x`` and ``y`` 

245 dimensions is by drawing a vector file, rasterizing it, and using it to 

246 select with ``xarray.where()``. 

247 

248 To get the flows per direction, pass ``netflow=False``. 

249 

250 >>> flowfront, flowlower, flowright = imod.evaluate.facebudget( 

251 >>> budgetzone=zone, front=front, lower=lower, right=right, netflow=False 

252 >>> ) 

253 

254 """ 

255 # Error handling 

256 if front is None and lower is None and right is None: 

257 raise ValueError("Atleast a single flow budget DataArray must be given") 

258 if tuple(budgetzone.dims) != ("layer", "y", "x"): 

259 raise ValueError('Dimensions of budgetzone must be exactly ("layer", "y", "x")') 

260 shape = budgetzone.shape 

261 for da, daname in zip((front, lower, right), ("front", "lower", "right")): 

262 if da is not None: 

263 dims = da.dims 

264 coords = da.coords 

265 if "time" in dims: 

266 da_shape = da.shape[1:] 

267 else: 

268 da_shape = da.shape 

269 if da_shape != shape: 

270 raise ValueError(f"Shape of {daname} does not match budgetzone") 

271 

272 # Determine control surface 

273 face = _outer_edge(budgetzone) 

274 # Convert indices array to dask array, otherwise garbage collection gets 

275 # rid of the array and we get a segfault. 

276 indices = dask.array.from_array(_face_indices(face.values, budgetzone.values)) 

277 # Make sure it returns NaNs if no zones are defined. 

278 if indices.size > 0: 

279 # Create dummy arrays for skipped values, allocate just once 

280 if front is None: 

281 F = xr.full_like(budgetzone, 0.0, dtype=np.float64) 

282 if lower is None: 

283 L = xr.full_like(budgetzone, 0.0, dtype=np.float64) 

284 if right is None: 

285 R = xr.full_like(budgetzone, 0.0, dtype=np.float64) 

286 

287 results_front = [] 

288 results_lower = [] 

289 results_right = [] 

290 

291 if "time" in dims: 

292 for itime in range(front.coords["time"].size): 

293 if front is not None: 

294 F = front.isel(time=itime) 

295 if lower is not None: 

296 L = lower.isel(time=itime) 

297 if right is not None: 

298 R = right.isel(time=itime) 

299 # collect dask arrays 

300 dF, dL, dR = delayed_collect(indices, F, L, R) 

301 # append 

302 results_front.append(dF) 

303 results_lower.append(dL) 

304 results_right.append(dR) 

305 

306 # Concatenate over time dimension 

307 dask_front = dask.array.stack(results_front, axis=0) 

308 dask_lower = dask.array.stack(results_lower, axis=0) 

309 dask_right = dask.array.stack(results_right, axis=0) 

310 else: 

311 if front is not None: 

312 F = front 

313 if lower is not None: 

314 L = lower 

315 if right is not None: 

316 R = right 

317 dask_front, dask_lower, dask_right = delayed_collect(indices, F, L, R) 

318 else: 

319 chunks = (1, *da_shape) 

320 dask_front = dask.array.full(front.shape, np.nan, chunks=chunks) 

321 dask_lower = dask.array.full(lower.shape, np.nan, chunks=chunks) 

322 dask_right = dask.array.full(right.shape, np.nan, chunks=chunks) 

323 

324 if netflow: 

325 return xr.DataArray(dask_front + dask_lower + dask_right, coords, dims) 

326 else: 

327 return ( 

328 xr.DataArray(dask_front, coords, dims), 

329 xr.DataArray(dask_lower, coords, dims), 

330 xr.DataArray(dask_right, coords, dims), 

331 ) 

332 

333 

334def flow_velocity( 

335 front: xr.DataArray, 

336 lower: xr.DataArray, 

337 right: xr.DataArray, 

338 top_bot: xr.DataArray, 

339 porosity: float = 0.3, 

340) -> Tuple[xr.DataArray, xr.DataArray, xr.DataArray]: 

341 """ 

342 Compute flow velocities (m/d) from budgets (m3/d). 

343 

344 Parameters 

345 ---------- 

346 front: xr.DataArray of floats, optional 

347 Dimensions must be exactly ``("layer", "y", "x")``. 

348 lower: xr.DataArray of floats, optional 

349 Dimensions must be exactly ``("layer", "y", "x")``. 

350 right: xr.DataArray of floats, optional 

351 Dimensions must be exactly ``("layer", "y", "x")``. 

352 top_bot: xr.Dataset of floats, containing 'top', 'bot' and optionally 

353 'dz' of layers. 

354 Dimensions must be exactly ``("layer", "y", "x")``. 

355 porosity: float or xr.DataArray of floats, optional (default 0.3) 

356 If xr.DataArray, dimensions must be exactly ``("layer", "y", "x")``. 

357 

358 Returns 

359 ------- 

360 vx, vy, vz: xr.DataArray of floats 

361 Velocity components in x, y, z direction. 

362 """ 

363 if "dz" not in top_bot: 

364 top_bot["dz"] = top_bot["top"] - top_bot["bot"] 

365 

366 # cell side area (m2) 

367 A_x = np.abs(top_bot.dz * right.dy) 

368 A_y = np.abs(top_bot.dz * front.dx) 

369 A_z = np.abs(lower.dx * lower.dy) 

370 

371 # Divide flux (m3/d) by area (m2) -> (m/d) 

372 # Flip direction around for x (right) 

373 return ( 

374 -right / (A_x * porosity), 

375 front / (A_y * porosity), 

376 lower / (A_z * porosity), 

377 )