Coverage for C:\src\imod-python\imod\visualize\cross_sections.py: 6%

170 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-08 10:26 +0200

1import matplotlib.pyplot as plt 

2import numpy as np 

3import xarray as xr 

4from mpl_toolkits.axes_grid1 import make_axes_locatable 

5 

6import imod 

7from imod.visualize import common 

8 

9 

10def _meshcoords(da, continuous=True): 

11 """ 

12 Generate coordinates for pcolormesh, or fill_between 

13 

14 Parameters 

15 ---------- 

16 da : xr.DataArray 

17 The array to plot 

18 continuous : bool, optional 

19 Whether the layers are connected, such that the bottom of layer N is 

20 the top of layer N + 1. 

21 

22 Returns 

23 ------- 

24 X : np.array 

25 x coordinates of mesh 

26 Y : np.array 

27 y coordinates of mesh 

28 C : np.array 

29 values of the mesh 

30 """ 

31 dims = tuple(da.dims) 

32 if not dims[0] == "layer": 

33 # Switch 'm around 

34 dims = (dims[1], dims[0]) 

35 

36 # Ensure dimensions are in the right order 

37 da = da.transpose(*dims, transpose_coords=True) 

38 data = da.values 

39 xcoord = da.dims[1] 

40 dx, xmin, xmax = imod.util.spatial.coord_reference(da[xcoord]) 

41 if isinstance(dx, (int, float)): 

42 if dx < 0.0: 

43 dx = abs(dx) 

44 data = data[:, ::-1] 

45 X = np.arange(xmin + dx, xmax, dx) 

46 else: # assuming dx is an array, non-equidistant dx 

47 if (dx < 0.0).all(): 

48 dx = abs(dx) 

49 data = data[:, ::-1] 

50 elif (dx > 0.0).all(): 

51 pass 

52 else: 

53 raise ValueError(f"{xcoord} is not monotonic") 

54 X = (xmin + dx[0]) + dx[1:].cumsum() 

55 

56 # if dimensions of top and bottom are 1D (voxels), promote to 2D 

57 if len(da["top"].dims) == 1 and len(da["bottom"].dims) == 1: 

58 top = np.repeat(np.atleast_2d(da.top), len(da[xcoord]), axis=0).T 

59 da = da.assign_coords(top=(("layer", xcoord), top)) 

60 bot = np.repeat(np.atleast_2d(da.bottom), len(da[xcoord]), axis=0).T 

61 da = da.assign_coords(bottom=(("layer", xcoord), bot)) 

62 

63 if continuous: 

64 Y = np.vstack([da["top"].isel(layer=0).values, da["bottom"].values]) 

65 else: 

66 nrow, ncol = da["top"].shape 

67 Y = np.empty((nrow * 2, ncol)) 

68 Y[0::2] = da["top"].values 

69 Y[1::2] = da["bottom"].values 

70 

71 Y = np.repeat(Y, 2, 1) 

72 nodata = np.isnan(Y) 

73 Y[nodata] = 0.0 

74 

75 X = np.hstack([xmin, np.repeat(X, 2), xmax]) 

76 X = np.full_like(Y, X) 

77 

78 if continuous: 

79 nrow, ncol = Y.shape 

80 C = np.full((nrow - 1, ncol - 1), np.nan) 

81 C[:, 0::2] = data 

82 else: 

83 _, ncol = Y.shape 

84 C = np.full((nrow, ncol - 1), np.nan) 

85 C[:, 0::2] = data 

86 

87 return X, Y, C, nodata 

88 

89 

90def _plot_aquitards(aquitards, ax, kwargs_aquitards): 

91 """ 

92 Overlay aquitards on ax 

93 

94 Parameters 

95 ---------- 

96 aquitards : xr.DataArray 

97 DataArray containing location of aquitards 

98 NaN's and zeros are treated as locations without aquitard 

99 ax : matplotlib.Axes object 

100 kwargs_aquitards : dict 

101 keyword arguments for ax.fill_between() 

102 """ 

