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.

Answer 1

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.)