python numpy: indexing broadcast

Question

suppose I have a 3-by-5 array:

  a=[[ 1.342,  2.244, -0.412, -1.456, -0.426],
     [ 1.884, -0.811,  0.193,  1.322,  0.76 ],
     [-0.654, -0.788,  1.264,  1.034,  0.356]]

and I want to select the 0th element from the first row, 2nd from the second row and 4th from the third row, I would use

a[range(3), [0, 2, 4]]

the result should be:

[1.342, 0.193, 0.356]

How to broadcast to more dimensions? suppose now I have a 2-by-3-by-5 tensor:

[[[-1.054,  0.068, -0.572,  1.535,  1.746],
  [-0.115,  0.356,  0.222, -0.391,  0.367],
  [-0.53 , -0.856,  0.58 ,  1.099,  0.605]],

 [[ 0.31 ,  0.037, -0.85 , -0.054, -0.75 ],
  [-0.097, -1.707, -0.702,  0.658,  0.548],
  [ 1.727, -0.326, -1.525, -0.656,  0.349]]]

For the first dimension a[0], I'd like to select [0,2,4]th element, and for a[1] I'd like to select [1,3,2]th element. Is there a way to do it? If I do it separately for each a[0] and a[1], the result should be:

print( a[0, range(3), [0,2,4]] )
print( a[1, range(3), [1,3,2]] )

>>>[-1.054  0.222  0.605]
   [ 0.037  0.658 -1.525]

The indexing arrays broadcast against in each other in much same as in addition (except without the actual sum). In the first case you have two (3,) arrays. In the second you want a set that broadcasts to (2,3), e.g. (2,1),(3,),(3,) — hpaulj
– hpaulj, Commented Sep 26, 2021 at 16:26

akuiper · Accepted Answer · 2021-09-26 16:02:26Z

3

You can do similar advanced indexing by providing an index for the 1st dimension (make sure it has the correct shape so it can broadcast correctly):

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

a[np.arange(2)[:,None], np.arange(3), idx]    
array([[-1.054,  0.222,  0.605],
       [ 0.037,  0.658, -1.525]])

answered Sep 26, 2021 at 16:02

akuiper

216k33 gold badges362 silver badges379 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

Sam-gege Over a year ago

Thanks @Psidom that works! actually I read the same link you provided before asking this question, but I stuck at the first index: np.arange(2)[:,None]. Could you please explain a bit on why this works?? I've tried using ..., [0, 1], : but they don't work

akuiper Over a year ago

np.arange(2)[:,None] creates an array of shape (2, 1) instead of (2,). We reshape it so it can broadcast with idx array which has shape (2, 3). You can read more about numpy broadcasting rule here.

Ivan · Accepted Answer · 2021-09-26 16:13:30Z

You could apply take_along_axis on the last dimension (axis=2). Using the following array of indices:

>>> indices = np.array([[0, 2, 4],
                        [1, 3, 2]])

However, you first need need to unsqueeeze an additional dimension on indices to match the number of dimensions of the indexed array a:

>>> indices = np.expand_dims(indices, -1)

The following np.take_along_axis call will gather a[i][j][indices[i][j]]:

>>> res = np.take_along_axis(a, indices, 2)
array([[[-1.054],
        [ 0.222],
        [ 0.605]],

       [[ 0.037],
        [ 0.658],
        [-1.525]]], dtype=float32)

You then have to unsqueeze the last dimension:

>>> np.squeeze(res, -1)
array([[-1.054,  0.222,  0.605],
       [ 0.037,  0.658, -1.525]])

As a one-liner this would look like:

>>> np.take_along_axis(a, indices[..., None], -1)[..., 0]
array([[-1.054,  0.222,  0.605],
       [ 0.037,  0.658, -1.525]])

Collectives™ on Stack Overflow

python numpy: indexing broadcast

2 Answers 2

2 Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

2 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related