Broadcasting a function to a 3D array Python

Question

I tried understanding numpy broadcasting with 3d arrays but I think the OP there is asking something slightly different.

I have a 3D numpy array like so -

IQ = np.array([
    [[1,2],
    [3,4]],
    [[5,6],
    [7,8]]
], dtype = 'float64')

The shape of this array is (2,2,2). I want to apply a function to each 1x2 array in this 3D matrix like so -

def func(IQ):
   I = IQ[0]
   Q = IQ[1]
   amp = np.power((np.power(I,2) + np.power(Q, 2)),1/2)
   phase = math.atan(Q/I)
   return [amp, phase]

As you can see, I want to apply my function to each 1x2 array and replace it with the return value of my function. The output is a 3D array with the same dimensions. Is there a way to broadcast this function to each 1x2 array in my original 3D array? Currently I am using loops which becomes very slow as the 3D array increases in dimensions.

Currently I am doing this -

#IQ is defined from above

for i in range(IQ.shape[0]):
    for j in range(IQ.shape[1]):
        I = IQ[i,j,0]
        Q = IQ[i,j,1]
        amp = np.power((np.power(I,2) + np.power(Q, 2)),1/2)
        phase = math.atan(Q/I)
        IQ[i,j,0] = amp
        IQ[i,j,1] = phase

And the returned 3D array is -

 [[[ 2.23606798  1.10714872]
  [ 5.          0.92729522]]

 [[ 7.81024968  0.87605805]
  [10.63014581  0.85196633]]]

Answer 1

One way is to slice the arrays to extract the I and Q values, perform the computations using normal broadcasting, and then stick the values back together:

>>> Is, Qs = IQ[...,0], IQ[...,1]
>>> np.stack(((Is**2 + Qs**2) ** 0.5, np.arctan2(Qs, Is)), axis=-1)
array([[[ 2.23606798,  1.10714872],
        [ 5.        ,  0.92729522]],

       [[ 7.81024968,  0.87605805],
        [10.63014581,  0.85196633]]])