为什么添加 scipy.sparse.dia_matrix 的实例这么慢？

Question

我正在编写一个数字代码，我使用的是scipy.sparse.dia_matrix。我的矩阵非常大（大约 1000000 x 1000000），但非常稀疏。有时是三对角线，有时是更多的对角线。

出于各种原因，从编码的角度来看，将这些矩阵中的几个（当然大小相同）相加是非常方便和清晰的。但是，我发现添加这些稀疏矩阵非常慢。以下示例说明了我的意思：

import numpy as np
from scipy.sparse import diags, dia_matrix

N = 100000

M1 = diags(diagonals = [np.random.random(N-1), np.random.random(N), np.random.random(N-1)], offsets = [-1, 0, 1])
M2 = diags(diagonals = [np.random.random(N-1), np.random.random(N), np.random.random(N-1)], offsets = [-1, 0, 1])
M3 = diags(diagonals = [np.random.random(N-1), np.random.random(N), np.random.random(N-1)], offsets = [-1, 0, 1])

def simple_add():
    M = M1 + M2 + M3
    
def complicated_add():
    M_ = dia_matrix((N, N))
    for d in [-1, 0, 1]:
        M_.setdiag(M1.diagonal(d) + M2.diagonal(d) + M3.diagonal(d), d)

%timeit simple_add()

%timeit complicated_add()

计时的输出是：

16.9 ms ± 730 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
959 µs ± 39.1 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

我不明白为什么将矩阵与 + 运算符相加比创建空对角矩阵并显式设置对角线要慢 17 倍。我能做些什么来加快速度吗？我更喜欢使用 + 运算符保留更简单的表达式，因为它更具可读性，但不会以计算时间增加一个数量级为代价。

更新：

我提议对 Scipy 进行更改，以更快地添加两个 dia_matrix 实例，经过一番讨论后，我向 Scipy 提交了一个拉取请求，该请求现已合并。所以在未来，添加两个dia_matrix 实例将不再转换为csr_matrix。

https://github.com/scipy/scipy/pull/14004

Answer 1

_add_sparse(self, other) 的实现是返回self.tocsr()._add_sparse(other)。额外的时间是把它变成一个CSR矩阵（它有一个C扩展用于加法）。

你能创建一个稀疏矩阵来满足你的需求吗？大概吧。

from scipy.sparse import dia_matrix, isspmatrix_dia

class dia_matrix_adder(dia_matrix):

    def __add__(self, other):

        if not isspmatrix_dia(other):
            return super(dia_matrix_adder, self).__add__(other)

        M_ = dia_matrix((self.shape[0], self.shape[1]))

        for d in [-1, 0, 1]:
            M_.setdiag(self.diagonal(d) + other.diagonal(d), d)

        return M_

我可能不会那样做，只是给自己写一个函数：

def add_dia_matrix(*mats):

    if len(mats) == 1:
        return mats[0]

    M_ = dia_matrix((mats[0].shape[0], mats[0].shape[1]))

    for d in [-1, 0, 1]:
        M_diag = mats[0].diagonal(d).copy()

        for i in range(1, len(mats)):
            M_diag += mats[i].diagonal(d)

        M_.setdiag(M_diag, d)

    return M_

这应该像一堆+ 一样可读，而不必处理新的类。

%timeit simple_add()
30.3 ms ± 218 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit complicated_add()
1.28 ms ± 2.63 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit add_dia_matrix(M1, M2, M3)
1.22 ms ± 4.14 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Answer 2

diags 从输入列表中生成dia_matrix：

In [84]: M=sparse.diags([np.arange(1,4),np.arange(1,5),np.arange(1,4)], offsets=[-1,0,1])
In [85]: M
Out[85]: 
<4x4 sparse matrix of type '<class 'numpy.float64'>'
    with 10 stored elements (3 diagonals) in DIAgonal format>
In [86]: M.offsets
Out[86]: array([-1,  0,  1], dtype=int32)
In [87]: M.data
Out[87]: 
array([[1., 2., 3., 0.],
       [1., 2., 3., 4.],
       [0., 1., 2., 3.]])

对角线列表（不同长度）已转换为 2 数组，带有偏移量。这主要用作输入格式。大多数（如果不是全部）数学以csr 格式实现。即使在那里，matrix_multiplication 也是相对强项。逐元素数学明显不如 numpy 数组等效项。

In [89]: Mr=M.tocsr()
In [90]: Mr
Out[90]: 
<4x4 sparse matrix of type '<class 'numpy.float64'>'
    with 10 stored elements in Compressed Sparse Row format>
In [91]: Mr.data
Out[91]: array([1., 1., 1., 2., 2., 2., 3., 3., 3., 4.])
In [92]: Mr.indices
Out[92]: array([0, 1, 0, 1, 2, 1, 2, 3, 2, 3], dtype=int32)
In [93]: Mr.indptr
Out[93]: array([ 0,  2,  5,  8, 10], dtype=int32)

如果偏移量和形状都相同，dia 格式建议更快的添加。

In [94]: M.data += M.data + M.data
In [95]: M.data
Out[95]: 
array([[ 3.,  6.,  9.,  0.],
       [ 3.,  6.,  9., 12.],
       [ 0.,  3.,  6.,  9.]])
In [96]: M.A
Out[96]: 
array([[ 3.,  3.,  0.,  0.],
       [ 3.,  6.,  6.,  0.],
       [ 0.,  6.,  9.,  9.],
       [ 0.,  0.,  9., 12.]])

对于任何稀疏格式，如果所有参数和输出的稀疏度都相同，您通常可以直接对 data 属性进行数学运算，保持隐含的 0 不变。