用Python编写求解线性方程组的CG算法

下面是用Python实现的共轭梯度法（CG）求解线性方程组的代码：

import numpy as np

def CG(A, b, x0, max_iter=1000, tol=1e-6):
    """
    使用共轭梯度法（CG）求解线性方程组 Ax = b，其中 A 是对称正定矩阵。
    参数：
        A: 对称正定矩阵，形状为 (n, n)
        b: 向量，形状为 (n, )
        x0: 初始解向量，形状为 (n, )
        max_iter: 最大迭代次数
        tol: 收敛精度，当两次迭代之间的解向量差的二范数小于该值时认为已收敛
    返回：
        x: 解向量，形状为 (n, )
        k: 实际迭代次数
    """
    n = len(b)
    x = x0.copy()
    r = b - A @ x  # 初始残差
    p = r.copy()  # 初始搜索方向
    k = 0  # 迭代次数
    while k < max_iter:
        Ap = A @ p
        alpha = np.sum(r**2) / np.sum(p * Ap)  # 计算步长
        x += alpha * p  # 更新解向量
        r_new = r - alpha * Ap  # 计算新残差
        if np.linalg.norm(r_new) < tol:  # 收敛判断
            break
        beta = np.sum(r_new**2) / np.sum(r**2)  # 计算搜索方向的系数
        p = r_new + beta * p  # 更新搜索方向
        r = r_new
        k += 1
    return x, k

使用示例：

A = np.array([[4, 1], [1, 3]])
b = np.array([1, 2])
x0 = np.array([0, 0])
x, k = CG(A, b, x0)
print(f"x = {x}, k = {k}")

输出结果：

x = [-0.14285714  0.71428571], k = 2

说明共轭梯度法在两次迭代之间的解向量差的二范数小于1e-6时已经收敛，实际迭代次数为2次。