CODEWITHSUNDEEP
pythonnumpyarray-broadcasting
Chris J Harris
Chris J Harris

Reputation: 1841

Defining numpy indexing arrays

I am having a point of confusion over numpy indexing. Let's say I have a three-dimensional array, like:

test_arr = np.arange(3*2*3).reshape(3,2,3)
test_arr
array([[[ 0,  1,  2],
        [ 3,  4,  5]],

       [[ 6,  7,  8],
        [ 9, 10, 11]],

       [[12, 13, 14],
        [15, 16, 17]]])

I would like to index this by a boolean array along dimension 1:

dim1_idx = np.array([True, False])
test_arr[:, dim1_idx, :]

which gives me 

array([[[ 0,  1,  2]],

       [[ 6,  7,  8]],

       [[12, 13, 14]]])

All good so far.

My question is, is there a way that I can define this boolean index array in advance - like (and this doesn't work):

all_dim_idx = dim1_idx[np.newaxis, :, np.newaxis]
test_arr[all_dim_idx]

I realize that the reason this doesn't is because it can't broadcast in a way to make the all_dim_idx array fit test_arr. I could use np.tile or np.reshape to make the index array fit onto the larger array, but (as well as not being then generalizable to other array shapes) I just get the impression that there's probably a better way. Can anyone enlighten me?

Thanks in advance!

Upvotes: 1

Views: 58

Answers (1)

hpaulj
hpaulj

Reputation: 231385

In [600]: test_arr = np.arange(3*2*3).reshape(3,2,3)                            
In [601]: dim1_idx = np.array([True, False])

Define an indexing tuple:

In [602]: idx = (slice(None), dim1_idx, slice(None))                            
In [603]: test_arr[idx]                                                         
Out[603]: 
array([[[ 0,  1,  2]],

       [[ 6,  7,  8]],

       [[12, 13, 14]]])

Upvotes: 3

Related Questions