import numpy as np
from matplotlib import pyplot as plt

def residual(h,u,f):
    n = len(u) - 1
    r = np.zeros(n)
    for i in range(1,n):
        r[i] = f[i] + (u[i-1]-2*u[i]+u[i+1])/h**2
    return np.linalg.norm(r)/np.linalg.norm(f)

xmin, xmax = 0.0, 1.0
n = 50
h = (xmax - xmin)/n

x = np.linspace(0.0, 1.0, n+1)
f = np.ones(n+1)
ue= 0.5*x*(1-x)

TOL   = 1.0e-6
itmax = 10000

uold = np.zeros(n+1)
unew = np.zeros(n+1)

for it in range(itmax):
    uold[:] = unew
    for i in range(1,n):
        unew[i] = 0.5*( h**2*f[i] + uold[i-1] + uold[i+1])
    change = residual(h,unew,f)
    if change < TOL:
        break

print "Number of iterations = %d" % it

plt.plot(x,ue,x,unew)
plt.legend(("Exact","Numerical"));

u = np.zeros(n+1)

for it in range(itmax):
    for i in range(1,n):
        u[i] = 0.5*( h**2*f[i] + u[i-1] + u[i+1])
    change = residual(h,u,f)
    if change < TOL:
        break

print "Number of iterations = %d" % it

plt.plot(x,ue,x,u)
plt.legend(("Exact","Numerical"));

omega = 1.5

u = np.zeros(n+1)

for it in range(itmax):
    for i in range(1,n):
        z = 0.5*( h**2*f[i] + u[i-1] + u[i+1])
        u[i] = (1-omega)*u[i] + omega*z
    change = residual(h,u,f)
    if change < TOL:
        break

print "Number of iterations = %d" % it

plt.plot(x,ue,x,u)
plt.legend(("Exact","Numerical"));