从多维numpy数组中选择

Question

我有一个多维数组a形状(nt,nz,ny,nx).尺寸是时间,z,y,x.对于x和y的每一次,我在具有形状(nt,ny,nx)的新索引数组J中选择了适当的z.J包含我想要选择的高度维度的索引.使用Python,我可以在循环中执行此操作:

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]

但我想更快地做到这一点,没有循环.这可能是微不足道的,但我无法理解它.任何人？

Answer 1

您可以numpy.indices()与高级索引一起使用:

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