python - Numpy 数组 : Efficient use of arrays containing indices

Question

我有一个有效的 python 代码，但我想知道它是否可以更快。

我们总共有四个形状相同的二维输入数组。其中三个数组包含整数(将用作索引)，另一个数组包含 double 。

要优化的函数需要将 double 组的值添加到 3D 数组中，位置由包含索引的数组定义。我的代码执行以下操作:

    array_3D[index_one, index_two, index_three] += array_with_doubles

所以问题是:这是一种有效的编程方式吗？我不确定，但我希望 [ ] 索引符号可以被有效的东西取代。这个函数被调用了很多，并且占用了我的执行时间的+- 50%(根据snakeviz)。

另一种策略可能是减少 3D 数组的维度，尽管我可以想象代码的可读性会损失很多。

Answer 1

更简单的二维案例:

In [48]: index1=np.array([1,1,2,2,3,3,4,4]);
     index2=np.array([0,2,1,2,3,4,4,5])
In [49]: data=np.arange(1,9)
In [50]: target=np.zeros((5,6))
In [53]: target[index1,index2]=data

In [54]: target
Out[54]: 
array([[ 0.,  0.,  0.,  0.,  0.,  0.],
       [ 1.,  0.,  2.,  0.,  0.,  0.],
       [ 0.,  3.,  4.,  0.,  0.,  0.],
       [ 0.,  0.,  0.,  5.,  6.,  0.],
       [ 0.,  0.,  0.,  0.,  7.,  8.]])

如果您“整理”索引，则可以使用 put 或 target.flat:

In [51]: flatindex=np.ravel_multi_index((index1,index2),target.shape)
In [52]: flatindex
Out[52]: array([ 6,  8, 13, 14, 21, 22, 28, 29], dtype=int32)
In [58]: np.put(target,flatindex,data)
In [61]: target.flat[flatindex]=data

一些快速时间比较(对于 =data，而不是 +=data):

In [63]: timeit target[index1,index2]=data
100000 loops, best of 3: 6.63 µs per loop

In [64]: timeit np.put(target,flatindex,data)
100000 loops, best of 3: 2.47 µs per loop

In [65]: timeit target.flat[flatindex]=data
100000 loops, best of 3: 2.77 µs per loop

In [66]: %%timeit
   ....: flatindex=np.ravel_multi_index((index1,index2),target.shape)
   ....: target.flat[flatindex]=data
   ....: 
100000 loops, best of 3: 7.34 µs per loop

target.flat[]= 是获胜者 - 如果已整理的索引已经可用。如果您对相同的索引数组重复应用此计算，则可能会出现这种情况。请记住，小型阵列上的时间测试可能与大型阵列上的时间测试不同。

使用 += 代替，put 不起作用。即使必须计算 ravel，flat 也具有速度优势:

In [78]: timeit target[index1,index2]+=data
100000 loops, best of 3: 16.2 µs per loop

In [79]: timeit target.flat[flatindex]+=data
100000 loops, best of 3: 7.45 µs per loop

In [80]: %%timeit                          
flatindex=np.ravel_multi_index((index1,index2),target.shape)
target.flat[flatindex]+=data
   ....: 
100000 loops, best of 3: 13.4 µs per loop

但是 - 如果索引有重复，并且您希望添加所有 data 值，则问题会发生显着变化。像这样的直接索引使用缓冲，因此只有点的最后添加适用。

请参阅最近的 SO 问题来讨论此缓冲问题和替代方案

Vector operations with numpy