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

Question

我有一个 M×N 数组 I，其中每一行都是一个 N 维数组 A 的索引。我想要一个矢量化表达式来从 A 中获取 M 索引值的一维数组。我发现 A[tuple(I.T)] 做的是正确的事情，但是分析表明尽管进行了矢量化，它还是非常昂贵。它也不是特别优雅或 "natural" 和 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 - 纯整数数组高级索引。

这总是比基本索引慢，returns 一个视图。基本索引选择连续的数据块，或者可以大步选择的值。你的索引是不可能的。

查看一些时间：

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

Answer 2

你可以用linear indexing另一种方式，像这样-

def ravel_einsum(A,I):

    # Get A's shape and calculate cummulative dimensions based on it
    shp = np.asarray(A.shape)
    cumdims = np.append(1,shp[::-1][:-1].cumprod())[::-1]

    # Use linear indexing of A to extract elements from A corresponding 
    # to linear indexing of it with I
    return A.ravel()[np.einsum('ij,j->i',I,cumdims)]

运行时测试

案例 #1：

In [84]: # Inputs
    ...: A = np.random.randint(0,10,(3,5,2,4,5,2,6,8,2,5,3,4,3))
    ...: I = np.mod(np.random.randint(0,10,(5,A.ndim)),A.shape)
    ...: 

In [85]: %timeit A[tuple(I.T)]
10000 loops, best of 3: 27.7 µs per loop

In [86]: %timeit ravel_einsum(A,I)
10000 loops, best of 3: 48.3 µs per loop

案例 #2：

In [87]: # Inputs
    ...: A = np.random.randint(0,10,(3,5,4,2))
    ...: I = np.mod(np.random.randint(0,5,(10000,A.ndim)),A.shape)
    ...: 

In [88]: %timeit A[tuple(I.T)]
1000 loops, best of 3: 357 µs per loop

In [89]: %timeit ravel_einsum(A,I)
1000 loops, best of 3: 240 µs per loop

案例 #3：

In [90]: # Inputs
    ...: A = np.random.randint(0,10,(30,50,40,20))
    ...: I = np.mod(np.random.randint(0,50,(5000,A.ndim)),A.shape)
    ...: 

In [91]: %timeit A[tuple(I.T)]
1000 loops, best of 3: 220 µs per loop

In [92]: %timeit ravel_einsum(A,I)
10000 loops, best of 3: 168 µs per loop

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

Faster way to do `A[tuple(B.T)]` in Numpy

python

arrays

numpy

vectorization

array-broadcasting