python - 从多维 numpy 数组中选择

Question

我有一个形状为 (nt, nz, ny, nx) 的多维数组 a。维度是时间、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]