使用numpy数组的坐标来索引numpy数组

Question

假设我们有：

- 一个n维的numpy.array A - 一个dtype为int且形状为(n, m)的numpy.array B

如何使用B对A进行索引，以便结果是一个形状为(m,)的数组，其中的值来自于B指定的列所示位置？

例如，考虑以下代码，当B是Python列表时，它可以实现我想要的功能：

>>> 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转换为元组而达到相同的效果？

- Sami Liedes

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的方法和@Kanak帖子中的基于map的方法。

情况＃1：dims = 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：dims = 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：dims = 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

因此，在处理高维输入和大数据时，寻找替代方案是明智的选择。

- Divakar

Answer 2

另一个符合您需求的选择是使用np.ravel函数。

>>> a[map(np.ravel, b)]
array([ 1, 10, 20])

然而并非完全基于numpy。

---

性能问题。 根据下面的评论进行了更新。

无论如何，你的方法比我的好，但不如@Divakar的任何方法。

import numpy as np
import timeit

a = np.arange(27).reshape(3,3,3)
bl = [[0, 1, 2], [0, 0, 0], [1, 1, 2]]
b = np.array(bl)

imps = "from __main__ import np,a,b"
reps = 100000

tup_cas_t = timeit.Timer("a[tuple(b)]", imps).timeit(reps)
map_rav_t = timeit.Timer("a[map(np.ravel, b)]", imps).timeit(reps)
fla_rp1_t = timeit.Timer("np.take(a,np.ravel_multi_index(b, a.shape))", imps).timeit(reps)
fla_rp2_t = timeit.Timer("a.flat[np.ravel_multi_index(b, a.shape)]", imps).timeit(reps)

print tup_cas_t/map_rav_t  ## 0.505382211881
print tup_cas_t/fla_rp1_t  ## 1.18185817386
print tup_cas_t/fla_rp2_t  ## 1.71288705886

- keepAlive

---

flat 和 np.take 是两个不同的替代方案。 - Divakar

此外，当处理如此微小的时间间隔时，我宁愿使用%timeit。你不需要这样做，我只是为了公平的基准测试而说。再次强调，OP似乎正在寻找一种惯用的方式。个人意见-无论哪种方式，tuple都不错。 - Divakar

很高兴知道 @Divakar。话虽如此，你不觉得100000次重复可以确保基准公正性的最低要求吗？ 另外，%timeit是[IPython]（http://ipython.org/）的特性/魔法，是吗？如果是这样，我不使用iPython。实际上，我没有不使用它的理由。 - keepAlive

好的。那么我们来试试更大一点的，比如说：a = np.random.randint(0,9,(10,10,10,10,10,10)); b = np.random.randint(0,10,(6,10000))？ - Divakar

在我的帖子中添加了一些大数据和维度的时间记录。 - Divakar

Answer 3

你正在寻找 numpy.ndarray.tolist() 吗？

>>> a = np.arange(27).reshape(3,3,3)
>>> bl = [[0, 1, 2], [0, 0, 0], [1, 1, 2]]
>>> b = np.array(bl)
>>> a[b.tolist()]
array([ 1, 10, 20])

或者对于数组索引数组，它与列表索引非常相似：

>>> a[np.array([0, 1, 2]), np.array([0, 0, 0]), np.array([1, 1, 2])]
array([ 1, 10, 20])

然而，正如您可以从前面的链接中看到的那样，直接使用数组b索引数组a意味着您仅使用整个b数组来索引a的第一个索引，这可能会导致混乱的输出。

- FollowK