import numpy as np
import matplotlib.pyplot as plt

#Parameters
T = 0.1 #Time horizon
K = 100   #Number of terms in the Fourier series
def error(N,h):
    M = int(np.ceil(h*T*(N+1)**2.))
    dt = float(T)/float(M)
    dx = 1./float(N+1)
    r = dt/(dx**2)
    print 'r: ', r
    x = np.linspace(0,1,N+2)
    t = np.linspace(0,T,M)
    v = np.zeros([M,N+2])

    #Initial value
    v[0] = 3.*np.sin(np.pi*x)+5.*np.sin(4.*np.pi*x)

    #Computing solution, explicit scheme
    for m in range(int(M-1)):
        for j in range(1,N+1):
            v[m+1][j] = r*v[m][j-1] + (1-2.*r)*v[m][j] + r*v[m][j+1]
    return dt,np.max(np.fabs(v[M-1]-u(x,T)))

#Fourier solution
def u(x,t):
    return 3.*np.exp(-t*np.pi**2)*np.sin(np.pi*x) + 5.*np.exp(-t*16.*np.pi**2)*np.sin(4.*np.pi*x)

#Checking error for different values of dt
for i in range(10,200,20):
    dt, e = error(i,2.)
   # print 'dt: ', dt, 'error: ', e
   # print 'ln(dt): ', np.log(dt), 'ln(e): ', np.log(e)
    plt.scatter(np.log(dt),np.log(e))
    dt, e = error(i,6.)
    plt.scatter(np.log(dt),np.log(e),c = u'r')
plt.show()