import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
from mpl_toolkits.axes_grid1 import make_axes_locatable
import imod
from imod.visualize import common
def _meshcoords(da, continuous=True):
"""
Generate coordinates for pcolormesh, or fill_between
Parameters
----------
da : xr.DataArray
The array to plot
continuous : bool, optional
Whether the layers are connected, such that the bottom of layer N is
the top of layer N + 1.
Returns
-------
X : np.array
x coordinates of mesh
Y : np.array
y coordinates of mesh
C : np.array
values of the mesh
"""
dims = tuple(da.dims)
if not dims[0] == "layer":
# Switch 'm around
dims = (dims[1], dims[0])
# Ensure dimensions are in the right order
da = da.transpose(*dims, transpose_coords=True)
data = da.values
xcoord = da.dims[1]
dx, xmin, xmax = imod.util.spatial.coord_reference(da[xcoord])
if isinstance(dx, (int, float)):
if dx < 0.0:
dx = abs(dx)
data = data[:, ::-1]
X = np.arange(xmin + dx, xmax, dx)
else: # assuming dx is an array, non-equidistant dx
if (dx < 0.0).all():
dx = abs(dx)
data = data[:, ::-1]
elif (dx > 0.0).all():
pass
else:
raise ValueError(f"{xcoord} is not monotonic")
X = (xmin + dx[0]) + dx[1:].cumsum()
# if dimensions of top and bottom are 1D (voxels), promote to 2D
if len(da["top"].dims) == 1 and len(da["bottom"].dims) == 1:
top = np.repeat(np.atleast_2d(da.top), len(da[xcoord]), axis=0).T
da = da.assign_coords(top=(("layer", xcoord), top))
bot = np.repeat(np.atleast_2d(da.bottom), len(da[xcoord]), axis=0).T
da = da.assign_coords(bottom=(("layer", xcoord), bot))
if continuous:
Y = np.vstack([da["top"].isel(layer=0).values, da["bottom"].values])
else:
nrow, ncol = da["top"].shape
Y = np.empty((nrow * 2, ncol))
Y[0::2] = da["top"].values
Y[1::2] = da["bottom"].values
Y = np.repeat(Y, 2, 1)
nodata = np.isnan(Y)
Y[nodata] = 0.0
X = np.hstack([xmin, np.repeat(X, 2), xmax])
X = np.full_like(Y, X)
if continuous:
nrow, ncol = Y.shape
C = np.full((nrow - 1, ncol - 1), np.nan)
C[:, 0::2] = data
else:
_, ncol = Y.shape
C = np.full((nrow, ncol - 1), np.nan)
C[:, 0::2] = data
return X, Y, C, nodata
def _plot_aquitards(aquitards, ax, kwargs_aquitards):
"""
Overlay aquitards on ax
Parameters
----------
aquitards : xr.DataArray
DataArray containing location of aquitards
NaN's and zeros are treated as locations without aquitard
ax : matplotlib.Axes object
kwargs_aquitards : dict
keyword arguments for ax.fill_between()
"""
if kwargs_aquitards is None:
kwargs_aquitards = {"alpha": 0.5, "facecolor": "grey"}
X_aq, Y_aq, C_aq, _ = _meshcoords(aquitards, continuous=False)
C_aq = C_aq.astype(np.float64)
for j, i in enumerate(range(0, X_aq.shape[0] - 1, 2)):
Y_i = Y_aq[i : i + 2]
C_i = C_aq[j]
C_i[C_i == 0.0] = np.nan
nodata = np.repeat(np.isnan(C_i[0::2]), 2)
Y_i[:, nodata] = np.nan
ax.fill_between(X_aq[0], Y_i[0], Y_i[1], **kwargs_aquitards)
[docs]
def cross_section(
da,
colors,
levels,
layers=False,
aquitards=None,
kwargs_pcolormesh={},
kwargs_colorbar={},
kwargs_aquitards=None,
return_cmap_norm=False,
fig=None,
ax=None,
):
"""
Wraps matplotlib.pcolormesh to draw cross-sections, drawing cell boundaries
accurately. Aquitards can be plotted on top of the cross-section, by providing
a DataArray with the aquitard location for `aquitards`.
Parameters
----------
da : xr.DataArray
Two dimensional DataArray containing data of the cross section. One
dimension must be "layer", and the second dimension will be used as the
x-axis for the cross-section.
Coordinates "top" and "bottom" must be present, and must have at least the
"layer" dimension (voxels) or both the "layer" and x-coordinate dimension.
*Use imod.select.cross_section_line() or cross_section_linestring() to obtain
the required DataArray.*
colors : list of str, or list of RGB tuples
Matplotlib acceptable list of colors. Length N.
Accepts both tuples of (R, G, B) and hexidecimal (e.g. "#7ec0ee").
Looking for good colormaps? Try: http://colorbrewer2.org/
Choose a colormap, and use the HEX JS array.
levels : listlike of floats or integers
Boundaries between the legend colors/classes. Length: N - 1.
layers : boolean, optional
Whether to draw lines separating the layers.
aquitards : xr.DataArray, optional
Datarray containing data on location of aquitard layers.
kwargs_pcolormesh : dict
Other optional keyword arguments for matplotlib.pcolormesh.
kwargs_colorbar : dict
If optional key ``plot_colorbar`` is set to False, no colorbar is drawn. Defaults to True.
Optional keyword argument ``whiten_triangles`` whitens respective colorbar triangle if
data is not larger/smaller than legend_levels-range. Defaults to True.
Other arguments are forwarded to fig.colorbar()
kwargs_aquitards: dict
These arguments are forwarded to matplotlib.fill_between to draw the
aquitards.
return_cmap_norm : boolean, optional
Return the cmap and norm of the plot, default False
fig : matplotlib Figure instance, optional
Figure to write plot to. If not supplied, a Figure instance is created
ax : matplotlib Axes instance, optional
Axes to write plot to. If not supplied, an Axes instance is created
Returns
-------
fig : matplotlib.figure
ax : matplotlig.ax
if return_cmap_norm == True:
cmap : matplotlib.colors.ListedColormap
norm : matplotlib.colors.BoundaryNorm
Examples
--------
Basic cross section:
>>> imod.visualize.cross_section(da, colors, levels)
Aquitards can be styled in multiple ways. For a transparent grey overlay
(the default):
>>> kwargs_aquitards = {"alpha": 0.5, "facecolor": "grey"}
>>> imod.visualize.cross_section(da, colors, levels, aquitards=aquitards, kwargs_aquitards)
For a hatched overlay:
>>> kwargs_aquitards = {"hatch": "/", "edgecolor": "k"}
>>> imod.visualize.cross_section(da, colors, levels, aquitards=aquitards, kwargs_aquitards)
"""
da = da.copy(deep=False)
if aquitards is not None:
aquitards = aquitards.copy(deep=False)
if len(da.dims) != 2:
raise ValueError("DataArray must be 2D")
if "layer" not in da.dims:
raise ValueError('DataArray must contain dimension "layer"')
if "top" not in da.coords:
raise ValueError('DataArray must contain coordinate "top"')
if "bottom" not in da.coords:
raise ValueError('DataArray must contain coordinate "bottom"')
if len(da["top"].dims) > 2:
raise ValueError('"top" coordinate be 1D or 2D')
if len(da["bottom"].dims) > 2:
raise ValueError('"bottom" coordinate be 1D or 2D')
# Read legend settings
cmap, norm = common._cmapnorm_from_colorslevels(colors, levels)
# cbar kwargs
settings_cbar = {"ticks": levels, "extend": "both"}
# Find a unit in the raster to use in the colorbar label
try:
settings_cbar["label"] = da.attrs["units"]
except (KeyError, AttributeError):
try:
settings_cbar["label"] = da.attrs["unit"]
except (KeyError, AttributeError):
pass
whiten_triangles = True
plot_colorbar = True
if kwargs_colorbar is not None:
whiten_triangles = kwargs_colorbar.pop("whiten_triangles", True)
plot_colorbar = kwargs_colorbar.pop("plot_colorbar", True)
settings_cbar.update(kwargs_colorbar)
# pcmesh kwargs
settings_pcmesh = {"cmap": cmap, "norm": norm}
if kwargs_pcolormesh is not None:
settings_pcmesh.update(kwargs_pcolormesh)
if fig is None or ax is None:
fig, ax = plt.subplots()
# Plot raster
X, Y, C, nodata = _meshcoords(da, continuous=True)
ax1 = ax.pcolormesh(X, Y, C, **settings_pcmesh)
# Plot aquitards if applicable
if aquitards is not None:
_plot_aquitards(aquitards, ax, kwargs_aquitards)
if layers:
Y[nodata] = np.nan
for y in Y:
ax.step(x=X[0], y=y)
# Make triangles white if data is not larger/smaller than legend_levels-range
if whiten_triangles:
if float(da.max().compute()) < levels[-1]:
ax1.cmap.set_over("#FFFFFF")
if float(da.min().compute()) > levels[0]:
ax1.cmap.set_under("#FFFFFF")
# Add colorbar
if plot_colorbar:
divider = make_axes_locatable(ax)
cbar_ax = divider.append_axes("right", size="5%", pad="5%")
fig.colorbar(ax1, cax=cbar_ax, **settings_cbar)
if not return_cmap_norm:
return fig, ax
else:
return fig, ax, cmap, norm
[docs]
def streamfunction(da, ax, n_streamlines=10, kwargs_contour={}):
"""
Wraps matplotlib.contour to draw stream lines. Function can be used to draw
stream lines on top of a cross-section.
Parameters
----------
da : xr.DataArray
Two dimensional DataArray containing data of the cross section. One
dimension must be "layer", and the second dimension will be used as the
x-axis for the cross-section.
Coordinates "top" and "bottom" must be present, and must have at least the
"layer" dimension (voxels) or both the "layer" and x-coordinate dimension.
*Use imod.evaluate.streamfunction_line() or streamfunction_linestring() to obtain
the required DataArray.*
ax : matplotlib Axes instance
Axes to write plot to.
n_streamlines : int or array_like
Determines the number and positions of the contour lines / regions.
If an int n, use n data intervals; i.e. draw n+1 contour lines. The level heights are automatically chosen.
If array-like, draw contour lines at the specified levels. The values must be in increasing order.
kwargs_contour : dict
Other optional keyword arguments for matplotlib.contour.
Returns
-------
matplotlib.contour.QuadContourSet
The drawn contour lines.
"""
da = da.copy(deep=False)
if len(da.dims) != 2:
raise ValueError("DataArray must be 2D")
if "layer" not in da.dims:
raise ValueError('DataArray must contain dimension "layer"')
if "top" not in da.coords:
raise ValueError('DataArray must contain coordinate "top"')
if "bottom" not in da.coords:
raise ValueError('DataArray must contain coordinate "bottom"')
if len(da["top"].dims) > 2:
raise ValueError('"top" coordinate be 1D or 2D')
if len(da["bottom"].dims) > 2:
raise ValueError('"bottom" coordinate be 1D or 2D')
# add a row of zeros at the bottom, with thickness zero
# moved here i/o imod.evaluate to not bother user with
# additional layer hampering coord assignment
zeros = xr.zeros_like(da.isel(layer=-1))
zeros.coords["layer"] += 1
zeros.coords["top"] = zeros.coords["bottom"] # layer of zero thickness
zeros.coords["bottom"] = zeros.coords["bottom"] - 0.1
da = xr.concat((da, zeros), dim="layer")
# _meshcoords returns mesh edges, but useful here to allow 1D/2D
# dimensions in da
# go back to cell centers in 2D
X = np.broadcast_to(da["s"].values, da.shape)
Y = (0.5 * da["top"] + 0.5 * da["bottom"]).values
if Y.ndim == 1:
Y = np.broadcast_to(Y[:, np.newaxis], da.shape)
settings_contour = {
"colors": "w",
"linestyles": "solid",
"linewidths": 0.8,
"zorder": 10,
}
settings_contour.update(kwargs_contour)
CS = ax.contour(X, Y, da.values, n_streamlines, **settings_contour)
return CS
[docs]
def quiver(u, v, ax, kwargs_quiver={}):
"""
Wraps matplotlib.quiver to draw quiver plots. Function can be used to draw
flow quivers on top of a cross-section.
Parameters
----------
u : xr.DataArray
Two dimensional DataArray containing u component of quivers. One
dimension must be "layer", and the second dimension will be used as the
x-axis for the cross-section.
Coordinates "top" and "bottom" must be present, and must have at least the
"layer" dimension (voxels) or both the "layer" and x-coordinate dimension.
*Use imod.evaluate.quiver_line() or quiver_linestring() to obtain
the required DataArray.*
v : xr.DataArray
Two dimensional DataArray containing v component of quivers. One
dimension must be "layer", and the second dimension will be used as the
x-axis for the cross-section.
Coordinates "top" and "bottom" must be present, and must have at least the
"layer" dimension (voxels) or both the "layer" and x-coordinate dimension.
*Use imod.evaluate.quiver_line() or quiver_linestring() to obtain
the required DataArray.*
ax : matplotlib Axes instance
Axes to write plot to.
kwargs_quiver : dict
Other optional keyword arguments for matplotlib.quiver.
Returns
-------
matplotlib.quiver.Quiver
The drawn quivers.
Examples
--------
First: apply evaluate.quiver_line to get the u and v components of the quivers
from a three dimensional flow field. Assign top and bottom coordinates if these are
not already present in the flow field data arrays.
>>> u, v = imod.evaluate.quiver_line(right, front, lower, start, end)
>>> u.assign_coords(top=top, bottom=bottom)
>>> v.assign_coords(top=top, bottom=bottom)
The quivers can then be plotted over a cross section created by imod.visualize.cross_section():
>>> imod.visualize.quiver(u, v, ax)
Quivers can easily overwhelm your plot, so it is a good idea to 'thin out' some of the quivers:
>>> # Only plot quivers at every 5th cell and every 3rd layer
>>> thin = {"s": slice(0, None, 5), "layer": slice(0, None, 3)}
>>> imod.visualize.quiver(u.isel(**thin), v.isel(**thin), ax)
"""
for da, name in zip([u, v], ["u", "v"]):
if len(da.dims) != 2:
raise ValueError(f"DataArray {name} must be 2D")
if "layer" not in da.dims:
raise ValueError(f'DataArray {name} must contain dimension "layer"')
if "top" not in da.coords:
raise ValueError(f'DataArray {name} must contain coordinate "top"')
if "bottom" not in da.coords:
raise ValueError(f'DataArray {name} must contain coordinate "bottom"')
if len(da["top"].dims) > 2:
raise ValueError(f'"top" coordinate in dataarray {name} must be 1D or 2D')
if len(da["bottom"].dims) > 2:
raise ValueError(
f'"bottom" coordinate in dataarray {name} must be 1D or 2D'
)
# do not draw quivers for cells that are only thinly sliced (ds > 25% of cellsize)
dsmin = 0.25 * u.dx
u = u.where(u.ds > dsmin)
v = v.where(v.ds > dsmin)
# _meshcoords returns mesh edges, but useful here to allow 1D/2D
# dimensions in da
# go back to cell centers in 2D
X, _ = np.meshgrid(u.s.values, u.layer.values)
Y = 0.5 * u.top + 0.5 * u.bottom
if Y.ndim == 1:
# promote to 2D
Y = xr.concat(u.shape[1] * [Y], dim=u.s)
settings_quiver = {
"color": "w",
"scale": None, # autoscaling
"linestyle": "-",
"zorder": 10,
}
settings_quiver.update(kwargs_quiver)
Q = ax.quiver(X, Y, u.values, v.values, **settings_quiver)
return Q
