from numpy import linspace, zeros, exp, sin, pi, arange
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=I(x)

    i = arange(1,n-1);

    for l in range(m):
        t = (l+1)*dt
        # inner nodes
        u[i] = u[i] + alpha*(u[i+1]-2*u[i]+u[i-1]) + 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)

pylab.plot(x, exact(T,x), "b", x, u,"r--")
pylab.xlabel('x')
pylab.ylabel('u')
pylab.show()