优化使用numpsin、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基函数（例如使用np.总和与sum（）相比，速度大大提高。现在已经到了使用numpy数组和numpy方法的地步，我不确定还能在哪里找到额外的收益。我想解决的公式是：

我相信一定有更好的方法来做这件事，所以任何关于更好解决办法的指导都将不胜感激：）。在

Answer 1

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

如果你可以放入内存，你可以利用数组广播来计算你的整个相关函数。通过将计算矢量化，可以最大限度地提高numpy的速度。其思想是使用2d数组索引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检查了我的版本中返回的数组是否与原始数组一致（由于矢量化，算术运算的顺序被移动，因此我们不能期望精确一致，只有一个在机器精度范围内）。在这个虚拟的例子中，我将求和维度定义为第0个还是第一个并不重要，但根据输入数组的形状，这可能会导致一些性能差异。在

此外，通常不鼓励使用全局变量，全局名称空间查找比本地名称空间查找稍微慢一些（尽管我不希望这在函数占用大量CPU的情况下无关紧要）。不管怎样，如果不使用多处理池，只需扩展calc_cc函数的定义，使其获得所有输入作为参数，而不是全局参数。在

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