超松弛迭代法

import numpy as np
def sor_method(A, b, omega, x0=None, tol=1e-6, max_iter=1000):
    """
    超松弛迭代法 (SOR) 求解线性方程组 Ax = b

    参数：
    A : np.ndarray, n x n 的系数矩阵
    b : np.ndarray, n x 1 的常数向量
    omega : float, 超松弛因子 (0 < omega < 2)
    x0 : np.ndarray, 初始解向量，默认全零
    tol : float, 收敛容差，默认1e-6
    max_iter : int, 最大迭代次数，默认1000

    返回：
    x : np.ndarray, 方程组的近似解
    k : int, 迭代次数
    """
    n = len(b)
    # 如果未提供初始解，则初始化为零向量
    if x0 is None:
        x = np.zeros(n)
    else:
        x = x0.copy()
    # 检查输入合法性
    if A.shape[0] != A.shape[1] or A.shape[0] != n:
        raise ValueError("矩阵A和向量b的维度不匹配")
    if not (0 < omega < 2):
        raise ValueError("超松弛因子omega必须在(0, 2)范围内")

    # 迭代过程
    for k in range(max_iter):
        x_old = x.copy()  # 保存上一轮迭代结果
        print(f"\n第{k+1}次迭代开始:")
        print(f"上一轮解向量 x_old = {x_old}")
        # 对每个变量进行更新
        for i in range(n):
            # 计算 sigma1: 已更新的变量部分 (j < i)
            sigma1 = np.dot(A[i, :i], x[:i])
            # 计算 sigma2: 未更新的变量部分 (j > i)
            sigma2 = np.dot(A[i, i + 1:], x_old[i + 1:])
            # SOR公式
            x[i] = (1 - omega) * x_old[i] + (omega / A[i, i]) * (b[i] - sigma1 - sigma2)
            print(f"更新 x[{i}] = {x[i]} (sigma1 = {sigma1}, sigma2 = {sigma2})")

        print(f"当前解向量 x = {x}")
        # 检查收敛性
        if np.linalg.norm(x - x_old, ord=np.inf) < tol:
            print(f"收敛，迭代次数 = {k + 1}")
            return x, k + 1  # 返回解和迭代次数

    # 如果超过最大迭代次数仍未收敛
    print("警告：达到最大迭代次数，未完全收敛")
    return x, max_iter

# 测试代码
if __name__ == "__main__":
    A = np.array([[4, 3, 0],
                  [3, 4, -1],
                  [0, -1, 4]], dtype=float)
    b = np.array([24, 30, -24], dtype=float)
    # 参数设置
    omega = 1.2  # 超松弛因子
    x0 = np.zeros(3)  # 初始值
    # 调用SOR方法
    solution, iterations = sor_method(A, b, omega, x0)
    # 输出结果
    print(f"解向量 x = {solution}")
    print(f"迭代次数 = {iterations}")
    # 验证结果
    print(f"验证 Ax - b = {np.dot(A, solution) - b}")