使用Python中的NumPy对循环进行向量化或优化

Question

我正在编写一个脚本，处理从signal_gen函数中表示的传感器数据。如你在testing函数中所看到的那样，它相当于循环的中心。由于该函数被调用多次，使得其速度有点慢，因此需要优化。

我阅读过可以使用向量化数组来替换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

- Runsiv

---

测试中的S是什么？ - Ohjeah

一个笔误，应该是“信号”。 - Runsiv

max_index 是另一个打错的地方。您测试过这段代码吗？ - Ohjeah

我做了，但我猜当我改变了一些变量名时有些马虎。 - Runsiv

Answer 1

使用矩阵乘法的向量化方法，利用 numpy.dot 替换嵌套循环，得到 I_avg, Q_avg，同时结合 NumPy广播 实现更高效的解决方案，代码如下 -

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

- Divakar

---

真的很有效，几乎将计算时间减半了，而且非常易读。谢谢！ - Runsiv

我们也可以使用 np.newaxis 而不是在那里硬编码 None。这只是一个建议，虽然两者本质上是相同的。 - kmario23

@Runsiv 考虑接受其中一个解决方案吗？在这里阅读更多信息：https://meta.stackexchange.com/questions/5234/how-does-accepting-an-answer-work - Divakar

谢谢你提醒我！ - Runsiv

Answer 2

您可以使用np.einsum 进行广播：

yx = 2*np.pi*np.einsum("i,j->ij", y, x)
I_avg = np.sin(yx) @ signal    
Q_avg = np.cos(yx) @ signal 

- Ohjeah