反应扩散算法中的Numba或Cython加速

时间:2014-11-08 23:41:00

标签: python numpy cython pde numba

我想加速用Python和NumPy编写的代码。我使用Gray-Skott算法(http://mrob.com/pub/comp/xmorphia/index.html)作为反应扩散模型,但是使用Numba和Cython它甚至更慢!有可能加快速度吗?提前谢谢!

的Python + NumPy的

def GrayScott(counts, Du, Dv, F, k):
    n = 300
    U = np.zeros((n+2,n+2), dtype=np.float_)
    V = np.zeros((n+2,n+2), dtype=np.float_)
    u, v = U[1:-1,1:-1], V[1:-1,1:-1]

    r = 20
    u[:] = 1.0
    U[n/2-r:n/2+r,n/2-r:n/2+r] = 0.50
    V[n/2-r:n/2+r,n/2-r:n/2+r] = 0.25
    u += 0.15*np.random.random((n,n))
    v += 0.15*np.random.random((n,n))

    for i in range(counts):
        Lu = (                 U[0:-2,1:-1] +
              U[1:-1,0:-2] - 4*U[1:-1,1:-1] + U[1:-1,2:] +
                               U[2:  ,1:-1] )
        Lv = (                 V[0:-2,1:-1] +
              V[1:-1,0:-2] - 4*V[1:-1,1:-1] + V[1:-1,2:] +
                               V[2:  ,1:-1] )
        uvv = u*v*v
        u += Du*Lu - uvv + F*(1 - u)
        v += Dv*Lv + uvv - (F + k)*v

    return V

Numba

from numba import jit, autojit

@autojit
def numbaGrayScott(counts, Du, Dv, F, k):
    n = 300
    U = np.zeros((n+2,n+2), dtype=np.float_)
    V = np.zeros((n+2,n+2), dtype=np.float_)
    u, v = U[1:-1,1:-1], V[1:-1,1:-1]

    r = 20
    u[:] = 1.0
    U[n/2-r:n/2+r,n/2-r:n/2+r] = 0.50
    V[n/2-r:n/2+r,n/2-r:n/2+r] = 0.25
    u += 0.15*np.random.random((n,n))
    v += 0.15*np.random.random((n,n))

    Lu = np.zeros_like(u)
    Lv = np.zeros_like(v)

    for i in range(counts):
        for row in range(n):
            for col in range(n):
                Lu[row,col] = U[row+1,col+2] + U[row+1,col] + U[row+2,col+1] + U[row,col+1] - 4*U[row+1,col+1]
                Lv[row,col] = V[row+1,col+2] + V[row+1,col] + V[row+2,col+1] + V[row,col+1] - 4*V[row+1,col+1]

        uvv = u*v*v
        u += Du*Lu - uvv + F*(1 - u)
        v += Dv*Lv + uvv - (F + k)*v

    return V

用Cython

%%cython
cimport cython
import numpy as np
cimport numpy as np

cpdef cythonGrayScott(int counts, double Du, double Dv, double F, double k):
    cdef int n = 300
    cdef np.ndarray U = np.zeros((n+2,n+2), dtype=np.float_)
    cdef np.ndarray V = np.zeros((n+2,n+2), dtype=np.float_)
    cdef np.ndarray u = U[1:-1,1:-1]
    cdef np.ndarray v = V[1:-1,1:-1]

    cdef int r = 20
    u[:] = 1.0
    U[n/2-r:n/2+r,n/2-r:n/2+r] = 0.50
    V[n/2-r:n/2+r,n/2-r:n/2+r] = 0.25
    u += 0.15*np.random.random((n,n))
    v += 0.15*np.random.random((n,n))

    cdef np.ndarray Lu = np.zeros_like(u)
    cdef np.ndarray Lv = np.zeros_like(v)
    cdef int i, row, col
    cdef np.ndarray uvv

    for i in range(counts):
        for row in range(n):
            for col in range(n):
                Lu[row,col] = U[row+1,col+2] + U[row+1,col] + U[row+2,col+1] + U[row,col+1] - 4*U[row+1,col+1]
                Lv[row,col] = V[row+1,col+2] + V[row+1,col] + V[row+2,col+1] + V[row,col+1] - 4*V[row+1,col+1]

        uvv = u*v*v
        u += Du*Lu - uvv + F*(1 - u)
        v += Dv*Lv + uvv - (F + k)*v

    return V

用法示例:

GrayScott(4000, 0.16, 0.08, 0.04, 0.06)

2 个答案:

答案 0 :(得分:7)

以下是加速cython版本的步骤:

  • cdef np.ndarray不会让元素访问速度更快,您需要在cython中使用memoryview:cdef double[:, ::1] bU = U

  • 关闭boundscheckwraparound

  • 在for-loop中进行所有计算。

这是修改过的cython代码:

%%cython
#cython: boundscheck=False
#cython: wraparound=False
cimport cython
import numpy as np
cimport numpy as np

cpdef cythonGrayScott(int counts, double Du, double Dv, double F, double k):
    cdef int n = 300
    cdef np.ndarray U = np.zeros((n+2,n+2), dtype=np.float_)
    cdef np.ndarray V = np.zeros((n+2,n+2), dtype=np.float_)
    cdef np.ndarray u = U[1:-1,1:-1]
    cdef np.ndarray v = V[1:-1,1:-1]

    cdef int r = 20
    u[:] = 1.0
    U[n/2-r:n/2+r,n/2-r:n/2+r] = 0.50
    V[n/2-r:n/2+r,n/2-r:n/2+r] = 0.25
    u += 0.15*np.random.random((n,n))
    v += 0.15*np.random.random((n,n))

    cdef np.ndarray Lu = np.zeros_like(u)
    cdef np.ndarray Lv = np.zeros_like(v)
    cdef int i, c, r1, c1, r2, c2
    cdef double uvv

    cdef double[:, ::1] bU = U
    cdef double[:, ::1] bV = V
    cdef double[:, ::1] bLu = Lu
    cdef double[:, ::1] bLv = Lv

    for i in range(counts):
        for r in range(n):
            r1 = r + 1
            r2 = r + 2
            for c in range(n):
                c1 = c + 1
                c2 = c + 2
                bLu[r,c] = bU[r1,c2] + bU[r1,c] + bU[r2,c1] + bU[r,c1] - 4*bU[r1,c1]
                bLv[r,c] = bV[r1,c2] + bV[r1,c] + bV[r2,c1] + bV[r,c1] - 4*bV[r1,c1]

        for r in range(n):
            r1 = r + 1
            for c in range(n):
                c1 = c + 1
                uvv = bU[r1,c1]*bV[r1,c1]*bV[r1,c1]
                bU[r1,c1] += Du*bLu[r,c] - uvv + F*(1 - bU[r1,c1])
                bV[r1,c1] += Dv*bLv[r,c] + uvv - (F + k)*bV[r1,c1]

    return V

它比numpy版本快11倍。

答案 1 :(得分:5)

除了涉及的循环和大量操作之外,在您的情况下最有可能的杀戮性能是数组分配。我不知道为什么你的Numba和Cython版本达不到你的期望,但是你可以通过就地完成所有操作,即更换你的numpy代码,速度提高2倍(以牺牲一些可读性为代价)电流循环:

Lu, Lv, uvv = np.empty_like(u), np.empty_like(v), np.empty_like(u)

for i in range(counts):
    Lu[:] = u
    Lu *= -4
    Lu += U[:-2,1:-1]
    Lu += U[1:-1,:-2]
    Lu += U[1:-1,2:]
    Lu += U[2:,1:-1]
    Lu *= Du

    Lv[:] = v
    Lv *= -4
    Lv += V[:-2,1:-1]
    Lv += V[1:-1,:-2]
    Lv += V[1:-1,2:]
    Lv += V[2:,1:-1]
    Lv *= Dv

    uvv[:] = u
    uvv *= v
    uvv *= v
    Lu -= uvv
    Lv += uvv

    u *= 1 - F
    u += F
    u += Lu

    v *= 1 - F - k
    v += Lv