CODEWITHSUNDEEP
pythonnumpy
amaurea
amaurea

Reputation: 5067

Slice numpy array with other array

I am trying to extract a subset of a a numpy array y specified by a set of indices contained in x, while still leaving some indices of y free. For a concrete example. Let y have shape (10,10,10,3) while x has shape (7,7,3). The last dimension of x corresponds to indices info the first three dimensions of y. That is, I would like an efficient slicing operation with the same result as this:

for i in x.shape[0]:
    for j in x.shape[1]:
        z[i,j,:] = y[x[i,j,0],x[i,j,1],x[i,j,2],:]

Ideally the answer would work regardless of the number of dimensions of x.

In general, y would be N+1-dimensional, with shape (...,N), while x would be Q+1-dimensional with shape (...,N), and the result would have the same shape as x.

The motivation for this is extracting a subset of vectors from a vector field.

Upvotes: 3

Views: 294

Answers (1)

ecatmur
ecatmur

Reputation: 157484

This should work reasonably well:

y[x[..., 0], x[..., 1], x[..., 2]]

In general:

y[tuple(np.rollaxis(x, -1))]

Upvotes: 4

Related Questions