numpy.add.at 比就地添加慢？

Question

从我的开始，我注意到操作 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, indices, 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 的文档指出：

Performs unbuffered in place operation on operand ‘a’ for elements specified by ‘indices’. For addition ufunc, this method is equivalent to a[indices] += b, except that results are accumulated for elements that are indexed more than once.

为什么 np.add.at 这么慢？

这个函数的用例是什么？

Answer 1

add.at 适用于索引包含重复且 += 未产生所需结果的情况

In [44]: A = np.zeros(5,int); idx = np.array([0,1,1,2,2,2,3,3,3,3])
In [45]: A[idx]+=1
In [46]: A
Out[46]: array([1, 1, 1, 1, 0])    # the duplicates in idx are ignored

与add.at:

In [47]: A = np.zeros(5,int); idx = np.array([0,1,1,2,2,2,3,3,3,3])
In [48]: np.add.at(A, idx, 1)
In [49]: A
Out[49]: array([1, 2, 3, 4, 0])

与显式迭代相同的结果：

In [50]: A = np.zeros(5,int); idx = np.array([0,1,1,2,2,2,3,3,3,3])
In [51]: for i in idx: A[i]+=1
In [52]: A
Out[52]: array([1, 2, 3, 4, 0])

一些时间：

In [53]: %%timeit A = np.zeros(5,int); idx = np.array([0,1,1,2,2,2,3,3,3,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,int); idx = np.array([0,1,1,2,2,2,3,3,3,3])
    ...: np.add.at(A, idx, 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,int); idx = np.array([0,1,1,2,2,2,3,3,3,3])
    ...: np.add.at(A, idx, 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.,  9.,  8.,  7.,
        6.,  5.,  4.,  3.,  2.,  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.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10.,  9.,  8.,  7.,
        6.,  5.,  4.,  3.,  2.,  1.])

无需添加您的 'slow' 案例，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., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
       1., 1.])

add.at 解决了：

In [115]: y=np.zeros(2*10-1)
     ...: np.add.at(y, idx, 1)
In [116]: y
Out[116]: 
array([ 1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10.,  9.,  8.,  7.,
        6.,  5.,  4.,  3.,  2.,  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.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10.,  9.,  8.,  7.,
        6.,  5.,  4.,  3.,  2.,  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 迭代与循环排列情况大致相同：

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, idx,1)
     ...: 
     ...: 
29.4 µs ± 21.2 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

在我的小型测试中，您的 slower 表现最好。但它的扩展性可能不如 base/fast。您的 indices 要大得多。通常对于非常大的数组，由于内存管理负载减少，一维迭代速度更快。

numpy.add.at 比就地添加慢？

numpy.add.at slower than in-place add?

arrays

numpy

add

indices

numpy-ufunc

编辑