python - 使用另一个数组的索引索引一个 numpy 数组

Question

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

我想从 idx 索引的数据中提取值。我做到了成功使用循环(下)。我读过几本 SO threads和 numpy's indexing docs但我没能使它更像 pythonic/矢量化。

有没有一种简单的方法我只是不太正确？或者循环无论如何，这是一种更清晰的方法......

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?