如何在Python中实现类似R中inla.qinv的稀疏精度矩阵逆对角线计算？

针对大型稀疏精度矩阵Q，计算其逆矩阵的对角线元素（即协方差矩阵的对角线，对应边际方差），以下是几种高效的实现方案，避免直接求逆或低效的线性系统求解：

一、优化现有sksparse.cholmod代码 #

你当前直接求解单位矩阵的方法等价于执行N次线性系统求解，时间复杂度极高。利用Rue & Held (2005)中提出的递归公式，可基于Cholesky分解的下三角因子L直接计算，时间复杂度仅为O(n + nnz(L))：

实现代码 #

import numpy as np
import scipy.sparse as sparse
from sksparse.cholmod import cholesky

# 执行Cholesky分解，获取下三角因子L
chol = cholesky(Q)
L = chol.L().tocsr()  # 转为CSR格式方便按行访问
n = L.shape[0]

# 预处理每个行的对角线元素
diag_vals = np.zeros(n)
for i in range(n):
    row_start, row_end = L.indptr[i], L.indptr[i+1]
    # 定位当前行的对角线元素位置
    diag_idx = np.argwhere(L.indices[row_start:row_end] == i)[0][0]
    diag_vals[i] = L.data[row_start + diag_idx]

# 从最后一行开始递归计算逆对角线元素
inv_diag = np.zeros(n)
inv_diag[-1] = 1.0 / (diag_vals[-1] ** 2)

for i in range(n-2, -1, -1):
    sum_term = 0.0
    row_start, row_end = L.indptr[i], L.indptr[i+1]
    # 遍历当前行中列索引大于i的非零元素
    for j in range(row_start, row_end):
        col = L.indices[j]
        if col > i:
            sum_term += (L.data[j] ** 2) * inv_diag[col]
    inv_diag[i] = (1.0 + sum_term) / (diag_vals[i] ** 2)

进一步加速：用Numba JIT编译 #

对递归循环部分用Numba编译，可大幅提升计算速度：

import numba

@numba.jit(nopython=True)
def compute_inv_diag(L_indptr, L_indices, L_data, diag_vals, n):
    inv_diag = np.zeros(n)
    inv_diag[-1] = 1.0 / (diag_vals[-1] ** 2)
    for i in range(n-2, -1, -1):
        sum_term = 0.0
        row_start = L_indptr[i]
        row_end = L_indptr[i+1]
        for j in range(row_start, row_end):
            col = L_indices[j]
            if col > i:
                sum_term += (L_data[j] ** 2) * inv_diag[col]
        inv_diag[i] = (1.0 + sum_term) / (diag_vals[i] ** 2)
    return inv_diag

# 调用加速函数
inv_diag = compute_inv_diag(L.indptr, L.indices, L.data, diag_vals, n)

二、直接调用GMRFLib（INLA底层库） #

R-INLA的inla.qinv基于GMRFLib实现，你可以直接调用GMRFLib的C函数，通过ctypes包装为Python可用接口：

步骤1：编译GMRFLib #

从INLA官方渠道获取GMRFLib源码，编译为动态库（如libgmrf.so或gmrf.dll）。

步骤2：Python调用示例 #

import ctypes
import numpy as np
import scipy.sparse as sparse

# 加载编译好的GMRFLib库
libgmrf = ctypes.CDLL('./libgmrf.so')

# 定义函数参数类型
libgmrf.gmrf_qinv.argtypes = [
    ctypes.c_int,
    ctypes.POINTER(ctypes.c_int),
    ctypes.POINTER(ctypes.c_int),
    ctypes.POINTER(ctypes.c_double),
    ctypes.c_int,
    ctypes.POINTER(ctypes.c_double)
]
libgmrf.gmrf_qinv.restype = ctypes.c_int

# 将稀疏矩阵转为COO格式
Q_coo = Q.tocoo()
n = Q.shape[0]
nnz = Q_coo.nnz

# 转换为ctypes兼容的数组
rows = np.array(Q_coo.row, dtype=np.int32)
cols = np.array(Q_coo.col, dtype=np.int32)
vals = np.array(Q_coo.data, dtype=np.float64)
inv_diag = np.zeros(n, dtype=np.float64)

# 调用GMRFLib的qinv函数
status = libgmrf.gmrf_qinv(
    ctypes.c_int(n),
    rows.ctypes.data_as(ctypes.POINTER(ctypes.c_int)),
    cols.ctypes.data_as(ctypes.POINTER(ctypes.c_int)),
    vals.ctypes.data_as(ctypes.POINTER(ctypes.c_double)),
    ctypes.c_int(nnz),
    inv_diag.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
)

if status != 0:
    raise RuntimeError("GMRFLib qinv计算失败")

三、其他可选工具 #

• PyAMG：代数多重网格库，可通过预处理矩阵快速估算逆对角线元素，适合超大规模矩阵的近似计算。
• sparse库：提供比scipy.sparse更高效的稀疏矩阵操作，结合其Cholesky分解工具可简化逆对角线计算流程。

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