Матрица выбора Numpy, заданная матрицей индексов, из многомерного массива

Question

У меня есть массив NumPy a размером 5x5x4x5x5.У меня есть другая матрица b размера 5x5.Я хочу получить a[i,j,b[i,j]] для i от 0 до 4 и для j от 0 до 4. Это даст мне матрицу 5x5x1x5x5.Есть ли способ сделать это без использования 2 for петель?

Answer 1

Есть np.take_along_axis именно для этой цели -

np.take_along_axis(a,b[:,:,None,None,None],axis=2)

Answer 2

Давайте подумаем о матрице a как о 100 (= 5 x 5 x 4) матрицах размера (5, 5).Итак, если вы можете получить индекс лайнера для каждого триплета - (i, j, b[i, j]) - все готово.Вот тут и приходит np.ravel_multi_index. Ниже приведен код.

import numpy as np
import itertools

# create some matrices
a = np.random.randint(0, 10, (5, 5, 4, 5, 5))
b = np.random(0, 4, (5, 5))

# creating all possible triplets - (ind1, ind2, ind3)
inds = list(itertools.product(range(5), range(5)))
(ind1, ind2), ind3 = zip(*inds), b.flatten()

allInds = np.array([ind1, ind2, ind3])
linearInds = np.ravel_multi_index(allInds, (5,5,4))

# reshaping the input array
a_reshaped = np.reshape(a, (100, 5, 5))

# selecting the appropriate indices
res1 = a_reshaped[linearInds, :, :]

# reshaping back into desired shape
res1 = np.reshape(res1, (5, 5, 1, 5, 5))

# verifying with the brute force method
res2 = np.empty((5, 5, 1, 5, 5))
for i in range(5):
    for j in range(5):
        res2[i, j, 0] = a[i, j, b[i, j], :, :]

print np.all(res1 == res2)  # should print True