如何获取Scipy中线性算子的对角线？CG法预条件需求

要获取LinearOperator对应的矩阵对角线，核心逻辑很直接：矩阵A的第i个对角线元素，就是A作用在第i个标准基向量（仅第i位为1，其余为0）上的结果的第i个分量。不需要显式构造完整矩阵，直接利用LinearOperator的matvec/matmat方法就能完成计算。

方法1：循环计算（适合中小规模问题） #

逐个生成标准基向量，调用matvec后提取对应分量：

import numpy as np
from scipy.sparse.linalg import LinearOperator

def get_operator_diag(operator):
    n_rows = operator.shape[0]
    diag = np.zeros(n_rows)
    for i in range(n_rows):
        # 生成第i个标准基向量
        e = np.zeros(n_rows)
        e[i] = 1.0
        # 计算A*e，取第i个分量作为对角线元素
        diag[i] = operator.matvec(e)[i]
    return diag

方法2：稀疏矩阵批量计算（适合大规模问题） #

用稀疏单位矩阵一次性生成所有标准基向量，通过matmat批量计算后提取对角线，效率比循环更高：

from scipy.sparse import eye

def get_operator_diag_sparse(operator):
    n_rows = operator.shape[0]
    # 生成稀疏单位矩阵，每一列对应一个标准基向量
    sparse_eye = eye(n_rows, format='csr')
    # 计算A乘以单位矩阵，结果就是A的所有列
    A_columns = operator.matmat(sparse_eye)
    # 直接提取结果的对角线
    return A_columns.diagonal()

构造预条件子 #

拿到对角线后，不用显式构造逆矩阵，直接用LinearOperator定义预条件子（作用是逐元素除以对角线值）：

# 获取A的对角线
diag_A = get_operator_diag(A_operator)
# 构造预条件子M，M(x) = x ./ diag_A
precond = LinearOperator(A_operator.shape, matvec=lambda x: x / diag_A)

之后调用CG时传入预条件子即可：

from scipy.sparse.linalg import cg

x, info = cg(A_operator, b, M=precond)

额外优化建议 #

如果你的算子A有特定结构（比如有限差分、有限元离散得到的矩阵），可以直接通过数学推导写出对角线元素的表达式，这比数值计算更快，也更节省内存。

内容的提问来源于stack exchange，提问作者lrisley