【后篇】Python高性能计算库——Numba

接下来我们使用Numba来实现相同的功能。

@jit

def abc_model_numba(a, b, c, rain):

   outflow = np.zeros((rain.size), dtype=np.float64)

   state_in = 0

   state_out = 0

   for i in range(rain.size):

       state_out = (1 - c) * state_in + a * rain[i]

       outflow[i] = (1 - a - b) * rain[i] + c * state_out

       state_in = state_out

   return outflow

我用随机数字作为输入来运行这些模块，这只是为了比较计算时间，而且也比较了针对fortran实现的时间（详见here）。我们来看看数字：

py_time = %timeit -r 5 -n 10 -o abc_model_py(0.2, 0.6, 0.1, rain)

>> 6.75 s ± 11.6 ms per loop (mean ± std. dev. of 5 runs, 10 loops each)

# Measure the execution time of the Numba implementation

numba_time = %timeit -r 5 -n 10 -o abc_model_numba(0.2, 0.6, 0.1, rain)

>> 30.6 ms ± 498 µs per loop (mean ± std. dev. of 5 runs, 10 loops each)

# Measure the execution time of the Fortran implementation

fortran_time = %timeit -r 5 -n 10 -o abc_model_fortran(0.2, 0.6, 0.1, rain)