Choosing from a multidimensional numpy array

Question

I have a multidimensional array a with the form (nt, nz, ny, nx). Dimensions are time, z, y, x. For each x and y, I selected the corresponding z in the new array of indexes J with the form (nt, ny, nx). J contains the height-size indexes I would like to select. Using Python, I could do this in a loop:

b=J.copy() for t in range(nt): for y in range(ny): for x in range(nx): z=J[t,y,x] b[t,y,x]=a[t,z,y,x] 

But I want to do it faster, without cycles. This is probably trivial, but I can't think it over. Is anyone

Answer 1

You can use numpy.indices() along with extended indexing:

 t, y, x = numpy.indices(J.shape) b = a[t, J, y, x]