3

I want to find the location of minima along a given axis in a rank-3 numpy array. I have obtained these locations with np.argmin, however I'm not sure how to "apply" this to the original matrix to get the actual minima.

For example:

import numpy as np

a = np.random.randn(10, 5, 2)
min_loc = a.argmin(axis = 0)   # this gives an array of shape (5, 2)

Now, the problem is how do I get the actual minima using min_loc? I have tried a[min_loc], which gives me a shape (5, 2, 5, 2). What's the logic for this shape? How can I use this auxiliary matrix to get a sensible solution of shape (5, 2)

Note that a.min(axis = 0) is not the solution I'm looking for. I need a solution via argmin.

Pythonist
  • 1,937
  • 1
  • 14
  • 25
  • 1
    [Documentation suggests np.unravel_index](https://numpy.org/doc/stable/reference/generated/numpy.argmin.html) – Mephy Feb 27 '21 at 19:30
  • 1
    Look at `take_along_axis`. `min_loc` is a (5,2) array of numbers in the 0-9 range. It's a valid index for the 1st dimension, but you have index the other dimensions correctly. `take_along_axis` facilitates this. – hpaulj Feb 27 '21 at 19:32
  • 1
    The other index arrays together also need to index a (5,2) block (using broadcasting). – hpaulj Feb 27 '21 at 20:24
  • I add you to a long and illustrious chain. This is mentioned in the docs and beaten to death here too, but only if you have the right keywords in your search. Apparently they're hard to intuit until after you've missed them. – Mad Physicist Feb 28 '21 at 04:41
  • @Mephy bad advice. `unravel_index` does not work for this case (I need minima along a given axis, not absolute minimum of the array as it's the task tackled in the doc). It took me a while to realise what was happening... – Pythonist Feb 28 '21 at 11:38

1 Answers1

1

a[min_loc] does integer array indexing on the first dimension, i.e. it will pick up (5, 2) shaped array for each index in min_loc. Since min_loc itself is (5, 2) shaped, and for each integer in min_loc, it picks up another (5, 2) shaped array. You end up with a (5, 2, 5, 2) array. Same reason a[np.array([0, 3])] has a shape of (2, 5, 2) and a[np.array([[0], [3]])] has a shape of (2, 1, 5, 2), since you only provide the index for the 1st dimension.

For your usecase, you do not want to pick up a subarray for each index in min_loc but rather you need an element. For instance, if you have min_loc = [[5, ...], ...], the first element should have a full indice of 5, 0, 0 instead of 5, :, :. This is exactly what advanced indexing does. Basically by providing an integer array as index for each dimension, you can pick up the element corresponding to the specific positions. And you can construct indices for the 2nd and 3rd dimensions from a (5, 2) shape with np.indices:

j, k = np.indices(min_loc.shape)
a[min_loc, j, k]

# [[-1.82762089 -0.80927253]
#  [-1.06147046 -1.70961507]
#  [-0.59913623 -1.10963768]
#  [-2.57382762 -0.77081778]
#  [-1.6918745  -1.99800825]]

where j, k are coordinates for the 2nd and 3rd dimensions:

j
#[[0 0]
# [1 1]
# [2 2]
# [3 3]
# [4 4]]

k  
#[[0 1]
# [0 1]
# [0 1]
# [0 1]
# [0 1]]

Or as @hpaulj commented, use np.take_along_axis method:

np.take_along_axis(a, min_loc[None], axis=0)

# [[[-0.93515242 -2.29665325]
#   [-1.30864779 -1.483428  ]
#   [-1.24262879 -0.71030707]
#   [-1.40322789 -1.35580273]
#   [-2.10997209 -2.81922197]]]
Psidom
  • 209,562
  • 33
  • 339
  • 356
  • 1
    Thanks. The `take_along_axis` is certainly the most elegant solution. But the "hand made" solution with explicit indices certainly helps to clarify why my approach was not working and what was missing. – Pythonist Feb 28 '21 at 11:41