Cythonize 偏微分方程积分器

Question

我正在尝试使用 Cython 为偏微分方程加速有限差分积分器。我不确定我需要做什么才能让 Cython 正确使用 numpy 数组。

我使用的扩散项函数是

def laplacian(var, dh2):
    """ (1D array, dx^2) -> laplacian(1D array)
    periodic_laplacian_1D_4th_order
    Implementing the 4th order 1D laplacian with periodic condition
    """
    lap = numpy.zeros_like(var)
    lap[1:]    = (4.0/3.0)*var[:-1]
    lap[0]     = (4.0/3.0)*var[1]
    lap[:-1]  += (4.0/3.0)*var[1:]
    lap[-1]   += (4.0/3.0)*var[0]
    lap       += (-5.0/2.0)*var

    lap[2:]   += (-1.0/12.0)*var[:-2]
    lap[:2]   += (-1.0/12.0)*var[-2:]
    lap[:-2]  += (-1.0/12.0)*var[2:]
    lap[-2:]  += (-1.0/12.0)*var[:2]

    return lap / dh2

模型方程的rhs为

from derivatives import laplacian

def dbdt(b,w,p,m,d,dx2):
    """ db/dt of Modified Klausmeier """
    return w*b**2 - m*b + laplacian(b,dx2)

def dwdt(b,w,p,m,d,dx2):
    """ dw/dt of Modified Klausmeier """
    return p - w - w*b**2 + d*laplacian(b,dx2)

如何使用 Cython 优化这些功能？

我的工作代码在 Github 上有一个存储库，它集成了 Gray-Scott 模型 - Gray-Scott model integrator。

Answer 1

要有效地使用 Cython，您应该明确所有循环并确保 cython -a 显示尽可能少的 Python 调用。第一次尝试是：

import numpy as np
cimport numpy as np
cimport cython
@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
def laplacian(double [::1] var, double dh2):
    """ (1D array, dx^2) -> laplacian(1D array)
    periodic_laplacian_1D_4th_order
    Implementing the 4th order 1D laplacian with periodic condition
    """
    cdef int n = var.shape[0]
    cdef double[::1] lap = np.zeros(n)
    cdef int i
    for i in range(0, n-1):
        lap[1+i] = (4.0/3.0)*var[i]
    lap[0]     = (4.0/3.0)*var[1]
    for i in range(0, n-1):
        lap[i]  += (4.0/3.0)*var[1+i]
    lap[n-1]   += (4.0/3.0)*var[0]
    for i in range(0, n):
        lap[i]       += (-5.0/2.0)*var[i]

    for i in range(0, n-2):
        lap[2+i]   += (-1.0/12.0)*var[i]
    for i in range(0, 2):
        lap[i]   += (-1.0/12.0)*var[n - 2 + i]
    for i in range(0, n-2):
        lap[i]   += (-1.0/12.0)*var[i+2]
    for i in range(0, 2):
        lap[n-2+i]  += (-1.0/12.0)*var[i]
    for i in range(0, n):
        lap[i]  /= dh2
    return lap

现在这给了你：

$ python -m timeit -s 'import numpy as np; from lap import laplacian; var = np.random.rand(1000000); dh2 = .01' 'laplacian(var, dh2)'
100 loops, best of 3: 11.5 msec per loop

而 NumPy 代码给出：

100 loops, best of 3: 18.5 msec per loop

请注意，Cython 可以通过合并循环等进一步优化。

我还尝试了Pythran 的自定义（即未在主控中提交）版本，并且在不更改原始 Python 代码的情况下，我获得了与 Cython 版本相同的加速，而无需转换代码：

#pythran export laplacian(float [], float)
import numpy
def laplacian(var, dh2):
    """ (1D array, dx^2) -> laplacian(1D array)
    periodic_laplacian_1D_4th_order
    Implementing the 4th order 1D laplacian with periodic condition
    """
    lap = numpy.zeros_like(var)
    lap[1:]    = (4.0/3.0)*var[:-1]
    lap[0]     = (4.0/3.0)*var[1]
    lap[:-1]  += (4.0/3.0)*var[1:]
    lap[-1]   += (4.0/3.0)*var[0]
    lap       += (-5.0/2.0)*var

    lap[2:]   += (-1.0/12.0)*var[:-2]
    lap[:2]   += (-1.0/12.0)*var[-2:]
    lap[:-2]  += (-1.0/12.0)*var[2:]
    lap[-2:]  += (-1.0/12.0)*var[:2]

    return lap / dh2

转换为：

$ pythran lap.py -O3

我得到：

100 loops, best of 3: 11.6 msec per loop

Answer 2

所以我想我已经想通了，虽然我不确定这是最优化的方法：

import numpy as np
cimport numpy as np
cdef laplacian(np.ndarray[np.float64_t, ndim=1] var,np.float64_t dh2):
    """ (1D array, dx^2) -> laplacian(1D array)
    periodic_laplacian_1D_4th_order
    Implementing the 4th order 1D laplacian with periodic condition
    """
    lap = np.zeros_like(var)
    lap[1:]    = (4.0/3.0)*var[:-1]
    lap[0]     = (4.0/3.0)*var[1]
    lap[:-1]  += (4.0/3.0)*var[1:]
    lap[-1]   += (4.0/3.0)*var[0]
    lap       += (-5.0/2.0)*var

    lap[2:]   += (-1.0/12.0)*var[:-2]
    lap[:2]   += (-1.0/12.0)*var[-2:]
    lap[:-2]  += (-1.0/12.0)*var[2:]
    lap[-2:]  += (-1.0/12.0)*var[:2]

    return lap / dh2

我使用了以下 setup.py

from distutils.core import setup
from Cython.Build import cythonize

setup(
    ext_modules = cythonize("derivatives_c.pyx")
)

欢迎任何关于改进它的建议..