加速Python/Cython循环。

我试图让python中的循环尽可能快地运行。所以我就跳进了NumPy和Cython。 以下是原始Python代码：

def calculate_bsf_u_loop(uvel,dy,dz):
   """
   Calculate barotropic stream function from zonal velocity

   uvel (t,z,y,x)
   dy   (y,x)
   dz   (t,z,y,x)

   bsf  (t,y,x)
   """

   nt = uvel.shape[0]
   nz = uvel.shape[1]
   ny = uvel.shape[2]
   nx = uvel.shape[3]

   bsf = np.zeros((nt,ny,nx))

   for jn in range(0,nt):
      for jk in range(0,nz):
         for jj in range(0,ny):
            for ji in range(0,nx):
               bsf[jn,jj,ji] = bsf[jn,jj,ji] + uvel[jn,jk,jj,ji] * dz[jn,jk,jj,ji] * dy[jj,ji] 

   return bsf

这只是k个指数的和。数组大小为nt=12、nz=75、ny=559、nx=1442，因此约有7.25亿个元素。 这花了68秒。现在，我已经在cython做过了

import numpy as np
cimport numpy as np
cimport cython

@cython.boundscheck(False) # turn off bounds-checking for entire function
@cython.wraparound(False)  # turn off negative index wrapping for entire function

## Use cpdef instead of def
## Define types for arrays
cpdef calculate_bsf_u_loop(np.ndarray[np.float64_t, ndim=4] uvel, np.ndarray[np.float64_t, ndim=2] dy, np.ndarray[np.float64_t, ndim=4] dz):
   """
   Calculate barotropic stream function from zonal velocity

   uvel (t,z,y,x)
   dy   (y,x)
   dz   (t,z,y,x)

   bsf  (t,y,x)
   """

   ## cdef the constants
   cdef int nt = uvel.shape[0]
   cdef int nz = uvel.shape[1]
   cdef int ny = uvel.shape[2]
   cdef int nx = uvel.shape[3]

   ## cdef loop indices
   cdef ji,jj,jk,jn

   ## cdef. Note that the cdef is followed by cython type
   ## but the np.zeros function as python (numpy) type
   cdef np.ndarray[np.float64_t, ndim=3] bsf = np.zeros([nt,ny,nx], dtype=np.float64)

   for jn in xrange(0,nt):
      for jk in xrange(0,nz):
         for jj in xrange(0,ny):
            for ji in xrange(0,nx):
               bsf[jn,jj,ji] += uvel[jn,jk,jj,ji] * dz[jn,jk,jj,ji] * dy[jj,ji] 

   return bsf

这花了49秒。 但是，将循环替换为

for jn in range(0,nt):
      for jk in range(0,nz):
         bsf[jn,:,:] = bsf[jn,:,:] + uvel[jn,jk,:,:] * dz[jn,jk,:,:] * dy[:,:]

只需要0.29秒！不幸的是，我不能在我的完整代码中这样做

为什么NumPy切片比Cython循环快得多？ 我以为NumPy跑得很快，因为它是引擎盖下的Cython。那么它们不应该有相似的速度吗

如您所见，我在cython中禁用了边界检查，并且还使用“快速数学”进行编译。然而，这只会带来很小的加速。 是否有任何方法可以使循环的速度与NumPy切片的速度相似，或者循环总是比切片的速度慢

非常感谢您的帮助！

---

/Joakim

鉴于您正在执行

元素乘法

，然后在

4D

乘积数组的第二个轴上执行

求和

，代码迫切需要用户的干预，这对于 ally

numpy.einsum

以高效的方式进行。要解决您的问题，您可以通过两种方式使用

numpy.einsum

-

bsf = np.einsum('ijkl,ijkl,kl->ikl',uvel,dz,dy)

bsf = np.einsum('ijkl,ijkl->ikl',uvel,dz)*dy

运行时测试和验证输出-

In [100]: # Take a (1/5)th of original input shapes
     ...: original_shape = [12,75, 559,1442]
     ...: m,n,p,q = (np.array(original_shape)/5).astype(int)
     ...: 
     ...: # Generate random arrays with given shapes
     ...: uvel = np.random.rand(m,n,p,q)
     ...: dy = np.random.rand(p,q)
     ...: dz = np.random.rand(m,n,p,q)
     ...: 

In [101]: bsf = calculate_bsf_u_loop(uvel,dy,dz)

In [102]: print(np.allclose(bsf,np.einsum('ijkl,ijkl,kl->ikl',uvel,dz,dy)))
True

In [103]: print(np.allclose(bsf,np.einsum('ijkl,ijkl->ikl',uvel,dz)*dy))
True

In [104]: %timeit calculate_bsf_u_loop(uvel,dy,dz)
1 loops, best of 3: 2.16 s per loop

In [105]: %timeit np.einsum('ijkl,ijkl,kl->ikl',uvel,dz,dy)
100 loops, best of 3: 3.94 ms per loop

In [106]: %timeit np.einsum('ijkl,ijkl->ikl',uvel,dz)*dy
100 loops, best of 3: 3.96 ms per loo

---

您忘记声明循环变量的类型。在当前Cython版本中不需要声明循环变量。@jakevdp在这种情况下，不声明它们似乎可以让它正确地推断类型，但如果不声明类型而将它们声明为

cdef ji，jj，jk，jn

，则会将其丢弃。对–我的意思是

cdef ji，jj，jk，jn

行可以完全删除，这可能会提高执行时间。非常感谢！将循环索引声明为

cdefintji、jj、jk、jn

产生了所有的不同！这是一个令人印象深刻的加速。我将研究einsum方法。谢谢