Defining multiple plots to be animated with a for loop in matplotlib

Question

Thanks to Jake Vanderplas, I know how to start to code an animated plot with matplotlib. Here is a sample code:

from matplotlib import pyplot as plt
from matplotlib import animation

fig = plt.figure()

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

line, = plt.plot([], [])

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

def animate(i):
    line.set_data([0, 2], [0,i])
    return line,

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

plt.show()

Suppose now I'd like to plot tons of functions (say four here), defined with the help of a loop. I did some voodoo programming, trying to understand how to mimic the comma following line and here is what I got (needless to say that it does not work: AttributeError: 'tuple' object has no attribute 'axes').

from matplotlib import pyplot as plt
from matplotlib import animation

fig = plt.figure()

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

line = []
N = 4

for j in range(N):
    temp, = plt.plot([], [])
    line.append(temp)

line = tuple(line)

def init():
    for j in range(N):
        line[j].set_data([], [])
    return line,

def animate(i):
    for j in range(N):
        line[j].set_data([0, 2], [10 * j,i])
    return line,

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

plt.show()

Some my question is: how can I make it work? Bonus (probably linked): what is the difference between line, = plt.plot([], []) and line = plt.plot([], [])?

Thanks

Answer 1

Here is a modified example which is more readable. This is just the code from the matplotlib website with another piece added. What threw me was I wasn't initially realizing the plt functions were returning lists, overlooking the trailing comma, or even how things were being tied together, which seems to be getting handled pretty implicitly by the library itself. But the key is to create several update-able objects and return them from the two key functions as part of the same list and they will be synced when the animation is run.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

fig, ax = plt.subplots()
xdata, ydata = [], []
xdata2,ydata2= [], []
ln, = plt.plot([], [], 'ro') 
ln2, = plt.plot([], [], 'ro')

def init():
    ax.set_xlim(0, 2*np.pi)
    ax.set_ylim(-1, 1)
    return ln,ln2

def update(frame):
    xdata.append(frame)
    ydata.append(np.sin(frame))
    ln.set_data(xdata, ydata)

    xdata2.append(frame)
    ydata2.append(np.cos(frame))
    ln2.set_data(xdata2,ydata2)
    
    return ln,ln2

ani = FuncAnimation(fig, update, frames=np.linspace(0, 2*np.pi, 128),
                    init_func=init, blit=True)


plt.show()

For more fun change the animation callback to the following:

def update(frame):
    xdata.append(frame)
    ydata.append(np.sin(frame))
    ln.set_data(xdata, ydata)

    delta = 2*np.pi/128
    xdata2.append([frame*2,frame*2+delta])
    ydata2.append([np.cos(frame*2), np.cos(frame*2+delta)])
    ln2.set_data(xdata2,ydata2)
    
    return ln,ln2

Answer 2

In the solution below I showcase a bigger example (with also bar plot) that may help people understand better what should be done for other cases. After the code I explain some details and answer the bonus question.

import matplotlib
matplotlib.use('Qt5Agg') #use Qt5 as backend, comment this line for default backend

from matplotlib import pyplot as plt
from matplotlib import animation

fig = plt.figure()

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

N = 4
lines = [plt.plot([], [])[0] for _ in range(N)] #lines to animate

rectangles = plt.bar([0.5,1,1.5],[50,40,90],width=0.1) #rectangles to animate

patches = lines + list(rectangles) #things to animate

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

    #init rectangles
    for rectangle in rectangles:
        rectangle.set_height(0)

    return patches #return everything that must be updated

def animate(i):
    #animate lines
    for j,line in enumerate(lines):
        line.set_data([0, 2], [10 * j,i])

    #animate rectangles
    for j,rectangle in enumerate(rectangles):
        rectangle.set_height(i/(j+1))

    return patches #return everything that must be updated

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

plt.show()

Explanation

The idea is to plot what you need and then reuse the artists (see more here) returned by matplotlib. This is done by first plotting a dummy sketch of what you want and keeping the objects matplotlib gives you. Then on your init and animate functions you can update the objects that need to be animated.

Note that in plt.plot([], [])[0] we get a single line artist, thus I collect them with [plt.plot([], [])[0] for _ in range(N)]. On the other hand plt.bar([0.5,1,1.5],[50,40,90],width=0.1) returns a container that can be iterated for the rectangle artists. list(rectangles) just convert this container into a list to be concatenated with lines.

I separate the lines from the rectangles because they are updated differently (and are different artists) but init and animate return all of them.

Answer to bonus question:

1. line, = plt.plot([], []) assign the first element of the list returned by plt.plot to the veriable line.
2. line = plt.plot([], []) just assign the whole list (of only one element).