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

问题描述

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

I have a numpy array "data" with dimensions [t, z, x, y]. Thedimensions represent time (t) and three spatial dimensions (x, y, z).I have a separate array "idx" of indices with dimensions [t, x, y]describing vertical coordinates in data: each value in idx describes asingle vertical level in data.

我想从idx索引的数据中提取值.我做完了成功使用循环(如下).我已经阅读了几个 SO线程和 numpy的索引文档，但我无法使其变得更pythonic/vectorized.

I want to pull out the values from data indexed by idx. I've done itsuccessfully using loops (below). I've read several SO threads and numpy's indexing docs but I haven't been able to make it more pythonic/vectorized.

有一种简单的方法让我变得不太正确吗?也许循环仍然是更清晰的方法...

Is there an easy way I'm just not getting quite right? Or maybe loopsare a clearer way to do this anyway...

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

推荐答案

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)相同的形状.

I,J,K broadcast together to the same shape as idx (4,4,4).

有关这种索引的更多详细信息，

More detail on this kind of indexing at

如何沿用元素给定的轴，由它们的索引给定?

这篇关于用另一个数组的索引来索引一个numpy数组的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持！