Выберите значения из одного массива, используя значения другого массива в numpy - PullRequest
0 голосов
/ 20 июня 2020

В numpy, у меня 3D массив. Вдоль оси 0 хранится несколько 2D-плоскостей. Мне нужно получить градиент каждой из этих плоскостей, выбрать среднюю величину градиента в каждой точке на этих плоскостях и, следовательно, изолировать соответствующие компоненты градиента x и y. Но у меня возникают проблемы с выполнением этого должным образом.

Пока что для получения градиента и медианы у меня есть: gradient, я использую:

pos=(mag_grad==med).argmax(axis=0) #This returns the first instance where the median element encountered along axis=0
G = np.stack([np.zeros(med.shape),np.zeros(med.shape)], axis=0) #Will store y and x median components of the gradient, respectively.
for i in range(med.shape[0]):
    for j in range(med.shape[1]):
        G[0,i,j], G[1,i,j] = grad[0][pos[i,j],i,j], grad[1][pos[i,j],i,j] #Manually select the median y and x components of the gradient, and save to G.

Я считаю, что второй блок кода работает правильно. Однако это очень неэлегантно, и поскольку я не смог найти способ сделать это в NumPy, мне пришлось использовать Python l oop, что добавляет много накладных расходов. Вдобавок, поскольку эта операция часто выполняется в NumPy, я подозреваю, что для этого должен быть встроенный способ.

Как я могу реализовать этот код более эффективно и элегантно?

1 Ответ

0 голосов
/ 20 июня 2020

Использование itertools для индексации вашего массива может сделать его более эффективным / элегантным.

import itertools

idxs = np.array(list(itertools.product(range(med.shape[0]), range(med.shape[1]))))
G[0,idxs], G[1,idxs] = grad[0][pos[idxs],idxs], grad[1][pos[idxs],idxs]
...