103 if kwargs_aquitards is None: 

104 kwargs_aquitards = {"alpha": 0.5, "facecolor": "grey"} 

105 X_aq, Y_aq, C_aq, _ = _meshcoords(aquitards, continuous=False) 

106 C_aq = C_aq.astype(np.float64) 

107 for j, i in enumerate(range(0, X_aq.shape[0] - 1, 2)): 

108 Y_i = Y_aq[i : i + 2] 

109 C_i = C_aq[j] 

110 C_i[C_i == 0.0] = np.nan 

111 nodata = np.repeat(np.isnan(C_i[0::2]), 2) 

112 Y_i[:, nodata] = np.nan 

113 ax.fill_between(X_aq[0], Y_i[0], Y_i[1], **kwargs_aquitards) 

114 

115 

116def cross_section( 

117 da, 

118 colors, 

119 levels, 

120 layers=False, 

121 aquitards=None, 

122 kwargs_pcolormesh={}, 

123 kwargs_colorbar={}, 

124 kwargs_aquitards=None, 

125 return_cmap_norm=False, 

126 fig=None, 

127 ax=None, 

128): 

129 """ 

130 Wraps matplotlib.pcolormesh to draw cross-sections, drawing cell boundaries 

131 accurately. Aquitards can be plotted on top of the cross-section, by providing 

132 a DataArray with the aquitard location for `aquitards`. 

133 

134 Parameters 

135 ---------- 

136 da : xr.DataArray 

137 Two dimensional DataArray containing data of the cross section. One 

138 dimension must be "layer", and the second dimension will be used as the 

139 x-axis for the cross-section. 

140 

141 Coordinates "top" and "bottom" must be present, and must have at least the 

142 "layer" dimension (voxels) or both the "layer" and x-coordinate dimension. 

143 

144 *Use imod.select.cross_section_line() or cross_section_linestring() to obtain 

145 the required DataArray.* 

146 colors : list of str, or list of RGB tuples 

147 Matplotlib acceptable list of colors. Length N. 

148 Accepts both tuples of (R, G, B) and hexidecimal (e.g. "#7ec0ee"). 

149 

150 Looking for good colormaps? Try: http://colorbrewer2.org/ 

151 Choose a colormap, and use the HEX JS array. 

152 levels : listlike of floats or integers 

153 Boundaries between the legend colors/classes. Length: N - 1. 

154 layers : boolean, optional 

155 Whether to draw lines separating the layers. 

156 aquitards : xr.DataArray, optional 

157 Datarray containing data on location of aquitard layers. 

158 kwargs_pcolormesh : dict 

159 Other optional keyword arguments for matplotlib.pcolormesh. 

160 kwargs_colorbar : dict 

161 If optional key ``plot_colorbar`` is set to False, no colorbar is drawn. Defaults to True. 

162 Optional keyword argument ``whiten_triangles`` whitens respective colorbar triangle if 

163 data is not larger/smaller than legend_levels-range. Defaults to True. 

164 Other arguments are forwarded to fig.colorbar() 

165 kwargs_aquitards: dict 

166 These arguments are forwarded to matplotlib.fill_between to draw the 

167 aquitards. 

168 return_cmap_norm : boolean, optional 

169 Return the cmap and norm of the plot, default False 

170 fig : matplotlib Figure instance, optional 

171 Figure to write plot to. If not supplied, a Figure instance is created 

172 ax : matplotlib Axes instance, optional 

173 Axes to write plot to. If not supplied, an Axes instance is created 

174 

175 Returns 

176 ------- 

177 fig : matplotlib.figure 

178 ax : matplotlig.ax 

179 if return_cmap_norm == True: 

180 cmap : matplotlib.colors.ListedColormap 

181 norm : matplotlib.colors.BoundaryNorm 

182 

183 Examples 

184 -------- 

185 

186 Basic cross section: 

187 

188 >>> imod.visualize.cross_section(da, colors, levels) 

189 

190 Aquitards can be styled in multiple ways. For a transparent grey overlay 

191 (the default): 

192 

193 >>> kwargs_aquitards = {"alpha": 0.5, "facecolor": "grey"} 

194 >>> imod.visualize.cross_section(da, colors, levels, aquitards=aquitards, kwargs_aquitards) 

195 

196 For a hatched overlay: 

197 

198 >>> kwargs_aquitards = {"hatch": "/", "edgecolor": "k"} 

199 >>> imod.visualize.cross_section(da, colors, levels, aquitards=aquitards, kwargs_aquitards) 

200 """ 

