Jeśli mam ndarray w ten sposób:

>>> 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]]])

Wiem, że mogę uzyskać maksimum wzdłuż określonej osi za pomocąnp.max(axis=...):

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

Alternatywnie, mógłbym dostaćindeksy wzdłuż tej osi, które odpowiadają maksymalnym wartościom z:

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

Moje pytanie - Biorąc pod uwagę tablicęindices i tablicaa, czy istnieje elegancki sposób na odtworzenie tablicy zwróconej przez tablicęa.max(axis=2)?

Prawdopodobnie to zadziała:

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

Ale wydaje się dość hacky / nieefektywny. Nie pozwala mi także używać tej osi z inną osią niż „ostatnia” oś. Czy istnieje funkcja numpy (lub użycie magicznego indeksowania liczb), aby to działało? Naiwnya[:,:,a.argmax(axis=2)] nie działa.

AKTUALIZACJA:

Wygląda na to, że również działa (i jest trochę ładniejszy):

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

Chciałbym to zrobić, ponieważ chciałbym wyodrębnić indeksy w oparciu o dane w 1 tablicy (zwykle przy użyciuargmax(axis=...)) i używaj tych indeksów do pobierania danych z szeregu innych (równoważnie ukształtowanych) tablic. Jestem otwarty na alternatywne sposoby osiągnięcia tego celu (np. Przy użyciu boolowskich tablic zamaskowanych). Podoba mi się jednak „bezpieczeństwo”, które otrzymuję za pomocą tablic „indeksowych”. Dzięki temu mam gwarancję posiadania odpowiedniej liczby elementów, aby utworzyć nową tablicę, która wygląda jak 2d „wycinek” w polu 3d.