Conjugate Gradient Method(CG法) 共役勾配法 を Python と numpy でやってみた

cgm.py

import numpy as np
#loop max: 200
#convergence condition: norm() < 1.0e-10

###
# Solve Ax = b
# A : real, symmetric, positive-definite matrix
# b : A*x_init
# x_init : Initial vector
###
def cgm(A, b, x_init):
    x = x_init
    r0 = b - np.dot(A,x)
    p = r0
    k = 0
    for i in range(200):
        a = float( np.dot(r0.T,r0) / np.dot(np.dot(p.T, A),p) )
        x = x + p*a
        r1 = r0 - np.dot(A*a, p)
        print np.linalg.norm(r1)
        if np.linalg.norm(r1) < 1.0e-10:
            return x
        b = float( np.dot(r1.T, r1) / np.dot(r0.T, r0) )
        p = r1 + b * p
        r0 = r1
    return x

if __name__ == "__main__":
    #I want b = A "x"
    A = np.matrix((
        (1,0,3,0),
        (0,0,0,1),
        (3,0,0,2),
        (0,1,2,1)
        ))
    b = np.matrix([[6],[1],[4],[1]])
    x0 = np.matrix([[0],[0],[0],[0]]) #init x
    ans = cgm(A, b, x0)
    print ans
    b_a = np.dot(A,ans)
    print b_a