Add legends to LineCollection plot

Question

This is a derivative question related to the answer given in Set line colors according to colormap where a great solution was suggested to plot several lines with colors according to a colorbar (see code and output image below).

I have a list that stores a string associated with each plotted line, like so:

legend_list = ['line_1', 'line_2', 'line_3', 'line_4']

and I'd like to add these strings as legends in a box (where the first string corresponds to the first plotted line and so on) in the upper right corner of the plot. How could I do this?

I'd be open to not use LineCollection if it was necessary, but I need to keep the colorbar and the colors of each line associated to it.

Code and output

import numpy
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection

# 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 = [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, cmap=plt.cm.rainbow, linewidths=5)
ax.add_collection(lines)
fig.colorbar(lines)

# Manually adding artists doesn't rescale the plot, so we need to autoscale
ax.autoscale()

plt.show()

enter image description here

I'm not sure it's possible to add legends to LineCollection because a LineCollection is not iterable. I'm not sure though. — , Nov 09 '13 at 16:40
Please see updated answer, I'd be open to drop `LineCollection` but I need to keep the colorbar. — Gabriel, Nov 09 '13 at 17:07
I assume it's a new arrival(?), but set_label("") seems to work for me. — Ian, Jul 11 '17 at 21:32
@Ian if you found an alternative easier way, you should post an answer. — Gabriel, Jul 11 '17 at 23:10

score 14 · Accepted Answer · edited Aug 08 '21 at 19:17

@unutbu's answer is the 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()

enter image description here

Thanks for showing how to use proxy artists. Perhaps change `make_proxy(item, ...)` to `make_proxy(item/z.max(),...)` so the legend colors match the LineCollection colors? — unutbu, Nov 09 '13 at 20:24

unutbu · Answer 2 · 2013-11-10T15:25:25.237

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')

enter image description here

This is a fantastic answer @unubtu. Joe's answer seems more fitted since the original question was about adding legends to a LineCollection, so I'll mark that one as accepted. Thank you very much! — Gabriel, Nov 10 '13 at 13:19
Is it possible that this is only working because the min value for `z` is close to zero? I tried it with my real data and I had to do a linear transformation instead of just using `color/z.max()`. Here's the code I used to get the colors right: `m, h = 1./(ages.max()-ages.min()), ages.min()/(ages.min()-ages.max())` (new line) `col_transf = m*color+h`. It does a simple linear transformation between 0 and 1. — Gabriel, Nov 10 '13 at 15:03
It depends on what you want. If you want purple to be associated with z-value 0, then `color/z.max()` works. But if you want purple to always be associated with the line with the lowest z-value, then indeed you need to rescale as you have done. I'll modify my answer to show another way to do it. — unutbu, Nov 10 '13 at 15:25

Add legends to LineCollection plot

2 Answers2

Linked

Related