python - 在 Numpy 中更快地执行 `A[tuple(B.T)]`

Question

我有一个 M×N 数组 I，其中每一行都是一个 N 维数组 A 的索引。我想要一个矢量化表达式来从 A 中获取 M 索引值的一维数组。我发现 A[tuple(I.T)] 做了正确的事情，但是分析表明它非常昂贵尽管被矢量化了。它也不是特别优雅或“自然”，A[I] 和 A[I.T] 做的事情完全不同

正确的做法是什么？它也应该适用于像

这样的分配

A[tuple(I.T)] = 1

Answer 1

我认为你在谈论类似这样的事情:

In [398]: A=np.arange(24).reshape(4,6)
In [401]: I=np.array([[0,1],[1,2],[3,4],[0,0],[2,5]])

In [402]: tuple(I.T)
Out[402]: (array([0, 1, 3, 0, 2]), array([1, 2, 4, 0, 5]))

In [403]: A[tuple(I.T)]
Out[403]: array([ 1,  8, 22,  0, 17])

这是 http://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#purely-integer-array-indexing - 纯整数数组高级索引。

这总是比返回 View 的基本索引慢。基本索引选择连续的数据 block ，或者可以大步选择的值。你的索引是不可能的。

看看一些时间:

In [404]: timeit tuple(I.T)
100000 loops, best of 3: 3.4 µs per loop

In [405]: timeit A[tuple(I.T)]
100000 loops, best of 3: 10 µs per loop

In [406]: %%timeit i,j=tuple(I.T)
   .....: A[i,j]
   .....: 
100000 loops, best of 3: 4.86 µs per loop

构造元组大约需要 1/3 的时间。 i,j=I.T 只是快了一点。但索引是最大的部分。

A[i,j] 与 A[(i,j)] 相同(A.__getitem__((i,j) ))。因此，将 I.T 包装在 tuple 中只会生成 2 个索引数组，每个维度一个。

在数组的扁平化版本上索引更快:

In [420]: J= np.ravel_multi_index(tuple(I.T),A.shape)
In [421]: J
Out[421]: array([ 1,  8, 22,  0, 17], dtype=int32)

In [422]: A.flat[J]
Out[422]: array([ 1,  8, 22,  0, 17])

In [425]: timeit A.flat[J]
1000000 loops, best of 3: 1.56 µs per loop

In [426]: %%timeit
    .....: J= np.ravel_multi_index(tuple(I.T),A.shape)
    .....: A.flat[J]
    .....: 
100000 loops, best of 3: 11.2 µs per loop

因此，能够预先计算和重用索引会节省您的时间，但无法了解从 A 中选择一堆单独的值会花费额外时间这一事实。

只是为了好玩，比较索引 A 和 I 的每一行所花费的时间:

In [442]: timeit np.array([A[tuple(i)] for i in I])
100000 loops, best of 3: 17.3 µs per loop
In [443]: timeit np.array([A[i,j] for i,j in I])
100000 loops, best of 3: 15.7 µs per loop