201 da = da.copy(deep=False) 

202 if aquitards is not None: 

203 aquitards = aquitards.copy(deep=False) 

204 

205 if len(da.dims) != 2: 

206 raise ValueError("DataArray must be 2D") 

207 if "layer" not in da.dims: 

208 raise ValueError('DataArray must contain dimension "layer"') 

209 if "top" not in da.coords: 

210 raise ValueError('DataArray must contain coordinate "top"') 

211 if "bottom" not in da.coords: 

212 raise ValueError('DataArray must contain coordinate "bottom"') 

213 if len(da["top"].dims) > 2: 

214 raise ValueError('"top" coordinate be 1D or 2D') 

215 if len(da["bottom"].dims) > 2: 

216 raise ValueError('"bottom" coordinate be 1D or 2D') 

217 

218 # Read legend settings 

219 cmap, norm = common._cmapnorm_from_colorslevels(colors, levels) 

220 

221 # cbar kwargs 

222 settings_cbar = {"ticks": levels, "extend": "both"} 

223 

224 # Find a unit in the raster to use in the colorbar label 

225 try: 

226 settings_cbar["label"] = da.attrs["units"] 

227 except (KeyError, AttributeError): 

228 try: 

229 settings_cbar["label"] = da.attrs["unit"] 

230 except (KeyError, AttributeError): 

231 pass 

232 

233 whiten_triangles = True 

234 plot_colorbar = True 

235 if kwargs_colorbar is not None: 

236 whiten_triangles = kwargs_colorbar.pop("whiten_triangles", True) 

237 plot_colorbar = kwargs_colorbar.pop("plot_colorbar", True) 

238 settings_cbar.update(kwargs_colorbar) 

239 

240 # pcmesh kwargs 

241 settings_pcmesh = {"cmap": cmap, "norm": norm} 

242 if kwargs_pcolormesh is not None: 

243 settings_pcmesh.update(kwargs_pcolormesh) 

244 

245 if fig is None or ax is None: 

246 fig, ax = plt.subplots() 

247 

248 # Plot raster 

249 X, Y, C, nodata = _meshcoords(da, continuous=True) 

250 ax1 = ax.pcolormesh(X, Y, C, **settings_pcmesh) 

251 # Plot aquitards if applicable 

252 if aquitards is not None: 

253 _plot_aquitards(aquitards, ax, kwargs_aquitards) 

254 

255 if layers: 

256 Y[nodata] = np.nan 

257 for y in Y: 

258 ax.step(x=X[0], y=y) 

259 

260 # Make triangles white if data is not larger/smaller than legend_levels-range 

261 if whiten_triangles: 

262 if float(da.max().compute()) < levels[-1]: 

263 ax1.cmap.set_over("#FFFFFF") 

264 if float(da.min().compute()) > levels[0]: 

265 ax1.cmap.set_under("#FFFFFF") 

266 

267 # Add colorbar 

268 if plot_colorbar: 

269 divider = make_axes_locatable(ax) 

270 cbar_ax = divider.append_axes("right", size="5%", pad="5%") 

271 fig.colorbar(ax1, cax=cbar_ax, **settings_cbar) 

272 

273 if not return_cmap_norm: 

274 return fig, ax 

275 else: 

276 return fig, ax, cmap, norm 

277 

