Numpy - индексировать последнее измерение массива индексным массивом

Question

Я пытаюсь проиндексировать последнее измерение трехмерной матрицы с помощью матрицы, состоящей из индексов, которые я хочу сохранить.

У меня есть матрица значений тяги с формой:

(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.