對部分微分方程積分器進行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模型集成器 。

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")
)

歡迎提供任何改進建議。