Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Slice numpy array with other array

Tags:

python

numpy

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.

like image 536
amaurea Avatar asked Jul 18 '26 22:07

amaurea


1 Answers

This should work reasonably well:

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

In general:

y[tuple(np.rollaxis(x, -1))]
like image 173
ecatmur Avatar answered Jul 20 '26 11:07

ecatmur



Donate For Us

If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!