Source code for landlab.plot.drainage_plot

"""Plot drainage network."""
import matplotlib.pylab as plt
import numpy as np

from landlab.plot.imshow import imshow_grid

# KRB, FEB 2017.


[docs]def drainage_plot( mg, surface="topographic__elevation", receivers=None, proportions=None, surf_cmap="gray", quiver_cmap="viridis", title="Drainage Plot", ): if isinstance(surface, str): colorbar_label = surface else: colorbar_label = "topographic_elevation" imshow_grid(mg, surface, cmap=surf_cmap, colorbar_label=colorbar_label) if receivers is None: receivers = mg.at_node["flow__receiver_node"] if proportions is None and "flow__receiver_proportions" in mg.at_node: proportions = mg.at_node["flow__receiver_proportions"] else: receivers = np.asarray(receivers) if receivers.ndim == 1: receivers = np.expand_dims(receivers, axis=1) nreceievers = receivers.shape[-1] propColor = plt.get_cmap(quiver_cmap) for j in range(nreceievers): rec = receivers[:, j] is_bad = rec == -1 xdist = -0.8 * (mg.x_of_node - mg.x_of_node[rec]) ydist = -0.8 * (mg.y_of_node - mg.y_of_node[rec]) if proportions is None: proportions = np.ones_like(receivers, dtype=float) if proportions.ndim == 1: proportions = np.expand_dims(proportions, axis=1) is_bad[proportions[:, j] == 0.0] = True xdist[is_bad] = np.nan ydist[is_bad] = np.nan prop = proportions[:, j] * 256.0 lu = np.floor(prop) colors = propColor(lu.astype(int)) shape = (mg.number_of_nodes, 1) plt.quiver( mg.x_of_node.reshape(shape), mg.y_of_node.reshape(shape), xdist.reshape(shape), ydist.reshape(shape), color=colors, angles="xy", scale_units="xy", scale=1, zorder=3, ) # Plot different types of nodes: (o,) = plt.plot( mg.x_of_node[mg.status_at_node == mg.BC_NODE_IS_CORE], mg.y_of_node[mg.status_at_node == mg.BC_NODE_IS_CORE], "b.", label="Core Nodes", zorder=4, ) (fg,) = plt.plot( mg.x_of_node[mg.status_at_node == mg.BC_NODE_IS_FIXED_VALUE], mg.y_of_node[mg.status_at_node == mg.BC_NODE_IS_FIXED_VALUE], "c.", label="Fixed Gradient Nodes", zorder=5, ) (fv,) = plt.plot( mg.x_of_node[mg.status_at_node == mg.BC_NODE_IS_FIXED_GRADIENT], mg.y_of_node[mg.status_at_node == mg.BC_NODE_IS_FIXED_GRADIENT], "g.", label="Fixed Value Nodes", zorder=6, ) (c,) = plt.plot( mg.x_of_node[mg.status_at_node == mg.BC_NODE_IS_CLOSED], mg.y_of_node[mg.status_at_node == mg.BC_NODE_IS_CLOSED], "r.", label="Closed Nodes", zorder=7, ) node_id = np.arange(mg.number_of_nodes) flow_to_self = receivers[:, 0] == node_id (fts,) = plt.plot( mg.x_of_node[flow_to_self], mg.y_of_node[flow_to_self], "kx", markersize=6, label="Flows To Self", zorder=8, ) ax = plt.gca() ax.legend( labels=[ "Core Nodes", "Fixed Gradient Nodes", "Fixed Value Nodes", "Closed Nodes", "Flows To Self", ], handles=[o, fg, fv, c, fts], numpoints=1, loc="center left", bbox_to_anchor=(1.7, 0.5), ) sm = plt.cm.ScalarMappable(cmap=propColor, norm=plt.Normalize(vmin=0, vmax=1)) sm._A = [] cx = plt.colorbar(sm, ax=ax) cx.set_label("Proportion of Flow") plt.title(title)