Python: basic conjugate gradient linear solver

Visit Wikipedia:Conjugate gradient method for a clear and succinct explanation of the conjugate gradient method for solving linear equations. We based our Python code on the algorithm described in this article. We are pleased that we were able to get the same results for the simple example involving a 2 by 2 system,

# -*- coding: utf-8 -*-

"""
File     conjgrad.py
Author   Ernesto P. Adorio
         UPDEPP Clarfield
email    ernesto.adorio@gmail.com          
ref.     http://en.wikipedia.org/wiki/Conjugate_gradient_method
"""

from matlib import *

def conjgrad(A, b, x = None, ztol = 1.0e-4, maxiter = None):
    """
    Solves the equation Ax = b by iteration using conjugate gradients.
    Arguments 
      A- symmetric positive definite matrix.

    """
    print "input A=", A
    print "input b=", b
    print "input x=", x
    if x is None:
       x = [0] * len(b)
    m, n = matdim(A)
    if m != n or m!= len(b):
       raise ValueError
    if maxiter is None:
       maxiter = m    
    rk = vecsub(b, matvec(A, x))
    print "r_0 =", rk
    pk = rk[:]
    num    = dot(rk, rk)

    for k in range(1,maxiter):
        print "iteration #", k
        Apk    = matvec(A, pk)
        den    = dot(pk, Apk)
        alphak = num/float(den)
        print "num/den=", num, den, alphak

        #update x.
        x  = [e  + alphak * p for (e,p) in zip(x, pk)]
        print "x = ", x

        #update rk.
        rk    = [r - alphak * apk for r, apk in zip(rk, Apk)]
        dotrk = dot(rk, rk)
        if num < ztol:
           return x, k
        print "rkp1=", rk

        #update pk
        betak = dotrk/num
        num   = dotrk
        pk    = [r + betak* p for r, p in zip(rk, pk)] 
        print "betak=", betak
        print "pk=", pk
        print
    return x, maxiter  

def vec2mat(x, nrows, ncols):
    if nrows * ncols != len(x):
       raise ValueError,"incompatible dimensions in vec2mat()"
    M = []
    for i in range(nrows):
       M.append(x[i*ncols:(i+1)*ncols])

    return M   

     
if __name__ == "__main__":
   A = vec2mat([4,1,1,3],2,2)
   b = [1,2]
   x = [2,1]
   sol = conjgrad(A,b, x,ztol = 1.0e-4, maxiter=5)
   print "sol=", sol

We will soon add the vec2mat() function included in the code above in the matlib.py function. Please note that this author recommends using the numpy and scipy Python modules. They are well tested and provide a wrapper to the linear algebra routines developed in EISPACK and LINPACK. scientific libraries which were coded in C. Here are the results of running the above from the terminal:

toto@toto-Aspire-4520:~/Blogs/lm$ python conjgrad.py 
input A= [[4, 1], [1, 3]]
input b= [1, 2]
input x= [2, 1]
r_0 = [-8, -3]
iteration # 1
num/den= 73 331 0.220543806647
x =  [0.2356495468277946, 0.33836858006042303]
rkp1= [-0.2809667673716012, 0.7492447129909365]
betak= 0.00877136937414
pk= [-0.3511377223647101, 0.7229306048685207]

iteration # 2
num/den= 0.640309964312 1.55338036659 0.412204234122
x =  [0.09090909090909094, 0.6363636363636365]
rkp1= [5.551115123125783e-17, 0.0]
betak= 4.81249407751e-33
pk= [5.551115123125783e-17, 3.4790992543769425e-33]

iteration # 3
num/den= 3.08148791102e-33 1.23259516441e-32 0.25
x =  [0.09090909090909095, 0.6363636363636365]
sol= ([0.09090909090909095, 0.6363636363636365], 3)
toto@toto-Aspire-4520:~/Blogs/lm$         
 

Of course remove the print statements to avoid the computational progress reports. Our aim here is mainly pedagogical. The Python code is easier to read and one can use a debugger (as in Eric4) to walk through the code!