from numpy import linspace, zeros, exp, sin, pi
import pylab

def solve(I, f, g0, g1, T, m, L, n):
    dx = L/(n-1.) # n unknowns, n-1 intervals of length dx.
    dt = 1.*T/m;
    alpha = dt/dx**2;
    print alpha
    x = linspace(0, L, n);
    u_new=zeros(n)
    u=I(x)
    
    im = range(0,n-2);
    i = range(1,n-1);
    ip = range(2,n);
    
    for l in range(m):
        t = (l+1)*dt
        # inner nodes
        u_new[i] = u[i] + alpha*(u[ip]-2*u[i]+u[im]) + dt*f(t,x[i])
        # boundary conditions
        u_new[0] = g0(t)
        u_new[n-1] = g1(t)
        # swap pointers
	tmp = u
	u = u_new
	u_new = tmp


    return x, u


# Version 2, no swapping called for due to vectorized operations
def solve2(I, f, g0, g1, T, m, L, n):
    dx = L/(n-1.) # n unknowns, n-1 intervals of length dx.
    dt = 1.*T/m;
    alpha = dt/dx**2;
    print alpha
    x = linspace(0, L, n);
    u=I(x)
    
    im = range(0,n-2);
    i = range(1,n-1);
    ip = range(2,n);
    
    for l in range(m):
        t = (l+1)*dt
        # inner nodes
        u[i] = u[i] + alpha*(u[ip]-2*u[i]+u[im]) + dt*f(t,x[i])
        # boundary conditions
        u[0] = g0(t)
        u[n-1] = g1(t)
	

    return x, u


Q = 4

def I(x):
    return sin(Q*pi*x)

def g0(t): 
    return 0.0

def g1(t):
    return 0.0

def f(t,x):
    return 0.0*x

def exact(t,x):
    return exp(-(Q*pi)**2*t)*sin(Q*pi*x)

T = 0.1;
x, u = solve(I, f, g0, g1, T, m = 2000, L = 1, n = 100)
x, v = solve2(I, f, g0, g1, T, m = 2000, L = 1, n = 100)
pylab.plot(x, exact(T,x), "b", x, u,"r--", x, v, "g--")
pylab.xlabel('x')
pylab.ylabel('u')
pylab.show()

pylab.plot(x, v-u,"g-")
pylab.show()