Question

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

    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]]

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

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

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

Answer 1

听起来你需要做两件事：

查找数据数组中的所有索引和
根据其他数组subindex翻译列索引。

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

使用numpy，没有必要在for循环中执行此操作 - 我们可以立即传入所有索引：

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

Answer 2

你真正想要的是：

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"

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

2 个答案: