问题描述
从我的 earlier posts 开始，我注意到操作 np.add.at(A,indices,B) 比 A[indices] += B 慢很多。
def fast(A):
n = A.shape[0]
retval = np.zeros(2*n-1)
for i in range(n):
retval[slice(i,i+n)] += A[i,:]
return retval
def slow(A):
n = A.shape[0]
retval = np.zeros(2*n-1)
for i in range(n):
np.add.at(retval,slice(i,i+n),A[i,:])
return retval
def slower(A):
n = A.shape[0]
retval = np.zeros(2*n-1)
indices = np.arange(n)
indices = indices[:,None] + indices
np.add.at(retval,A) # bottleneck here
return retval
我的时间：
A = np.random.randn(10000,10000)
timeit(lambda: fast(A),number=10) # 0.8852798199995959
timeit(lambda: slow(A),number=10) # 56.633683917999406
timeit(lambda: slower(A),number=10) # 57.763389584000834
显然，使用 __iadd__ 要快得多。但是，np.add.at 的文档指出：
对“索引”指定的元素对操作数“a”执行无缓冲的就地操作。对于加法 ufunc，此方法等价于 a[indices] += b，只是对被索引多次的元素进行累加。
为什么 np.add.at 这么慢？
这个函数的用例是什么？
解决方法
add.at 适用于索引包含重复项且 += 不会产生预期结果的情况
In [44]: A = np.zeros(5,int); idx = np.array([0,1,2,3,3])
In [45]: A[idx]+=1
In [46]: A
Out[46]: array([1,0]) # the duplicates in idx are ignored
使用add.at：
In [47]: A = np.zeros(5,3])
In [48]: np.add.at(A,idx,1)
In [49]: A
Out[49]: array([1,4,0])
与显式迭代相同的结果：
In [50]: A = np.zeros(5,3])
In [51]: for i in idx: A[i]+=1
In [52]: A
Out[52]: array([1,0])
一些时间：
In [53]: %%timeit A = np.zeros(5,3])
...: A[idx]+=1
3.65 µs ± 13.2 ns per loop (mean ± std. dev. of 7 runs,100000 loops each)
In [54]: %%timeit A = np.zeros(5,3])
...: np.add.at(A,1)
6.47 µs ± 24.8 ns per loop (mean ± std. dev. of 7 runs,100000 loops each)
In [55]: %%timeit A = np.zeros(5,1)
...: for i in idx: A[i]+=1
15.6 µs ± 41.5 ns per loop (mean ± std. dev. of 7 runs,100000 loops each)
add.at 比 += 慢，但比 python 迭代好。
我们可以测试 A[:4]+1、A[:4]+=1 等变体
无论如何，如果您不需要，请不要使用 add.at 变体。
编辑
你的例子，简化了一点：
In [108]: x = np.zeros(2*10-1)
...: for i in range(10):
...: x[i:i+10] += 1
...:
In [109]: x
Out[109]:
array([ 1.,2.,3.,4.,5.,6.,7.,8.,9.,10.,1.])
因此，您正在向重叠切片添加值。我们可以用数组替换切片：
In [110]: x = np.zeros(2*10-1)
...: for i in range(10):
...: x[np.arange(i,i+10)] += 1
...:
In [111]: x
Out[111]:
array([ 1.,1.])
无需添加带有切片的“慢”案例 add.at，因为索引没有重复项。
创建所有索引。 += 因缓冲无法工作
In [112]: idx=np.arange(10); idx=(idx[:,None]+idx).ravel()
In [113]: y=np.zeros(2*10-1)
...: y[idx]+=1
In [114]: y
Out[114]:
array([1.,1.,1.])
add.at 解决了这个问题：
In [115]: y=np.zeros(2*10-1)
...: np.add.at(y,1)
In [116]: y
Out[116]:
array([ 1.,1.])
以及完整的python迭代：
In [117]: y=np.zeros(2*10-1)
...: for i in idx: y[i]+=1
In [118]:
In [118]: y
Out[118]:
array([ 1.,1.])
现在有些时间。
基线：
In [119]: %%timeit
...: x = np.zeros(2*10-1)
...: for i in range(10):
...: x[i:i+10] += 1
...:
50.5 µs ± 177 ns per loop (mean ± std. dev. of 7 runs,10000 loops each)
高级索引会减慢一些速度：
In [120]: %%timeit
...: x = np.zeros(2*10-1)
...: for i in range(10):
...: x[np.arange(i,i+10)] += 1
...:
75.2 µs ± 79.9 ns per loop (mean ± std. dev. of 7 runs,10000 loops each)
如果它有效，一个高级索引 += 将是最快的：
In [121]: %%timeit
...: idx=np.arange(10); idx=(idx[:,None]+idx).ravel()
...: y=np.zeros(2*10-1)
...: y[idx]+=1
...:
...:
17.5 µs ± 693 ns per loop (mean ± std. dev. of 7 runs,10000 loops each)
完整的python迭代与循环的arange情况大致相同：
In [122]: %%timeit
...: idx=np.arange(10); idx=(idx[:,None]+idx).ravel()
...: y=np.zeros(2*10-1)
...: for i in idx: y[i]+=1
...:
...:
76.3 µs ± 2.51 µs per loop (mean ± std. dev. of 7 runs,10000 loops each)
add.at 比 flat += 慢，但仍然比基线好：
In [123]: %%timeit
...: idx=np.arange(10); idx=(idx[:,None]+idx).ravel()
...: y=np.zeros(2*10-1)
...: np.add.at(y,1)
...:
...:
29.4 µs ± 21.2 ns per loop (mean ± std. dev. of 7 runs,10000 loops each)
在我较小的测试中，您的 slower 表现最好。但它的扩展性可能不如基本/快速。您的 indices 大得多。通常对于非常大的数组，由于内存管理负载减少，一维迭代速度更快。