PAYATHON 2026

如何为 N-body simulation 正确实现 RK45 算法

技术老齐

结论

RK45 可以用于 N-body system。当前实现的问题是没有满足 RK 方法的两个基本条件:

  1. 每个 Runge-Kutta stage 都必须基于同一个完整的系统中间状态,计算所有天体的导数。不能逐个推进天体,却让其余天体停留在旧状态。
  2. 误差估计系数写错了。-2197/4104 应为 -2197/75240,四阶解中的 1408/2625 * k2 应为 1408/2565 * k3。使用错误的误差公式后,误差无法随步长正常下降。即使步长已经降到 min_step,程序仍会反复重试,看起来就像陷入了无限循环。

接受一个积分步后,还要分别返回本次实际使用的步长和下一步建议步长。更新当前时间时必须使用前者。

为什么必须整体推进系统

N-body system 的状态可以写成:

[ Y= \begin{bmatrix} x_1 & y_1 & v_{x1} & v_{y1}\ x_2 & y_2 & v_{x2} & v_{y2}\ \vdots & \vdots & \vdots & \vdots\ x_N & y_N & v_{xN} & v_{yN} \end{bmatrix} ]

对应的微分方程是:

[ \frac{dY}{dt}=F(Y) ]

第 (i) 个天体的加速度取决于所有天体在同一时刻的位置:

[ \mathbf a_i = -G\sum_{j\ne i} m_j \frac{\mathbf r_i-\mathbf r_j} {\left(\lVert\mathbf r_i-\mathbf r_j\rVert^2+\epsilon^2\right)^{3/2}} ]

计算 k2 时,要先把所有天体移动到 Y + h*k1/4,再用这个完整状态计算加速度。原代码只修改当前天体的中间状态,而 arr_calc() 会通过全局 state_arr 读取其他天体的旧位置,所以算出的结果并不是 N-body ODE 的 RK stage。

原实现中的具体错误

1. RKF45 系数写错

经典 Runge-Kutta-Fehlberg 4(5) 方法的四阶增量是:

y4 = y + h * (
    25/216 * k1
    + 1408/2565 * k3
    + 2197/4104 * k4
    - 1/5 * k5
)

原实现写成了:

k1*25/216 + k2*1408/2625 + k4*2197/4104 - k5*1/5

其中有两处错误:

  • k2 应为 k3
  • 2625 应为 2565

五阶解和四阶解之差为:

error = h * (
    1/360 * k1
    - 128/4275 * k3
    - 2197/75240 * k4
    + 1/50 * k5
    + 2/55 * k6
)

原代码把 75240 写成了 4104。两者相差很大,会直接导致误差估计失效。

2. 修改步长后仍使用旧步长算出的增量

当前流程是:

  1. 用旧 dt 计算 k_list
  2. 根据误差修改 dt
  3. 接受旧 k_list
  4. 返回修改后的 dt
  5. 外层执行 time_val += delta_t

状态按旧步长推进,时间却按新步长推进,两者从第一步开始就可能不一致。

3. 到达 min_step 后没有失败出口

如果步长已经是 min_step,误差却仍不合格,下面的循环不会再改变任何条件:

while True:
    ...
    dt = min_h
    ...
    if r_val <= err_tol:
        return ...

程序此时应抛出异常,说明给定的最小步长无法满足误差要求,否则循环会一直继续。

4. Python list 不适合这里的切片更新

state_arr 是嵌套 list,而这段代码:

out_arr[i][1:] += k_list[i]

不能可靠地完成数值数组加法。更合适的做法是从一开始就把质量和动态状态分别保存为 NumPy array。

5. out_arr = inp_arr 没有复制数据

下面的赋值只会创建另一个引用:

out_arr = inp_arr

修改 out_arr 时,inp_arr 也会随之改变。被拒绝的试探步不应修改原状态。

一个完整的 RKF45 实现

下面的实现使用经典 Fehlberg 4(5) 系数。质量单独保存,动态状态使用 (N, 4) 数组。

import numpy as np


