Waves Equation Animation in Python

I use matplotlib module to do the animation.

The main code is in def waves(y0,y1,cb) that use finite difference that solved initial value problem and boundary value problem simultaneously.    

code
from pylab import *
import matplotlib.animation as animation

fig,ax = subplots()

def waves(y0,y1,cb):
    y2 = y0
    for i in range(1,len(y0)-1):
        y2[i] = 2*y1[i]-y0[i]+cb*(y1[i+1]-2*y1[i]+y1[i-1])
    return y2

x   = linspace(0.,1.,20)
dx  = 1./(len(x))
y0  = sin(2*pi*x)
vy0 = 12.

b   = 1./32.  #dt2/dx2
dt  = sqrt(b*dx*dx)
print dt

c   = 1.

cb  = c*b

y1  = y0 + vy0*dt

print y0
print y1

line,   = ax.plot(x,y0)
def animate(i):
    global y0,y1,cb
    y2  = waves(y0,y1,cb)
    y0  = y1
    y1  = y2
    line.set_ydata(y0)
    return line,


#plot (x,y0)

ani =  animation.FuncAnimation(fig, animate, np.arange(1,200), interval =  25, blit=False)

grid(True)
ylim(-10,10)
show()

.