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 

Для моего приложения это, к сожалению, слишком медленно, я думаю, из-за использования двойных для петель. Поэтому я попытался собрать все в 4-мерном массиве (скажем, BigMatrix), где первые две координаты - это индексы p1, p2, а остальные 2 - координаты матриц.

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