Selecting from multidimensional numpy array

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?

Sven Marnach · Accepted Answer · 2011-04-18 11:06:06Z

8

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

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

answered Apr 18, 2011 at 11:06

Sven Marnach

608k123 gold badges966 silver badges865 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Selecting from multidimensional numpy array

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related