Numpy 3d array indexing

Question

I have a 3d numpy array (n_samples x num_components x 2) in the example below n_samples = 5 and num_components = 7.

I have another array (indices) which is the selected component for each sample which is of shape (n_samples,).

I want to select from the data array given the indices so that the resulting array is n_samples x 2.

The code is below:

import numpy as np
np.random.seed(77)
data=np.random.randint(low=0, high=10, size=(5, 7, 2))
indices = np.array([0, 1, 6, 4, 5])
#how can I select indices from the data array?

For example for data 0, the selected component should be the 0th and for data 1 the selected component should be 1.

Note that I can't use any for loops because I'm using it in Theano and the solution should be solely based on numpy.

hpaulj · Accepted Answer · 2017-03-15 06:35:37Z

5

Is this what you are looking for?

In [36]: data[np.arange(data.shape[0]),indices,:]
Out[36]: 
array([[7, 4],
       [7, 3],
       [4, 5],
       [8, 2],
       [5, 8]])

answered Mar 15, 2017 at 6:35

hpaulj

233k14 gold badges260 silver badges392 bronze badges

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

1 Comment

Ash Over a year ago

yeah, that's it. I think this is the best answer.

kennytm · Accepted Answer · 2017-03-15 06:36:58Z

4

To get component #0, use

data[:, 0]

i.e. we get every entry on axis 0 (samples), and only entry #0 on axis 1 (components), and implicitly everything on the remaining axes.

This can be easily generalized to

data[:, indices]

to select all relevant components.

But what OP really wants is just the diagonal of this array, i.e. (data[0, indices[0]], (data[1, indices[1]]), ...) The diagonal of a high-dimensional array can be extracted using the diagonal function:

>>> np.diagonal(data[:, indices])
array([[7, 7, 4, 8, 5],
       [4, 3, 5, 2, 8]])

(You may need to transpose the result.)

edited Mar 15, 2017 at 6:36

answered Mar 15, 2017 at 6:11

kennytm

526k110 gold badges1.1k silver badges1k bronze badges

8 Comments

Ash Over a year ago

data[:, indices] results in shape (5, 5, 2) while I need a (5, 2) shape.

kennytm Over a year ago

@Afshin data[:, 0] is (5, 2), for a single component. What (5, 2) array you want when you have got five components?

jyalim Over a year ago

data[:,indices][np.arange(n_samples),np.arange(n_samples)] works but is obtuse.

kennytm Over a year ago

@Afshin What about np.diagonal(data[:, indices]).T

kennytm Over a year ago

@Afshin Updated. I guess hpaulj's one is faster since it doesn't need to build the intermediate #Samples×#Indices×2 array.

|

jyalim · Accepted Answer · 2017-03-15 06:37:41Z

2

You have a variety of ways to do so, but this is my loop recommendation:

selection = np.array([ datum[indices[k]] for k,datum in enumerate(data)])

The resulting array, selection, has the desired shape.

edited Mar 15, 2017 at 6:37

answered Mar 15, 2017 at 6:11

jyalim

3,4491 gold badge17 silver badges22 bronze badges

2 Comments

Ash Over a year ago

I can't use for loops because I'm using Theano, need a solely numpy solution.

jyalim Over a year ago

I think you should add that constraint to your original post.

Collectives™ on Stack Overflow

Numpy 3d array indexing

3 Answers 3

1 Comment

8 Comments

2 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

1 Comment

8 Comments

2 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related