用另一个数组的索引索引一个 numpy 数组

Question

我有一个尺寸为 [t, z, x, y] 的 numpy 数组“数据”。这 维度代表时间 (t) 和三个空间维度 (x, y, z)。 我有一个单独的索引数组“idx”，维度为 [t, x, y] 描述数据中的垂直坐标：idx 中的每个值描述一个 数据中的单个垂直级别。

我想从 idx 索引的数据中提取值。我已经做到了 成功使用循环（下）。我已经阅读了几篇SO threads 和numpy's indexing docs，但我无法让它更加pythonic/vectorized。

有没有一种简单的方法让我不太对劲？或者可能是循环 无论如何都是一种更清晰的方法......

import numpy as np

dim = (4, 4, 4, 4)  # dimensions time, Z, X, Y

data = np.random.randint(0, 10, dim)
idx = np.random.randint(0, 3, dim[0:3])

# extract vertical indices in idx from data using loops
foo = np.zeros(dim[0:3])
for this_t in range(dim[0]):
    for this_x in range(dim[2]):
        for this_y in range(dim[3]):
            foo[this_t, this_x, this_y] = data[this_t,
                                               idx[this_t, this_x, this_y],
                                               this_x,
                                               this_y]

# surely there's a better way to do this with fancy indexing
# data[idx] gives me an array with dimensions (4, 4, 4, 4, 4, 4)
# data[idx[:, np.newaxis, ...]] is a little closer
# data[tuple(idx[:, np.newaxis, ...])] doesn't quite get it either
# I tried lots of variations on those ideas but no luck yet

Answer 1

In [7]: I,J,K = np.ogrid[:4,:4,:4]
In [8]: data[I,idx,J,K].shape
Out[8]: (4, 4, 4)
In [9]: np.allclose(foo, data[I,idx,J,K])
Out[9]: True

I,J,K 一起广播到与idx (4,4,4) 相同的形状。

有关此类索引的更多详细信息，请参见

How to take elements along a given axis, given by their indices?