Чтение элементов из массива 4D с использованием другого массива в качестве индекса

Question

У меня есть массив dataX, имеющий форму (4, 9359, 10852), и у меня есть второй массив dataY, имеющий форму (9359, 10852). dataY может иметь значения от 1 до 6. Я хочу сделать следующее.

1. Получить местоположение 5.0s в dataY.
2. Используйте это как индекс в dataX и получите значения во всех 4 измерениях в виде меньшего массива. И я сделал это

---

idx = np.where(dataY == 5.0)    
element = dataX[:,idx[0],idx[1]]

idx[0].size = 860498

element.shape = (4, 860498)

Но я хочу, чтобы элемент был трехмерным массивом, а не двухмерным списком. Как мне это сделать?

Answer 1

idx идентифицирует те элементы dataY, которые равны 5,0. Было найдено много таких, но они не могут (вообще) быть расположены в 2d массиве. Таким образом, dataY[idx] вернет (860498,) массив из 5-х.

Аналогичным образом, когда idx используется для выбора элементов из dataX, он «выравнивает» последние 2 измерения.

Это может быть легче визуализировать с помощью небольших массивов, которые можно распечатать полностью.

In [104]: y=np.arange(12).reshape(4,3)
In [105]: y
Out[105]: 
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11]])
In [106]: idx = np.where(y%2)
In [107]: idx
Out[107]: (array([0, 1, 1, 2, 3, 3]), array([1, 0, 2, 1, 0, 2]))
In [108]: y[idx]
Out[108]: array([ 1,  3,  5,  7,  9, 11])
In [110]: np.ma.masked_array(y, y%2==0)
Out[110]: 
masked_array(
  data=[[--, 1, --],
        [3, --, 5],
        [--, 7, --],
        [9, --, 11]],
  mask=[[ True, False,  True],
        [False,  True, False],
        [ True, False,  True],
        [False,  True, False]],
  fill_value=999999)