具有重复索引的最快添加：np.add.at/sparse.csr_matrix？

说我有一个num_indices * n指数矩阵（在range(m)) and a num_indices * n值矩阵，即

m, n = 100, 50
num_indices = 100000
indices = np.random.randint(0, m, (num_indices, n))
values = np.random.uniform(0, 1, (num_indices, n))

我想得到一个形状的结果矩阵m * n，这样每一列都会被索引并在索引和值矩阵中的相应列上求和。以下是一些实现：

• 基本的for循环

tic = time.time()
res1 = np.zeros((m, n))
for i in range(num_indices):
    for j in range(n):
        res1[indices[i, j], j] += values[i, j]
print(f'element-wise for loop used {time.time() - tic:.5f}s')

• np.add.at，多维

tic = time.time()
res2 = np.zeros((m, n))
np.add.at(res2, (indices, np.arange(n)[None, :]), values)
print(f'np.add.at used {time.time() - tic:.5f}s')

• np.add.at，带有 for 循环

tic = time.time()
reslst = []
for colid in range(n):
    rescol = np.zeros((m, 1))
    np.add.at(rescol, (indices[:, colid], np.zeros(num_indices, dtype=int)), values[:, colid])
    reslst.append(rescol)
res3 = np.hstack(reslst)
print(f'np.add.at with loop used {time.time() - tic:.5f}s')

• scipy.sparse，多维

from scipy import sparse
tic = time.time()
res4 = sparse.csr_matrix((values.ravel(), (indices.ravel(), np.tile(np.arange(n), num_indices))), shape=(m, n)).toarray()
print(f'sparse.csr used {time.time() - tic:.5f}s')

• scipy.sparse，带有 for 循环

tic = time.time()
reslst = []
for colid in range(n):
    rescol = sparse.csr_matrix((values[:, colid], (indices[:, colid], np.zeros(num_indices, dtype=int))), shape=(m, 1)).toarray()
    reslst.append(rescol)
res5 = np.hstack(reslst)
print(f'sparse.csr with loop used {time.time() - tic:.5f}s')

结果都是一样的：

assert np.isclose(res2, res1).all()
assert np.isclose(res3, res1).all()
assert np.isclose(res4, res1).all()
assert np.isclose(res5, res1).all()

然而，速度很奇怪，而且不令人满意：

for loop used 1.69889s
np.add.at used 0.17287s
np.add.at with loop used 0.47767s
sparse.csr used 0.16847s
sparse.csr with loop used 0.09845s

我的基本问题是：

• Why is np.add.at这么慢 - 甚至比scipy.sparse？我认为 numpy 会受益匪浅，尤其是在多维情况下。
• For scipy.sparse，为什么多维比for循环还慢？没有使用并发吗？ （为什么对于 numpy 来说，多维更快？）
• 总的来说，有没有更快的方法来实现我想要的？

令人惊讶的是，事实证明np.add.at在 Numpy 中实现效率非常低.

关于scipy.sparse，我无法在我的机器上使用 Scipy 1.7.1 重现相同的性能改进：它只是快一点。就像 Numpy 的np.add.at，还远远谈不上快。

您可以使用以下命令有效地实现此操作Numba:

import numba as nb

@nb.njit('(int_[:,::1], float64[:,::1], int_, int_)')
def compute(indices, values, n, m):
    res = np.zeros((m, n), dtype=np.float64)
    for i in range(num_indices):
        for j in range(n):
            #assert 0 <= indices[i, j] < m
            res[indices[i, j], j] += values[i, j]
    return res

tic = time.time()
result = compute(indices, values, n, m)
print(f'Numba used {time.time() - tic:.5f}s')

请注意，该函数假设indices包含有效值（即没有越界）。否则，该函数可能会崩溃或默默地损坏程序。如果您不确定，您可以启用断言，但代价是计算速度较慢（在我的机器上慢两倍）。

请注意，只要满足以下条件，此实现就会很快8 * n * m足够小，所以输出数组适合L1缓存（一般为 16 KB 到 64 KB）。否则，由于随机访问模式效率低下，它可能会慢一些。如果输出数组不适合 L2 缓存（通常为几百 KB），那么速度会明显变慢。

Results

element-wise for loop used 2.45158s
np.add.at used 0.28600s
sparse.csr used 0.19600s
sparse.csr with loop used 0.18900s
Numba used 0.00500s

Thus, Numba 大约快 40 倍比最快的实现快，比最慢的实现快约 500 倍。 numba 代码速度要快得多，因为与 Numpy 和 Scipy 相比，代码被编译为优化的本机二进制文件，没有额外的开销（即边界检查、临时数组、函数调用）。