Numpy max along axis

Question

Am I missing something here? I would expect that np.max in the following snippet would return [0, 4]...

>>> a
array([[1, 2],
       [0, 4]])

>>> np.max(a, axis=0)
array([1, 4])

Thanks for any pointers.

Answer 1

Looks like you want the row that contains the maximum value, right?

max(axis=0) returns the maximum of [1,0] and [2,4] independently.

argmax without axis parameter finds the maximum over the whole array - in flattened form. To turn that index into row number we have to use unravel_index:

In [464]: a.argmax()
Out[464]: 3
In [465]: np.unravel_index(3,(2,2))
Out[465]: (1, 1)
In [466]: a[1,:]
Out[466]: array([0, 4])

or in one expression:

In [467]: a[np.unravel_index(a.argmax(), a.shape)[0], :]
Out[467]: array([0, 4])

As you can see from the length of the answer it's not the usual definition of maximum along/over an axis.

Sum along axis in numpy array may give more insight into the meaning of 'along axis'. The same definitions apply to the sum, mean and max operations.

===================

To pick row with the largest norm, first calculate the norm. norm uses the axis parameter in the same way.

In [537]: np.linalg.norm(a,axis=1)
Out[537]: array([ 2.23606798,  4.        ])
In [538]: np.argmax(_)
Out[538]: 1
In [539]: a[_,:]
Out[539]: array([0, 4])