需要帮助的矢量化代码或优化

Question

我想通过第一次内插数据来做一个双积分做一个表面。我正在使用numba来加速这个过程，但这只是需要很长时间。

Here is my code,带有运行位于here和here的代码所需的图像。

Answer 1

注意到你的代码有一个for循环的四重嵌套集，我专注于优化内对。这里的旧代码：

for i in xrange(K.shape[0]): 
    for j in xrange(K.shape[1]): 

     print(i,j) 
     '''create an r vector ''' 
     r=(i*distX,j*distY,z) 

     for x in xrange(img.shape[0]): 
      for y in xrange(img.shape[1]): 
       '''create an ksi vector, then calculate 
        it's norm, and the dot product of r and ksi''' 
       ksi=(x*distX,y*distY,z) 
       ksiNorm=np.linalg.norm(ksi) 
       ksiDotR=float(np.dot(ksi,r)) 

       '''calculate the integrand''' 
       temp[x,y]=img[x,y]*np.exp(1j*k*ksiDotR/ksiNorm) 

     '''interpolate so that we can do the integral and take the integral''' 
     temp2=rbs(a,b,temp.real) 
     K[i,j]=temp2.integral(0,n,0,m) 

由于K和IMG每个约为2000×2000，最里面的语句需要执行160000亿次。这对于使用Python来说并不实用，但我们可以使用NumPy将工作转换为C和/或Fortran以进行矢量化。我一次只做了一步，试图确保结果一致;这里是我结束了：

'''create all r vectors''' 
R = np.empty((K.shape[0], K.shape[1], 3)) 
R[:,:,0] = np.repeat(np.arange(K.shape[0]), K.shape[1]).reshape(K.shape) * distX 
R[:,:,1] = np.arange(K.shape[1]) * distY 
R[:,:,2] = z 

'''create all ksi vectors''' 
KSI = np.empty((img.shape[0], img.shape[1], 3)) 
KSI[:,:,0] = np.repeat(np.arange(img.shape[0]), img.shape[1]).reshape(img.shape) * distX 
KSI[:,:,1] = np.arange(img.shape[1]) * distY 
KSI[:,:,2] = z 

# vectorized 2-norm; see http://stackoverflow.com/a/7741976/4323              
KSInorm = np.sum(np.abs(KSI)**2,axis=-1)**(1./2) 

# loop over entire K, which is same shape as img, rows first               
# this loop populates K, one pixel at a time (so can be parallelized)            
for i in xrange(K.shape[0]):                      
    for j in xrange(K.shape[1]):                     

     print(i, j) 

     KSIdotR = np.dot(KSI, R[i,j]) 
     temp = img * np.exp(1j * k * KSIdotR/KSInorm) 

     '''interpolate so that we can do the integral and take the integral''' 
     temp2 = rbs(a, b, temp.real) 
     K[i,j] = temp2.integral(0, n, 0, m) 

内部对循环现在已经完全消失了，提前完成矢量操作替代（在输入端的大小的空间成本直线）。

这样可以在不使用Numba的情况下，将我的Macbook Air 1.6 GHz i5上的外部两个循环的每次迭代的时间从340秒减少到1.3秒。在每次迭代1.3秒中，0.68秒用于rbs函数，即scipy.interpolate.RectBivariateSpline。有可能有进一步优化的空间 - 这里有一些想法：

1. Reenable Numba。我没有在我的系统上。它在这一点上可能没有太大的区别，但很容易让你测试。
2. 做更多特定领域的优化，例如试图简化正在完成的基本计算。我的优化旨在无损，并且我不知道您的问题域，因此无法尽可能深入地进行优化。
3. 尝试矢量化剩余的循环。这可能很难，除非你愿意用每次调用多次计算的东西来替换scipy RBS函数。
4. 获得更快的CPU。我的速度很慢;通过使用比我的微型笔记本电脑更好的计算机，你可以获得至少2倍的加速比。
5. 对您的数据进行降采样。您的测试图像是2000x2000像素，但包含的细节很少。如果你将他们的线性尺寸减少2-10倍，你会得到巨大的加速。

这就是我现在的情况。这在哪里离开你？假设一台稍微好一点的计算机并且没有进一步的优化工作，即使优化的代码也需要大约一个月的时间来处理你的测试图像。如果你只需要做一次，也许没关系。如果您需要更频繁地执行此操作，或者在尝试不同的操作时需要迭代代码，那么您可能需要继续优化 - 从现在消耗一半以上时间的RBS函数开始。

特别提示：您的代码会更容易处理，如果它不具有几乎相同的变量名称，如k和K，也没用过j作为变量名，也可以作为一个复杂的数字后缀很多（0j） 。