如何找到最后 2 个轴的 argmax

Question

嘿，我已经看到这个问题了 - 但输出不是我想要的形状。因此，例如，如果我给函数一个维度数组：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]]]])

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

然后我们可以在最后一个维度上取最大值和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

但我看不到这样的堆栈的价值。 linked 问题确实要求这样的数组，但没有说明如何使用这样的数组。

如何找到最后 2 个轴的 argmax

How to find argmax of last 2 axes

python

numpy

argmax

编辑