import numpy as np
import matplotlib.pyplot as plt

n = 3 #Number of interior points
h = 1./float(n+1) #size of partition
x = np.linspace(0,1,n+2) #x_i = ih

#Linear system to be solved: Av = b, A is (n x n). and tridiagonal
b = (h**2)*(3*x + x**2)*np.exp(x)
print 'b: '
print(b)
v = np.empty(n+2)


#Boundary values
v[0] = 0
v[n+1] = 0

#Solving tridiagonal system (Algorithm 2.1)
d = np.empty(n+2)
c = np.empty(n+1)
d[1] = 2
c[1] = b[1]
for k in range(2,n+1):
    m = -1./d[k-1]
    d[k] = 2.+m
    c[k] = b[k]-m*c[k-1]
v[n] = c[n]/d[n]
for k in range(n-1,0,-1):
    v[k] = (c[k] + v[k+1])/d[k]

print 'v: '
print(v)
#Exact solution
x1 = np.linspace(0,1,100)
u = x1*(1-x1)*np.exp(x1)

#Plotting
plt.plot(x,v,'--')
plt.plot(x1,u)
plt.show()