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

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循環外，我想使用numpy.ix_使用numpy.ix_為data numpy.ix_索引，我想我會有類似的東西：

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

但我不知道第二維參數應該是什么。 我猜應該是涉及numpy.choose東西？

Answer 1

聽起來您需要做兩件事：

• 找到數據數組中的所有索引，然后
• 根據其他數組，即子索引轉換列索引。

因此，下面的代碼為所有數組位置生成索引（使用np.indices ），並將其np.indices為(..., 2) -表示數組中每個位置的二維坐標列表。 然后，對於每個坐標(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])

只是data的第一維“ y軸”的索引。 和

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

向該數組添加一個新軸 （在現有軸之后）並有效地對其進行轉置。 現在，當您在索引表達式中將此數組與subindex數組一起使用時，前者將廣播為后者的形狀。 因此y_idx變得有效了：

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

現在，對於每對y_idx和子subindex您都會在data數組中查找一個元素。

在這里您可以找到有關“花式索引”的更多信息