使用numpy中另一个数组的值从一个数组中选择值

Question

在 numpy 中，我有一个 3D 数组。沿 0 轴，它存储多个 2D 平面。我需要获取每个平面的梯度，选择这些平面上每个点的中值梯度幅度，从而隔离相应的 x 和 y 梯度分量。但我很难正确执行此操作。

到目前为止，为了得到梯度和中位数，我有：

img_l = #My 3D array of 2D planes
grad = np.gradient(img_l,axis=[1,2]) #Get gradient of each image. This is a list with 2 elements.
mag_grad = np.sqrt(grad[0]**2 + grad[1]**2) #Get magnitude of gradient in each case
med = np.median(mag_grad, axis=0) #Get median value at each point in the planes

然后为了选择正确的渐变 x & y 分量，我使用：

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 循环，这会增加大量开销。另外，由于这个操作在 NumPy 中经常发生，我怀疑应该有一个内置的方法来做到这一点。

我怎样才能更有效、更优雅地实现这段代码？

Answer 1

使用 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]