N体运动模拟

给定n个质量为，坐标为，初速度为的质点，求其轨迹方程。对于质点i，有：

对于每个质点的方程，可转换为6个一阶微分方程，然后使用龙格库塔等方法进行数值求解。

代码如下：

from matplotlib.animation import FuncAnimation
import numpy as np
import matplotlib.pyplot as plt

G = 6.67430e-11  # 引力常数 [m^3 kg^(-1) s^(-2)]


# 定义天体类
class Body:
    def __init__(self, name, mass, position, velocity):
        self.name = name  # 天体名称
        self.mass = mass  # 天体质量 (kg)
        self.position = np.array(position, dtype='float64')  # 位置 (m)
        self.velocity = np.array(velocity, dtype='float64')  # 速度 (m/s)


# 计算所有天体之间的引力
def compute_forces(bodies):
    positions = np.array([body.position for body in bodies])
    masses = np.array([body.mass for body in bodies])

    # 计算位置的差值
    r = positions[:, np.newaxis, :] - positions[np.newaxis, :, :]
    distances = np.linalg.norm(r, axis=2)
    np.fill_diagonal(distances, np.inf)  # 避免自引力的影响

    # 计算引力大小
    force_magnitudes = G * masses[:, np.newaxis] * masses / (distances ** 2)
    forces = -np.sum((force_magnitudes[:, :, np.newaxis] * r / distances[:, :, np.newaxis]), axis=1)

    return forces


# 四阶龙格-库塔法
def rk4_step(bodies, dt):
    num_bodies = len(bodies)
    forces = compute_forces(bodies)  # 计算当前步骤的引力

    positions = np.array([body.position for body in bodies])
    velocities = np.array([body.velocity for body in bodies])
    masses = np.array([body.mass for body in bodies])

    # 定义四个中间步骤 k1, k2, k3, k4
    k1_positions = velocities * dt
    k1_velocities = forces / masses[:, np.newaxis] * dt

    # k2 步骤
    temp_positions = positions + k1_positions / 2
    temp_velocities = velocities + k1_velocities / 2
    forces_temp = compute_forces(
        [Body(body.name, body.mass, temp_positions[i], temp_velocities[i]) for i, body in enumerate(bodies)])
    k2_positions = temp_velocities * dt
    k2_velocities = forces_temp / masses[:, np.newaxis] * dt

    # k3 步骤
    temp_positions = positions + k2_positions / 2
    temp_velocities = velocities + k2_velocities / 2
    forces_temp = compute_forces(
        [Body(body.name, body.mass, temp_positions[i], temp_velocities[i]) for i, body in enumerate(bodies)])
    k3_positions = temp_velocities * dt
    k3_velocities = forces_temp / masses[:, np.newaxis] * dt

    # k4 步骤
    temp_positions = positions + k3_positions
    temp_velocities = velocities + k3_velocities
    forces_temp = compute_forces(
        [Body(body.name, body.mass, temp_positions[i], temp_velocities[i]) for i, body in enumerate(bodies)])
    k4_positions = temp_velocities * dt
    k4_velocities = forces_temp / masses[:, np.newaxis] * dt

    # 更新每个天体的位置和速度
    for i in range(num_bodies):
        bodies[i].position += (k1_positions[i] + 2 * k2_positions[i] + 2 * k3_positions[i] + k4_positions[i]) / 6
        bodies[i].velocity += (k1_velocities[i] + 2 * k2_velocities[i] + 2 * k3_velocities[i] + k4_velocities[i]) / 6


# 初始化天体
bodies = [
    Body('PlanetA', 2e29, [-1e11, 0, 0], [0, 1e4, 0]),
    Body('PlanetB', 2e29, [0, 0, 0], [0, -1e4, 0]),
    Body('PlanetC', 2e29, [1e11, 0, 0], [0, 0, 0]),
]

dt = 1000  # 每步时长 (秒)，越低精度越高，但模拟速度越慢
n = 200  # 每帧步数，越低动画越流畅，但渲染速度越慢
num_steps = 100000  # 总帧数
freme_time = 1  # 帧间隔 (毫秒)
scope = 1.5e11  # 视野范围
delay = 100  # 轨迹保留长度（帧数）

# 存储每一步的轨迹
positions = np.zeros((num_steps * n, len(bodies), 3))

# 设置绘图
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

# 初始化轨迹和散点图
lines = [ax.plot([], [], [], lw=1, label=f'{body.name}')[0] for body in bodies]

# 设置轴标签
ax.set_xlabel('X (m)')
ax.set_ylabel('Y (m)')
ax.set_zlabel('Z (m)')
ax.legend()

# 设置坐标范围
ax.set_xlim(-scope, scope)
ax.set_ylim(-scope, scope)
ax.set_zlim(-scope, scope)


# 更新动画帧的函数
def update(frame):
    # 每帧进行多次迭代
    for i in range(n):
        rk4_step(bodies, dt)

        # 存储位置
        for j in range(len(bodies)):
            positions[frame * n + i, j] = bodies[j].position

    # 更新绘图数据
    for i in range(len(bodies)):
        pos = positions[max((frame - delay) * n + i, 0):frame * n + i, i]
        lines[i].set_data(pos[:, 0], pos[:, 1])
        lines[i].set_3d_properties(pos[:, 2])

    return lines


# 创建动画
ani = FuncAnimation(fig, update, frames=num_steps, blit=True, interval=freme_time)
plt.show()

在此，给出一些其他模拟数据，注意手动调整视野范围

'''
bodies = [
    Body('Sun', 1.989e30, [0, 0, 0], [0, 0, 0]),  # 太阳
    Body('Mercury', 3.3011e23, [5.791e10, 0, 0], [0, 47400, 0]),  # 水星
    Body('Venus', 4.867e24, [1.082e11, 0, 0], [0, 35020, 0]),  # 金星
    Body('Earth', 5.972e24, [1.496e11, 0, 0], [0, 29780, 0]),  # 地球
    Body('Mars', 6.4171e23, [2.279e11, 0, 0], [0, 24077, 0]),  # 火星
    Body('Jupiter', 1.898e27, [7.785e11, 0, 0], [0, 13070, 0]),  # 木星
    Body('Saturn', 5.683e26, [1.433e12, 0, 0], [0, 9690, 0]),  # 土星
]
'''

'''
bodies = [
    Body('PlanetA', 1e30, [0, 2 * np.sqrt(3) * 1e11 / 3, 0], [1.82e4, 0, 0]),
    Body('PlanetB', 1e30, [-1e11, - np.sqrt(3) * 1e11 / 3, 0], [-1.82e4 / 2, np.sqrt(3) * 1.82e4 / 2, 0]),
    Body('PlanetC', 1e30, [1e11, - np.sqrt(3) * 1e11 / 3, 0], [-1.82e4 / 2, - np.sqrt(3) * 1.82e4 / 2, 0]),
]
'''

动画效果截图如下，效率问题不渲染球体：