2

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

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

A[tuple(I.T)] = 1
4

2 回答 2

4

我想你在谈论类似的事情：

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 - 纯整数数组高级索引。

这总是比返回视图的基本索引慢。基本索引选择连续的数据块，或可以跨步选择的值。您的索引无法做到这一点。

看一些时间：

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)](as is A.__getitem__((i,j))) 相同。所以包装只I.T产生tuple2个索引数组，每个维度一个。

在数组的扁平版本上建立索引会更快：

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
于 2015-08-15T20:28:45.667 回答
1

你可以使用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
于 2015-08-16T08:19:09.970 回答