数值分析-线性方程组的迭代法

数值积分上机课写的 Python 代码，实现线性方程组的迭代法。

import numpy as np
from numpy.linalg import inv


def prep(A, b, X_0):
    if A.shape[0] != A.shape[1]:
        raise ValueError('A must be a square matirx', A)
    if b.shape[1] != 1:
        raise ValueError('B must be a column vector', b)
    if A.shape[0] != b.shape[0]:
        raise ValueError('A and b must have the same number of rows', A, b)
    if X_0.shape[1] != 1:
        raise ValueError('X_0 must be a column vector', X_0)
    if A.shape[0] != X_0.shape[0]:
        raise ValueError('X_0 and b must have the same number of rows', X_0, b)
        
    D = np.zeros_like(A)
    for i in range(A.shape[0]):
        D[i, i] = A[i, i]
    L = -np.tril(A, -1)
    U = -np.triu(A, 1)
    return(D, L, U)

def jacobi(A, b, X_0, err):
    D, L, U = prep(A, b, X_0)
    B = inv(D).dot(L + U)
    f = inv(D).dot(b)
    X_k = X_0
    while True:
        X = B.dot(X_k) + f
        if max(abs(X - X_k)) <= err:
            return X
        else:
            X_k = X
            
            
def gauss_seidel(A, b, X_0, err):
    D, L, U = prep(A, b, X_0)
    B = inv(D - L).dot(U)
    f = inv(D - L).dot(b)
    X_k = X_0
    while True:
        X = B.dot(X_k) + f
        if max(abs(X - X_k)) <= err:
            return X
        else:
            X_k = X


def sor(A, b, omega, true_X, X_0, err):
    D, L, U = prep(A, b, X_0)
    L_omega = np.identity(A.shape[0]) - omega * inv(D - omega * L).dot(A)
    f = omega * inv(D - omega * L).dot(b)
    X_k = X_0
    while True:
        X = L_omega.dot(X_k) + f
        if max(abs(true_X - X)) <= err:
            return X
        else:
            X_k = X