Adding a legend to a matplotlib plot with a multicolored line

Question

Following the example on how to draw multicolored lines I can draw lines that change color along their length based on some color map. Trying to add a legend to the plot I added this code:

plt.legend([lc], ["test"],\
    handler_map={lc: matplotlib.legend_handler.HandlerLineCollection()})

This adds a legend to the plot (figure below) but the color of the icon in the legend does not relate at all to the colors of the line. Is this the wrong way to try to add a legend to this plot, or is this a limitation of matplotlib?

The handler for a LineCollection is a line. So this is expected. It's not a priori clear how a legend entry for a linecollection should look like instead, right? But if you want to tell how you envision the legend to look like instead, one may find a way to accomplish this. — ImportanceOfBeingErnest, Mar 11 '18 at 18:48
Yes, I totally agree. Especially in the other example case where the line has discontinuous color jumps, it's tricky to imagine what the icon should be. In this case I was picturing something like a horizontal colorbar (obviously with no numerical scale attached) associated with the colormap. — RGWinston, Mar 11 '18 at 18:54

ImportanceOfBeingErnest · Accepted Answer · 2018-03-11T19:53:50.387

The idea would be to show a line collection in the legend as well. There is no inbuilt way to do that but one may subclass HandlerLineCollection and create the respective LineCollection within its create_artists method.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.legend_handler import HandlerLineCollection
from matplotlib.collections import LineCollection

class HandlerColorLineCollection(HandlerLineCollection):
    def create_artists(self, legend, artist ,xdescent, ydescent,
                        width, height, fontsize,trans):
        x = np.linspace(0,width,self.get_numpoints(legend)+1)
        y = np.zeros(self.get_numpoints(legend)+1)+height/2.-ydescent
        points = np.array([x, y]).T.reshape(-1, 1, 2)
        segments = np.concatenate([points[:-1], points[1:]], axis=1)
        lc = LineCollection(segments, cmap=artist.cmap,
                     transform=trans)
        lc.set_array(x)
        lc.set_linewidth(artist.get_linewidth())
        return [lc]

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), linewidth=3)
lc.set_array(t)

fig, ax = plt.subplots()
ax.add_collection(lc)

plt.legend([lc], ["test"],\
    handler_map={lc: HandlerColorLineCollection(numpoints=4)}, framealpha=1)

ax.autoscale_view()
plt.show()

Adding a legend to a matplotlib plot with a multicolored line

1 Answers1

Linked

Related