Есть ли способ применить numpy .extract к 3d-массиву?

Question

У меня есть следующее:

X = np.ndarray (324,349,24)
y = np.ndarray (324,349)

Я хотел бы создать словарь для условного извлечения, например:

myDict = {'keyA':cond,'keyB':cond,'keyC':cond,'keyD':cond}

Каждый ключ в myDict.keys () - это численно представлены в пределах y . Я хотел бы применить маску и извлечь только те индексы в пределах X , которые соответствуют заданной маске.

Например,

condA = y == 0 ... condD = y == 3

Как бы я go применил это условие к X ? Я думал что-то вроде этого:

for k in range(1, X.shape[2]):
  myDict['keyA'] = np.extract(condA,k)
  myDict['keyB'] = np.extract(condB,k)
  myDict['keyC'] = np.extract(condC,k)
  myDict['keyD'] = np.extract(condD,k)

Однако я получаю ошибку:

IndexError: index is out of bounds for size

Ожидаемый результат: Словарь:

myDict{'keyA':ndarray [n,n,24],'keyB':ndarray[n,n,24],'keyC':ndarray[n,n,24],'keyD':ndarray[n,n,24]}

Answer 1

Функция numpy where может быть полезна здесь:

myDict['keyA'] = X.where(condA)