How to add legend/label in python animation

Question

I want to add a legend in a python animation, like the line.set_label() below. It is similar to plt.plot(x,y,label='%d' %*variable*).

However, I find that codes do not work here. The animation only shows lines changing but no label or legend available. How can I fix this problem?

 from matplotlib import pyplot as plt 
 from matplotlib import animation

fig = plt.figure()

ax = plt.axes(xlim=(0, 2), ylim=(0, 100))

N = 3
lines = [plt.plot([], [])[0] for _ in range(N)]


def init():    
    for line in lines:
        line.set_data([], [])
    return lines

def animate(i):
    for j,line in enumerate(lines):
        line.set_data([0, 2], [10*j,i])
        line.set_label('line %d, stage %d'%(j,i))
    return lines

anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=100, interval=20, blit=True)

plt.show()

Answer 1

I'm no expert on matplotlib at all, but in the Double Pendulum animation they display texts which changes, and this leads to some variations which can help you.

To get legends with the actual color of the lines, you can either change the initial setting lines to:

lines = [plt.plot([], [], label = 'line {}'.format(i))[0] for i in range(N)]

or add a line.set_label() to the for loop in the init() function. Both these seem to work as expected. At least if you add plt.legend(loc="upper left") right before plt.show().

However the set_label doesn't work within the animate() function, but according to the linked animation you can use specific text fields added to the animation, and that seems to work nicely. Add the following code after initialisation of lines:

texts = [ax.text(0.80, 0.95-i*0.05,  '', transform=ax.transAxes) for i in range(N)]

And change animate() to be:

def animate(i):
    for j in range(N):
        lines[j].set_data([0, 2], [10*j,i]) #, label="hei {}".format(i))
        texts[j].set_text('line %d, stage %d'%(j,i))
    return lines

This places the text close to the upper right corner, and is updated for each animation step. Since the lines still have their legend displayed, you possibly simplify into one text only displaying the stage. But I leave the fine tuning of messages to your discretion.

Addendum: Extend Line2D

Another alternative could possibly be to extend lines.Line2D and use these lines in your animation, something similar to this article. Not sure if this would work with animation, but if you can't get the above to work, this might be worth a try.

Answer 2

you must return the legend in your animation function for it to be rendered. Try this, instead :

legend = plt.legend()
def animate(i):
    for j,line in enumerate(lines):
        line.set_data([0, 2], [10*j,i])
        line.set_label('line %d, stage %d'%(j,i))
    legend.remove()
    legend = plt.legend()
    return lines + [legend]

You should also include the same code in your init function, init is used when resizing the window, otherwise the legend will disappear when resizing