Coverage for C:\src\imod-python\imod\visualize\cross_sections.py: 49%
170 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +0200
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-08 14:15 +0200
1import matplotlib.pyplot as plt
2import numpy as np
3import xarray as xr
4from mpl_toolkits.axes_grid1 import make_axes_locatable
6import imod
7from imod.visualize import common
10def _meshcoords(da, continuous=True):
11 """
12 Generate coordinates for pcolormesh, or fill_between
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.
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])
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()
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))
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
71 Y = np.repeat(Y, 2, 1)
72 nodata = np.isnan(Y)
73 Y[nodata] = 0.0
75 X = np.hstack([xmin, np.repeat(X, 2), xmax])
76 X = np.full_like(Y, X)
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
87 return X, Y, C, nodata
90def _plot_aquitards(aquitards, ax, kwargs_aquitards):
91 """
92 Overlay aquitards on ax
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)
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`.
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.
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.
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").
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
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
183 Examples
184 --------
186 Basic cross section:
188 >>> imod.visualize.cross_section(da, colors, levels)
190 Aquitards can be styled in multiple ways. For a transparent grey overlay
191 (the default):
193 >>> kwargs_aquitards = {"alpha": 0.5, "facecolor": "grey"}
194 >>> imod.visualize.cross_section(da, colors, levels, aquitards=aquitards, kwargs_aquitards)
196 For a hatched overlay:
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)
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')
218 # Read legend settings
219 cmap, norm = common._cmapnorm_from_colorslevels(colors, levels)
221 # cbar kwargs
222 settings_cbar = {"ticks": levels, "extend": "both"}
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
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)
240 # pcmesh kwargs
241 settings_pcmesh = {"cmap": cmap, "norm": norm}
242 if kwargs_pcolormesh is not None:
243 settings_pcmesh.update(kwargs_pcolormesh)
245 if fig is None or ax is None:
246 fig, ax = plt.subplots()
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)
255 if layers:
256 Y[nodata] = np.nan
257 for y in Y:
258 ax.step(x=X[0], y=y)
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")
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)
273 if not return_cmap_norm:
274 return fig, ax
275 else:
276 return fig, ax, cmap, norm
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.
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.
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.
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.
301 If an int n, use n data intervals; i.e. draw n+1 contour lines. The level heights are automatically chosen.
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.
307 Returns
308 -------
309 matplotlib.contour.QuadContourSet
310 The drawn contour lines.
311 """
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')
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")
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)
344 settings_contour = {
345 "colors": "w",
346 "linestyles": "solid",
347 "linewidths": 0.8,
348 "zorder": 10,
349 }
350 settings_contour.update(kwargs_contour)
352 CS = ax.contour(X, Y, da.values, n_streamlines, **settings_contour)
353 return CS
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.
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.
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.
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.
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.
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.
388 Returns
389 -------
390 matplotlib.quiver.Quiver
391 The drawn quivers.
393 Examples
394 --------
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.
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)
404 The quivers can then be plotted over a cross section created by imod.visualize.cross_section():
406 >>> imod.visualize.quiver(u, v, ax)
408 Quivers can easily overwhelm your plot, so it is a good idea to 'thin out' some of the quivers:
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 """
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 )
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)
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)
445 settings_quiver = {
446 "color": "w",
447 "scale": None, # autoscaling
448 "linestyle": "-",
449 "zorder": 10,
450 }
451 settings_quiver.update(kwargs_quiver)
453 Q = ax.quiver(X, Y, u.values, v.values, **settings_quiver)
454 return Q