从多维numpy数组中选择

Question

我有一个带有形状的多维数组（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]