#! /usr/bin/env python
"""Methods to plot data defined on Landlab grids.
Plotting functions
++++++++++++++++++
.. autosummary::
~imshow_grid
~imshow_grid_at_cell
~imshow_grid_at_node
"""
from warnings import warn
import numpy as np
from ..field import FieldError
from .event_handler import query_grid_on_button_press
try:
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.collections import PatchCollection
from matplotlib.patches import Polygon
except ImportError:
import warnings
warnings.warn("matplotlib not found", ImportWarning, stacklevel=2)
[docs]
class ModelGridPlotterMixIn:
"""MixIn that provides plotting functionality.
Inhert from this class to provide a ModelDataFields object with the
method function, ``imshow``, that plots a data field.
"""
[docs]
def imshow(self, *args, **kwds):
"""Plot a data field.
This is a wrapper for `plot.imshow_grid`, and can take the same
keywords. See that function for full documentation.
Parameters
----------
values : str, or array-like
Name of a field or an array of values to plot.
See Also
--------
landlab.plot.imshow_grid
LLCATS: GINF
"""
if len(args) == 1:
values = args[0]
elif len(args) == 2:
at, values = args
warn(f"use grid.imshow(values, at={at!r})", DeprecationWarning)
if at != kwds.get("at", at):
raise ValueError(f"multiple locations provided ({at}, {kwds['at']})")
kwds["at"] = at
else:
raise TypeError(f"imshow expected 1 or 2 arguments, got {len(args)}")
imshow_grid(self, values, **kwds)
[docs]
def imshow_grid_at_node(grid, values, **kwds):
"""Prepare a map view of data over all nodes in the grid.
Data is plotted as cells shaded with the value at the node at its center.
Outer edges of perimeter cells are extrapolated. Closed elements are
colored uniformly (default black, overridden with kwd 'color_for_closed');
other open boundary nodes get their actual values.
*values* can be a field name, a regular array, or a masked array. If a
masked array is provided, masked entries will be treated as if they were
Landlab BC_NODE_IS_CLOSED. Used together with the color_at_closed=None
keyword (i.e., "transparent"), this can allow for construction of overlay
layers in a figure (e.g., only defining values in a river network, and
overlaying it on another landscape).
Use matplotlib functions like xlim, ylim to modify your plot after calling
:func:`imshow_grid`, as desired.
Node coordinates are printed when a mouse button is pressed on a cell in
the plot.
This function happily works with both regular and irregular grids.
Parameters
----------
grid : ModelGrid
Grid containing the field to plot, or describing the geometry of the
provided array.
values : array_like, masked_array, or str
Node values, or a field name as a string from which to draw the data.
plot_name : str, optional
String to put as the plot title.
var_name : str, optional
Variable name, to use as a colorbar label.
var_units : str, optional
Units for the variable being plotted, for the colorbar.
grid_units : tuple of str, optional
Units for y, and x dimensions. If None, component will look to the
grid property `axis_units` for this information. If no units are
specified there, no entry is made.
symmetric_cbar : bool
Make the colormap symetric about 0.
cmap : str
Name of a colormap
limits : tuple of float
Minimum and maximum of the colorbar.
vmin, vmax: floats
Alternatives to limits.
allow_colorbar : bool
If True, include the colorbar.
colorbar_label : str or None
The string with which to label the colorbar.
norm : matplotlib.colors.Normalize
The normalizing object which scales data, typically into the interval
[0, 1]. Ignore in most cases.
shrink : float
Fraction by which to shrink the colorbar.
color_for_closed : str or None
Color to use for closed nodes (default 'black'). If None, closed
(or masked) nodes will be transparent.
color_for_background : color str or other color declaration, or None
Color to use for closed elements (default None). If None, the
background will be transparent, and appear white.
show_elements : bool
If True, and grid is a Voronoi, the faces will be plotted in black
along with just the colour of the cell, defining the cell outlines
(defaults False).
output : None, string, or bool
If None (or False), the image is sent to the imaging buffer to await
an explicit call to show() or savefig() from outside this function.
If a string, the string should be the path to a save location, and the
filename (with file extension). The function will then call
plt.savefig([string]) itself. If True, the function will call
plt.show() itself once plotting is complete.
"""
if isinstance(values, str):
values_at_node = grid.at_node[values]
else:
values_at_node = values.reshape((-1,))
if values_at_node.size != grid.number_of_nodes:
raise ValueError("number of values does not match number of nodes")
values_at_node = np.ma.masked_where(
grid.status_at_node == grid.BC_NODE_IS_CLOSED, values_at_node
)
_imshow_grid_values(grid, values_at_node, **kwds)
if isinstance(values, str):
plt.title(values)
plt.gcf().canvas.mpl_connect(
"button_press_event", lambda event: query_grid_on_button_press(event, grid)
)
[docs]
def imshow_grid_at_cell(grid, values, **kwds):
"""Map view of grid data over all grid cells.
Prepares a map view of data over all cells in the grid.
Method can take any of the same ``**kwds`` as :func:`imshow_grid_at_node`.
Parameters
----------
grid : ModelGrid
Grid containing the field to plot, or describing the geometry of the
provided array.
values : array_like, masked_array, or str
Values at the cells on the grid. Alternatively, can be a field name
(string) from which to draw the data from the grid.
plot_name : str, optional
String to put as the plot title.
var_name : str, optional
Variable name, to use as a colorbar label.
var_units : str, optional
Units for the variable being plotted, for the colorbar.
grid_units : tuple of str, optional
Units for y, and x dimensions. If None, component will look to the
gri property `axis_units` for this information. If no units are
specified there, no entry is made.
symmetric_cbar : bool
Make the colormap symetric about 0.
cmap : str
Name of a colormap
limits : tuple of float
Minimum and maximum of the colorbar.
vmin, vmax: floats
Alternatives to limits.
allow_colorbar : bool
If True, include the colorbar.
colorbar_label : str or None
The string with which to label the colorbar.
norm : matplotlib.colors.Normalize
The normalizing object which scales data, typically into the interval
[0, 1]. Ignore in most cases.
shrink : float
Fraction by which to shrink the colorbar.
color_for_closed : str or None
Color to use for closed elements (default 'black'). If None, closed
(or masked) elements will be transparent.
color_for_background : color str or other color declaration, or None
Color to use for closed elements (default None). If None, the
background will be transparent, and appear white.
show_elements : bool
If True, and grid is a Voronoi, the faces will be plotted in black
along with just the colour of the cell, defining the cell outlines
(defaults False).
output : None, string, or bool
If None (or False), the image is sent to the imaging buffer to await
an explicit call to show() or savefig() from outside this function.
If a string, the string should be the path to a save location, and the
filename (with file extension). The function will then call
plt.savefig([string]) itself. If True, the function will call
plt.show() itself once plotting is complete.
Raises
------
ValueError
If input grid is not uniform rectilinear.
"""
kwds.setdefault("color_for_closed", None)
if isinstance(values, str):
try:
values_at_cell = grid.at_cell[values]
except FieldError:
values_at_cell = grid.at_node[values]
else:
values_at_cell = values
if values_at_cell.size == grid.number_of_nodes:
values_at_cell = values_at_cell[grid.node_at_cell]
if values_at_cell.size != grid.number_of_cells:
raise ValueError(
"number of values must match number of cells or " "number of nodes"
)
values_at_node = np.ma.masked_array(grid.empty(at="node"))
values_at_node.mask = True
values_at_node[grid.node_at_cell] = values_at_cell
values_at_node.mask[grid.node_at_cell] = False
myimage = _imshow_grid_values(grid, values_at_node, **kwds)
if isinstance(values, str):
plt.title(values)
return myimage
def _imshow_grid_values(
grid,
values,
plot_name=None,
var_name=None,
var_units=None,
grid_units=(None, None),
symmetric_cbar=False,
cmap="pink",
limits=None,
colorbar_label=None,
allow_colorbar=True,
vmin=None,
vmax=None,
norm=None,
shrink=1.0,
color_for_closed="black",
color_for_background=None,
show_elements=False,
output=None,
alpha=1.0,
):
from ..grid.raster import RasterModelGrid
if isinstance(cmap, str):
cmap = plt.colormaps[cmap]
if color_for_closed is not None:
cmap.set_bad(color=color_for_closed)
else:
cmap.set_bad(alpha=0.0)
if isinstance(grid, RasterModelGrid):
values = values.reshape(grid.shape)
if values.ndim != 2:
raise ValueError("values must have ndim == 2")
y = (
np.arange(values.shape[0] + 1) * grid.dy
- grid.dy * 0.5
+ grid.xy_of_lower_left[1]
)
x = (
np.arange(values.shape[1] + 1) * grid.dx
- grid.dx * 0.5
+ grid.xy_of_lower_left[0]
)
kwds = {"cmap": cmap}
(kwds["vmin"], kwds["vmax"]) = (values.min(), values.max())
if (limits is None) and ((vmin is None) and (vmax is None)):
if symmetric_cbar:
(var_min, var_max) = (values.min(), values.max())
limit = max(abs(var_min), abs(var_max))
(kwds["vmin"], kwds["vmax"]) = (-limit, limit)
elif limits is not None:
(kwds["vmin"], kwds["vmax"]) = (limits[0], limits[1])
else:
if vmin is not None:
kwds["vmin"] = vmin
if vmax is not None:
kwds["vmax"] = vmax
kwds["alpha"] = alpha
myimage = plt.pcolormesh(x, y, values, **kwds)
myimage.set_rasterized(True)
myimage.axes.set_aspect("equal")
plt.autoscale(tight=True)
if allow_colorbar:
cb = plt.colorbar(norm=norm, shrink=shrink)
if colorbar_label:
cb.set_label(colorbar_label)
else:
import matplotlib.cm as cmx
import matplotlib.colors as colors
values = values.reshape(-1)
if limits is not None:
(vmin, vmax) = (limits[0], limits[1])
else:
if vmin is None:
vmin = values.min()
if vmax is None:
vmax = values.max()
if symmetric_cbar:
vmin, vmax = -max(abs(vmin), abs(vmax)), max(abs(vmin), abs(vmax))
cNorm = colors.Normalize(vmin, vmax)
scalarMap = cmx.ScalarMappable(norm=cNorm, cmap=cmap)
colorVal = scalarMap.to_rgba(values, alpha=alpha)[grid.node_at_cell]
patches = []
for corners in grid.corners_at_cell:
valid_corners = corners[corners != grid.BAD_INDEX]
closed_loop_corners = np.concatenate([valid_corners, [valid_corners[0]]])
x = grid.x_of_corner[closed_loop_corners]
y = grid.y_of_corner[closed_loop_corners]
xy = np.vstack((x, y)).T
patches.append(Polygon(xy, closed=True, fill=True))
patchcollection = PatchCollection(
patches, facecolor=colorVal, edgecolor=colorVal
)
ax = plt.gca()
ax.add_collection(patchcollection)
if show_elements:
x = grid.x_of_corner[grid.corners_at_face]
y = grid.y_of_corner[grid.corners_at_face]
segs = np.dstack((x, y))
line_segments = LineCollection(segs)
line_segments.set_color("black")
ax.add_collection(line_segments)
ax.set_aspect("equal")
ax.set_rasterized(True)
plt.xlim((np.min(grid.x_of_node), np.max(grid.x_of_node)))
plt.ylim((np.min(grid.y_of_node), np.max(grid.y_of_node)))
scalarMap.set_array(values)
if allow_colorbar:
cb = plt.colorbar(scalarMap, shrink=shrink, ax=ax)
if colorbar_label:
cb.set_label(colorbar_label)
if grid_units[1] is None and grid_units[0] is None:
grid_units = grid.axis_units
if grid_units[1] == "-" and grid_units[0] == "-":
plt.xlabel("X")
plt.ylabel("Y")
else:
plt.xlabel("X (%s)" % grid_units[1])
plt.ylabel("Y (%s)" % grid_units[0])
else:
plt.xlabel("X (%s)" % grid_units[1])
plt.ylabel("Y (%s)" % grid_units[0])
if plot_name is not None:
plt.title("%s" % (plot_name))
if var_name is not None or var_units is not None:
if var_name is not None:
assert type(var_name) is str
if var_units is not None:
assert type(var_units) is str
colorbar_label = var_name + " (" + var_units + ")"
else:
colorbar_label = var_name
else:
assert type(var_units) is str
colorbar_label = "(" + var_units + ")"
assert type(colorbar_label) is str
assert allow_colorbar
cb.set_label(colorbar_label)
if color_for_background is not None:
plt.gca().set_facecolor(color_for_background)
if output is not None:
if type(output) is str:
plt.savefig(output)
plt.clf()
elif output:
plt.show()
[docs]
def imshow_grid(grid, values, **kwds):
"""Prepare a map view of data over all nodes or cells in the grid.
Data is plotted as colored cells. If at='node', the surrounding cell is
shaded with the value at the node at its center. If at='cell', the cell
is shaded with its own value. Outer edges of perimeter cells are
extrapolated. Closed elements are colored uniformly (default black,
overridden with kwd 'color_for_closed'); other open boundary nodes get
their actual values.
*values* can be a field name, a regular array, or a masked array. If a
masked array is provided, masked entries will be treated as if they were
Landlab BC_NODE_IS_CLOSED. Used together with the color_for_closed=None
keyword (i.e., "transparent"), this can allow for construction of overlay
layers in a figure (e.g., only defining values in a river network, and
overlaying it on another landscape).
Use matplotlib functions like xlim, ylim to modify your plot after calling
:func:`imshow_grid`, as desired.
This function happily works with both regular and irregular grids.
Parameters
----------
grid : ModelGrid
Grid containing the field to plot, or describing the geometry of the
provided array.
values : array_like, masked_array, or str
Node or cell values, or a field name as a string from which to draw
the data.
at : str, {'node', 'cell'}
Tells plotter where values are defined.
plot_name : str, optional
String to put as the plot title.
var_name : str, optional
Variable name, to use as a colorbar label.
var_units : str, optional
Units for the variable being plotted, for the colorbar.
grid_units : tuple of str, optional
Units for y, and x dimensions. If None, component will look to the
gri property `axis_units` for this information. If no units are
specified there, no entry is made.
symmetric_cbar : bool
Make the colormap symetric about 0.
cmap : str
Name of a colormap
alpha : array-like or scalar or None, optional
Set the transparency.
limits : tuple of float
Minimum and maximum of the colorbar.
vmin, vmax: floats
Alternatives to limits.
allow_colorbar : bool
If True, include the colorbar.
colorbar_label : str or None
The string with which to label the colorbar.
norm : matplotlib.colors.Normalize
The normalizing object which scales data, typically into the interval
[0, 1]. Ignore in most cases.
shrink : float
Fraction by which to shrink the colorbar.
color_for_closed : str or None
Color to use for closed elements (default 'black'). If None, closed
(or masked) elements will be transparent.
color_for_background : color str or other color declaration, or None
Color to use for closed elements (default None). If None, the
background will be transparent, and appear white.
show_elements : bool
If True, and grid is a Voronoi, the faces will be plotted in black
along with just the colour of the cell, defining the cell outlines
(defaults False).
output : None, string, or bool
If None (or False), the image is sent to the imaging buffer to await
an explicit call to show() or savefig() from outside this function.
If a string, the string should be the path to a save location, and the
filename (with file extension). The function will then call
plt.savefig([string]) itself. If True, the function will call
plt.show() itself once plotting is complete.
"""
if "values_at" in kwds:
warn(
"the 'values_at' keyword is deprecated, use the 'at' keyword instead",
DeprecationWarning,
)
kwds.setdefault("at", kwds.pop("values_at"))
values_at = kwds.pop("at", None)
if values_at is None:
values_at = _guess_location(grid, values)
if values_at is None:
raise TypeError("unable to determine location of values, use 'at' keyword")
elif values_at not in {"node", "cell"}:
raise TypeError(
f"value location, {values_at!r}, is not supported (must be one of 'node', 'cell')"
)
if isinstance(values, str):
values = grid.field_values(values, at=values_at)
if values_at == "node":
imshow_grid_at_node(grid, values, **kwds)
elif values_at == "cell":
imshow_grid_at_cell(grid, values, **kwds)
def _guess_location(
grid, values, search_order=("node", "cell", "link", "patch", "corner", "face")
):
"""Make an educated guess as to where a field is located on a grid."""
if isinstance(values, str):
return _guess_location_from_name(grid, values)
else:
return _guess_location_from_size(grid, values)
def _guess_location_from_name(
grid, name, search_order=("node", "cell", "link", "patch", "corner", "face")
):
"""Given a name, make an educated guess as to where a field is located on a grid.
Parameters
----------
grid : ModelGrid
A landlab ModelGrid.
name : str
Name of a field.
Returns
-------
str or None
Grid element where the field is likely defined.
"""
for location in search_order:
if grid.has_field(name, at=location):
return location
return None
def _guess_location_from_size(
grid, values, search_order=("node", "cell", "link", "patch", "corner", "face")
):
"""Given an array, make an educated guess as to where a field is located on a grid.
Parameters
----------
grid : ModelGrid
A landlab ModelGrid.
values : array-like
An array of values.
Returns
-------
str or None
Grid element where the field is likely defined.
"""
for location in search_order:
if values.size == grid.number_of_elements(location):
return location
return None