从多维numpy数组中选择

Question

I have a multidimensional array a with shape (nt, nz, ny, nx). The dimensions are time, z, y, x. For each time, x and y, I've selected the appropriate z in a new index array J with shape (nt, ny, nx). J contains the indices along the height dimension that I'd 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 this faster, without the loops. This is probably trivial, but I can't get my head around it. Anyone?

Answer 1

You can use numpy.indices() together with advanced indexing:

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