numpy/scipy中牛顿力的最佳性能计算

Question

对于在大学里进行的练习,我们必须在Python中实现具有精确牛顿力的Leapfrog积分器.课程结束了,我们的解决方案已经足够好了,但我想知道是否/如何能够更好地提高力计算的性能.

瓶颈是计算所有力(即加速度):

对于大(1000和更大)数量的粒子N(i,j <N).

这里r ₁和r ₂是存储在形状(N,3)的ndarray中的粒子位置的三维向量,Gm是粒子质量乘以我在形状ndarray中保存的重力常数(N) .

我到目前为止发现的最快版本如下:

def a(self):
    sep = self.r[np.newaxis, :] - self.r[:, np.newaxis]
    # Calculate the distances between all particles with cdist
    # this is much faster than by hand
    dists = cdist(self.r, self.r)
    scale =dists*dists*dists
    # set diagonal elements of dist to something != 0, to avoid division by 0
    np.fill_diagonal(scale,1)
    Fsum = (sep/scale.reshape(self.particlenr,self.particlenr,1))*self.Gm[:,None]
    return np.add.reduce(Fsum,axis=1)

但它让我觉得这可能不是最快的版本.与cdist进行比较时,第一行似乎太慢,而cdist的计算基本相同.此外,该解决方案不使用在问题中切换r ₁和r ₂的对称性并且计算所有元素两次.

您是否知道任何性能改进(不将力计算更改为某种近似值或更改编程语言)？

Answer 1

我怀疑 numpy 实际上是双重计算距离（因为它总是对称的）。它可能正在执行一项计算并在两个地方分配相同的值。

不过，我确实想到了一些想法：

1. 您可以按照 numpy 源代码编写 cdist 的自定义版本。每次迭代时程序可能会解析许多选项。不多，但也许它可以给你带来一个小小的百分比提升。
2. 预分配。每次运行 a() 时，您可能会为所有中间矩阵值重新分配内存。你能持续生产这些数量吗？

我还没有完成计算，但如果你能以某种方式优雅地减少冗余对称计算，我不会感到惊讶。