如何使用一维的子索引生成多维 2D numpy 索引

Question

我想使用numpy.ix_ 为二维值空间生成多维索引。但是，我需要使用子索引来查找一维的索引。例如，

    assert subindex.shape == (ny, nx)

    data = np.random.random(size=(ny,nx))

    # Generator returning the index tuples 
    def get_idx(ny,nx,subindex):
      for y in range(ny):
        for x in range(nx):
           yi = y             # This is easy
           xi = subindex[y,x] # Get the second index value from the subindex

           yield (yi,xi)

    # Generator returning the data values
    def get_data_vals(ny,nx,data,subindex):
      for y in range(ny):
        for x in range(nx):
           yi = y             # This is easy
           xi = subindex[y,x] # Get the second index value from the subindex

           yield data[y,subindex[y,x]]

所以我想使用多维索引来索引data，而不是上面的for循环，使用numpy.ix_，我想我会有类似的东西：

    idx = numpy.ix_([np.arange(ny), ?])
    data[idx]

但我不知道第二维参数应该是什么。我猜应该是涉及到numpy.choose？

Answer 1

你真正想要的是：

y_idx = np.arange(ny)[:,np.newaxis]
data[y_idx, subindex]

顺便说一句，您可以使用y_idx = np.arange(ny).reshape((-1, 1)) 实现相同的效果。

我们来看一个小例子：

import numpy as np

ny, nx = 3, 5
data = np.random.rand(ny, nx)
subindex = np.random.randint(nx, size=(ny, nx))

现在

np.arange(ny)
# array([0, 1, 2])

只是“y 轴”的索引，data 的第一个维度。和

y_idx = np.arange(ny)[:,np.newaxis]
# array([[0],
#        [1],
#        [2]])

将new axis 添加到此数组（在现有轴之后）并有效地转置它。现在，当您在索引表达式中将此数组与subindex 数组一起使用时，前者将broadcasted 变为后者的形状。所以y_idx 变得有效：

# array([[0, 0, 0, 0, 0],
#        [1, 1, 1, 1, 1],
#        [2, 2, 2, 2, 2]])

现在，对于每对 y_idx 和 subindex，您在 data 数组中查找一个元素。

Here you can find out more about "fancy indexing"

Answer 2

听起来你需要做两件事：

• 找到数据数组中的所有索引并
• 根据其他数组子索引转换列索引。

因此，下面的代码会为所有数组位置生成索引（使用np.indices），并将其重塑为(..., 2)——代表数组中每个位置的二维坐标列表。对于每个坐标(i, j)，然后我们使用提供的子索引数组转换列坐标j，然后使用转换后的索引作为新的列索引。

使用 numpy，没有必要在 for 循环中这样做——我们可以简单地一次传入所有索引：

i, j = np.indices(data.shape).reshape((-1, 2)).T
data[i, subindex[i, j]]