矩阵过大无法分配内存时如何使用sparse matrix生成csr矩阵

核心思路是直接构造稀疏对角矩阵，全程不生成任何大尺寸密集矩阵，内存占用仅和非零元素数量成正比。

实现代码 #

import numpy as np
from scipy.sparse import diags

nx = 50
ny = 50
nz = 50
num1 = 123
num2 = 345

n_total = nx * ny * nz
# 定义两个对角的偏移量：主对角线偏移为0，第二个对角偏移为nx*ny
offsets = [0, nx * ny]
# 定义每个偏移对应的数值序列
diag_values = [
    [num1] * n_total,
    [num2] * (n_total - nx * ny)
]
# 直接生成CSR格式的稀疏矩阵，无中间密集矩阵内存占用
csrout = diags(diag_values, offsets, shape=(n_total, n_total), format="csr")

说明 #

• 原方案报错的原因是np.diag会直接生成完整的密集二维数组，125000阶的密集矩阵存储需要116GiB以上内存，完全超出常规硬件容量。
• 上述方案用scipy.sparse.diags构造稀疏矩阵时，仅存储非零元素，总非零元素数仅为247500个，内存占用不足1MB，可正常运行。
• 生成的csrout和你原逻辑得到的矩阵数值完全一致，可直接使用。

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