Multi-dimensional np.argmax?

Question

I've got a 3D array of shape (n, n, g) and I'd need for every (n, n) the argmax, i.e. the result should be two index vectors (x, y) of length g each.

The intuitive solution would be:

array = np.random.uniform(size=[5, 5, 1000])
np.argmax(array, axis=[0, 1])

However, numpy does not support multiple axes as argument.

Is there a solution to get this result anyways?

Answer 1

It seems you are looking to get two-dimensional (row,column) argmax indices for each flattened version of 2D slice with the first two axes merged/combined as to speak. So, the first block would be array[:,:,0] and so on and we need to need find argmax with that slice being flattened and retraced back to original 2D shape. So, to solve it, we can simply reshape to merge first two axes, get argmax along the first axis which is the merged one after reshaping and retrace back the original indices with np.unravel_index, like so -

array2D = array.reshape(-1,array.shape[-1])
r,c = np.unravel_index(array2D.argmax(0),array.shape[:2])

Sample run -

In [29]: array = np.random.uniform(size=[5, 5, 1000])

In [30]: array2D = array.reshape(-1,array.shape[-1])

In [31]: r,c = np.unravel_index(array2D.argmax(0),array.shape[:2])

In [32]: len(r), len(c)
Out[32]: (1000, 1000)

Let's verify results for the first 2D slice -

In [33]: array[:,:,0]
Out[33]: 
array([[0.81590174, 0.17919069, 0.22717883, 0.67863625, 0.97390595],
       [0.82096447, 0.05894774, 0.86379174, 0.13494354, 0.10003756],
       [0.37243189, 0.33714008, 0.21165031, 0.35910642, 0.15163255],
       [0.1376776 , 0.86866599, 0.43602004, 0.85421372, 0.77805012],
       [0.10519547, 0.7422571 , 0.35632275, 0.24168307, 0.76882613]])

In [34]: array[:,:,0].argmax() 
Out[34]: 4  # flattened index for 0.97390595 at (0,4) in the 2D slice

In [36]: r[0],c[0]
Out[36]: (0, 4)

Answer 2

Simple using vstack before argmax

np.argmax(np.vstack(array),0)//5
Out[61]: array([4, 0], dtype=int64)