np.argmax 在多维数组上，保持一些索引固定

Question

我有一组二维数组，取决于两个整数索引，比如 p1 和 p2，每个矩阵的形状相同。

然后我需要为每一对 (p1,p2) 找到矩阵的最大值和这些最大值的索引。一个简单但缓慢的方法是做这样的事情

import numpy as np
import itertools
range1=range(1,10)
range2=range(1,20)

for p1,p2 in itertools.product(range1,range1):
    mat=np.random.rand(10,10)
    index=np.unravel_index(mat.argmax(), mat.shape)
    m=mat[index]
    print m, index

不幸的是，对于我的应用程序，这太慢了，我猜是由于使用了双 for 循环。因此，我尝试将所有内容打包到一个 4 维数组（比如 BigMatrix）中，其中前两个坐标是索引 p1、p2，另外两个是矩阵的坐标。

np.amax命令

    >>res=np.amax(BigMatrix,axis=(2,3))
    >>res.shape
         (10,20)
    >>res[p1,p2]==np.amax(BigMatrix[p1,p2,:,:])
         True

按预期工作，因为它循环通过 2 和 3 轴。我怎样才能为 np.argmax 做同样的事情？请记住，速度很重要。

非常感谢您，

恩佐

Answer 1

这对我有用，其中 Mat 是大矩阵。

# flatten the 3 and 4 dimensions of Mat and obtain the 1d index for the maximum
# for each p1 and p2
index1d = np.argmax(Mat.reshape(Mat.shape[0],Mat.shape[1],-1),axis=2)

# compute the indices of the 3 and 4 dimensionality for all p1 and p2
index_x, index_y = np.unravel_index(index1d,Mat[0,0].shape)

# bring the indices into the right shape
index = np.array((index_x,index_y)).reshape(2,-1).transpose()

# get the maxima
max_val = np.amax(Mat,axis=(2,3)).reshape(-1)

# combine maxima and indices
sol = np.column_stack((max_val,index))

print sol

np.argmax 在多维数组上，保持一些索引固定

np.argmax on multidimensional arrays, keeping some indexes fixed

indexing

numpy

multidimensional-array

argmax