astromancer
astromancer

Reputation: 611

python matplotlib errorbar legend picking

I'm trying to code a legend picker for matplotlib errorbar plots similar to this example. I want to be able to click on the errorbars / datapoints in the legend to toggle their visibility in the axis. The problem is that the legend object returned by plt.legend() does not contain any data on the artists used in creating the legend. If I for eg. do:

import numpy as np
import matplotlib.pyplot as plt

fig, ax = plt.subplots()
x = np.linspace(0,10,100)
y = np.sin(x) + np.random.rand(100)
yerr = np.random.rand(100)

erbpl1 = ax.errorbar(x, y, yerr=yerr, fmt='o', label='A')
erbpl2 = ax.errorbar(x, 0.02*y, yerr=yerr, fmt='o', label='B')

leg = ax.legend()

From here it seems impossible to access the legend's artists by using the leg object. Normally, one can do this with simpler legends, eg:

plt.plot(x, y, label='whatever')
leg = plt.legend()
proxy_lines = leg.get_lines()

gives you the Line2D objects used in the legend. With errorbar plots, however, leg.get_lines() returns an empty list. This kind of makes sense because plt.errorbar returns a matplotlib.container.ErrorbarContainer object (which contains the data points, errorbar end caps, errorbar lines). I would expect the legend to have a similar data container, but i cant see this. The closest I could manage was leg.legendHandles which points to the errorbar lines, but not data points nor the end caps. If you can pick the legends, you can map them to the original plots with a dict and use the following functions to turn the errorbars on/off.

def toggle_errorbars(erb_pl):
    points, caps, bars = erb_pl
    vis = bars[0].get_visible()
    for line in caps:
        line.set_visible(not vis)
    for bar in bars:
        bar.set_visible(not vis)
    return vis

def onpick(event):
    # on the pick event, find the orig line corresponding to the
    # legend proxy line, and toggle the visibility
    legline = event.artist
    origline = lined[legline]

    vis = toggle_errorbars(origline)
    ## Change the alpha on the line in the legend so we can see what lines
    ## have been toggled
    if vis:
        legline.set_alpha(.2)
    else:
        legline.set_alpha(1.)
    fig.canvas.draw()

My question is, is there a workaround that could allow me to do event picking on an errorbar / other complex legend??

Upvotes: 4

Views: 4705

Answers (1)

tacaswell
tacaswell

Reputation: 87376

This makes the markers pickable:

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.legend_handler

from matplotlib.container import ErrorbarContainer

class re_order_errorbarHandler(matplotlib.legend_handler.HandlerErrorbar):
    """
    Sub-class the standard error-bar handler 
    """
    def create_artists(self, *args, **kwargs):
        #  call the parent class function
        a_list = matplotlib.legend_handler.HandlerErrorbar.create_artists(self, *args, **kwargs)
        # re-order the artist list, only the first artist is added to the
        # legend artist list, this is the one that corresponds to the markers
        a_list = a_list[-1:] + a_list[:-1]
        return a_list

my_handler_map = {ErrorbarContainer: re_order_errorbarHandler(numpoints=2)}

fig, ax = plt.subplots()
x = np.linspace(0,10,100)
y = np.sin(x) + np.random.rand(100)
yerr = np.random.rand(100)

erbpl1 = ax.errorbar(x, y, yerr=yerr, fmt='o', label='A')
erbpl2 = ax.errorbar(x, 0.02*y, yerr=yerr, fmt='o', label='B')

leg = ax.legend(handler_map=my_handler_map)

lines = [erbpl1, erbpl2]
lined = dict()
# not strictly sure about ordering, but 
for legline, origline in zip(leg.legendHandles, lines):
    legline.set_picker(5)  # 5 pts tolerance
    lined[legline] = origline


def onpick(event):
    # on the pick event, find the orig line corresponding to the
    # legend proxy line, and toggle the visibility
    legline = event.artist
    origline = lined[legline]
    for a in origline.get_children():
        vis = not a.get_visible()
        a.set_visible(vis)
    # Change the alpha on the line in the legend so we can see what lines
    # have been toggled
    if vis:
        legline.set_alpha(1.0)
    else:
        legline.set_alpha(0.2)
    fig.canvas.draw()

fig.canvas.mpl_connect('pick_event', onpick)

What is going on here is that the standard handler for ErrorbarContainers uses 4 artists to generate the legend entry (LineCollection for bars, LineCollection for caps, Line2D for the connecting line, and Line2D for the markers). The code that generate the artists only returns the first artists in the list of artists added to the legend (see matplotlib.legend_handler.HandlerBase.__call__). The first artist in the list for error bars happens to be the line collection that is the vertical lines, which is what ends up in leg.legendHandles. The reason the picking wasn't working seems to be that they were being hidden by the other artists (I think).

The solution is to make a local sub-class of HandlerErrorbar which re-orders the artist list so that the artist that gets saved in leg.legendHandles is the Line2D object for the markers.

I will probably open a PR to make this the default behavior.

Upvotes: 5

Related Questions