Using a miniature version of the plotted data as t

问题:

Is there any way to use the line plotted in a figure in matplotlib as the handle in the legend? I'm thinking for example in this basic code rather than a straight line in the legend have a small version of the sine wave I've plotted as the handle.

import matplotlib.pyplot as plt
import numpy as np

x = np.arange(0,15,0.1)
y = [np.sin(x1) for x1 in x]

plt.plot(x,y, label = 'sine wave')
plt.legend()
plt.show()

Thanks in advance!

回答1:

I thought this was a pretty fun question. Most of the information you need to solve your problem is provided in matplotlib's "Legend Guide". However, I definitely think the documentation is a little sparse.

The key to solving your question is to create a new "handler" that will be called when trying to draw the legend of your artist, and that returns an artists, which can have any properties or shape. In this case, we simply create a new Line2D object with the correct data to draw a sine wave.

from matplotlib.legend_handler import HandlerLine2D
import matplotlib.pyplot as plt


class HandlerSinWave(HandlerLine2D):
    def create_artists(self, legend, orig_handle,
                       xdescent, ydescent, width, height, fontsize,
                       trans):
        xdata, y1 = orig_handle.get_data()

        # scale the sin wave amplitude so that it fits neatly in the legend box
        # first, rescale so that it is centered around 0 and have an amplitude of 1
        min_val = np.min(y1)
        max_val = np.max(y1)
        y2 = (y1-(min_val+(max_val-min_val)/2))/((max_val-min_val)/2)
        # then rescale based on the dimension of the box
        y3 = height*y2 + xdescent+height/2
        legline = matplotlib.lines.Line2D(xdata, y3)

        self.update_prop(legline, orig_handle, legend)
        legline.set_transform(trans)
        return [legline]


fig, ax = plt.subplots()
x = np.arange(0,15,0.1)
y = np.sin(x)

sin1, = plt.plot(x,y, label='sine wave')

plt.legend(handler_map={sin1: HandlerSinWave()})

回答2:

The answer by @DizietAsahi gives the correct result for the simple example, but would fail for other x values. One may hence more generally use a transformed bbox such that one needs not care about the actual values of the data.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.legend_handler import HandlerLine2D
import matplotlib.path as mpath
from matplotlib.transforms import BboxTransformFrom, BboxTransformTo, Bbox


class HandlerMiniatureLine(HandlerLine2D):
    def create_artists(self, legend, orig_handle,
                       xdescent, ydescent, width, height, fontsize,
                       trans):

        legline, _ = HandlerLine2D.create_artists(self,legend, orig_handle,
                                xdescent, ydescent, width, height, fontsize, trans)

        legline.set_data(*orig_handle.get_data())

        ext = mpath.get_paths_extents([orig_handle.get_path()])
        if ext.width == 0:
            ext.x0 -= 0.1
            ext.x1 += 0.1
        bbox0 = BboxTransformFrom(ext)
        bbox1 = BboxTransformTo(Bbox.from_bounds(xdescent, ydescent, width, height))

        legline.set_transform(bbox0 + bbox1 + trans)
        return legline,


fig, ax = plt.subplots()
x = np.arange(0,15,0.1)
y = np.sin(x)

plt.plot(x-900,y+1500, label='sine wave')

plt.legend(handler_map={plt.Line2D: HandlerMiniatureLine()})

plt.show()