如何找到最后 2 个轴的 argmax

Question

嘿，我看到了这个问题 - Numpy: argmax over multiple axes without loop，但输出不是我想要的形状。例如，如果我给函数一个维度数组：10x20x12x12x2x2，它将输出一个维度数组：10x20x12x12，其中的值是索引

Answer 1

制作一个更简单，但我认为仍然相关的数组：

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

按照链接中的建议重塑：

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

那么我们可以在最后一个维度上取 max 和 argmax：

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

我们可以通过索引工作从 argmax 中恢复最大值：

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

使用unravel，我们可以将其应用于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]])

所以arr的最后2个轴上的max仍然是前2个的形状。

所以即使arr 是(4,1,3,2)，有用的argmax 也没有这个形状。相反，我们需要一个由 4 个数组组成的元组，每个维度对应一个 arr。像这样在超过 2 个维度上的高级索引是棘手的，并且难以可视化。我不得不玩这个很长一段时间。

根据您的尺寸：

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

编辑

我们也可以将此任务简化为使用二维数组：

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 与多维 idx1 错误，在轴 1 上连接。我们想在新的尾轴上连接，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

但我看不到这样一个堆栈的价值。链接的问题确实要求这样一个数组，但没有显示如何使用这样的数组。