[MRG+1] Simplify and clean multicolor_line example #8173

Merged
merged 5 commits into from Mar 8, 2017
@@ -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)
@tacaswell

tacaswell Feb 28, 2017

Owner

It is a shame to loose the boundary norm example here. Maybe do two axes, one with continuous, one with boundary norm?

@dstansby

dstansby Feb 28, 2017

Contributor

👍 That shouldn't be too hard to add back in

-
-# 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()