def nbody_rhs(y, masses, G=1.0, softening=0.0):
    """
    y.shape == (N, 4)
    每行依次为 [x, y, vx, vy]
    """
    positions = y[:, :2]
    velocities = y[:, 2:]

    # diff[i, j] = r_i - r_j
    diff = positions[:, None, :] - positions[None, :, :]

    distance_squared = np.sum(diff * diff, axis=2)
    distance_squared += softening**2

    inverse_r3 = distance_squared ** (-1.5)

    # 去掉天体对自身的作用
    np.fill_diagonal(inverse_r3, 0.0)

    accelerations = -G * np.sum(
        masses[None, :, None] * diff * inverse_r3[:, :, None],
        axis=1,
    )

    return np.column_stack((velocities, accelerations))


def rkf45_trial(y, h, masses, G=1.0, softening=0.0):
    """
    完成一次 RKF45 试探。
    返回五阶解以及四阶、五阶解之差。
    k1...k6 均为导数,不包含 h。
    """
    f = lambda stage: nbody_rhs(
        stage,
        masses,
        G=G,
        softening=softening,
    )

    k1 = f(y)

    k2 = f(
        y + h * (
            1 / 4 * k1
        )
    )

    k3 = f(
        y + h * (
            3 / 32 * k1
            + 9 / 32 * k2
        )
    )

    k4 = f(
        y + h * (
            1932 / 2197 * k1
            - 7200 / 2197 * k2
            + 7296 / 2197 * k3
        )
    )

    k5 = f(
        y + h * (
            439 / 216 * k1
            - 8 * k2
            + 3680 / 513 * k3
            - 845 / 4104 * k4
        )
    )

    k6 = f(
        y + h * (
            -8 / 27 * k1
            + 2 * k2
            - 3544 / 2565 * k3
            + 1859 / 4104 * k4
            - 11 / 40 * k5
        )
    )

    y4 = y + h * (
        25 / 216 * k1
        + 1408 / 2565 * k3
        + 2197 / 4104 * k4
        - 1 / 5 * k5
    )

    y5 = y + h * (
        16 / 135 * k1
        + 6656 / 12825 * k3
        + 28561 / 56430 * k4
        - 9 / 50 * k5
        + 2 / 55 * k6
    )

    error = y5 - y4
    return y5, error

所有 k 都根据完整的 (N, 4) 中间状态计算,因此每个 stage 中的天体位置是同步的。

自适应步长和误差归一化

只用一个绝对误差阈值通常不够。位置和速度的尺度可能不同,各天体的状态量也可能相差很大。通常可以为每个分量定义:

[ scale_i = atol_i + rtol \max(|y_i|, |y_i^{new}|) ]

然后计算无量纲误差:

[ E = \max_i \left|\frac{error_i}{scale_i}\right| ]

当 (E \le 1) 时接受该步。

def rkf45_step(
    y,
    h,
    masses,
    *,
    rtol=1e-8,
    atol=1e-10,
    min_h=1e-10,
    max_h=1e-2,
    G=1.0,
    softening=0.0,
    safety=0.9,
    max_rejections=100,
):
    """
    返回:
        y_new: 接受后的状态
        h_used: 本次实际使用的步长
        h_next: 下一次建议使用的步长
    """
    h = float(np.clip(h, min_h, max_h))

    for _ in range(max_rejections):
        y_new, error = rkf45_trial(
            y,
            h,
            masses,
            G=G,
            softening=softening,
        )

        if not (
            np.all(np.isfinite(y_new))
            and np.all(np.isfinite(error))
        ):
            raise FloatingPointError(
                "积分中出现 NaN 或无穷大;请检查碰撞、单位和初始条件"
            )

        scale = atol + rtol * np.maximum(
            np.abs(y),
            np.abs(y_new),
        )

        error_norm = np.max(np.abs(error) / scale)

        if error_norm <= 1.0:
            if error_norm == 0.0:
                factor = 5.0
            else:
                factor = safety * error_norm ** (-1.0 / 5.0)
                factor = np.clip(factor, 0.2, 5.0)

            h_next = float(np.clip(
                h * factor,
                min_h,
                max_h,
            ))

            return y_new, h, h_next

        if h <= min_h * (1.0 + 10.0 * np.finfo(float).eps):
            raise RuntimeError(
                "步长已达到 min_h,但局部误差仍超过容限;"
                "请减小 min_h、调整容限,或检查是否发生近碰撞"
            )

        factor = safety * error_norm ** (-1.0 / 5.0)
        factor = np.clip(factor, 0.1, 0.9)
        h = max(min_h, h * factor)

    raise RuntimeError(
        f"连续拒绝了 {max_rejections} 个试探步,积分无法继续"
    )

