您应该使用编译器指令。我用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 http://docs.cython.org/src/reference/compilation.html#compiler-directives)
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 倍。