source: inundation/ga/storm_surge/pyvolution/cg_solve.py @ 433

def conjugate_gradient(A,b,x0,imax=2000,tol=1.0e-8,iprint=0):
   #
   # Try to solve linear equation Ax = b using
   # conjugate gradient method
   #
   # Input
   # A: matrix or function which applies a matrix, assumed symmetric
   # b: right hand side
   # x0: inital guess
   # imax: max number of iterations
   # tol: tolerance used for residual
   #
   # Output
   # x: approximate solution
   # i: number of iterations
   #

   #if nargin<3
   #   x0 = zeros(b);
   #end

   from Numeric import dot
   #from operator import mod

   print "A shape",A.shape
   print "b shape",b.shape
   print "x0 shape",x0.shape
   i=0
   x = x0
   r = b - A*x
   d = r
   rTr = dot(r,r)
   rTr0 = rTr

   while (i<imax and rTr>tol**2*rTr0):
       q = A*d
       alpha = rTr/dot(d,q)
       x = x + alpha*d
       if i%50 :
           r = b - A*x
       else:
           r = r - alpha*q
       rTrOld = rTr
       rTr = dot(r,r)
       bt = rTr/rTrOld

       d = r + bt*d
       i = i+1
       print "i",i
       #if i%iprint :
       #   #print 'i = %g rTr = %20.15e \n': i,rTr
       #   pass
   if i==imax:
     print 'max number of iterations attained'

   return x