如何使我的 python 代码更快地求和？

Question

下面的代码是定义函数从下限 2 的总和，我改变总和的上限，因为我想找出总和收敛的位置。我发现上限越大，代码运行越慢。对于任何较大的上限值，如何使此代码快速运行？代码如下：

def h(i):
    x = (-1)**(i+1)
    y = 1000000000-(i-1)
    z = (log(i))**20
    return x*y*z

 gx = sum(h(i) for i in range (2, 1000000000+1))
 d_gx = gx/1000000000
 print(d_gx)

Answer 1

Numba是一个python库，用于对纯python代码进行即时优化，无需外部编译步骤。

这里我要介绍的两个重要函数是numba.njit 和numba.vectorize，它们都是装饰器。 njit 优化任意纯函数，vectorize 使函数同时在标量和 ndarray 上运行。

In [1]: from numba import vectorize, njit; from math import log

In [2]: def h(i):
   ...:     x = (-1)**(i+1)
   ...:     y = 1000000000-(i-1)
   ...:     z = (log(i))**20
   ...:     return x*y*z
   ...:

In [3]: %timeit sum(h(i) for i in range (2, 1000000+1))
646 ms ± 9.16 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

正如您在此处看到的，在减少输入空间的情况下，您的函数的简单实现平均需要 646 毫秒。我们可以通过 jitting 你的函数来改善这一点：

In [4]: jit_h = njit()(h)

In [5]: %timeit sum(jit_h(i) for i in range (2, 1000000+1))

179 ms ± 3.88 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
    

我们已将其缩减至 179 毫秒，与原来的 646 毫秒相比有了巨大的改进。因为 for 循环很慢，我们可以尝试使用 numpy 数组作为输入来向量化操作：

In [6]: vectorize_h = vectorize()(h)

In [7]: %timeit sum(vectorize_h(i) for i in range (2, 1000000+1))
657 ms ± 4.55 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

正如预期的那样，向量化输入允许传递标量，但不会显着提高性能 - 事实上，它有点慢！如果我们对整个 numpy 数组进行操作会怎样？

In [8]: import numpy as np

In [9]: %timeit sum(vectorize_h(np.arange(2,1000000+1))) 

149 ms ± 1.78 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

最后，如果我们用 numpy ndarray sum 替换 sum 内置函数呢？

In [10]: %timeit vectorize_h(np.arange(2,1000000+1)).sum() 
17.2 ms ± 207 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

进一步减少到 17.2 毫秒 - 对幼稚实现的巨大改进。