[MRG+1] Simplify and clean multicolor_line example #8173
Merged
Commits
Show all changes
5 commits
Select commit
Hold shift + click to select a range
Jump to file
| @@ -0,0 +1,48 @@ | ||
| +''' | ||
| +================== | ||
| +Multicolored lines | ||
| +================== | ||
| + | ||
| +This example shows how to make a multi-colored line. In this example, the line | ||
| +is colored based on its derivative. | ||
| +''' | ||
| + | ||
| +import numpy as np | ||
| +import matplotlib.pyplot as plt | ||
| +from matplotlib.collections import LineCollection | ||
| +from matplotlib.colors import ListedColormap, BoundaryNorm | ||
| + | ||
| +x = np.linspace(0, 3 * np.pi, 500) | ||
| +y = np.sin(x) | ||
| +dydx = np.cos(0.5 * (x[:-1] + x[1:])) # first derivative | ||
| + | ||
| +# Create a set of line segments so that we can color them individually | ||
| +# This creates the points as a N x 1 x 2 array so that we can stack points | ||
| +# together easily to get the segments. The segments array for line collection | ||
| +# needs to be (numlines) x (points per line) x 2 (for x and y) | ||
| +points = np.array([x, y]).T.reshape(-1, 1, 2) | ||
| +segments = np.concatenate([points[:-1], points[1:]], axis=1) | ||
| + | ||
| +fig, axs = plt.subplots(2, 1, sharex=True, sharey=True) | ||
| + | ||
| +# Create a continuous norm to map from data points to colors | ||
| +norm = plt.Normalize(dydx.min(), dydx.max()) | ||
| +lc = LineCollection(segments, cmap='viridis', norm=norm) | ||
| +# Set the values used for colormapping | ||
| +lc.set_array(dydx) | ||
| +lc.set_linewidth(2) | ||
| +line = axs[0].add_collection(lc) | ||
| +fig.colorbar(line, ax=axs[0]) | ||
| + | ||
| +# Use a boundary norm instead | ||
| +cmap = ListedColormap(['r', 'g', 'b']) | ||
| +norm = BoundaryNorm([-1, -0.5, 0.5, 1], cmap.N) | ||
| +lc = LineCollection(segments, cmap=cmap, norm=norm) | ||
| +lc.set_array(dydx) | ||
| +lc.set_linewidth(2) | ||
| +line = axs[1].add_collection(lc) | ||
| +fig.colorbar(line, ax=axs[1]) | ||
| + | ||
| +axs[0].set_xlim(x.min(), x.max()) | ||
| +axs[0].set_ylim(-1.1, 1.1) | ||
| +plt.show() |
| @@ -1,53 +0,0 @@ | ||
| -''' | ||
| -Color parts of a line based on its properties, e.g., slope. | ||
| -''' | ||
| - | ||
| -import numpy as np | ||
| -import matplotlib.pyplot as plt | ||
| -from matplotlib.collections import LineCollection | ||
| -from matplotlib.colors import ListedColormap, BoundaryNorm | ||
| - | ||
| -x = np.linspace(0, 3 * np.pi, 500) | ||
| -y = np.sin(x) | ||
| -z = np.cos(0.5 * (x[:-1] + x[1:])) # first derivative | ||
| - | ||
| -# Create a colormap for red, green and blue and a norm to color | ||
| -# f' < -0.5 red, f' > 0.5 blue, and the rest green | ||
| -cmap = ListedColormap(['r', 'g', 'b']) | ||
| -norm = BoundaryNorm([-1, -0.5, 0.5, 1], cmap.N) | ||
|
|
||
| - | ||
| -# Create a set of line segments so that we can color them individually | ||
| -# This creates the points as a N x 1 x 2 array so that we can stack points | ||
| -# together easily to get the segments. The segments array for line collection | ||
| -# needs to be numlines x points per line x 2 (x and y) | ||
| -points = np.array([x, y]).T.reshape(-1, 1, 2) | ||
| -segments = np.concatenate([points[:-1], points[1:]], axis=1) | ||
| - | ||
| -# Create the line collection object, setting the colormapping parameters. | ||
| -# Have to set the actual values used for colormapping separately. | ||
| -lc = LineCollection(segments, cmap=cmap, norm=norm) | ||
| -lc.set_array(z) | ||
| -lc.set_linewidth(3) | ||
| - | ||
| -fig1 = plt.figure() | ||
| -plt.gca().add_collection(lc) | ||
| -plt.xlim(x.min(), x.max()) | ||
| -plt.ylim(-1.1, 1.1) | ||
| - | ||
| -# Now do a second plot coloring the curve using a continuous colormap | ||
| -t = np.linspace(0, 10, 200) | ||
| -x = np.cos(np.pi * t) | ||
| -y = np.sin(t) | ||
| -points = np.array([x, y]).T.reshape(-1, 1, 2) | ||
| -segments = np.concatenate([points[:-1], points[1:]], axis=1) | ||
| - | ||
| -lc = LineCollection(segments, cmap=plt.get_cmap('copper'), | ||
| - norm=plt.Normalize(0, 10)) | ||
| -lc.set_array(t) | ||
| -lc.set_linewidth(3) | ||
| - | ||
| -fig2 = plt.figure() | ||
| -plt.gca().add_collection(lc) | ||
| -plt.xlim(-1, 1) | ||
| -plt.ylim(-1, 1) | ||
| -plt.show() | ||
It is a shame to loose the boundary norm example here. Maybe do two axes, one with continuous, one with boundary norm?