Index n-th dimension of Numpy array with 2d index array

127 views Asked by At

I have the following 2-D Numpy arrays:

X # X.shape = (11688, 144)
Y # Y.shape = (2912, 1000)

The first array is populated with atmospheric data, and the second array is populated with random index values from 0 to X.shape[0]-1. I want to index the rows of X with each column of Y to yield a 3-D array result, where result.shape = (2912, 1000, 144), and I want to do this without looping.

My current approach is:

result = X[Y,:]

but this one line of code can take more than 10 seconds to execute depending on the shape of the 0th axis of Y.

Is there a more optimal way to perform this type of indexing in order to speed up its execution?

EDIT: Here's a more complete example of what I'm trying to accomplish.

X = np.random.rand(11688, 144) # Time-by-longitude array of atmospheric data
t = np.arange(X.shape[0])      # Time vector

# Populate array of randomly drawn time steps
Y = np.zeros((2912, 1000), dtype='i')
for i in xrange(1000):
    Y[:,i] = np.random.choice(t, 2912)

# Index X with each column of Y
result = X[Y,:]
0

There are 0 answers