python - 将函数广播到 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

一种方法是对数组进行切片以提取 I 和 Q 值，使用正常广播执行计算，然后将这些值重新粘在一起:

>>> 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]]])