Ottieni indici di elementi numpy.argmax su un asse

Question

Ho una matrice N-dimensionale che contiene i valori per una funzione con parametri N. Ogni parametro ha un numero discreto di valori. Ho bisogno di massimizzare la funzione su tutti i parametri tranne uno, risultando in un vettore unidimensionale di dimensione uguale al numero di valori del parametro non massimizzato. Devo anche salvare quali valori sono presi dagli altri parametri.

Per fare ciò, ho voluto applicare iterativamente numpy.max su diversi assi per ridurre la dimensionalità della matrice per trovare ciò di cui ho bisogno. Il vettore finale dipenderà quindi solo dal parametro che ho omesso.

Tuttavia, ho difficoltà a trovare gli indici originali degli elementi finali (che contengono le informazioni sui valori presi dagli altri parametri). Ho pensato di usare numpy.argmax allo stesso modo di numpy.max ma non riesco a ottenere gli indici originali.

Un esempio di quello che sto cercando è:

x = [[[1,2],[0,1]],[[3,4],[6,7]]] 
args = np.argmax(x, 0) 

Questo restituisce

[[1 1] 
[1 1]] 

Il che significa che argmax sta selezionando gli elementi (2,1,4,7) all'interno l'originale matrice. Ma come ottenere i loro indici? Ho provato unravel_index, usando lo args direttamente come indice per la matrice x, un sacco di funzioni da numpy a indice senza successo.

L'utilizzo di numpy.where non è una soluzione poiché la matrice di input può avere valori uguali all'interno, quindi non sarei in grado di discernere da valori originali diversi.

Answer 1

x.argmax(0) fornisce gli indici lungo il 1o asse per i valori massimi. Utilizzare np.indices per generare gli indici per l'altro asse.

x = np.array([[[1,2],[0,1]],[[3,4],[6,7]]]) 
x.argmax(0) 
    array([[1, 1], 
      [1, 1]]) 
a1, a2 = np.indices((2,2)) 
(x.argmax(0),a1,a2) 
    (array([[1, 1], 
      [1, 1]]), 
    array([[0, 0], 
      [1, 1]]), 
    array([[0, 1], 
      [0, 1]])) 


x[x.argmax(0),a1,a2] 
    array([[3, 4], 
      [6, 7]]) 

x[a1,x.argmax(1),a2] 
    array([[1, 2], 
      [6, 7]]) 

x[a1,a2,x.argmax(2)] 
    array([[2, 1], 
      [4, 7]]) 

Se x ha altre dimensioni, generare a1 e a2 appropriato.

La documentazione ufficiale non dice molto su come utilizzare argmax, ma i thread SO precedenti ne hanno discusso. Ho ricevuto questa idea generale da Using numpy.argmax() on multidimensional arrays