Numpy - 用索引數組索引數組的最后一維

Question

我正在嘗試使用由我希望保留的索引組成的矩陣來索引 3D 矩陣的最后一個維度。

我有一個形狀的推力值矩陣：

(3, 3, 5)

我想根據某些標准過濾最后一個索引，使其從大小 5 減少到大小 1。我已經在最后一個維度中找到了符合我的條件的索引：

[[0 0 1]
 [0 0 1]
 [1 4 4]]

我想要實現的目標：對於第一行和第一列，我想要最后一個維度的第 0 個索引。 對於第一行和第三列，我想要最后一個維度的第一個索引。 在保持最終矩陣的索引方面將成為這樣的(3, 3) 2D 矩陣：

[[0,0,0], [0,1,0], [0,2,1];
 [1,0,0], [1,1,0], [1,2,1];
 [2,0,1], [2,1,4], [2,2,4]]

我非常有信心 numpy 可以做到這一點，但我無法弄清楚究竟是如何做到的。 我寧願不構建帶有嵌套 for 循環的結構。

我已經嘗試過：

minValidPower = totalPower[:, :, tuple(indexMatrix)]

但這會導致(3, 3, 3, 3)矩陣，所以我不完全確定我應該如何處理這個問題。

Answer 1

使用a作為輸入數組和idx作為索引 -

np.take_along_axis(a,idx[...,None],axis=-1)[...,0]

或者，使用開放式網格 -

I,J = np.ogrid[:idx.shape[0],:idx.shape[1]]
out = a[I,J,idx]

Answer 2

您可以為前兩個維度構建相應的索引數組。 這些基本上是：

[0 1 2]
[0 1 2]
[0 1 2]

[0 0 0] 
[1 1 1] 
[2 2 2]

您可以使用 meshgrid 函數構建這些。 我在示例中將它們存儲為 m1 和 m2：

vals = np.arange(3*3*5).reshape(3, 3, 5) # test sample
m1, m2 = np.meshgrid(range(3), range(3), indexing='ij')
m3 = np.array([[0, 0, 1], 0, 0, 1], [1, 4, 4]])
sel_vals = vals[m1, m2, m3]

結果的形狀與索引數組 m1、m2 和 m3 的形狀相匹配。