Select values from one array using values of another array in numpy

Question

In numpy, I have a 3D array. Along the 0 axis, it stores multiple 2D planes. I need to get the gradient of each of these planes, select the median gradient magnitude at each point across these planes, and hence isolate the corresponding x and y gradient components. But I'm having difficulty carrying this out properly.

So far, to get the gradient and median, I have:

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

Then to select the correct x & y components of the gradient, I use:

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.

I believe the 2nd code block works correctly. However, it is very inelegant, and because I couldn't find a way to do this in NumPy, I had to use a Python loop which adds a large amount of overhead. In addition, since this operation occurs frequently in NumPy, I suspect there should be an in-built way to do this.

How can I implement this code more effectively and elegantly?

Answer 1

Using itertools to index your array can make it more efficient/elegant.

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]