python - 在 Python 中使用 numpy 对循环进行矢量化/优化

Question

我正在编写一个脚本来处理来自 signal_gen 函数中表示的传感器的一些数据。正如您在测试函数中看到的那样，它是完全循环发送的。由于这个函数被调用了很多次，所以它有点慢，如果朝着正确的方向优化它会很可爱。

我读到可以用 vectorizatid 数组交换 for 循环，但我无法理解应该如何编写 i_avg[i] 行，因为我们有单个元素 y[i] 相乘整个数组 x 都在 np.cos 中，所有这些又只是 i_avg 的一个刺激。

def testing(signal):
    y = np.arange(0.0108, 0.0135, 0.001) # this one changes over time, set 
    #to constant for easier reading
    x = np.arange(0, (len(signal)))
    I_avg = np.zeros(len(y))
    Q_avg = np.zeros_like(I_avg)
    for i in range(0, len(y)):
        I_avg[i] = np.array(signal * (np.cos(2 * np.pi * y[i] * x))).sum()
        Q_avg[i] = np.array(signal * (np.sin(2 * np.pi * y[i] * x))).sum()
    D = np.power(I_avg, 2) + np.power(Q_avg, 2)
    max_index = np.argmax(D)
    phaseOut = np.arctan2(Q_avg[max_index], I_avg[max_index])

#just a test signal
def signal_gen():
    signal = np.random.random(size=251)
    return signal

Answer 1

一种使用矩阵乘法 和 numpy.dot 的矢量化方法替换嵌套循环为我们提供 I_avg, Q_avg 并合并 NumPy broadcasting从而实现更有效的解决方案就像这样-

mult = 2*np.pi*y[:,None]*x
I_avg, Q_avg = np.cos(mult).dot(signal), np.sin(mult).dot(signal)

请注意，对于给定的示例，我们正在与循环版本竞争，循环版本只需迭代 3 次迭代(y 的长度为 3)。因此，我们不会在这里看到巨大的加速。

运行时测试-

In [9]: #just a test signal
   ...: signal = np.random.random(size=251)
   ...: y = np.arange(0.0108, 0.0135, 0.001)
   ...: x = np.arange(0, (len(signal)))
   ...: 

# Original approach
In [10]: %%timeit I_avg = np.zeros(len(y))
    ...: Q_avg = np.zeros_like(I_avg)
    ...: for i in range(0, len(y)):
    ...:     I_avg[i] = np.array(signal * (np.cos(2 * np.pi * y[i] * x))).sum()
    ...:     Q_avg[i] = np.array(signal * (np.sin(2 * np.pi * y[i] * x))).sum()
    ...: 
10000 loops, best of 3: 68 µs per loop

# Proposed approach
In [11]: %%timeit mult = 2*np.pi*y[:,None]*x
    ...: I_avg, Q_avg = np.cos(mult).dot(signal), np.sin(mult).dot(signal)
    ...: 
10000 loops, best of 3: 34.8 µs per loop