共轭梯度法求解线性方程组python代码

共轭梯度法是把求解线性方程组的问题转化为求解一个与之等价的二次函数极小值的问题。从任意给定的初始点出发，沿一组关于矩阵A的共轭方向进行线性搜索，在无舍入误差的假定下，最多迭代n次（其中n为矩阵A的阶数），就可求得二次函数的极小点，也就求得线性方程组Ax= b的解。其python代码：

import math
import numpy as np  
from numpy import * 
import matplotlib.pyplot as plt

def cg(A,b,x): #共轭梯度法
    r = b-np.dot(A,x)      #r=b-Ax         r也是是梯度方向
    p= np.copy(r)
    i=0
    while(max(abs(r))>1.e-10 and i < 100  ): 
        print('i',i)
        print('r',r)
        pap=np.dot(np.dot(A,p),p)
        if pap==0:                  #分母太小时跳出循环
            return x
        print('pap=',pap)
        alpha = np.dot(r,r)/pap   #直接套用公式
        x1 = x + alpha*p
        r1 = r-alpha*np.dot(A,p)
        beta = np.dot(r1,r1)/np.dot(r,r)
        p1 = r1 +beta*p
        r = r1
        x = x1
        p = p1
        i=i+1
    return x

   
n=10  #n 个解
A = np.array([                                         #系数矩阵，注意必须是正定对称矩阵
[1., 1., 0.,  0.,  0.,  0,  0.,  0.,  0.,  0], 
[1., 1., 1. , 0.,  0,   0., 0.,  0.,  0.,  0], 
[ 0., 1.,  1. , 1., 0.,  0., 0.,  0.,  0.,  0], 
[ 0.,  0., 1., 1., 1., 0., 0.,  0. , 0.,  0], 
[ 0.,  0.,  0., 1.,  1., 1.,0.,   0.,  0.,  0], 
[ 0.,  0.,  0.,  0.,  1., 1.,1., 0.,  0.,  0], 
[ 0.,  0. , 0. , 0. , 0., 1., 1., 1., 0., 0], 
[ 0.,  0.,  0. , 0. , 0., 0.,  1., 1.,  1.,0], 
[ 0.,  0. , 0.,  0,  0.,  0.,  0.,  1.,1., 1.], 
[ 0.,  0. , 0. , 0, 0.,   0.,  0.,  0.,  1.,1], 
])
b = np.array([2., 3., 3.,3., 3.,3., 3., 3., 3., 2.])       
x = np.zeros(n, dtype = float)
f=cg(A,b,x)                         #调用共轭斜量法
print('x=',f)      
'''
n = 6   #其他几个矩阵，测试用
A = np.array([
[-3., 1., 0. , 0. ,  0., 0.5 ], 
[1., -3., 1. , 0. , 0,   0. ], 
[ 0., 1.,-3., 1., 0. ,   0. ], 
[ 0.,  0. , 1., -3., 1.,   0.], 
[ 0.,  0. , 0. , 1.,-3.,  1.], 
[ 0.5,  0. , 0. , 0. , 1.,  -3.], 
]) 
b = np.array([2.5, 1.5, 1., 1., 1.5, 2.5])

n = 3
A = np.array([
[2.,0.,1.],
[0.,1.,0.],
[1.,0.,2.]
]) 
b = np.array([3.,1.,3.])
'''