用 numpy 优化这个循环

Question

递归生成 500 万个点r[i]：

import numpy as np
n, a, b, c = 5000000, 0.0000002, 0.5, 0.4
eps = np.random.normal(0, 1, n)
sigma = np.ones(n) * np.sqrt(a)
r = np.zeros(n)
for i in range(1,n):
    sigma[i] = np.sqrt(a + b * r[i-1] ** 2 + c * sigma[i-1] ** 2)
    r[i] = sigma[i] * eps[i]

在我的标准 i5 笔记本电脑上使用大约 17 秒。

我过去经常使用 Cython，我知道在这里使用它可能会优化 10 < k < 100。

但在这种情况下必须使用 Cython 之前，我想知道：是否会有一个我不知道的普通 Numpy/Python 方法可以优化这么多？

Answer 1

这里我已经排除了一些术语，所以它应该运行更快

sigma2 = np.ones(n)*a
eps2 = eps**2
eps2[0] = 0
abc = b*eps2+c
for i in range(1,n):
    #sigma2[i] = a + b * sigma2[i-1]*eps2[i-1] + c * sigma2[i-1]
    sigma2[i] = a + abc[i-1]*sigma2[i-1]
print np.sqrt(sigma2)
print np.allclose(sigma, np.sqrt(sigma2))
print np.sqrt(sigma2*eps2)  # r

我试图分解出 a+，但还没有完全匹配的结果。如果我能做到这一点，我认为可以用

替换循环

np.cumprod(abc)

Answer 2

只需将其更改为 math.sqrt 而不是 np.sqrt 即可在此处提供大约 40% 的加速。

因为我是一个 numba 狂热者，所以我尝试了 numba 版本与你的版本 (initial) 和数学版本 (normal)

import numpy as np
import math
import numba as nb
n, a, b, c = 500000, 0.0000002, 0.5, 0.4
eps = np.random.normal(0, 1, n)
sigma = np.ones(n) * np.sqrt(a)
r = np.zeros(n)

def initial(n, a, b, c, eps, sigma, r):
    for i in range(1,n):
        sigma[i] = np.sqrt(a + b * r[i-1] ** 2 + c * sigma[i-1] ** 2)
        r[i] = sigma[i] * eps[i]

def normal(n, a, b, c, eps, sigma, r):
    for i in range(1,n):
        sigma[i] = math.sqrt(a + b * r[i-1] ** 2 + c * sigma[i-1] ** 2)
        r[i] = sigma[i] * eps[i]

@nb.njit
def function(n, a, b, c, eps, sigma, r):
    for i in range(1,n):
        sigma[i] = math.sqrt(a + b * r[i-1] ** 2 + c * sigma[i-1] ** 2)
        r[i] = sigma[i] * eps[i]

然后只是为了验证结果是否相同：

sigma1 = sigma.copy()
sigma2 = sigma.copy()
sigma3 = sigma.copy()
r1 = r.copy()
r2 = r.copy()
r3 = r.copy()

initial(n, a, b, c, eps, sigma1, r1)   
normal(n, a, b, c, eps, sigma2, r2)       
function(n, a, b, c, eps, sigma3, r3)
np.testing.assert_array_almost_equal(sigma1, sigma2)
np.testing.assert_array_almost_equal(sigma1, sigma3)
np.testing.assert_array_almost_equal(r1, r2)
np.testing.assert_array_almost_equal(r1, r3)

那么速度呢（我使用 n=500000 来获得更快的 timeit 结果）：

%timeit initial(n, a, b, c, eps, sigma1, r1)   
1 loop, best of 3: 7.27 s per loop
%timeit normal(n, a, b, c, eps, sigma2, r2)   
1 loop, best of 3: 4.49 s per loop    
%timeit function(n, a, b, c, eps, sigma3, r3)
100 loops, best of 3: 17.7 ms per loop

我知道你不想要 cython，所以 numba 可能也不在考虑范围之内，但加速是惊人的（快 410 倍！）

用 numpy 优化这个循环

Optimize this loop with numpy

python

numpy

cython

numerical-methods