import numpy as np
from matplotlib import pyplot as plt

def ax(h,u):
    n = u.shape[0]
    res = np.zeros((n,n))
    for i in range(1,n-1):
        for j in range(1,n-1):
            res[i,j] = -(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
    return 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

u   = np.zeros((n,n))
p   = np.zeros((n,n))
res = np.array(f)

# First and last grid point, solution is fixed to zero.
# Hence we make residual zero, in which case solution
# will not change at these points.
res[0,:]   = 0.0
res[n-1,:] = 0.0
res[:,0]   = 0.0
res[:,n-1] = 0.0

f_norm = np.linalg.norm(f)
res_old, res_new = 0.0, 0.0
for it in range(itmax):
    res_new = np.linalg.norm(res)
    if res_new < TOL * f_norm:
        break
    if it==0:
        beta = 0.0
    else:
        beta = res_new**2 / res_old**2
    p = res + beta * p
    ap= ax(h,p)
    alpha = res_new**2 / np.sum(p*ap)
    u += alpha * p
    res -= alpha * ap
    res_old = res_new
    
print "Number of iterations = %d" % it
cs = plt.contourf(X,Y,u)
plt.colorbar(cs);