如何提升Python粒子n体模拟的运行速度？

我正在开发一个粒子模拟项目，粒子间存在随距离缩放的引力作用，还加入了进食、能量消耗、移动等功能。当前模拟100个粒子时只能跑到约5帧/秒，远达不到预期。我已经做了不少计算优化：比如移除三角函数，只保留平方根和基础运算；用JIT和向量化加速数学运算，实现多粒子并行处理。我用Python开发，JIT应该已经把密集运算转成更快的底层代码执行了。

为了提速，我已经设置只对合理距离内的粒子进行计算；也了解过分块模拟方案，但觉得和现有优化思路类似，而且重写成本高、收益有限。

请问还有什么方法能进一步提升运行效率？我希望能跑到60帧/秒，现在差距很大，但我见过同类模拟跑得更快的。

核心代码片段 #

向量化计算函数 #

@guvectorize([(float64[:], float64[:],  int32, float64[:], float64[:])],'(m),(n),(),(p)->(n)', nopython=True)
def mathStuff(ipos, j, interact_range, matrix, vels2):
    vels2[2] = 0 
    vels2[3] = 0
    if (ipos[0] != j[0] and ipos[1] != j[1]):
        delta1 = (ipos[0]-j[0])
        delta2 = (ipos[1]-j[1])
        dist = (math.sqrt(delta1**2 + delta2**2))
        if in_range(dist, interact_range): # if j is in range 
            dx,sector = distance_calc(dist, interact_range) # get sector the particles are in
            if sector <= 4: # if close by, calculate remainder of the sector and repel based on that + inverse square law 
                #dy/dx = 3
                dx = (dx + sector +1)/5
                vel_mag = -(1/dx)**2
            elif sector > 4 and sector <= 16: #medium distance: scale from 0 attraction to attraction matrrix attraction
                dx = dx + sector-5
                grad = matrix[int(j[4])]/12
                vel_mag = grad*dx 
            elif sector >16 and sector <= 26: # large distance: scale from attraction matrix attraction down to 0 
                dx = dx + sector-16
                grad = matrix[int(j[4])]
                vel_mag = grad -(grad/10)*dx
            else: # too far - 0 velocity
                vel_mag= 0 
            vels = trig(delta2,delta1, vel_mag) # get what the velocity should be 
            vels2[0] = vels[0]
            vels2[1] = vels[1]
        else:
            vels2[0] = 0.0
            vels2[1] = 0.0
        if in_range(dist, interact_range/4):
            vels2[2] = int(j[4])+1
    else:
        vels2[0] = 0.0
        vels2[1] = 0.0

主循环逻辑 #

for c1,i in enumerate(particles): # for each particle
        i.timestep(positions[c1][0],positions[c1][1]) # update position  
        positions[c1][2] = 0 
        positions[c1][3] = 0
         
        matrixes = [i.attraction_matrix[key] for key in i.attraction_matrix]
        expand_pos = np.concatenate((positions,np.asarray([types]).T), axis=1)
        v = np.asarray(mathStuff([(y) for y in positions[c1]],expand_pos, int(interact_range), matrixes))

辅助工具函数 #

@jit(nopython=True)
def in_range(dist, interact_range):
    if dist < math.sqrt(2*interact_range**2):
        return True
    else:
        return False 

# calculate the sector that a particle is in (based on distance split into 32 sectors)
@jit(nopython=True)
def distance_calc(other_dist, interact_range):
    other_dist = other_dist/(math.sqrt(2*interact_range**2)/32)
    dx = other_dist%1 
    sector = other_dist-dx 
    return dx,sector

# get velx and vely from velTotal and posx and posy 
@jit(nopython=True)
def trig(a,b, vel_mag):
    angle = a/b
    sign = np.sign(b)
    sqrt = math.sqrt(1+angle**2)
    self_vel0 = (-vel_mag/sqrt)*sign
    self_vel1 = self_vel0*angle
    return self_vel0,self_vel1

以下是针对你的代码和场景的具体优化方向，均为可快速落地且收益明显的方案：

1. 消除主循环内的重复内存分配与计算 #

主循环中每轮执行的np.concatenate、np.asarray([types]).T、[i.attraction_matrix[key] for key...]都会重复分配内存，是核心性能开销点：

• 提前缓存固定不变的计算结果：比如预计算types_col = np.asarray([types]).T，主循环直接复用；若attraction_matrix结构固定，提前将所有粒子的矩阵数据转为二维数组，避免每轮重新推导。
• 预初始化expand_pos数组，每轮仅更新positions部分，避免重复拼接操作。

2. 优化距离判断，避免冗余开根号 #

平方根是高耗时浮点操作，可通过平方比较替代：

• 提前预计算并缓存最大交互距离的平方：max_interact_dist_sq = 2 * interact_range ** 2，判断粒子是否在范围内时，直接用delta1**2 + delta2**2 < max_interact_dist_sq，完全跳过开根号步骤。
• in_range函数中的math.sqrt(2*interact_range**2)是固定值，提前缓存为max_interact_dist，避免重复计算。

3. 简化distance_calc的固定计算项 #

distance_calc中的分母math.sqrt(2*interact_range**2)/32是固定值，提前预计算为sector_scale，函数内直接用other_dist / sector_scale即可。

4. 重写trig函数，消除除法与角度计算 #

原函数的角度计算存在除法开销与零值风险，改用向量归一化方案：

@jit(nopython=True)
def trig(a, b, vel_mag):
    dist_sq = a**2 + b**2
    if dist_sq == 0:
        return 0.0, 0.0
    dist_inv = 1.0 / math.sqrt(dist_sq)
    self_vel0 = -vel_mag * b * dist_inv
    self_vel1 = -vel_mag * a * dist_inv
    return self_vel0, self_vel1

该方案无需计算角度，避免除法与符号判断，更快且更稳定。

5. 优化guvectorize内核的分支与调用 #

• 将in_range逻辑直接内联到mathStuff中，消除函数调用开销。
• 用delta1**2 + delta2**2 == 0代替ipos[0] != j[0] and ipos[1] != j[1]，更高效判断粒子是否为自身。
• 调整sector判断分支的顺序，将最常出现的情况放在前面，减少分支跳转的性能损耗。

6. 引入空间索引减少无效计算 #

用scipy.spatial.cKDTree构建粒子空间索引，每轮仅查询当前粒子周围有效范围内的邻居，而非遍历所有粒子：

# 提前构建KD树
tree = cKDTree(positions[:, :2])
for c1, i in enumerate(particles):
    i.timestep(positions[c1][0], positions[c1][1])
    positions[c1][2] = 0
    positions[c1][3] = 0
    # 查询当前粒子的邻居索引
    neighbor_indices = tree.query_ball_point(positions[c1][:2], max_interact_dist)
    # 仅处理邻居粒子
    neighbor_positions = expand_pos[neighbor_indices]
    v = np.asarray(mathStuff(positions[c1], neighbor_positions, int(interact_range), matrixes))
    # 合并速度结果到总数组

若每个粒子平均仅有10个邻居，计算量可直接降至原有的10%。

7. 优化JIT与向量化参数 #

• 给guvectorize添加target='parallel'参数，利用CPU多核心并行计算（需确保CPU支持）。
• 检查mathStuff的维度签名是否匹配实际数据，避免不必要的维度转换开销。

8. 减少Python循环开销 #

主循环的Python层循环本身存在性能损耗，可尝试：

• 将主循环逻辑用@jit(nopython=True)装饰，让JIT编译整个循环逻辑。
• 用NumPy向量化操作替代Python循环，将更多逻辑转移到底层执行。

内容的提问来源于stack exchange，提问作者Mercury