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.
This should work reasonably well:
y[x[..., 0], x[..., 1], x[..., 2]]
In general:
y[tuple(np.rollaxis(x, -1))]
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With