优化使用 numpy sin、cos、sum 和 abs 的 python 代码

Question

我有一些 Python 代码目前运行速度太慢而无法使用。进行了一些速度测试后，大部分时间似乎都花在了计算 corr 的数学运算上（参见下面的代码）。

import numpy as np
from multiprocessing import Pool
from contextlib import closing

def calc_cc(tlag):
    global phi_t, psi_t, T
    tstart = tlag-T/2
    x = np.arange(tstart, tstart+T, dtype=int)

    #bulk of time spent running line below calculating corr
    corr = np.sum(np.abs(np.cos(phi_t[tlag+x]-psi_t[x]))-np.abs(np.sin(phi_t[tlag+x]-psi_t[x])))


def main():
    global phi_t, psi_t, T

    phi_t = #some large numpy array
    psi_t = #another numpy array
    tlag = #another numpy array

    with closing(Pool(processes=5)) as p:
        corr = p.map(calc_cc,tlag)
        p.terminate()

if __name__ == "__main__":
    main()

我花了一些时间在网上四处寻找优化我的 python 代码的方法，更具体地说是数学运算，但很多建议都是围绕使用 numpy 等包而不是在需要时使用 python 基本函数（例如使用与 sum() 相比，np.sum 大大提高了速度。现在已经到了使用 numpy 数组和 numpy 方法的地步，我不确定我还能在哪里找到额外的收益。我正在有效地尝试的公式解决方法是：

我确信必须有更好的方法来做到这一点，因此任何关于更好解决方案的指导将不胜感激:)。

Answer 1

首先，您可以通过对两个三角函数as @hpaulj noted 计算一次phi_t(...)-psi_t(...) 来稍微改进您的代码。但是，我看到的最大问题是您使用多处理池进行相对较轻的计算。我的猜测是，您运行时的很大一部分来自使用池的开销。

如果你能适应内存，你可以通过使用数组广播来计算你的整个相关函数。通过矢量化您的计算，您将最大限度地提高您可以通过 numpy.这个想法是使用二维数组索引您的phi/psi 数组：列对应于原始x 索引，行对应于单个tlag 移位。您使用的内置 numpy 函数自然会使用广播。

这就是我的意思，加上一些虚拟数据：

import numpy as np

def calc_cc_new(tlags):
    global phi_t, psi_t, T

    tstarts = tlags - T//2
    xs = tstarts + np.arange(T)[:,None]
    dphipsi = phi_t[tlags+xs] - psi_t[xs]

    corr = np.sum(np.abs(np.cos(dphipsi)) - np.abs(np.sin(dphipsi)),axis=0)
    return corr

def main_new():
    global phi_t, psi_t, tlags

    return calc_cc_new(tlags)

N = 10000
T = 100
phi_t = np.random.rand(N)*2*np.pi
psi_t = np.random.rand(N)*2*np.pi
tlags = np.arange(T//2,3*T)

if __name__ == "__main__":
    print(main_new())

我没有检查您的多处理版本，但您原始的串行版本花费了 8.1 毫秒，而矢量化版本花费了 2.1 毫秒（无论我是否将总和的两项分离为对 np.sum() 的单独调用）。我使用np.allclose 检查了从我的版本返回的数组是否与原始数组签出（由于矢量化，算术运算的顺序被移动了，所以我们不能期望 exact 一致，只有一个在机器精度）。对于虚拟示例，我将求和维度定义为第零还是第一个并不重要，但根据您的输入数组形状，这可能会导致一些性能差异。

此外，通常不鼓励使用全局变量，全局名称空间查找比本地名称空间稍慢（尽管如果您的函数占用大量 CPU 资源，我认为这并不重要）。无论如何，如果您放弃多处理池，您可以简单地扩展您的 calc_cc 函数的定义，以便它获取所有输入作为参数，而不是全局变量。

最后，完成后不要忘记 1/N 前置因子。