最快添加重复索引：np.add.at / sparse.csr_matrix？
Fastest add with repeated indices: np.add.at / sparse.csr_matrix?
假设我有一个 num_indices * n 索引矩阵（在 range(m) 中）和一个 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
我的基本问题是：
- 为什么
np.add.at 这么慢 - 甚至比 scipy.sparse 还慢？我虽然 numpy 会受益很多，尤其是在多维情况下。
- 对于
scipy.sparse，为什么多维比for循环还要慢？没有使用并发吗？ （以及为什么对于 numpy，多维更快？）
- 总的来说，有没有更快的方法来实现我想要的？
复制这个以更好地理解正在发生的事情：
In [48]: 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))
add.at案例
In [49]: res2 = np.zeros((m, n))
...: np.add.at(res2, (indices, np.arange(n)[None, :]), values)
In [50]: res2
Out[50]:
array([[472.25600004, 526.36528387, 493.10320355, ..., 548.43850955,
497.3412013 , 508.51132267],
...])
In [51]: %%timeit
...: res2 = np.zeros((m, n))
...: np.add.at(res2, (indices, np.arange(n)[None, :]), values)
613 ms ± 11.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
稀疏方法
In [52]: from scipy import sparse
In [53]: res4 = sparse.csr_matrix((values.ravel(), (indices.ravel(), np.tile(np.
...: arange(n), num_indices))), shape=(m, n)).toarray()
In [54]: res4
Out[54]:
array([[472.25600004, 526.36528387, 493.10320355, ..., 548.43850955,
497.3412013 , 508.51132267],
[....)
值是相同的，并且它们没有任何病态（即不全是 0。只是很多浮点数的总和。）
确实，稀疏的时间好一点：
In [55]: timeit res4 = sparse.csr_matrix((values.ravel(), (indices.ravel(), np.t
...: ile(np.arange(n), num_indices))), shape=(m, n)).toarray()
402 ms ± 22.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
通常我发现创建稀疏矩阵比创建密集矩阵花费的时间更长，尤其是当大小仅为 (100,50) 时。从某种意义上说，它确实需要很长时间，半秒。但它比 add.at.
短
用 1 替换 values，我确认每个数组元素平均有 1000 个重复项（我可以从 res2 值及其 (0,1) 随机分布中推断出这一点).
In [57]: res5 = sparse.csr_matrix((np.ones_like(indices.ravel()), (indices.ravel
...: (), np.tile(np.arange(n), num_indices))), shape=(m, n)).toarray()
In [58]: res5
Out[58]:
array([[ 971, 1038, 1004, ..., 1056, 988, 1030],
[1016, 992, 949, ..., 994, 982, 1011],
[ 984, 1014, 980, ..., 1001, 1053, 1057],
...,])
如果我正确地回忆起过去的稀疏代码探索，使用这些输入，csr_matrix 首先创建一个 coo 矩阵，然后按词法对索引进行排序 - 即按行排序，并在列中排序。这将重复的索引放在一起，可以很容易地求和。考虑到大量重复项，这相对有效也就不足为奇了。
随机性就在行索引中。 sparse_with_loop 通过迭代 n, 50 来利用这一点。因此创建每个 (100,1) 稀疏矩阵更容易；再次对这些行索引进行排序并求和。
In [61]: %%timeit
...: reslst = []
...: for colid in range(n):
...: rescol = sparse.csr_matrix((values[:, colid], (indices[:, colid], n
...: p.zeros(num_indices, dtype=int))), shape=(m, 1)).toarray()
...: reslst.append(rescol)
...: res5 = np.hstack(reslst)
258 ms ± 452 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
我可能会补充说 np.add.at，虽然比完整的 python 循环更快，但通常比它更正的缓冲 += 慢。
In [72]: %%timeit
...: res10 = np.zeros((m,n))
...: res10[indices, np.arange(n)] += values
103 ms ± 55.1 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
为了正确看待这些时间，生成 values 数组的时间是：
In [75]: timeit values = np.random.uniform(0, 1, (num_indices, n))
97.5 ms ± 90.4 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
大约是 add.at 时间的 1/6。以及从中创建稀疏矩阵的时间：
In [76]: timeit sparse.csr_matrix(np.random.uniform(0, 1, (num_indices, n)))
379 ms ± 4.51 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
因此执行索引加法大约 300 毫秒的时间并不少见。以任何方式处理 (m,n) 形数组都需要时间。
res5 的情况下，重复计数可以用 bincount:
In [88]: np.array([np.bincount(indices[:,i]) for i in range(n)]).T
Out[88]:
array([[ 971, 1038, 1004, ..., 1056, 988, 1030],
[1016, 992, 949, ..., 994, 982, 1011],
[ 984, 1014, 980, ..., 1001, 1053, 1057],
...,])
In [89]: timeit np.array([np.bincount(indices[:,i]) for i in range(n)]).T
64.7 ms ± 100 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
编辑
有趣的是，如果我做 coo_matrix 时间会快很多：
In [95]: timeit res4 = sparse.coo_matrix((values.ravel(), (indices.ravel(), np.t
...: ile(np.arange(n), num_indices))), shape=(m, n)).toarray()
78.3 ms ± 46.3 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
它仍在执行重复项总和步骤，但作为 to_array() 步骤的一部分。但它不必做创建 csr 矩阵的额外工作。
编辑
我只记得之前的 add.at SO 使用了 bincount 和 weights:
In [105]: res0 = np.array([np.bincount(indices[:,i],weights=values[:,i]) for i i
...: n range(n)]).T
In [106]: np.allclose(res2,res0)
Out[106]: True
In [107]: timeit res0 = np.array([np.bincount(indices[:,i],weights=values[:,i])
...: for i in range(n)]).T
142 ms ± 133 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
令人惊讶的是，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），那么这将显着变慢。
结果
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
因此，Numba 比最快的实现快 40 倍，比最慢的实现快约 500 倍。与 Numpy 和 Scipy.[=17= 相比，numba 代码更快，因为代码被编译为优化的本机二进制文件，没有额外开销（即 bound-checking、临时数组、函数调用） ]
假设我有一个 num_indices * n 索引矩阵（在 range(m) 中）和一个 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
我的基本问题是：
- 为什么
np.add.at这么慢 - 甚至比scipy.sparse还慢？我虽然 numpy 会受益很多，尤其是在多维情况下。 - 对于
scipy.sparse，为什么多维比for循环还要慢？没有使用并发吗？ （以及为什么对于 numpy，多维更快？） - 总的来说，有没有更快的方法来实现我想要的？
复制这个以更好地理解正在发生的事情：
In [48]: 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))
add.at案例
In [49]: res2 = np.zeros((m, n))
...: np.add.at(res2, (indices, np.arange(n)[None, :]), values)
In [50]: res2
Out[50]:
array([[472.25600004, 526.36528387, 493.10320355, ..., 548.43850955,
497.3412013 , 508.51132267],
...])
In [51]: %%timeit
...: res2 = np.zeros((m, n))
...: np.add.at(res2, (indices, np.arange(n)[None, :]), values)
613 ms ± 11.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
稀疏方法
In [52]: from scipy import sparse
In [53]: res4 = sparse.csr_matrix((values.ravel(), (indices.ravel(), np.tile(np.
...: arange(n), num_indices))), shape=(m, n)).toarray()
In [54]: res4
Out[54]:
array([[472.25600004, 526.36528387, 493.10320355, ..., 548.43850955,
497.3412013 , 508.51132267],
[....)
值是相同的，并且它们没有任何病态（即不全是 0。只是很多浮点数的总和。）
确实，稀疏的时间好一点：
In [55]: timeit res4 = sparse.csr_matrix((values.ravel(), (indices.ravel(), np.t
...: ile(np.arange(n), num_indices))), shape=(m, n)).toarray()
402 ms ± 22.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
通常我发现创建稀疏矩阵比创建密集矩阵花费的时间更长，尤其是当大小仅为 (100,50) 时。从某种意义上说，它确实需要很长时间，半秒。但它比 add.at.
用 1 替换 values，我确认每个数组元素平均有 1000 个重复项（我可以从 res2 值及其 (0,1) 随机分布中推断出这一点).
In [57]: res5 = sparse.csr_matrix((np.ones_like(indices.ravel()), (indices.ravel
...: (), np.tile(np.arange(n), num_indices))), shape=(m, n)).toarray()
In [58]: res5
Out[58]:
array([[ 971, 1038, 1004, ..., 1056, 988, 1030],
[1016, 992, 949, ..., 994, 982, 1011],
[ 984, 1014, 980, ..., 1001, 1053, 1057],
...,])
如果我正确地回忆起过去的稀疏代码探索，使用这些输入，csr_matrix 首先创建一个 coo 矩阵，然后按词法对索引进行排序 - 即按行排序，并在列中排序。这将重复的索引放在一起，可以很容易地求和。考虑到大量重复项，这相对有效也就不足为奇了。
随机性就在行索引中。 sparse_with_loop 通过迭代 n, 50 来利用这一点。因此创建每个 (100,1) 稀疏矩阵更容易；再次对这些行索引进行排序并求和。
In [61]: %%timeit
...: reslst = []
...: for colid in range(n):
...: rescol = sparse.csr_matrix((values[:, colid], (indices[:, colid], n
...: p.zeros(num_indices, dtype=int))), shape=(m, 1)).toarray()
...: reslst.append(rescol)
...: res5 = np.hstack(reslst)
258 ms ± 452 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
我可能会补充说 np.add.at，虽然比完整的 python 循环更快，但通常比它更正的缓冲 += 慢。
In [72]: %%timeit
...: res10 = np.zeros((m,n))
...: res10[indices, np.arange(n)] += values
103 ms ± 55.1 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
为了正确看待这些时间，生成 values 数组的时间是：
In [75]: timeit values = np.random.uniform(0, 1, (num_indices, n))
97.5 ms ± 90.4 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
大约是 add.at 时间的 1/6。以及从中创建稀疏矩阵的时间：
In [76]: timeit sparse.csr_matrix(np.random.uniform(0, 1, (num_indices, n)))
379 ms ± 4.51 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
因此执行索引加法大约 300 毫秒的时间并不少见。以任何方式处理 (m,n) 形数组都需要时间。
res5 的情况下，重复计数可以用 bincount:
In [88]: np.array([np.bincount(indices[:,i]) for i in range(n)]).T
Out[88]:
array([[ 971, 1038, 1004, ..., 1056, 988, 1030],
[1016, 992, 949, ..., 994, 982, 1011],
[ 984, 1014, 980, ..., 1001, 1053, 1057],
...,])
In [89]: timeit np.array([np.bincount(indices[:,i]) for i in range(n)]).T
64.7 ms ± 100 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
编辑
有趣的是，如果我做 coo_matrix 时间会快很多：
In [95]: timeit res4 = sparse.coo_matrix((values.ravel(), (indices.ravel(), np.t
...: ile(np.arange(n), num_indices))), shape=(m, n)).toarray()
78.3 ms ± 46.3 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
它仍在执行重复项总和步骤，但作为 to_array() 步骤的一部分。但它不必做创建 csr 矩阵的额外工作。
编辑
我只记得之前的 add.at SO 使用了 bincount 和 weights:
In [105]: res0 = np.array([np.bincount(indices[:,i],weights=values[:,i]) for i i
...: n range(n)]).T
In [106]: np.allclose(res2,res0)
Out[106]: True
In [107]: timeit res0 = np.array([np.bincount(indices[:,i],weights=values[:,i])
...: for i in range(n)]).T
142 ms ± 133 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
令人惊讶的是，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），那么这将显着变慢。
结果
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
因此，Numba 比最快的实现快 40 倍，比最慢的实现快约 500 倍。与 Numpy 和 Scipy.[=17= 相比，numba 代码更快，因为代码被编译为优化的本机二进制文件，没有额外开销（即 bound-checking、临时数组、函数调用） ]