import numpy as np
from matplotlib import pyplot as plt

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

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

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

TOL   = 1.0e-6
itmax = 100

u   = np.zeros(n+1)
p   = np.zeros(n+1)
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] = 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)
    print it, res_new
    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 / p.dot(ap)
    u += alpha * p
    res -= alpha * ap
    res_old = res_new

print "Number of iterations = %d" % it
plt.plot(x,ue,x,u)
plt.legend(("Exact","Numerical"));