在零重力二维空间中优化粒子引力计算

Question

我用Python做了一个小的粒子可视化项目。
我在一个没有重力的二维空间里计算粒子的运动。每个粒子会根据它的质量和与其他粒子的距离来吸引其他粒子。

我在pygame中做了可视化，所有的计算都按计划进行，效果不错。不过，我需要极大地优化这些计算。目前这个系统大约能在合理的帧率下计算100到150个粒子。我把所有的计算放在一个单独的线程里，这样虽然有些提升，但还是远远不够。

我看过scipy和numpy，但因为我不是科学家或数学高手，所以看得一头雾水。感觉我走在正确的路上，但就是不知道该怎么做。

我需要计算所有粒子之间的吸引力，这就需要在一个循环里再嵌套一个循环。
而且因为我还需要检查粒子之间是否发生了碰撞，所以得重复一遍这个过程。

写这种代码真让我心痛……

Numpy可以实现数组与数组之间的计算，但我还没找到办法让一个数组里的所有项与另一个数组里的所有项进行计算。有没有这样的功能？
如果有的话，我可以创建几个数组，计算速度会快很多，而且应该有个函数可以找到两个数组中值相同的项的索引（也就是碰撞检测）。

这是我今天的吸引力/碰撞计算代码：

class Particle:
    def __init__(self):
        self.x = random.randint(10,790)
        self.y = random.randint(10,590)
        self.speedx = 0.0
        self.speedy = 0.0
        self.mass = 4

#Attraction    
for p in Particles:
    for p2 in Particles:
        if p != p2:
            xdiff = P.x - P2.x
            ydiff = P.y - P2.y
            dist = math.sqrt((xdiff**2)+(ydiff**2))
            force = 0.125*(p.mass*p2.mass)/(dist**2)
            acceleration = force / p.mass
            xc = xdiff/dist
            yc = ydiff/dist
            P.speedx -= acceleration * xc
            P.speedy -= acceleration * yc
for p in Particles:
    p.x += p.speedx
    p.y += p.speedy

#Collision
for P in Particles:
   for P2 in Particles:
        if p != P2:
            Distance = math.sqrt(  ((p.x-P2.x)**2)  +  ((p.y-P2.y)**2)  )
            if Distance < (p.radius+P2.radius):
                p.speedx = ((p.mass*p.speedx)+(P2.mass*P2.speedx))/(p.mass+P2.mass)
                p.speedy = ((p.mass*p.speedy)+(P2.mass*P2.speedy))/(p.mass+P2.mass)
                p.x = ((p.mass*p.x)+(P2.mass*P2.x))/(p.mass+P2.mass)
                p.y = ((p.mass*p.y)+(P2.mass*P2.y))/(p.mass+P2.mass)
                p.mass += P2.mass
                p.radius = math.sqrt(p.mass)
                Particles.remove(P2)

Answer 1

为了快速计算，你需要把x、y、speedx、speedy和m这些数据存储在numpy数组里。比如：

import numpy as np

p = np.array([
    (0,0),
    (1,0),
    (0,1),
    (1,1),
    (2,2),
], dtype = np.float)

p是一个5行2列的数组，用来存储粒子的x和y位置。要计算每一对粒子之间的距离，你可以使用：

print np.sqrt(np.sum((p[:, np.newaxis] - p[np.newaxis, :])**2, axis=-1))

输出结果是：

[[ 0.          1.          1.          1.41421356  2.82842712]
 [ 1.          0.          1.41421356  1.          2.23606798]
 [ 1.          1.41421356  0.          1.          2.23606798]
 [ 1.41421356  1.          1.          0.          1.41421356]
 [ 2.82842712  2.23606798  2.23606798  1.41421356  0.        ]]

或者你也可以使用scipy里的cdist功能：

from scipy.spatial.distance import cdist
print cdist(p, p)

Answer 2

我之前做过这个，发现加速碰撞计算的一个方法是存储一个附近粒子的列表。

简单来说，在计算重力的时候，你可以这样做：

for (int i = 0; i < n; i++)
{
    for (int j = i + 1; j < n; j++)
    {
        DoGravity(Particle[i], Particle[j]);
        if (IsClose(Particle[i], Particle[j]))
        {
            Particle[i].AddNeighbor(Particle[j]);
            Particle[j].AddNeighbor(Particle[i]);
        }
    }
}

然后，你就可以遍历所有粒子，逐个进行碰撞检测。通常情况下，这样的操作在最好的情况下是 O(n)，但在最坏的情况下可能会变成 O(n^2)。

另一种选择是把粒子放进一个八叉树里。构建一个八叉树的复杂度大约是 O(n)，然后你可以查询看看有没有粒子靠得很近。到那时，你只需要对这些成对的粒子进行碰撞检测。我认为这样做的复杂度是 O(n log n)。

不仅如此，你还可以利用八叉树来加速重力计算。这样一来，复杂度从 O(n^2) 降到 O(n log n)。大多数八叉树的实现都有一个“开启参数”，可以控制速度和准确性之间的权衡。因此，八叉树通常比直接的成对计算准确性低，而且编写起来也比较复杂，但它们让大规模模拟成为可能。

如果你以这种方式使用八叉树，你就会进行一种叫做巴恩斯-哈特模拟。

注意：因为你是在做二维的工作，二维的八叉树叫做 四叉树。想了解更多信息，可以查看这个维基百科文章：http://en.wikipedia.org/wiki/Quadtree

Answer 3

你可以先尝试使用复数来进行计算：在这种方式下，相关的引力和动力学公式非常简单，而且计算速度也很快（因为NumPy可以内部处理计算，而不是让你单独处理x和y坐标）。比如，两个粒子在z和z'之间的力可以简单表示为：

(z-z')/abs(z-z')**3

NumPy可以非常快速地计算所有z/z'的组合。举个例子，所有z-z'值的矩阵可以通过一维数组Z的坐标来获得，方法是Z-Z[:, numpy.newaxis]（对角线上的项[z=z']在计算1/abs(z-z')**3时需要特别处理：它们应该设为零）。

至于时间演变，你当然可以使用SciPy的快速微分方程例程：它们比逐步的欧拉积分要精确得多。

无论如何，深入了解NumPy会非常有帮助，特别是如果你打算进行科学计算，因为NumPy的速度非常快。