共轭梯度法求解线性方程组-python实现

共轭梯度法求解线性方程组 Ax = b 的子程序，其中，A为系数矩阵，x为待求未知数向量，b为方程右端向量。
具体算法见
https://wenku.baidu.com/view/f5157b0c9a6648d7c1c708a1284ac850ac02047c.html?rec_flag=default
29页
这里需要注意的是共轭梯度法要求系数矩阵A是对称正定矩阵，也就意味着A不对称情况下没有稳定解。

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.])
'''
                      #输出结果