将函数广播到3D数组Python

Question

我曾尝试理解numpy broadcasting with 3d arrays，但我认为那里的OP要求有所不同。

我有一个像这样的3D numpy数组-

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

此数组的形状为（2,2,2）。我想对这个3D矩阵中的每个1x2数组应用一个函数，就像这样-

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]

如您所见，我想将我的函数应用于每个1x2数组，并将其替换为函数的返回值。输出是具有相同尺寸的3D阵列。有没有办法将此功能广播到原始3D阵列中的每个1x2阵列？当前，我使用的循环随着3D阵列尺寸的增加而变得非常慢。

当前我正在这样做-

#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

返回的3D数组是-

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

 [[ 7.81024968  0.87605805]
  [10.63014581  0.85196633]]]

非常感谢！附言我是新手，所以请让我知道如何改善我的问题！

Answer 1

可以使用数组来完成：

# sort of sum of squares along axis 0, ie (IQ[0, ...]**2 + IQ[1, ...]**2 + ...)**0.5
amp = np.sqrt(np.square(IQ).sum(axis=0))

# and phase is arctan for each component in each matrix
phase = np.arctan2(IQ[1], IQ[0])