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

1 import matplotlib.pyplot as plt

2 import numpy as np

3 import xarray as xr

4 from mpl_toolkits.axes_grid1 import make_axes_locatable

5

6 import imod

7 from imod.visualize import common

8

9

10 def _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

90 def _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

116 def 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

279 def 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

356 def 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