Imagine you have a 3d array like e.g.

A = np.array([[[1,2], [2,3] ], [[3,4], [1,2]], [[2,3], [3,4]]])

A.shape (3, 2, 2)

array([[[1, 2],
        [2, 3]],

       [[3, 4],
        [1, 2]],

       [[2, 3],
        [3, 4]]])

Now, if I want to get the indices of maxima along the first dimension this is easy as

A.argmax(axis=0)

array([[1, 1],
       [2, 2]])

How do I do the same for a 4d array finding the maximum along the first two dimensions?

Example 4d array

np.random.randint(0, 9,[2,3,2,2])

B = array([[[[5, 4],
         [3, 8]],

        [[3, 5],
         [0, 4]],

        [[0, 1],
         [3, 0]]],


       [[[0, 2],
         [7, 3]],

        [[7, 3],
         [8, 0]],

        [[8, 3],
         [2, 7]]]])

B.shape (2, 3, 2, 2)

In that case, the output should still be a (2, 2) matrix, but each cell containing a tuple of the maximum index of dimension 0 and dimension 1, i.e.

Example output

array([[(1, 2), (0, 1)],
       [(1, 1), (0, 0)]])

Solution

As correctly pointed out by @hpaulj the solution to this problem lies in calling

idx = B.reshape(6,2,2).argmax(axis=0)
2d_indices = np.unravel_index(idx,(2,3))

However, this way the ordering is a bit different to what I expected (and outlined in my question above). To get the output exactly as above, just add

list_of_tuples = [a for a in zip(2d_indices[0].flatten(), 2d_indices[1].flatten())]
np.array(list_of_tuples).reshape(3, 3, -1)

It looks a bit cumbersome to me and there may be a more direct way of getting there, but I still want to post it for the sake of completeness :)

CodePudding user response：

argmax only accepts scalar axis values (some other functions allow tuples). But we can reshape B:

In [18]: B.reshape(6,2,2).argmax(axis=0)
Out[18]: 
array([[5, 1],
       [4, 0]])

and convert those back to 2d indices:

In [21]: np.unravel_index(_18,(2,3))
Out[21]: 
(array([[1, 0],
        [1, 0]]),
 array([[2, 1],
        [1, 0]]))

those values could be reordered in various ways, for example:

In [23]: np.transpose(_21)
Out[23]: 
array([[[1, 2],
        [1, 1]],

       [[0, 1],
        [0, 0]]])