Question

我有一个numpy数组＆＃34;数据＆＃34;尺寸为[t，z，x，y]。该尺寸表示时间（t）和三个空间尺寸（x，y，z）。我有一个单独的阵列＆＃34; idx＆＃34;尺寸为[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?

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

1 个答案: