import numpy as np
from matplotlib import pyplot as plt

def residual(u,f):
    n = u.shape[0]
    res = 0.0
    for i in range(1,n-1):
        for j in range(1,n-1):
            r = -(u[i-1,j]-2*u[i,j]+u[i+1,j])/h**2 \
                - (u[i,j-1]-2*u[i,j]+u[i,j+1])/h**2 \
                - f[i,j]
            res += r**2
    return np.sqrt(res)

xmin,xmax = 0.0,1.0
n = 21
h = (xmax - xmin)/(n-1)
x = np.linspace(xmin,xmax,n)
X,Y = np.meshgrid(x,x)
f    = np.ones((n,n))

TOL   = 1.0e-6
itmax = 2000

uold = np.zeros((n,n))
unew = np.zeros((n,n))
res0 = residual(unew,f)

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

print "Number of iterations = %d" % it
cs = plt.contourf(X,Y,unew)
plt.colorbar(cs);

u    = np.zeros((n,n))
res0 = residual(u,f)

for it in range(itmax):
    for i in range(1,n-1):
        for j in range(1,n-1):
            u[i,j] = 0.25*(h**2 * f[i,j] \
                              + (u[i-1,j] + u[i+1,j] + u[i,j-1] + u[i,j+1]))
    res = residual(u,f)
    if res < TOL * res0:
        break

print "Number of iterations = %d" % it
cs = plt.contourf(X,Y,u)
plt.colorbar(cs);

omega= 1.5
u    = np.zeros((n,n))
res0 = residual(u,f)

for it in range(itmax):
    for i in range(1,n-1):
        for j in range(1,n-1):
            z = 0.25*(h**2 * f[i,j] \
                       + (u[i-1,j] + u[i+1,j] + u[i,j-1] + u[i,j+1]))
            u[i,j] = (1.0-omega)*u[i,j] + omega*z
    res = residual(u,f)
    if res < TOL * res0:
        break

print "Number of iterations = %d" % it
cs = plt.contourf(X,Y,u)
plt.colorbar(cs);