import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

#Parameters
N = 5
M = 50
T = 3.0 #Time horizon

#Implicit/Explicit
expl = True
impl = False

dt = float(T)/float(M)
dx = 5./float(N+1)
r = 4.*dt/(dx**2)

x = np.linspace(-2,3,N+2)
t = np.linspace(0,T,M)
v = np.zeros([M,N+2]) #Explicit
w = np.zeros([M,N+2]) #Implicit

#Boundary values
v[0] = np.ones(N+2) + x
v[:,0] = np.exp(t)-2.*np.ones(M)
v[:,N+1] = np.exp(t)+3.*np.ones(M)

#Source function
def q(x,t):
    return 11.*np.exp(t) + 10.*x

#Computing solution, explicit scheme
if expl:
    for m in range(M-1):
        for j in range(1,N+1):
            v[m+1][j] = r*v[m][j-1] + (1-2.*r-10.*dt)*v[m][j] + r*v[m][j+1] + dt*q(x[j],t[m])

#Computing solution, implicit scheme
if impl:
    A = (1.+2.*r+10.*dt)*np.eye(N)
    for i in range(N-1):
        A[i+1][i] = -r
        A[i][i+1] = -r
    print np.linalg.cond(A)
    B = np.linalg.inv(A)
    for m in range(M-1):
        b = -dt*q(x[1:N+1],t[m+1])
        b[0] = b[0]-r*v[m+1,0]
        b[N-1] = b[N-1]-r*v[m+1,N+1]
        v[m+1][1:N+1] = np.dot(B,v[m][1:N+1]-b)

#Plotting solution
fig = plt.figure()
ax = fig.gca(projection='3d')
x, t = np.meshgrid(x, t)
ax.plot_wireframe(x,t,v)
plt.show()