|
| 1 | +''' |
| 2 | +================== |
| 3 | +Multicolored lines |
| 4 | +================== |
| 5 | +
|
| 6 | +This example shows how to make a multi-colored line. In this example, the line |
| 7 | +is colored based on its derivative. |
| 8 | +''' |
| 9 | + |
| 10 | +import numpy as np |
| 11 | +import matplotlib.pyplot as plt |
| 12 | +from matplotlib.collections import LineCollection |
| 13 | +from matplotlib.colors import ListedColormap, BoundaryNorm |
| 14 | + |
| 15 | +x = np.linspace(0, 3 * np.pi, 500) |
| 16 | +y = np.sin(x) |
| 17 | +dydx = np.cos(0.5 * (x[:-1] + x[1:])) # first derivative |
| 18 | + |
| 19 | +# Create a set of line segments so that we can color them individually |
| 20 | +# This creates the points as a N x 1 x 2 array so that we can stack points |
| 21 | +# together easily to get the segments. The segments array for line collection |
| 22 | +# needs to be (numlines) x (points per line) x 2 (for x and y) |
| 23 | +points = np.array([x, y]).T.reshape(-1, 1, 2) |
| 24 | +segments = np.concatenate([points[:-1], points[1:]], axis=1) |
| 25 | + |
| 26 | +fig, axs = plt.subplots(2, 1, sharex=True, sharey=True) |
| 27 | + |
| 28 | +# Create a continuous norm to map from data points to colors |
| 29 | +norm = plt.Normalize(dydx.min(), dydx.max()) |
| 30 | +lc = LineCollection(segments, cmap='viridis', norm=norm) |
| 31 | +# Set the values used for colormapping |
| 32 | +lc.set_array(dydx) |
| 33 | +lc.set_linewidth(2) |
| 34 | +line = axs[0].add_collection(lc) |
| 35 | +fig.colorbar(line, ax=axs[0]) |
| 36 | + |
| 37 | +# Use a boundary norm instead |
| 38 | +cmap = ListedColormap(['r', 'g', 'b']) |
| 39 | +norm = BoundaryNorm([-1, -0.5, 0.5, 1], cmap.N) |
| 40 | +lc = LineCollection(segments, cmap=cmap, norm=norm) |
| 41 | +lc.set_array(dydx) |
| 42 | +lc.set_linewidth(2) |
| 43 | +line = axs[1].add_collection(lc) |
| 44 | +fig.colorbar(line, ax=axs[1]) |
| 45 | + |
| 46 | +axs[0].set_xlim(x.min(), x.max()) |
| 47 | +axs[0].set_ylim(-1.1, 1.1) |
| 48 | +plt.show() |
0 commit comments