Indexing numpy array by a numpy array of coordinates

Question

Suppose we have

• an n-dimensional numpy.array A
• a numpy.array B with dtype=int and shape of (n, m)

How do I index A by B so that the result is an array of shape (m,), with values taken from the positions indicated by the columns of B?

For example, consider this code that does what I want when B is a python list:

>>> a = np.arange(27).reshape(3,3,3)
>>> a[[0, 1, 2], [0, 0, 0], [1, 1, 2]]
array([ 1, 10, 20])    # the result we're after
>>> bl = [[0, 1, 2], [0, 0, 0], [1, 1, 2]]
>>> a[bl]
array([ 1, 10, 20])   # also works when indexing with a python list
>>> a[bl].shape
(3,)

但是，当 B 是一个 numpy 数组时，结果就不同了：

>>> b = np.array(bl)
>>> a[b].shape
(3, 3, 3, 3)

现在，我可以通过将 B 转换为元组来获得所需的结果，但这肯定不是正确/惯用的方法吗？

>>> a[tuple(b)]
array([ 1, 10, 20])

是否有一个 numpy 函数可以在不将 B 转换为元组的情况下实现相同的功能？

Answer 1

一种替代方法是转换为线性索引，然后使用索引np.take或索引到其扁平版本 -

np.take(a,np.ravel_multi_index(b, a.shape))
a.flat[np.ravel_multi_index(b, a.shape)]

定制np.ravel_multi_index以提高性能

我们可以实现一个自定义版本来模拟行为np.ravel_multi_index以提高性能，如下所示 -

def ravel_index(b, shp):
    return np.concatenate((np.asarray(shp[1:])[::-1].cumprod()[::-1],[1])).dot(b)

使用它，可以通过两种方式找到所需的输出 -

np.take(a,ravel_index(b, a.shape))
a.flat[ravel_index(b, a.shape)]

标杆管理

另外，还结合了tuple问题中的方法和map@Kanak 帖子中的方法。

情况 #1：暗淡 = 3

In [23]: a = np.random.randint(0,9,([20]*3))

In [24]: b = np.random.randint(0,20,(a.ndim,1000000))

In [25]: %timeit a[tuple(b)]
    ...: %timeit a[map(np.ravel, b)]  
    ...: %timeit np.take(a,np.ravel_multi_index(b, a.shape))
    ...: %timeit a.flat[np.ravel_multi_index(b, a.shape)]
    ...: %timeit np.take(a,ravel_index(b, a.shape))
    ...: %timeit a.flat[ravel_index(b, a.shape)]
100 loops, best of 3: 6.56 ms per loop
100 loops, best of 3: 6.58 ms per loop
100 loops, best of 3: 6.95 ms per loop
100 loops, best of 3: 9.17 ms per loop
100 loops, best of 3: 6.31 ms per loop
100 loops, best of 3: 8.52 ms per loop

情况 #2：暗淡 = 6

In [29]: a = np.random.randint(0,9,([10]*6))

In [30]: b = np.random.randint(0,10,(a.ndim,1000000))

In [31]: %timeit a[tuple(b)]
    ...: %timeit a[map(np.ravel, b)]  
    ...: %timeit np.take(a,np.ravel_multi_index(b, a.shape))
    ...: %timeit a.flat[np.ravel_multi_index(b, a.shape)]
    ...: %timeit np.take(a,ravel_index(b, a.shape))
    ...: %timeit a.flat[ravel_index(b, a.shape)]
10 loops, best of 3: 40.9 ms per loop
10 loops, best of 3: 40 ms per loop
10 loops, best of 3: 20 ms per loop
10 loops, best of 3: 29.9 ms per loop
100 loops, best of 3: 15.7 ms per loop
10 loops, best of 3: 25.8 ms per loop

情况 #3：变暗 = 10

In [32]: a = np.random.randint(0,9,([4]*10))

In [33]: b = np.random.randint(0,4,(a.ndim,1000000))

In [34]: %timeit a[tuple(b)]
    ...: %timeit a[map(np.ravel, b)]  
    ...: %timeit np.take(a,np.ravel_multi_index(b, a.shape))
    ...: %timeit a.flat[np.ravel_multi_index(b, a.shape)]
    ...: %timeit np.take(a,ravel_index(b, a.shape))
    ...: %timeit a.flat[ravel_index(b, a.shape)]
10 loops, best of 3: 60.7 ms per loop
10 loops, best of 3: 60.1 ms per loop
10 loops, best of 3: 27.8 ms per loop
10 loops, best of 3: 38 ms per loop
100 loops, best of 3: 18.7 ms per loop
10 loops, best of 3: 29.3 ms per loop

因此，在处理高维输入和大数据时寻找替代方案是有意义的。