Is there a way to conditionally index 3D-numpy array?

Question

Having an array A with the shape (2,6, 60), is it possible to index it based on a binary array B of shape (6,)?

The 6 and 60 is quite arbitrary, they are simply the 2D data I wish to access.

The underlying thing I am trying to do is to calculate two variants of the 2D data (in this case, (6,60)) and then efficiently select the ones with the lowest total sum - that is where the binary (6,) array comes from.

Example: For B = [1,0,1,0,1,0] what I wish to receive is equal to stacking

A[1,0,:]
A[0,1,:]
A[1,2,:]
A[0,3,:]
A[1,4,:]
A[0,5,:]

but I would like to do it by direct indexing and not a for-loop.

I have tried A[B], A[:,B,:], A[B,:,:] A[:,:,B] with none of them providing the desired (6,60) matrix.

import numpy as np
A = np.array([[4, 4, 4, 4, 4, 4], [1, 1, 1, 1, 1, 1]])
A = np.atleast_3d(A)
A = np.tile(A, (1,1,60)
B = np.array([1, 0, 1, 0, 1, 0])
A[B]

Expected results are a (6,60) array containing the elements from A as described above, the received is either (2,6,60) or (6,6,60).

Thank you in advance, Linus

blubberdiblub · Accepted Answer · 2019-04-15 14:02:35Z

1

You can generate a range of the indices you want to iterate over, in your case from 0 to 5:

count = A.shape[1]

indices = np.arange(count)  # np.arange(6) for your particular case

>>> print(indices)
array([0, 1, 2, 3, 4, 5])

And then you can use that to do your advanced indexing:

result_array = A[B[indices], indices, :]

If you always use the full range from 0 to length - 1 (i.e. 0 to 5 in your case) of the second axis of A in increasing order, you can simplify that to:

result_array = A[B, indices, :]
# or the ugly result_array = A[B, np.arange(A.shape[1]), :]

Or even this if it's always 6:

result_array = A[B, np.arange(6), :]

edited Apr 15, 2019 at 14:02

answered Apr 15, 2019 at 13:55

blubberdiblub

4,1451 gold badge31 silver badges32 bronze badges

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

Comments

Tarifazo · Accepted Answer · 2019-04-15 14:03:31Z

0

An alternative solution using np.take_along_axis (from version 1.15 - docs)

import numpy as np
x = np.arange(2*6*6).reshape((2,6,6))
m = np.zeros(6, int)
m[0] = 1
#example: [1, 0, 0, 0, 0, 0]

np.take_along_axis(x, m[None, :, None], 0)   #add dimensions to mask to match array dimensions

>>array([[[36, 37, 38, 39, 40, 41],
        [ 6,  7,  8,  9, 10, 11],
        [12, 13, 14, 15, 16, 17],
        [18, 19, 20, 21, 22, 23],
        [24, 25, 26, 27, 28, 29],
        [30, 31, 32, 33, 34, 35]]])

answered Apr 15, 2019 at 14:03

Tarifazo

4,3531 gold badge11 silver badges24 bronze badges

1 Comment

Linus Eriksson Over a year ago

Thank you, I will check this solution out and compare to blubberdiblub's!

Collectives™ on Stack Overflow

Is there a way to conditionally index 3D-numpy array?

2 Answers 2

Comments

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related