278 

279def streamfunction(da, ax, n_streamlines=10, kwargs_contour={}): 

280 """ 

281 Wraps matplotlib.contour to draw stream lines. Function can be used to draw 

282 stream lines on top of a cross-section. 

283 

284 Parameters 

285 ---------- 

286 da : xr.DataArray 

287 Two dimensional DataArray containing data of the cross section. One 

288 dimension must be "layer", and the second dimension will be used as the 

289 x-axis for the cross-section. 

290 

291 Coordinates "top" and "bottom" must be present, and must have at least the 

292 "layer" dimension (voxels) or both the "layer" and x-coordinate dimension. 

293 

294 *Use imod.evaluate.streamfunction_line() or streamfunction_linestring() to obtain 

295 the required DataArray.* 

296 ax : matplotlib Axes instance 

297 Axes to write plot to. 

298 n_streamlines : int or array_like 

299 Determines the number and positions of the contour lines / regions. 

300 

301 If an int n, use n data intervals; i.e. draw n+1 contour lines. The level heights are automatically chosen. 

302 

303 If array-like, draw contour lines at the specified levels. The values must be in increasing order. 

304 kwargs_contour : dict 

305 Other optional keyword arguments for matplotlib.contour. 

306 

307 Returns 

308 ------- 

309 matplotlib.contour.QuadContourSet 

310 The drawn contour lines. 

311 """ 

312 

313 da = da.copy(deep=False) 

314 if len(da.dims) != 2: 

315 raise ValueError("DataArray must be 2D") 

316 if "layer" not in da.dims: 

317 raise ValueError('DataArray must contain dimension "layer"') 

318 if "top" not in da.coords: 

319 raise ValueError('DataArray must contain coordinate "top"') 

320 if "bottom" not in da.coords: 

321 raise ValueError('DataArray must contain coordinate "bottom"') 

322 if len(da["top"].dims) > 2: 

323 raise ValueError('"top" coordinate be 1D or 2D') 

324 if len(da["bottom"].dims) > 2: 

325 raise ValueError('"bottom" coordinate be 1D or 2D') 

326 

327 # add a row of zeros at the bottom, with thickness zero 

328 # moved here i/o imod.evaluate to not bother user with 

329 # additional layer hampering coord assignment 

330 zeros = xr.zeros_like(da.isel(layer=-1)) 

331 zeros.coords["layer"] += 1 

332 zeros.coords["top"] = zeros.coords["bottom"] # layer of zero thickness 

333 zeros.coords["bottom"] = zeros.coords["bottom"] - 0.1 

334 da = xr.concat((da, zeros), dim="layer") 

335 

336 # _meshcoords returns mesh edges, but useful here to allow 1D/2D 

337 # dimensions in da 

338 # go back to cell centers in 2D 

339 X = np.broadcast_to(da["s"].values, da.shape) 

340 Y = (0.5 * da["top"] + 0.5 * da["bottom"]).values 

341 if Y.ndim == 1: 

342 Y = np.broadcast_to(Y[:, np.newaxis], da.shape) 

343 

344 settings_contour = { 

345 "colors": "w", 

346 "linestyles": "solid", 

347 "linewidths": 0.8, 

348 "zorder": 10, 

349 } 

350 settings_contour.update(kwargs_contour) 

351 

352 CS = ax.contour(X, Y, da.values, n_streamlines, **settings_contour) 

353 return CS 

354 

355 

356def quiver(u, v, ax, kwargs_quiver={}): 

