numpy 3d à la transformation 2d basée sur la matrice de masque 2d

Question

Si j'ai un ndarray comme ceci:

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

Je sais que je peux obtenir le maximum le long d'un certain axe en utilisant np.max(axis=...):

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

Alternativement , Je pourrais obtenir les indices le long de cet axe qui correspondent aux valeurs maximales de:

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

Ma question - Étant donné le tableau indices et le tableau a, existe-t-il un moyen élégant de reproduire le tableau le tableau retourné par a.max(axis=2)?

Ce serait probablement:

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 

Mais, il semble assez hacky/inefficace. Cela ne me permet pas non plus de l'utiliser avec un autre axe que le "dernier" axe. Y a-t-il une fonction numpy (ou une utilisation de l'indexation numpy magique) pour que cela fonctionne? Le naïf a[:,:,a.argmax(axis=2)] ne fonctionne pas.

MISE À JOUR:

Il semble que le suivant fonctionne aussi (et est un peu plus agréable):

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 

---

Je voudrais faire cela parce que je voudrais extraire les indices basé sur les données dans 1 tableau (en utilisant généralement argmax(axis=...)) et utiliser ces indices pour extraire des données d'un tas d'autres tableaux (de forme équivalente). Je suis ouvert à d'autres moyens d'accomplir ceci (par exemple en utilisant des tableaux masqués booléens). Cependant, j'aime la "sécurité" que j'obtiens en utilisant ces tableaux "index". Avec cela, je suis sûr d'avoir le bon nombre d'éléments pour créer un nouveau tableau qui ressemble à une "tranche" 2d à travers le champ 3D.

Answer 1

Voici une indexation magique qui fera ce que vous voulez, mais malheureusement c'est assez illisible.

def apply_mask(a, indices, axis): 
    magic_index = [np.arange(i) for i in indices.shape] 
    magic_index = np.ix_(*magic_index) 
    magic_index = magic_index[:axis] + (indices,) + magic_index[axis:] 
    return a[magic_index] 

ou tout aussi illisible:

def apply_mask(a, indices, axis): 
    magic_index = np.ogrid[tuple(slice(i) for i in indices.shape)] 
    magic_index.insert(axis, indices) 
    return a[magic_index] 

Answer 2

J'utilise index_at() pour créer l'index complet:

import numpy as np 

def index_at(idx, shape, axis=-1): 
    if axis<0: 
     axis += len(shape) 
    shape = shape[:axis] + shape[axis+1:] 
    index = list(np.ix_(*[np.arange(n) for n in shape])) 
    index.insert(axis, idx) 
    return tuple(index) 

a = np.random.randint(0, 10, (3, 4, 5)) 

axis = 1 
idx = np.argmax(a, axis=axis) 
print a[index_at(idx, a.shape, axis=axis)] 
print np.max(a, axis=axis)