如何加速这个 Python 函数？答案

【问题标题】：How can I speed up this Python function?如何加速这个 Python 函数？
【发布时间】：2021-10-31 12:13:27
【问题描述】：

有没有办法向量化这个表达式？

import numpy as np
def phase(f, dt=1):
    c = [f[0]] + [fi*dt for fi in f[1:]]
    
    s=0
    x = []
    for ci in c:
        s = (s+ci)%(2*np.pi)
        x.append(s)
    return x

f 是一个任意长度的数组，dt 是一个常数。代码：

将整个向量 f 但第一个实例乘以 dt。
返回部分和，即output[i] = sum(c[0:i])模2pi。

关于如何加快速度的任何想法？谢谢！

【问题讨论】：

你试过numba吗？
您可以通过使用列表理解而不是 for 循环：x=[(s + ci) % (2 * np.pi) for ci in c] 来获得较小的性能提升。你也同意s 在所有迭代中都不同这一事实吗？
就我个人而言，我会使用 for idx, ci in enumerate(c) 编写这个 for 循环，这样您就可以跳过单独定义 s 并可能在调试中节省一些神经细胞。

标签： python performance loops vectorization

【解决方案1】：

只需要一个for循环。

import numpy as np


def phase(f, dt=1):
    s = 0
    x = []
    if f:
        s = (s + f[0]) % (2 * np.pi)
        x.append(s)
    
        for fi in f[1:]:
            s = (s + fi * dt) % (2 * np.pi)
            x.append(s)
    return x

【讨论】：

【解决方案2】：

您可以使用 Numpy 的 cumsum() 函数以矢量化方式执行此操作。这应该会快一些：

import numpy as np

def phase(f, dt=1):
    pi_2 = 2 * np.pi
    c = np.concatenate((f[:1], f[1:] * dt))
    
    return np.cumsum(c) % pi_2

时间安排：

n = np.arange(5000)

%timeit original_phase(n, 3)
# 3.2 ms ± 120 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit phase(n, 3)
# 272 µs ± 12.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

np.testing.assert_allclose(original_phase(n, 3), phase(n, 3))

【讨论】：