How can i get sub array depending on condition?

Question

For example, i have an complex array of shape (3,3,3) and i want to find angle of the element when magnitude of the array along axis 2 is maximum. That is,

a = np.random.random((3, 3, 3)) + 1j * np.random.random((3, 3, 3))

amp3d = np.abs(a)
amp2d = np.max(amp3d, axis=2)

angle2d =  np.zeros_like(amp2d)

for i in range(a.shape[0]):
    for j in range(a.shape[1]):
        for k in range(a.shape[2]):
            if amp3d[i, j, k] == amp2d[i,j]:
                angle2d[i, j] = np.angle(a[i, j, k])

Of course i can loop, but i guess there is exist a pretty pythonic way to accomplish this task

In Matlab it would be look similar to np.angle(a[np.argmax(a, axis=2)])

Thanks!

Answer 1

Create a mask/boolean array for the relevant values by expanding the dimension of amp2d to match the dimension of amp3d, and multiply that with a and select the maximum value for axis=2.

a = np.random.random((3, 3, 3)) + 1j * np.random.random((3, 3, 3))

amp3d = np.abs(a)
amp2d = np.max(amp3d, axis=2)

mask = amp3d == amp2d[...,np.newaxis]
angle2d = np.angle((a*mask).max(axis=2))

Check the numpy documentation on indexing for reference what np.newaxis does, or check this stackoverflow question.