Cython 优化

Question

我正在用 Python 编写一个相当大的模拟，并希望从 Cython 获得一些额外的性能。但是，对于下面的代码，我似乎并没有得到那么多，即使它包含一个相当大的循环。大约 10 万次迭代。

我是否犯了一些初学者的错误，或者这个循环大小只是太小而产生很大的影响？ （在我的测试中，Cython 代码只快了大约 2 倍）。

import numpy as np;
cimport numpy as np;
import math

ctypedef np.complex64_t cpl_t
cpl = np.complex64

def example(double a, np.ndarray[cpl_t,ndim=2] A):

    cdef int N = 100

    cdef np.ndarray[cpl_t,ndim=3] B = np.zeros((3,N,N),dtype = cpl)

    cdef Py_ssize_t n, m;
    for n in range(N):
        for m in range(N):

            if np.sqrt(A[0,n]) > 1:
                B[0,n,m] = A[0,n] + 1j * A[0,m]

    return B;

Answer 1

您应该使用编译器指令。我用 Python 编写了你的​​函数

import numpy as np

def example_python(a, A):
    N = 100
    B = np.zeros((3,N,N),dtype = np.complex)
    aux = np.sqrt(A[0])
    for n in range(N):
        if aux[n] > 1:
            for m in range(N):
                B[0,n,m] = A[0,n] + 1j * A[0,m]
return B

在 Cython 中（您可以了解编译器指令here）

import cython
import numpy as np
cimport numpy as np

ctypedef np.complex64_t cpl_t
cpl = np.complex64

@cython.boundscheck(False) # compiler directive
@cython.wraparound(False) # compiler directive
def example_cython(double a, np.ndarray[cpl_t,ndim=2] A):

    cdef int N = 100
    cdef np.ndarray[cpl_t,ndim=3] B = np.zeros((3,N,N),dtype = cpl)
    cdef np.ndarray[float, ndim=1] aux
    cdef Py_ssize_t n, m
    aux = np.sqrt(A[0,:]).real
    for n in range(N):
        if aux[n] > 1.:
            for m in range(N):
                B[0,n,m] = A[0,n] + 1j * A[0,m]
    return B

我比较两个函数

c = np.array(np.random.rand(100,100)+1.5+1j*np.random.rand(100,100), dtype=np.complex64)

%timeit example_python(100, c)
10 loops, best of 3: 61.8 ms per loop

%timeit example_cython(100, c)
10000 loops, best of 3: 134 µs per loop

在这种情况下，Cython 比 Python 快约 450 倍。