加速Python/Cython循环。

Question

我试图让python中的一个循环尽可能快地运行。所以我开始研究纽比和赛顿。 以下是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秒。现在，我已经在赛顿做过了

^{pr2}$

那花了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循环快得多？ 我以为纽比跑得快是因为它是引擎盖下的赛顿。所以它们的速度不是应该差不多吗？

如您所见，我在cython中禁用了边界检查，并且我还使用“快速数学”进行了编译。然而，这只是一个微小的加速。 有没有办法让一个循环的速度与纽比切片相似，或者循环总是比切片慢？

非常感谢任何帮助！ /约金

Answer 1

这段代码要求^{}的干预，假设您正在elementwise-multiplication，然后sum-reduction在{}乘积数组的第二个轴上执行，这意味着 ally numpy.einsum以一种高效的方式进行。要解决您的问题，您可以使用numpy.einsum两种方法-

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

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

运行时测试和验证输出-

^{pr2}$