Slice numpy array with other array

Question

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.

ecatmur · Accepted Answer · 2012-11-08 13:02:33Z

4

This should work reasonably well:

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

In general:

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

answered Nov 8, 2012 at 13:02

ecatmur

158k28 gold badges311 silver badges387 bronze badges

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

1 Comment

amaurea Over a year ago

The first version does what I want, but only when the number of dimensions of y is known beforehand. The second is general, but sadly seems to be very slow (20 times slower in one of my tests).

Collectives™ on Stack Overflow

Slice numpy array with other array

1 Answer 1

1 Comment

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Related