357 """ 

358 Wraps matplotlib.quiver to draw quiver plots. Function can be used to draw 

359 flow quivers on top of a cross-section. 

360 

361 Parameters 

362 ---------- 

363 u : xr.DataArray 

364 Two dimensional DataArray containing u component of quivers. One 

365 dimension must be "layer", and the second dimension will be used as the 

366 x-axis for the cross-section. 

367 

368 Coordinates "top" and "bottom" must be present, and must have at least the 

369 "layer" dimension (voxels) or both the "layer" and x-coordinate dimension. 

370 

371 *Use imod.evaluate.quiver_line() or quiver_linestring() to obtain 

372 the required DataArray.* 

373 v : xr.DataArray 

374 Two dimensional DataArray containing v component of quivers. One 

375 dimension must be "layer", and the second dimension will be used as the 

376 x-axis for the cross-section. 

377 

378 Coordinates "top" and "bottom" must be present, and must have at least the 

379 "layer" dimension (voxels) or both the "layer" and x-coordinate dimension. 

380 

381 *Use imod.evaluate.quiver_line() or quiver_linestring() to obtain 

382 the required DataArray.* 

383 ax : matplotlib Axes instance 

384 Axes to write plot to. 

385 kwargs_quiver : dict 

386 Other optional keyword arguments for matplotlib.quiver. 

387 

388 Returns 

389 ------- 

390 matplotlib.quiver.Quiver 

391 The drawn quivers. 

392 

393 Examples 

394 -------- 

395 

396 First: apply evaluate.quiver_line to get the u and v components of the quivers 

397 from a three dimensional flow field. Assign top and bottom coordinates if these are 

398 not already present in the flow field data arrays. 

399 

400 >>> u, v = imod.evaluate.quiver_line(right, front, lower, start, end) 

401 >>> u.assign_coords(top=top, bottom=bottom) 

402 >>> v.assign_coords(top=top, bottom=bottom) 

403 

404 The quivers can then be plotted over a cross section created by imod.visualize.cross_section(): 

405 

406 >>> imod.visualize.quiver(u, v, ax) 

407 

408 Quivers can easily overwhelm your plot, so it is a good idea to 'thin out' some of the quivers: 

409 

410 >>> # Only plot quivers at every 5th cell and every 3rd layer 

411 >>> thin = {"s": slice(0, None, 5), "layer": slice(0, None, 3)} 

412 >>> imod.visualize.quiver(u.isel(**thin), v.isel(**thin), ax) 

413 """ 

414 

415 for da, name in zip([u, v], ["u", "v"]): 

416 if len(da.dims) != 2: 

417 raise ValueError(f"DataArray {name} must be 2D") 

418 if "layer" not in da.dims: 

419 raise ValueError(f'DataArray {name} must contain dimension "layer"') 

420 if "top" not in da.coords: 

421 raise ValueError(f'DataArray {name} must contain coordinate "top"') 

422 if "bottom" not in da.coords: 

423 raise ValueError(f'DataArray {name} must contain coordinate "bottom"') 

424 if len(da["top"].dims) > 2: 

425 raise ValueError(f'"top" coordinate in dataarray {name} must be 1D or 2D') 

426 if len(da["bottom"].dims) > 2: 

427 raise ValueError( 

428 f'"bottom" coordinate in dataarray {name} must be 1D or 2D' 

429 ) 

430 

431 # do not draw quivers for cells that are only thinly sliced (ds > 25% of cellsize) 

432 dsmin = 0.25 * u.dx 

433 u = u.where(u.ds > dsmin) 

434 v = v.where(v.ds > dsmin) 

435 

436 # _meshcoords returns mesh edges, but useful here to allow 1D/2D 

437 # dimensions in da 

438 # go back to cell centers in 2D 

439 X, _ = np.meshgrid(u.s.values, u.layer.values) 

440 Y = 0.5 * u.top + 0.5 * u.bottom 

441 if Y.ndim == 1: 

442 # promote to 2D 

443 Y = xr.concat(u.shape[1] * [Y], dim=u.s) 

444 

445 settings_quiver = { 

446 "color": "w", 

447 "scale": None, # autoscaling 

448 "linestyle": "-", 

449 "zorder": 10, 

450 } 

451 settings_quiver.update(kwargs_quiver) 

452 

453 Q = ax.quiver(X, Y, u.values, v.values, **settings_quiver) 

454 return Q