如何找到最后 2 个轴的 argmax

Question

Hey I have seen this question - Numpy: argmax over multiple axes without loop but the output is not the shape I desire. So for example, if I give the function an array of dimension: 10x20x12x12x2x2, it will output an array of dimension: 10x20x12x12, which values are the indices

Answer 1

Make a simpler, but I think still relevant array:

In [268]: arr = np.random.randint(0,20,(4,1,3,2))                                        
In [269]: arr                                                                            
Out[269]: 
array([[[[16,  1],
         [13, 17],
         [19,  0]]],


       [[[ 2, 13],
         [12,  9],
         [ 6,  6]]],


       [[[13,  2],
         [18, 10],
         [ 7, 10]]],


       [[[ 8, 19],
         [ 6, 17],
         [ 2,  6]]]])

reshape as suggested in the link:

In [270]: arr1 = arr.reshape(arr.shape[:-2]+(-1,))                                       
In [271]: arr1                                                                           
Out[271]: 
array([[[16,  1, 13, 17, 19,  0]],

       [[ 2, 13, 12,  9,  6,  6]],

       [[13,  2, 18, 10,  7, 10]],

       [[ 8, 19,  6, 17,  2,  6]]])

then we can take the max and argmax on the last dimension:

In [272]: np.max(arr1, -1)                                                               
Out[272]: 
array([[19],
       [13],
       [18],
       [19]])
In [273]: idx = np.argmax(arr1, -1)                                                      
In [274]: idx                                                                            
Out[274]: 
array([[4],
       [1],
       [2],
       [1]])

we can recover the max from the argmax with indexing work:

In [282]: ij = np.ix_(np.arange(4),np.arange(1))                                         
In [283]: ij+(idx,)                                                                      
Out[283]: 
(array([[0],
        [1],
        [2],
        [3]]),
 array([[0]]),
 array([[4],
        [1],
        [2],
        [1]]))
In [284]: arr1[ij+(idx,)]                                                                
Out[284]: 
array([[19],
       [13],
       [18],
       [19]])

With unravel we can apply this to arr :

In [285]: idx1 = np.unravel_index(idx, (3,2))                                            
In [286]: idx1                                                                           
Out[286]: 
(array([[2],
        [0],
        [1],
        [0]]),
 array([[0],
        [1],
        [0],
        [1]]))
In [287]: arr[ij+idx1]       # tuple concatenate                                                            
Out[287]: 
array([[19],
       [13],
       [18],
       [19]])

So the max on the last 2 axes of arr is still the shape of the first 2.

So even though arr is (4,1,3,2), the useful argmax does not have this shape. Instead we need a tuple of 4 arrays, one for each dimension of arr . Advanced indexing like this on more than 2 dimensions is tricky, and hard to visualize. I had to play around with this quite a while.

With your dimensions:

In [322]: barr = np.random.randint(0,100,(10,20,12,12,2,2))                              
In [323]: barr1 = barr.reshape(barr.shape[:-2]+(-1,))                                    
In [324]: ms = np.max(barr1, axis=-1)                                                    
In [325]: idx = np.argmax(barr1,-1)                                                      
In [326]: idx1 = np.unravel_index(idx, barr.shape[-2:])                                  
In [327]: ij = np.ix_(*[np.arange(i) for i in barr.shape[:-2]])                          
In [328]: np.allclose(barr[ij+idx1], ms)                                                 
Out[328]: True

edit

We could just as well simplify this task to working with a 2d array:

In [65]: barr2 = barr.reshape(-1,4)                                             
In [66]: idx2 = np.argmax(barr2, axis=1)                                        
In [67]: idx2.shape                                                             
Out[67]: (28800,)
In [68]: np.allclose(idx.ravel(), idx2)                                         
Out[68]: True
In [69]: ms2 = barr2[np.arange(barr2.shape[0]),idx2]                            
In [70]: ms2.shape                                                              
Out[70]: (28800,)
In [72]: np.allclose(ms2.reshape(barr.shape[:-2]), ms)                          
Out[72]: True

column_stack is wrong with the multidimensional idx1 , joining on axis 1. We want to join on a new trailing axis, with stack :

In [77]: np.column_stack(idx1).shape                                            
Out[77]: (10, 40, 12, 12)
In [78]: np.stack(idx1,axis=-1).shape                                           
Out[78]: (10, 20, 12, 12, 2)
In [79]: np.allclose(x, np.stack(idx1,-1).reshape(-1,2))                        
Out[79]: True

But I don't see the value of such a stack. The linked question does ask for such an array, but doesn't show how such an array might be used.