Source code for landlab.plot.layers

from functools import partial
from itertools import tee

import matplotlib.pyplot as plt
import numpy as np
from matplotlib import cm
from matplotlib.patches import Patch
from scipy.interpolate import interp1d


def pairwise(iterable):
    "s -> (s0,s1), (s1,s2), (s2, s3), ..."
    a, b = tee(iterable)
    next(b, None)
    return zip(a, b)


[docs]def plot_layers( elevation_at_layer, x=None, sea_level=0.0, color_water=(0.8, 1.0, 1.0), color_bedrock=(0.8, 0.8, 0.8), color_layer=None, layer_line_width=0.5, layer_line_color="k", title=None, x_label="Distance", y_label="Elevation", legend_location="lower left", ): """Plot a stack of sediment layers as a cross section. Create a plot of the elevation sediment layers, including surfaces for sea level and bedrock. Parameters ---------- elevation_at_layer : array-like of shape *(n_layers, n_stacks)* Elevation to each layer along the profile. Layers are provided row-by-row, with the bottom-most layer being the first row. x : array-like, optional Distance to each stack along the cross-section. If not provided, stack number will be used. sea_level : float, optional Elevation of sea level. color_water : tuple of float, optional Tuple of *(red, green, blue)* values for water. color_bedrock : tuple of float, optional Tuple of *(red, green, blue)* values for bedrock. color_layer : string, optional Colormap to use to color in the layers. layer_line_width : float, optional Width of line used to plot layer surfaces. layer_line_color : string, optional Color of the line used to plot layer surfaces. title : string, optional Text to be used for the graph's title. The default is to not include a title. x_label : string, optional Text to be used for the x (horizontal) axis label. y_label : string, optional Text to be used for the y (vertical) axis label. legend_location : string, optional Where to put the legend. """ elevation_at_layer = np.asarray(elevation_at_layer) elevation_at_layer = np.expand_dims( elevation_at_layer, axis=tuple(np.arange(2 - elevation_at_layer.ndim)), ) if len(elevation_at_layer) == 0: raise ValueError( f"no layers to plot (elevation_at_layer.shape is {np.shape(elevation_at_layer)}" ) if x is None: x = np.arange(elevation_at_layer.shape[1]) top_surface = elevation_at_layer[-1] bottom_surface = elevation_at_layer[0] if len(elevation_at_layer) > 0: _plot_layers( x, elevation_at_layer, # [layers_to_plot], color=color_layer, lc=layer_line_color, lw=layer_line_width, ) _plot_water(x, top_surface, sea_level=sea_level, fc=color_water) _plot_bedrock(x, bottom_surface, fc=color_bedrock) _plot_surface(x, top_surface, sea_level=sea_level) legend_location and _plot_legend( legend_location, color_water=color_water, color_bedrock=color_bedrock ) plt.xlabel(x_label) plt.ylabel(y_label) title and plt.title(title) plt.show()
def _plot_water(x, y, sea_level=0.0, fc=(0.8, 1.0, 1.0)): x_water, y_water = _insert_shorelines(x, y, sea_level=sea_level) if fc is not None: plt.fill_between(x_water, y_water, sea_level, where=y_water <= sea_level, fc=fc) water_surface = np.full_like(x_water, sea_level, dtype=float) water_surface[y_water > sea_level] = np.nan plt.plot(x_water, water_surface, color="b") def _plot_bedrock(x, y, fc=(0.8, 0.8, 0.8)): if fc is not None: plt.fill_between( x, y, np.full_like(y, y.min()), color=fc, ) plt.plot(x, y, color="k") def _plot_surface(x, y, sea_level=0.0): under_water = y <= sea_level plt.plot(x[~under_water], y[~under_water], color="g") plt.plot(x[under_water], y[under_water], color="b") def _plot_layers(x, layers, color=None, lc="k", lw=0.5): if color is not None: cmap = cm.get_cmap(color) for layer, (lower, upper) in enumerate(pairwise(layers)): plt.fill_between( x, lower, upper, fc=cmap(layer * 256 // len(layers)), ) plt.plot( x, layers.T, color=lc, linewidth=lw, ) def _plot_legend(legend_location, color_water=None, color_bedrock=None): legend_item = partial(Patch, edgecolor="k", linewidth=0.5) items = [ ("Ocean", color_water), ("Bedrock", color_bedrock), ] legend = [legend_item(label=label, fc=color) for label, color in items if color] legend and plt.legend(handles=legend, loc=legend_location) def _insert_shorelines(x, y, sea_level=0.0): """Insert shorelines into x-y arrays. Examples -------- >>> from landlab.plot.layers import _insert_shorelines >>> _insert_shorelines([0, 1, 2],[2, 1, -1]) (array([ 0. , 1. , 1.5, 2. ]), array([ 2., 1., 0., -1.])) """ x, y = np.asarray(x, dtype=float), np.asarray(y, dtype=float) y_relative_to_sea_level = y - sea_level shorelines = _search_zero_crossings(y_relative_to_sea_level) x_of_shoreline = _interp_zero_crossings(x, y_relative_to_sea_level, shorelines) return ( np.insert(x, shorelines + 1, x_of_shoreline), np.insert(y, shorelines + 1, sea_level), ) def _search_zero_crossings(y): """Search an array for changes in sign between elements. Parameters ---------- y : array-like Input array to check for sign changes. Returns ------- int Indices into *y* where a sign has changed. Examples -------- >>> from landlab.plot.layers import _search_zero_crossings The returned index is to the element before the zero-crossing. >>> list(_search_zero_crossings([2, 1, -1])) [1] >>> list(_search_zero_crossings([-2, -2, 1, 2])) [1] >>> list(_search_zero_crossings([-2, -2, 1, 2, -1])) [1, 3] These are not zero-crossings. >>> len(_search_zero_crossings([2, 0, 0, -2])) == 0 True >>> len(_search_zero_crossings([2, 0, 1])) == 0 True >>> len(_search_zero_crossings([2, 3, 4])) == 0 True >>> len(_search_zero_crossings([0, 0, 0])) == 0 True """ sign = np.sign(y) # zeros = sign == 0 # if not np.all(zeros): # while np.any(zeros): # sign[zeros] = np.roll(sign, 1)[zeros] # zeros = sign == 0 # return np.where(sign[1:] != sign[:-1])[0] return np.where(sign[1:] * sign[:-1] < 0)[0] def _interp_zero_crossings(x, y, shorelines): """Interpolate between adjacent elements to find x-locations of zero-crossings. Parameters ---------- x : array-like Distances. y : array-like Elevations. shorelines : array-like of int Indices to shoreline elements. Returns ------- array of float Distances to interpolated shorelines. Examples -------- >>> from landlab.plot.layers import _interp_zero_crossings >>> _interp_zero_crossings([0, 1, 2], [1, -1, -1], [0]) array([ 0.5]) >>> _interp_zero_crossings([0, 1, 2, 3], [1, -1, -1, 4], [0, 2]) array([ 0.5, 2.2]) """ x_of_shoreline = [] for shoreline in shorelines: coast = slice(shoreline, shoreline + 2) # for scipy<1.10 interp1d requires x and y to have at least two elements, # which is not the case if theshoreline is the last element. x_of_shoreline.append( interp1d(np.broadcast_to(y[coast], 2), np.broadcast_to(x[coast], 2))(0.0) ) return np.asarray(x_of_shoreline)