Source code for landlab.plot.network_sediment_transporter.plot_network_and_parcels

"""
This code was designed to plot outputs of the NetworkSedimentTransporter
landlab component.

This code plots:
    - the network, with option to color each link according to a link attribute.
    - the parcels, with option to color and size each parcel according to
    parcel attributes.

Authors: Katy Barnhart, Jon Czuba, Allison Pfeiffer
"""
import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.collections import LineCollection
from matplotlib.colors import Normalize

from landlab.plot.network_sediment_transporter.locate_parcel_xy import locate_parcel_xy
from landlab.utils.return_array import return_array_at_link


[docs]def plot_network_and_parcels( grid, parcels, parcel_time_index=None, map_buffer=0.1, parcel_filter=None, network_color=None, link_attribute=None, link_attribute_title=None, network_cmap="cividis", network_norm=None, network_linewidth=None, parcel_color=None, parcel_color_attribute=None, parcel_color_attribute_title=None, parcel_color_cmap="plasma", parcel_color_norm=None, parcel_size=None, parcel_size_attribute=None, parcel_size_attribute_title=None, parcel_size_norm=None, parcel_size_min=5, parcel_size_max=40, parcel_alpha=0.5, fig=None, **kwargs ): """Plot a river network and parcels on the river network. Intended to display the results of the NetworkSedimentTransporter component. The river network (an instance of NetworkModelGrid) is plotted either as straight links between grid nodes, or (if the network was created using a shapefile to set network topology) as sinuous lines representing the actual link geometry. The parcels (an instance of DataRecord) are represented as dot markers along the links, with the marker location set by parcel attribute `location_at_link`. The default is to plot the parcel locations at the last timestep in DataRecord, though any time index may be specified. Use of this plotting tool is described in detail in a landlab tutorial. Parameters ---------- grid : NetworkModelGrid Instance of NetworkModelGrid. parcels : DataRecord Instance of Landlab DataRecord, with the same attribute requirements as NetworkSedimentTransporter. parcel_time_index : int, time index of parcels DataRecord Parcel time index to plot. Default is last timestep in parcels DataRecord. map_buffer : 0.1 Increase the plot extent by at least this much (default 0.1). Note, b/c of axis equal, may be more. parcel_filter : boolean array of shape (number_of_parcels, ) Filter to plot only a selection of the parcels. ######################## ## Part 1: Network. To set the link colors provide either: network_color="k" Uniform color for network links. # or link_attribute : array or field name at link Value used to set link color. Categorical options not supported. Must be continuous. link_attribute_title : str String to use as the title, if link_attribute is a string, it is used as the default. network_cmap : "cividis" Name of colormap for network. network_norm : matplotlib color normalizer. https://matplotlib.org/3.1.1/tutorials/colors/colormapnorms.html Default is linear between min and max of link_attribute. # linewidth will be recognized by either link coloring option. network_linewidth : float Width of network lines (default 0.5). ######################## ## Part 2: Parcels. To set the parcel color, provide either: parcel_color : color str Constant color used for parcel markers (default "k"). # or parcel_color_attribute : parcel attribute name. Categorical options not supported. Must be continuous. parcel_color_attribute_title : str String to use as the legend title. If parcel_color_attribute is a string, it is used as the default. parcel_color_cmap : cmap str Name of colormap for variable parcel color (default "plasma"). parcel_color_norm : matplotlib color normalizer. https://matplotlib.org/3.1.1/tutorials/colors/colormapnorms.html Default is linear between min and max of parcel_color_attribute. # for parcel size use either: parcel_size : float Marker size, in points. # or parcel_size_attribute: parcel atribute name. Categorical options not supported. Must be continuous. parcel_size_attribute_title : str string to use as the title, if parcel_size_attribute is a string, it is used as the default. parcel_size_norm : par matplotlib color normalizer. https://matplotlib.org/3.1.1/tutorials/colors/colormapnorms.html Default is linear between min and max of parcel_size_attribute., parcel_size_min : float Specify the smallest size of the dot markers plotted, in units of points (default 5). Use with parcel_size_max. They will be aligned with the limits of parcel_size_norm. parcel_size_max : float Specify the largest size of the dot markers plotted, in units of points (default 40). Use with parcel_size_min. They will be aligned with the limits of parcel_size_norm. # with constant or attribute, can set parcel transparency parcel_alpha : float, between 0 and 1 Specify parcel marker transparency (default 0.5). ################### Miscellaneous fig : figure object Default is to create a new figure object. **kwargs Anything else to pass to figure creation. Returns ------- fig : Figure object. """ # part 0 checking and default setting. # only network color/linewidth provided OR link attribute. if (link_attribute is not None) and (network_color is not None): msg = "Only one of link_attribute and network_color can be provided." raise ValueError(msg) if link_attribute is None: network_color = network_color or "c" network_linewidth = network_linewidth or 0.5 legend_link = False else: legend_link = True if link_attribute_title is None: if isinstance(link_attribute, str): link_attribute_title = link_attribute else: link_attribute_title = "" # only parcel color OR parcel_color_attribute. if (parcel_color_attribute is not None) and (parcel_color is not None): msg = "Only one of parcel_color_attribute and parcel_color can be provided." raise ValueError(msg) if parcel_color_attribute is None: parcel_color = parcel_color or "k" legend_parcel_color = False else: legend_parcel_color = True if parcel_color_attribute_title is None: parcel_color_attribute_title = parcel_color_attribute # only parcel size or parcel_size_attribute if (parcel_size_attribute is not None) and (parcel_size is not None): msg = "Only one of parcel_size_attribute and parcel_size can be provided." raise ValueError(msg) if parcel_size_attribute is None: parcel_size = parcel_size or 1.0 legend_parcel_size = False else: legend_parcel_size = True if parcel_size_attribute_title is None: parcel_size_attribute_title = parcel_size_attribute # parcel time: # cant use standard value or default because a value of 0 is valid. if parcel_time_index is None: parcel_time_index = -1 # Figure out whether the legend will have one, two, or three # parts (linewidth, parcel size, parcel color) n_legends = legend_link + legend_parcel_size + legend_parcel_color # set up figure, label and legend gridspecs. if fig is None: fig = plt.figure(**kwargs) spec = gridspec.GridSpec( ncols=1, nrows=3, left=0, right=1, top=1, bottom=0, figure=fig, height_ratios=[1, 0.1, 0.2], ) ax = fig.add_subplot(spec[0, 0]) if n_legends > 0: label_spec = spec[1, 0].subgridspec( ncols=2 * n_legends - 1, nrows=1, wspace=0.1, hspace=1, width_ratios=[1] + (n_legends - 1) * [0.2, 1], ) legend_spec = spec[2, 0].subgridspec( ncols=2 * n_legends - 1, nrows=1, wspace=0.1, hspace=1, width_ratios=[1] + (n_legends - 1) * [0.2, 1], ) legend_idx = 0 # SET UP LINESEGMENTS FOR NETWORK. If polylines exist use, otherwise use # endpoints. Also get the ranges so a buffer can be placed around the # network. if "x_of_polyline" in grid.at_link: xy_of_polylines = _get_xy_of_polylines( grid.at_link["x_of_polyline"], grid.at_link["y_of_polyline"] ) else: xy_of_polylines = grid.xy_of_node[grid.nodes_at_link] xlim, ylim = _calc_xy_limits(xy_of_polylines, buffer_frac=map_buffer) # Add Linesegments and Configure. # if there is a link attribute. if link_attribute is not None: line_segments = LineCollection( xy_of_polylines, cmap=network_cmap, norm=network_norm, linewidth=network_linewidth, zorder=1, ) line_segments.set_array(return_array_at_link(grid, link_attribute)) ax.add_collection(line_segments) # create label lax = fig.add_subplot(label_spec[0, legend_idx]) lax.text( 0.5, 0.0, "Line Color", transform=lax.transAxes, color="k", ha="center", va="top", size=plt.rcParams["font.size"] + 2, ) lax.axis("off") # add legend. lax = fig.add_subplot(legend_spec[0, legend_idx]) legend_idx += 2 fig.colorbar( line_segments, cax=lax, orientation="horizontal", label=link_attribute_title ) # if link values are constant. else: line_segments = LineCollection( xy_of_polylines, colors=network_color, linewidth=network_linewidth, zorder=1 ) ax.add_collection(line_segments) # Part 2: Add Parcels. X = np.empty(len(parcels.dataset.element_id)) Y = np.empty(len(parcels.dataset.element_id)) # Locate parcel XY for each parcel at a particular time # some aspects of this may be possible to speed up, but at minimum # locate_parcel_xy must be called for each link (since calculating location) # along link requires interpoloation. # if this occurs we must also ensure the parcel order is maintained b/c of # color and shape formatting. for parcel_idx in range(parcels.dataset.item_id.size): XY = locate_parcel_xy(grid, parcels, parcel_time_index, parcel_idx) X[parcel_idx] = XY[0] Y[parcel_idx] = XY[1] # plot X,Y point on delineated network and color/size point according to a # certain attribute of the parcel or the link in which the parcel resides # if a parcel color attribute is provided. if parcel_color_attribute is not None: # if this is true, then instead of none we need to get the right # values from the parcels and scale/normalize them correctly. At present # plan to support only continuous values. Can be extended to strs as # categorical. if parcel_color_attribute in parcels.dataset: if "time" in parcels.dataset[parcel_color_attribute].dims: parcel_color = parcels.dataset[parcel_color_attribute].values[ :, parcel_time_index ] else: parcel_color = parcels.dataset[parcel_color_attribute].values else: msg = "Parcel color attribute {attribute} not present in parcels.".format( attribute=parcel_color_attribute ) raise ValueError(msg) if parcel_filter is not None: parcel_color = parcel_color[parcel_filter] if parcel_size_attribute is not None: # if this is true, then instead of none we need to get the right # values from the parcels and scale/normalize them correctly. At present # plan to support only continuous values. Can be extended to strs as # categorical. if parcel_size_attribute in parcels.dataset: if "time" in parcels.dataset[parcel_size_attribute].dims: parcel_size_values = parcels.dataset[parcel_size_attribute].values[ :, parcel_time_index ] else: parcel_size_values = parcels.dataset[parcel_size_attribute].values if parcel_size_norm is None: parcel_size_norm = Normalize( vmin=parcel_size_values.min(), vmax=parcel_size_values.max() ) parcel_size = parcel_size_min + ( parcel_size_max - parcel_size_min ) * parcel_size_norm(parcel_size_values) else: msg = "Parcel size attribute {attribute} not present in parcels.".format( attribute=parcel_size_attribute ) raise ValueError(msg) if parcel_filter is not None: parcel_size = parcel_size[parcel_filter] # add scatter, filter x and y if necessary. if parcel_filter is not None: X = X[parcel_filter] Y = Y[parcel_filter] scatter = ax.scatter( X, Y, s=parcel_size, c=parcel_color, alpha=parcel_alpha, cmap=parcel_color_cmap, norm=parcel_color_norm, zorder=2, ) # create legends. if legend_parcel_color: lax = fig.add_subplot(label_spec[0, legend_idx]) lax.text( 0.5, 0.0, "Parcel Color", transform=lax.transAxes, color="k", ha="center", va="top", size=plt.rcParams["font.size"] + 2, ) lax.axis("off") lax = fig.add_subplot(legend_spec[0, legend_idx]) legend_idx += 2 fig.colorbar( scatter, cax=lax, orientation="horizontal", label=parcel_color_attribute_title, ) if legend_parcel_size: lax = fig.add_subplot(label_spec[0, legend_idx]) lax.text( 0.5, 0.0, "Parcel Size", transform=lax.transAxes, color="k", ha="center", va="top", size=plt.rcParams["font.size"] + 2, ) lax.axis("off") lax = fig.add_subplot(legend_spec[0, legend_idx]) handles, _ = scatter.legend_elements(prop="sizes", alpha=0.6) if len(handles) - 1 == 0: han = handles lab = [parcel_size_values.min()] else: han = handles[:: len(handles) - 1] lab = [parcel_size_values.min(), parcel_size_values.max()] lax.legend( han, lab, title=parcel_size_attribute_title, loc="center", frameon=False, ) plt.axis("off") # Set the plot limits ax.set_xlim(xlim) ax.set_ylim(ylim) # make axes equal ax.axis("equal") return fig
def _get_xy_of_polylines(x_of_polylines, y_of_polylines): """Zip together x and y coordinate arrays. Parameters ---------- x_of_polylines : ndarray x coordinates of a series of polyline segments. y_of_polylines : ndarray y coordinates of a series of polyline segments. Returns ------- ndarray An ndarray of zipped polyline coordinates. Examples -------- >>> x = [[0, 1, 2], [3, 4], [4, 3, 2, 1]] >>> y = [[5, 7, 6], [9, 8], [4, 5, 6, 7]] >>> xy_of_polylines = _get_xy_of_polylines(x, y) >>> len(xy_of_polylines) 3 >>> xy_of_polylines[0] array([[0, 5], [1, 7], [2, 6]]) >>> xy_of_polylines[1] array([[3, 9], [4, 8]]) >>> xy_of_polylines[2] array([[4, 4], [3, 5], [2, 6], [1, 7]]) """ return [np.stack(xy, axis=1) for xy in zip(x_of_polylines, y_of_polylines)] def _calc_xy_limits(xy_of_segment, buffer_frac=0.0): """Calculate xy limits with an optional buffer. Parameters ---------- xy_of_segment : iterable of ndarray xy coordinates of each segment. buffer_frac : float, optional Size of buffer as a fraction of the size of the bounding box of all the segments. A value of zero mean the limits will be 'tight'. Returns ------- x_limits, y_limits x and y limits. Examples -------- >>> xy_of_segments = ([[0, 1], [1, 2]], [[4, 5], [2, 3], [6, 6]], [[2, 9]]) >>> _calc_xy_limits(xy_of_segments) ((0.0, 6.0), (1.0, 9.0)) >>> _calc_xy_limits(xy_of_segments, buffer_frac=0.5) ((-3.0, 9.0), (-3.0, 13.0)) """ segments = np.concatenate(xy_of_segment) x_limits = _calc_limits(segments[:, 0], buffer_frac=buffer_frac) y_limits = _calc_limits(segments[:, 1], buffer_frac=buffer_frac) return x_limits, y_limits def _calc_limits(values, buffer_frac=0.0): """Calculate min and max limits with a buffer on each end. Parameters ---------- values : iterable Values to find limits of. buffer_frac : float, optional Size of buffer, as a fraction of peak-to-peak, to add to the upper and lower limit. Returns ------- limits : tuple Lower and upper limits. Examples -------- >>> _calc_limits([1, 3, 5, 4, 2]) (1.0, 5.0) >>> _calc_limits([1, 3, 5, 4, 2], buffer_frac=0.25) (0.0, 6.0) """ values = np.asarray(values) min_value, max_value = values.min(), values.max() buffer_width = buffer_frac * (max_value - min_value) return min_value - buffer_width, max_value + buffer_width