Выбор из многомерного массива NumPy

Question

У меня есть многомерный массив a с формой (nt, nz, ny, nx).Размеры - время, z, y, x.Для каждого времени x и y я выбрал соответствующий z в новом индексном массиве J с формой (nt, ny, nx).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]

Но я хочу сделать это быстрее, без циклов.Это, вероятно, тривиально, но я не могу разобраться с этим.Кто-нибудь? * * 1004

Answer 1

Вы можете использовать numpy.indices() вместе с расширенной индексацией:

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