获取轴上numpy.argmax元素的索引

Question

我有N维矩阵,其中包含具有N个参数的函数的值.每个参数都有一个离散数量的值.我需要最大化除了一个参数之外的所有参数的函数,从而产生大小等于非最大化参数的值的一维向量.我还需要保存其他参数所采用的值.

为此,我想迭代地应用numpy.max不同的轴来减少矩阵的维数以找到我需要的东西.最终的矢量将取决于我遗漏的参数.

然而,我无法找到最终元素的原始索引(其中包含有关其他参数所采用值的信息).我虽然numpy.argmax以同样的方式使用,numpy.max但我无法获得原始索引.

我正在尝试的一个例子是:

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

这回来了

[[1 1]
 [1 1]]

这意味着argmax正在选择原始矩阵中的元素(2,1,4,7).但如何获得他们的指数？我尝试unravel_index使用args直接作为矩阵的索引x,从numpy到索引的一堆函数没有成功.

使用numpy.where不是解决方案,因为输入矩阵内部可能具有相等的值,因此我无法从不同的原始值中辨别出来.

Answer 1

x.argmax(0)给出沿第一轴的索引的最大值.使用np.indices产生的其他轴的索引.

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

如果x有其他维度,生成a1和a2适当.

官方文档没有说明如何使用argmax,但早期的SO线程已经讨论过它.我从多维数组上使用numpy.argmax()得到了这个一般性的想法