Add legends to LineCollection plot
@ubuntu's answer is right approach if you have a small number of lines. (And if you're wanting to add a legend, you presumably do!)
Just to show the other option, though, you can still use a LineCollection
, you just need to use "proxy artists" for the legend:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.lines import Line2D
# The line format you curently have:
lines = [[(0, 1, 2, 3, 4), (4, 5, 6, 7, 8)],
[(0, 1, 2, 3, 4), (0, 1, 2, 3, 4)],
[(0, 1, 2, 3, 4), (8, 7, 6, 5, 4)],
[(4, 5, 6, 7, 8), (0, 1, 2, 3, 4)]]
# Reformat it to what `LineCollection` expects:
lines = [tuple(zip(x, y)) for x, y in lines]
z = np.array([0.1, 9.4, 3.8, 2.0])
fig, ax = plt.subplots()
lines = LineCollection(lines, array=z, linewidths=5,
cmap=plt.cm.rainbow, norm=plt.Normalize(z.min(), z.max()))
ax.add_collection(lines)
fig.colorbar(lines)
# Manually adding artists doesn't rescale the plot, so we need to autoscale
ax.autoscale()
def make_proxy(zvalue, scalar_mappable, **kwargs):
color = scalar_mappable.cmap(scalar_mappable.norm(zvalue))
return Line2D([0, 1], [0, 1], color=color, **kwargs)
proxies = [make_proxy(item, lines, linewidth=5) for item in z]
ax.legend(proxies, ['Line 1', 'Line 2', 'Line 3', 'Line 4'])
plt.show()
Using a LineCollection is faster than using plt.plot
if you have a large number of lines, but I haven't been able to figure out how to add a legend if using LineCollection. The legend guide says to use a proxy artist, but if you have to create a different proxy artist for each line segment in the LineCollection, it might be better to bite the bullet and just use plt.plot
.
And since you want a legend, it seems plausible that you have a small number of lines. Indeed, that would be fortunate, since trying to plot thousands of lines with plt.plot
is a recipe for slowness.
So, if you have a small number of lines, the following should work fine:
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
lines = [[(0, 1, 2, 3, 4), (4, 5, 6, 7, 8)],
[(0, 1, 2, 3, 4), (0, 1, 2, 3, 4)],
[(0, 1, 2, 3, 4), (8, 7, 6, 5, 4)],
[(4, 5, 6, 7, 8), (0, 1, 2, 3, 4)]]
z = np.array([0.1, 9.4, 3.8, 2.0])
legend_list = ['line_1', 'line_2', 'line_3', 'line_4']
fig, ax = plt.subplots()
cmap = plt.get_cmap('rainbow')
def normalize(z):
z = z.copy()
z -= z.min()
z /= z.max()
return z
for (x, y), color, label in zip(lines, normalize(z), legend_list):
plt.plot(x, y, label=label, color=cmap(color), lw=5)
m = cm.ScalarMappable(cmap=cmap)
m.set_array(z)
plt.colorbar(m)
ax.legend()
plt.savefig('/tmp/test.png')