從多維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]