I have a collection of 2D narrays, depending on two integer indexes, say p1 and p2, with each matrix of the same shape.
Then I need to find, for each pair (p1,p2), the maximum value of the matrix and the indexes of these maxima. A trivial, albeit slow, way to do this would would be to do something like this
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 my application this is unfortunately too slow, I guess due to the usage of double for loops. Therefore I tried to pack everything in a 4-dimensional array (say BigMatrix), where the first two coordinates are the indexes p1,p2, and the other 2 are the coordinates of the matrices.
The np.amax command
>>res=np.amax(BigMatrix,axis=(2,3)) >>res.shape (10,20) >>res[p1,p2]==np.amax(BigMatrix[p1,p2,:,:]) True
works as expected, as it loops through the 2 and 3 axis. How can I do the same for np.argmax? Please keep in mind that speed is important.
Thank you very much in advance,