Python在for循环和数百个属性查找上变慢。使用Numba?

时间:2016-08-08 01:11:47

标签: python jit numba kdtree

我正在研究一个简单的展示SPH(平滑粒子流体动力学,虽然这里不相关)在python中的实现。代码有效,但执行有点迟钝。我经常要将单个粒子与一定数量的邻居进行比较。在早期的实现中,我将所有粒子位置和所有距离 - 每个存在的粒子保持在大的numpy阵列中 - >到某一点,这是非常快的。但在视觉上并不令人愉悦和n ** 2。现在我想用类+ kdTree来干净简单地加速邻居搜索。

这一切都发生在我的全球模拟课程中。此外,还有一个名为" particle"包含所有个人信息。我之前创建了数百个实例并循环遍历它们。

 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"展示他们如何加速数学重要的个人功能。我不知道选择哪个。在一个侧面节点,我甚至无法让Numba在这种情况下工作。我得到一个looong错误消息。我希望它像打耳机一样简单" @ jit"在它面前。

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

1 个答案:

答案 0 :(得分:2)

这种循环密集型计算在Python中很慢。在这些情况下,您要做的第一件事就是看看是否可以对这些操作进行矢量化并摆脱循环。然后,实际的计算将在C或Fortran库中完成,您将获得大量的加速。如果你能做到这一点通常就是这样,因为维护代码要容易得多。

然而,某些操作本质上是循环密集型的。在这些情况下使用Cython会对你有很大帮助 - 当你循环循环时,你通常可以期望加速60X +。我也有过与numba相似的经历 - 当我的功能变得复杂时,它无法让它更快,所以通常我只是使用Cython。

Cython中的编码不是太糟糕 - 比在C中实际编码容易得多,因为您可以通过内存视图轻松访问numpy数组。另一个优点是可以很容易地将循环与openMP并行化,这可以为您提供额外的4倍+加速(当然,这取决于您机器中的核心数量),因此您的代码可以快上百倍。

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