Индекс Numpy - Использование распутанного индекса для базовой индексации

Question

Если у меня есть следующий массив 4D:

mat = np.array(np.arange(27)).reshape((3,3,3))
[[[ 0  1  2]
  [ 3  4  5]
  [ 6  7  8]]

 [[ 9 10 11]
  [12 13 14]
  [15 16 17]]

 [[18 19 20]
  [21 22 23]
  [24 25 26]]]

и следующий развернутый индекс:

ind = np.unravel_index([7], mat.shape[1:])
(array([2], dtype=int64), array([1], dtype=int64))

, каков наилучший способ доступа

mat[:, 2, 1]
[ 7 16 25]

используя распутанный индекс?Я ищу общее решение этой проблемы, где число измерений mat может варьироваться.

Я знаю, что мог бы сделать что-то вроде этого:

new_ind = (np.arange(mat.shape[0]),) +  ind
mat[new_ind]
[ 7 16 25]

но яИнтересно, есть ли способ сделать это, не требующий явного построения нового индекса?

Answer 1

Вам необходимо создать новый индексный кортеж:

In [8]: ind=np.unravel_index([7,8],(3,3))                                       
In [9]: ind                                                                     
Out[9]: (array([2, 2]), array([1, 2]))
In [10]: (slice(None),*ind)                                                     
Out[10]: (slice(None, None, None), array([2, 2]), array([1, 2]))
In [11]: np.arange(27).reshape(3,3,3)[_]                                        
Out[11]: 
array([[ 7,  8],
       [16, 17],
       [25, 26]])

Out[10] эквивалентно добавлению : к вашим развернутым индексам:

In [12]: np.s_[:,[2,2],[1,2]]                                                   
Out[12]: (slice(None, None, None), [2, 2], [1, 2])