matplotlib - multiple legend styles in an entry

In a current project, I had to create a lot of plots. In one plot type, several lines should be connected with the same text entry in the legend. To implement this I looked around the internet and stumbled upon a StackOverflow answer from 2015. It guided me to implement a solution, that worked better for my use-case.

The result can be seen in the plot below.

Plot to showcase stacked legend entry.

To paint the symbol that is shown in front of the legend text, matplotlib uses legend handlers defined in matplotlib.legend_handler. You can implement a style by implementing your own handler.

Let us start with how to create a legend in matplotlib.

plt.legend([(line1, line2), line3], ['text', 'more text', 'even more'])

were line1 to line3 are the unpacked return values of plt.plot(…), so called artists. If two artists are passed as a tuple, the legend command uses the legend_handler.HandlerTuple handler. This handler draws both symbols on top of each other.

I used this behaviour and implemented a new handler that works similar, but draws the two symbols as a stack.

from matplotlib.legend_handler import HandlerTuple

class HandlerTupleVertical(HandlerTuple):
    """Plots all the given Lines vertical stacked."""

    def __init__(self, **kwargs):
        """Run Base Handler."""
        HandlerTuple.__init__(self, **kwargs)

    def create_artists(self, legend, orig_handle,
                       xdescent, ydescent, width, height, fontsize, trans):
        """Create artists (the symbol) for legend entry."""
        # How many lines are there.
        numlines = len(orig_handle)
        handler_map = legend.get_legend_handler_map()

        # divide the vertical space where the lines will go
        # into equal parts based on the number of lines
        height_y = (height / numlines)

        leglines = []
        for i, handle in enumerate(orig_handle):
            handler = legend.get_legend_handler(handler_map, handle)

            legline = handler.create_artists(legend, handle,
                                             xdescent,
                                             (2*i + 1)*height_y,
                                             width,
                                             2*height,
                                             fontsize, trans)
            leglines.extend(legline)

        return leglines

This new handler can now be set to be used instead of the original HandlerTuple by modifying the command to create a legend to

line1, = plt.plot(xy,xy, c='k', label='solid')
line2, = plt.plot(xy,xy+1, c='k', ls='dashed', label='dashed')
line3, = plt.plot(xy,xy-1, c='k', ls='dashed', label='dashed')

plt.legend([(line1, line2), line3], ['text', 'more text', 'even more'],
           handler_map = {tuple : HandlerTupleVertical()})

Here it is important to get the first item from the list returned by plt.plot.

Using this handler, nicer legends can be made using matplotlib.

Changes:

Tags: