头文件

import sys
import copy
import numpy as np
import networkx as nx
from scipy.sparse import spdiags, tril, triu, coo_matrix, csr_matrix

gauss_seidel

# A csr_matrix 
def gauss_seidel(A, b, x0 = None, maxiter=PRE_SMOOTH_ITER):
    # x0 = spsolve(A, b)
    if x0 is None:
        x0 =  np.array( [0.0]*A.shape[0]  )
    
    for _ in range(maxiter):
        n, _, indptr, indices, data = A.shape[0], A.nnz, A.indptr, A.indices, A.data
        for row in range(n):
            aii, sumi = 0.0, 0.0
            for j in range(indptr[row], indptr[row+1]):
                if row == indices[j]:
                    aii = 1.0/data[j]
                else:
                    sumi += data[j] * x0[ indices[j] ]
            x0[row] = (aii * (b[row] - sumi) )
    return x0

jacobi

# A csr_matrix
def jacobi(A, b, x0 = None, maxiter=PRE_SMOOTH_ITER):
    if x0 is None:
        x0 =  np.array( [0.0]*A.shape[0]  )

    D = spdiags(A.diagonal(), np.array([0]), A.shape[0], A.shape[1]) # default with <class 'scipy.sparse.dia.dia_matrix'> 

    L = tril(A) - D
    U = triu(A) - D 

    dia = generate_dia_inv(D)
    B = -( dia ).dot(L+U)
    f = ( dia ).dot(b)

    for _ in range(maxiter):
        x0 = B.dot(x0) + f
    return x0
def generate_dia_inv(mat):
    res_dia = mat.diagonal()
    res_dia = 1.0/res_dia
    res_mat = csr_matrix((res_dia.shape[0], res_dia.shape[0]))
    res_mat.setdiag(res_dia)
    return res_mat