带有函数调用的Python嵌套循环优化

时间:2019-02-09 18:48:39

标签: python optimization

我正在python3.5中构建一个简单的N体集成器,该集成器实现了leapfrog timestepping as a position verlet。 本质上,它来回更新两个浮点变量 x_tmp v_tmp ,在这里我需要对 self.forces()进行函数调用更新 v_tmp

此函数调用使我的运行速度极其慢(我已经介绍了)。该调用没有任何异常,只有几个平方根以及一些加减数字。

   for t in range(self.max_timesteps):

        #For all objects do the position verlet with generator expression / list comprehensions
        x_tmp     = [x_tmp[j] + 0.5*self.timestep*v_tmp[j] for j in range(self.num_objects)]
        v_tmp     = [v_tmp[j] + self.timestep*self.forces(x_tmp[j]) for j in range(self.num_objects)]
        x_tmp     = [x_tmp[j] + 0.5*self.timestep*v_tmp[j] for j in range(self.num_objects)]

        if(t % self.outputtime == 0):
            self.x_list[outputcounter] = x_tmp
            self.v_list[outputcounter] = v_tmp

和函数 self.forces()

def forces(self,x):

    r = np.sqrt((x[0])**2+(x[1])**2+(x[2])**2) # spherical radius
    R = math.hypot(x[0], x[1])   # cylindrical radius

    def _f1(r,x,y,z):
        f = -G*self.Mb/(r*(r+self.rb)**2) 
        return np.array([f*x, f*y, f*z])

    def _f2(R,x,y,z): 
        rr = math.hypot(self.disc_b,z)
        arr = (self.disc_a+rr)
        arrR = math.hypot(arr,R)
        f = -G*self.Md/ arrR**3.
        fz = f*(arr/rr)
        return np.array([f*x,f*y,fz*z])

    def _f3(r,x,y,z):
        f = -self.Vh**2/(self.rh**2+r**2)
        return np.array([f*x,f*y,f*z])

    a = _f1(r,x[0],x[1],x[2]) + _f2(R,x[0],x[1],x[2]) + _f3(r,x[0],x[1],x[2])

    return np.array((a[0], a[1], a[2]))

现在上方代码块中带有 x_tmp = ... 的两行都可以与 num_objects 很好地缩放(它们几乎没有),但是具有 v_tmp 及其中的函数调用随 num_objects 线性缩放。

这很糟糕。在 max_timestes = 10 ^ 6 的情况下,使用此代码可以获得 num_objects 秒的运行时间,因此,如果要使用此代码计算200个对象,则需要200秒。这是完全不能接受的。

但是我对这里的操作有点不知所措,因为我已经使用 math.hypot()和其他一些方法优化了一些2D平方根。但是forces()调用仍然非常慢,这种情况在C或等效语言中永远不会发生。

所以现在我要寻求帮助,在优化那些函数调用时,有什么明显的我被忽略了的东西吗?或者我可以快速构建一个C函数来调用它来加快速度。

任何想法都值得赞赏。

0 个答案:

没有答案