How to configure FuncAnimation (dynamic plotting from matplotlib.animation) to numpy arrays?

Question

I have some issues with the FuncAnimation fonction of MatPlotLib. I can't configure it to my code... I hope someone could help me !!

This is a diffusion equation, I need to plot it for each step of the time. At each step, the result of the calculation is a numpy array. I manage to plot it in a dynamic way with pyplot.interactive(True) but it is very laggy. I read that FuncAnimation can deal with that problem but I did not manage to have it working with results in lists or arrays.

Here is the code with a classic slow plot :
It produces a vector of vectors (U) wich are ploted after all calculations

import numpy as np
import scipy
from scipy.linalg import solve_banded
from matplotlib import pyplot as plt
import matplotlib.animation as animation


def DrawRecord(U):
    plt.interactive(True)
    plt.figure(1)
    for i in range(0,len(U)):
        plt.clf()
        plt.plot(U[i])
        plt.ylim([0,1])
        plt.draw()

J=350.0 
dt=0.01 
T=3.0 
t=np.arange(dt,T,dt) 
dx=1.0/J 

D=0.005 
c=0.5 
r=0.1 

mu=c*dt/(2.0*dx) 
lambd=D*dt/(dx**2.0) 

K_x=50.0*np.ones(J-1) 
alpha_t=0.5*np.ones(len(t)) 

#initial conditions
u=np.zeros(J) 
u[J/5*1:J/5*2]=1 
U=u

espace=np.linspace(0,1,J) 

#Matrix
A=np.diag(-lambd*np.ones(J-2),1)+np.diag((1+2*lambd)*np.ones(J-1),0)+np.diag(-lambd*np.ones(J-2),-1)  
AA=scipy.linalg.inv(A)


for i in t:

    u[1:J]=scipy.dot(AA,u[1:J]+(r-alpha_t[i/dt])*dt*(u[1:J]-u[1:J]/K_x))
    u[0]=0 
    u[J-1]=0  

    U=np.vstack([U,u])

DrawRecord(U)

And here is my try of making turn the FuncAnimation with the previous code (big fail) :
nb : U contents the arrays of results calculated for each steps

global U

fig = plt.figure()
window = fig.add_subplot(111)
line, = window.plot(list(U[1,:]))

def init():
    line=list(U[1,:])
    return line

def animate(i):
    line.set_ydata(list(U[i,:]))
    return line

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

plt.show()

That produces a lot of errors ... Maybe someone can set it up for the previous code !
I hope I'm clear (sorry for my english) and thank you for your help.

Answer 1

Your init function isn't right, it should be:

def init():
    line.set_ydata(U[1,:])
    return line

You also don't need to convert the numpy slice to a list.

Answer 2

FuncAnimation expects an iterable to return from init and update

try this:

def init():
    line.set_ydata(U[1,:])
    return line,

Answer 3

I ran into a similar error, and I fixed it by adding a comma.

def init():
    line.set_ydata(U[1,:])
    return line,

For some reason this fixed it. You don't need the trailing comma if you have more than one object you're returning. Bloody cryptic error message, if you ask me.