Question

我试图让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切片速度相似，或者循环总是慢于切片？

非常感谢任何帮助！ /乔金姆

Answer 1

考虑到您在elementwise-multiplication产品数组的第二轴上sum-reduction然后4D，该代码会为numpy.einsum's的干预而尖叫，essenti 盟友numpy.einsum以高效的方式开展工作。要解决您的问题，您可以通过两种方式使用numpy.einsum -

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

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

运行时测试＆amp;验证输出 -

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

加速Python / Cython循环。

1 个答案: