使用 Numba 改进有限差分拉普拉斯算子

Question

我正在使用 Python 求解反应扩散方程组 (Fitz-Hugh-Nagumo model)。我想学习如何使用 Numba 来加速计算。我目前正在将以下 laplacian.py 模块导入到我的集成脚本中：

def neumann_laplacian_1d(u,dx2):
    """Return finite difference Laplacian approximation of 2d array.
    Uses Neumann boundary conditions and a 2nd order approximation.
    """
    laplacian = np.zeros(u.shape)
    laplacian[1:-1] =  ((1.0)*u[2:] 
                       +(1.0)*u[:-2]
                       -(2.0)*u[1:-1])
    # Neumann boundary conditions
    # edges
    laplacian[0]  =  ((2.0)*u[1]-(2.0)*u[0])
    laplacian[-1] =  ((2.0)*u[-2]-(2.0)*u[-1])

    return laplacian/ dx2

其中 u 是代表其中一个字段的 NumPy 一维数组。我尝试在导入 from numba import autojit 后添加装饰器 @autojit(target="cpu")。我没有看到计算有任何改进。谁能告诉我如何在这种情况下正确使用 Numba？

我这里使用的输入数组是

a = random.random(252)

所以我将性能与以下行进行了比较：

%timeit(neumann_laplacian_1d(a,1.0))

使用 Numba 我得到了：

%timeit(neumann_laplacian_1d(a,1.0))
The slowest run took 22071.00 times longer than the fastest. This could mean that an intermediate result is being cached 
1 loops, best of 3: 14.1 µs per loop

没有 Numba 我得到了 (!!):

%timeit(neumann_laplacian_1d(a,1.0))
The slowest run took 11.84 times longer than the fastest. This could mean that an intermediate result is being cached 
100000 loops, best of 3: 9.12 µs per loop

Numba 实际上让它变慢了..

Answer 1

我无法复制您的结果。

Python 版本：3.4.4 |Anaconda 2.4.1（64 位）| （默认，2016 年 1 月 19 日，12:10:59）[MSC v.1600 64 位 (AMD64)]

numba 版本：0.23.1

import numba as nb
import numpy as np

def neumann_laplacian_1d(u,dx2):
    """Return finite difference Laplacian approximation of 2d array.
    Uses Neumann boundary conditions and a 2nd order approximation.
    """
    laplacian = np.zeros(u.shape)
    laplacian[1:-1] =  ((1.0)*u[2:] 
                       +(1.0)*u[:-2]
                       -(2.0)*u[1:-1])
    # Neumann boundary conditions
    # edges
    laplacian[0]  =  ((2.0)*u[1]-(2.0)*u[0])
    laplacian[-1] =  ((2.0)*u[-2]-(2.0)*u[-1])

    return laplacian/ dx2

@nb.autojit(nopython=True)
def neumann_laplacian_1d_numba(u,dx2):
    """Return finite difference Laplacian approximation of 2d array.
    Uses Neumann boundary conditions and a 2nd order approximation.
    """
    laplacian = np.zeros(u.shape)
    laplacian[1:-1] =  ((1.0)*u[2:] 
                       +(1.0)*u[:-2]
                       -(2.0)*u[1:-1])
    # Neumann boundary conditions
    # edges
    laplacian[0]  =  ((2.0)*u[1]-(2.0)*u[0])
    laplacian[-1] =  ((2.0)*u[-2]-(2.0)*u[-1])

    return laplacian/ dx2

a = np.random.random(252)
#run once to make the JIT do it's work before timing
neumann_laplacian_1d_numba(a, 1.0)


%timeit neumann_laplacian_1d(a, 1.0)
%timeit neumann_laplacian_1d_numba(a, 1.0)

>>10000 loops, best of 3: 21.5 µs per loop
>>The slowest run took 4.49 times longer than the fastest. This could mean that an intermediate result is being cached 
>>100000 loops, best of 3: 3.53 µs per loop

我看到 python 2.7.11 和 numba 0.23

的结果相似

>>100000 loops, best of 3: 19.1 µs per loop
>>The slowest run took 8.55 times longer than the fastest. This could mean that an intermediate result is being cached 
>>100000 loops, best of 3: 2.4 µs per loop

使用 Numba 改进有限差分拉普拉斯算子

Using Numba to improve finite-differences laplacian

python

numpy

numba