Python在for循环和数百个属性查找上运行缓慢。使用Numba？

Question

我正在研究一个简单的展示SPH（平滑粒子流体力学，这里虽然没有相关），但在python中实现。代码有效，但执行有点迟缓。我经常需要比较单个粒子与一定数量的邻居。在早期的实现中，我保留了所有粒子位置和所有距离到每个存在粒子的大块数组 - >到某个点，这个速度非常快。但在视觉上不令人愉悦，并且n ** 2。现在我想用class + kdTree来简化它，以加速邻居搜索。

这一切都发生在我的全球模拟类。此外还有一个名为“粒子”的类，其中包含所有个人信息。我之前创建了数百个实例，并通过它们循环。

def calculate_density(self): 
    #Using scipys advanced nearest neighbour seach magic 
    tree = scipy.spatial.KDTree(self.particle_positions) 

    #here we go... loop through all existing particles. set attributes.. 
    for particle in self.my_particles: 

     #get the indexes for the nearest neighbours 
     particle.index_neighbours = tree.query_ball_point(particle.position,self.h,p=2) 


     #now loop through the list of neighbours and perform some additional math 
     particle.density = 0 
     for neighbour in particle.index_neighbours: 


      r = np.linalg.norm(particle.position - self.my_particles[neighbour].position) 
      particle.density += particle.mass * (315/(64*math.pi*self.h**9)) *(self.h**2-r**2)**3 

我只为216颗粒计时0.2717630863189697s。

现在我想知道：如何加快速度？ 像“Numba”这样的大多数在线工具都展示了他们如何加快数学繁重的个人功能。我不知道该选哪个。在一个sidenode上，我甚至无法让Numba在这种情况下工作。我收到一条looong错误消息。我希望它像在它前面打“@jit”一样简单。

我知道它的循环带有属性调用，无论如何都会破坏我的性能 - 而不是数学或邻居搜索。可悲的是，我是编程的新手，我喜欢干净的方法，我在这里工作:(任何想法？

Answer 1

这些类型的循环密集型计算在Python中很慢，在这些情况下，你想要做的第一件事是看看你是否可以矢量化这些操作并摆脱循环，然后在C或Fortran库中完成实际的计算，并且你会得到很多的加速，如果你能做到这一点，通常这是要走的路，因为在这些情况下，使用Cython将会对你有很大的帮助 - 当你循环你的循环时，你通常会希望60X +的速度加快。也有类似的经验与numba - 当我的功能变得复杂，它没有成功所以通常我只是使用Cython。

在Cython中进行编码并不算太差 - 比C中的实际代码容易得多，因为您可以通过内存视图轻松访问numpy数组。另一个优点是，使用openMP并行化循环非常简单，它可以为您提供额外的4倍加速（当然，取决于您机器中的核心数量），因此您的代码可以快上百倍。

一个问题是，要获得最佳速度，你必须删除你的循环内的所有python调用，这意味着你不能调用numpy/scipy函数。因此，您必须将tree.query_ball_point和np.linalg.norm部分转换为Cython以获得最佳速度。