Numpy: Index 3D array with index of last axis stored in 2D array

Question

I have a ndarray of shape(z,y,x) containing values. I am trying to index this array with another ndarray of shape(y,x) th

Answer 1

If you have numpy >= 1.15.0 you could use numpy.take_along_axis. In your case:

result_array = numpy.take_along_axis(val_arr, z_indices.reshape((3,3,1)), axis=2)

That should give you the result you want in one neat line of code. Note the size of the indices array. It needs to have the same number of dimensions as your val_arr (and the same size in the first two dimensions).