from scipy import sparse,linsolve
import scipy
from pylab import *
from time import sleep
from math import sin,cos,exp,pi as PI

n = 100
xmin = 0.
xmax = 1.
dx = (xmax-xmin)/(n-1)
xi = [xmin+dx*i for i in range(n)]
lam = 1.0
dt = lam*dx**2

K = 10.
F = lambda x: 10.
Fx = lambda x: 0.
A = sparse.lil_matrix((n,n))

A.setdiag([-2.*lam*K - 1. - K*dt*Fx(x) for x in xi])
A.setdiag([lam*K*(1.+dx*F(x)) for x in xi[:-1]],1)
A.setdiag([lam*K*(1.-dx*F(x)) for x in xi[1:]],-1)

A[0,n-1] = lam*K*(1.-dx*F(xi[n-1]))
A[n-1,0] = lam*K*(1.+dx*F(xi[1]))

a = 20.
f = lambda x: a/scipy.pi**0.5 * scipy.exp(-(x-0.5)**2*a**2) 
b = scipy.array([f(x) for x in xi])

ion()
line, = plot(xi,b)
#savefig("out/%6.5f.png"%0.0)
t = 0.
while t < .1:
    bt = b[:]
    x = linsolve.spsolve(A,b)
    b = -x
    b[0] = -x[n-1]
    b[n-1] = -x[0]
    line.set_ydata(b)
    draw()
    t += dt
    #savefig("out/%6.5f.png"%t)