Numpy многомерная расширенная индексация

Question

У меня есть массив a с формой [3,x,y,z,n] (три 4d-изображения).И второй массив b с формой [x,y,z], который содержит индексы, которые я хочу выбрать из первого измерения a (поэтому значения b находятся в диапазоне от 0 до 2).
РезультатыЯ хотел бы иметь форму [x,y,z,n].Как я могу сделать это в NumPy?

Answer 1

Использование advanced-indexing -

a[b,np.arange(x)[:,None,None],np.arange(y)[:,None],np.arange(z)]

Более короткий способ выразить это -

a[tuple([b] + np.ogrid[:x,:y,:z])]

Использование встроенной функции NumPy np.take_along_axis для выполнения расширенной индексации, выполняя грязную работу под капотами -

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