从多维numpy数组中选择

Question

我有一个多维数组a，具有形状（nt，nz，ny，nx）。尺寸是时间，z，y，x。对于每一次x和y，我已经在shape（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]