python 加速循环_使用numba加速循环

是的，这是Numba真正解决的问题。我更改了dk的值，因为对于简单的演示来说这是不明智的。代码如下：import numpy as np

import numba as nb

def f_big(A, k, std_A, std_k, mean_A=10, mean_k=0.2, hh=100):

return ( 1 / (std_A * std_k * 2 * np.pi) ) * A * (hh/50) ** k * np.exp( -1*(k - mean_k)**2 / (2 * std_k **2 ) - (A - mean_A)**2 / (2 * std_A**2))

def func():

outer_sum = 0

dk = 0.01 #0.000001

for k in np.arange(dk, 0.4, dk):

inner_sum = 0

for A in np.arange(dk, 20, dk):

inner_sum += dk * f_big(A, k, 1e-5, 1e-5)

outer_sum += inner_sum * dk

return outer_sum

@nb.jit(nopython=True)

def f_big_nb(A, k, std_A, std_k, mean_A=10, mean_k=0.2, hh=100):

return ( 1 / (std_A * std_k * 2 * np.pi) ) * A * (hh/50) ** k * np.exp( -1*(k - mean_k)**2 / (2 * std_k **2 ) - (A - mean_A)**2 / (2 * std_A**2))

@nb.jit(nopython=True)

def func_nb():

outer_sum = 0

dk = 0.01 #0.000001

X = np.arange(dk, 0.4, dk)

Y = np.arange(dk, 20, dk)

for i in xrange(X.shape[0]):

k = X[i] # faster to do lookup than iterate over an array directly

inner_sum = 0

for j in xrange(Y.shape[0]):

A = Y[j]

inner_sum += dk * f_big_nb(A, k, 1e-5, 1e-5)

outer_sum += inner_sum * dk

return outer_sum

然后计时：In [7]: np.allclose(func(), func_nb())

Out[7]: True

In [8]: %timeit func()

1 loops, best of 3: 222 ms per loop

In [9]: %timeit func_nb()

The slowest run took 419.10 times longer than the fastest. This could mean that an intermediate result is being cached

1000 loops, best of 3: 362 µs per loop

所以在我的笔记本电脑上，numba版本要快600倍。