这里使用 -1/5 次幂,因为直接控制的是四阶解与五阶解之差,其主导项随 (h^5) 缩放。有些教材把误差定义为 error / h,此时控制器会使用 -1/4 次幂。这两种写法都可以成立,但误差定义必须与指数匹配。

推进 Golden Pancake Drama 初始状态

import numpy as np
from matplotlib import pyplot as plt


initial_state = np.array([
    [1.0,                -1.0, 0.0,  0.151837333844,  0.139577137808],
    [0.6123560888279949,  0.0, 0.0, -0.495911893796, -0.455869192302],
    [1.0,                 1.0, 0.0,  0.151837333844,  0.139577137808],
], dtype=float)

masses = initial_state[:, 0].copy()
state = initial_state[:, 1:].copy()

G = 1.0
softening = 5e-5

t = 0.0
t_end = 8.8457131110
h = 1e-3

min_step = 1e-10
max_step = 1e-4

times = [t]
states = [state.copy()]

accepted_steps = 0

while t < t_end:
    # 避免最后一步越过终点
    h_try = min(h, t_end - t)

    # 最后剩余时间可能小于通常设置的 min_step
    call_min_h = min(min_step, h_try)

    state, h_used, h_next = rkf45_step(
        state,
        h_try,
        masses,
        rtol=1e-8,
        atol=1e-10,
        min_h=call_min_h,
        max_h=max_step,
        G=G,
        softening=softening,
    )

    # 必须按本次实际使用的步长更新时间
    t += h_used
    h = h_next
    accepted_steps += 1

    if accepted_steps % 500 == 0 or t >= t_end:
        times.append(t)
        states.append(state.copy())
        print(f"TRAJECTORY CALCULATION: {100 * t / t_end:.3f}%")

states = np.asarray(states)

fig = plt.figure(layout="constrained", figsize=(5, 5))

for i, color in enumerate(("blue", "red", "green")):
    plt.plot(
        states[:, i, 0],
        states[:, i, 1],
        color=color,
        linewidth=0.5,
        label=f"m{i + 1}",
    )

plt.axis("equal")
plt.legend()
plt.show()

这里的 rtol=1e-8 和 atol=1e-10 只是示例值,并不表示这组轨道一定需要或适合这些容限。实际使用时,应根据单位、期望精度和守恒量误差调整。

近距离交会时的注意事项

步长变得很小不一定是程序错误

近距离交会时,加速度及其变化率会迅速增大。为了满足误差要求,自适应积分器可能连续缩短步长。只要误差确实随步长下降,这种行为就是正常的。

可以记录以下数据,判断程序是否真的卡住:

print(t, h_used, h_next)

如果 h 已降到 min_h,误差仍不合格,程序应停止并报告原因,不能继续重试。

max_step=1e-4 本身会产生大量步骤

积分区间约为 8.8457,最大步长只有 1e-4。即使每次都采用最大步长,也至少需要约 8.8 万个接受步,而且每一步都要计算六次完整的 N-body 导数。因此,运行时间较长并不一定表示程序陷入了死循环。

softening 会改变所求解的物理模型

加入:

distance_squared += softening**2

可以避免距离趋近于零时产生奇异加速度,但求解结果对应的是软化后的引力系统。如果要复现某个未软化模型的周期轨道,softening 会改变轨道和周期。是否保留它,要看初始条件对应的是哪种模型。

用守恒量检查实现

单看轨迹图很难判断积分是否正确。对于没有外力的封闭系统,还应监控:

  • 总线动量
  • 质心运动
  • 总角动量
  • 总能量

这些量不应出现明显的单向漂移。RK45 不是辛积分器,长时间积分时无法保证能量误差始终严格受限。如果目标是长期模拟太阳系动力学,还需要单独评估这一点,但这不影响使用 RK45 验证当前实现。

修复时,需要把所有天体合并为一个 ODE 状态,在每个 RK stage 中同步计算整个系统,同时严格区分本次接受的步长和下一步建议步长。处理好这两点后,无限重试的问题会消失,误差控制也能正确反映 N-body system 的局部截断误差。

备注:内容仅供参考。