Se eu tenho um ndarray assim:

>>> a = np.arange(27).reshape(3,3,3)
>>> a
array([[[ 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]]])

Eu sei que posso obter o máximo ao longo de um determinado eixo usandonp.max(axis=...):

>>> a.max(axis=2)
array([[ 2,  5,  8],
       [11, 14, 17],
       [20, 23, 26]])

Alternativamente, eu poderia obter oíndices ao longo desse eixo que correspondem aos valores máximos de:

>>> indices = a.argmax(axis=2)
>>> indices
array([[2, 2, 2],
       [2, 2, 2],
       [2, 2, 2]])

Minha pergunta - dada a matrizindices e o arrayaexiste uma maneira elegante de reproduzir a matriz da matriz retornada pora.max(axis=2)?

Isso provavelmente funcionaria:

import itertools as it
import numpy as np
def apply_mask(field,indices):
    data = np.empty(indices.shape)

    #It seems highly likely that there is a more numpy-approved way to do this.
    idx = [range(i) for i in indices.shape]
    for idx_tup,zidx in zip(it.product(*idx),indices.flat):
        data[idx_tup] = field[idx_tup+(zidx,)]
    return data

Mas, parece bastante hacky / ineficiente. Ele também não permite que eu use isso com qualquer eixo diferente do "último" eixo. Existe uma função numpy (ou algum uso de indexação mágica numpy) para fazer este trabalho? O ingênuoa[:,:,a.argmax(axis=2)] não funciona.

ATUALIZAR:

Parece que o seguinte também funciona (e é um pouco melhor):

import numpy as np
def apply_mask(field,indices):
    data = np.empty(indices.shape)

    for idx_tup,zidx in np.ndenumerate(indices):
        data[idx_tup] = field[idx_tup+(zidx,)]

    return data

Eu gostaria de fazer isso porque eu gostaria de extrair os índices com base nos dados em 1 array (normalmente usandoargmax(axis=...)) e use esses índices para extrair dados de um monte de outros arrays (de formato equivalente). Estou aberto a formas alternativas de realizar isso (por exemplo, usando matrizes booleanas mascaradas). No entanto, eu gosto da "segurança" que eu uso esses arrays "index". Com isso, tenho a garantia de ter o número certo de elementos para criar uma nova matriz que se parece com uma "fatia" 2D através do campo 3D.