I have a three dimensional numpy array. What is the fastest way to get a 3D array that has the largest item of each of final axis of the array without writing a loop.(I will later use CuPy with the same syntax, and loops would take away the GPU parallelism and the speed which is the most important factor here.)
Getting the indexes for largest items is easy:
>>> arr = np.array(
[[[ 6, -2, -6, -5],
[ 1, 12, 3, 9],
[21, 7, 9, 8]],
[[15, 12, 20, 12],
[17, 15, 17, 23],
[22, 18, 27, 32]]])
>>> indexes = arr.argmax(axis=2, keepdims=True)
>>> indexes
array([[[0],
[1],
[0]],
[[2],
[3],
[3]]])
but how to use those indexes to get the chosen values from arr? All the ways I have tried either produce errors (such as arr[indexes]), or wrong results. What I would like to get in this example is
array([[[6],
[12],
[21]],
[[20],
[23],
[32]]])
As pointed by hpaulj, you can use take_along_axis
>>>arr = np.array(
[[[ 6, -2, -6, -5],
[ 1, 12, 3, 9],
[21, 7, 9, 8]],
[[15, 12, 20, 12],
[17, 15, 17, 23],
[22, 18, 27, 32]]])
>>>indexes = arr.argmax(axis=2, keepdims=True)
>>>np.take_along_axis(arr , indexes , axis=2)
array([[[6],
[12],
[21]],
[[20],
[23],
[32]]])