如何在每n个时间步长保存数组状态（行星轨道模拟优化）

问题背景 #

我正在开发8颗行星绕日的简易测试粒子轨道模拟，当前代码在设置较大t_end时，存储位置（pos）和速度（vel）的数组会占用过多内存。需求是每10000个时间步长保存一次pos和vel，而非存储所有时间步数据，可选择两种实现方式：

• 仅保留当前计算所需的最小数据，每n步存储目标数据
• 直接将数据写入文件，每n步追加新行

当前代码通过新建数组存储每n步数据，但仍保留了全量时间步的pos/vel数组，内存占用问题未解决。以下是优化后的两种实现方案：

方案1：内存极致优化（仅保留必要计算数据） #

Leapfrog积分仅依赖上一步的位置和速度即可计算下一步，因此不需要创建存储全量时间步的巨大数组，只需要两个数组分别存储当前步和上一步的pos/vel，每10000步将数据存入结果列表。

修改后代码 #

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 常量定义
M_Sun = 1.989e30  # 太阳质量
G = 6.67430e-11   # 引力常数 m^3 kg^(-1) s^(-2)
yr = 365 * 24 * 60 * 60  # 一年的秒数

# 粒子数量
num_particles = 8

# 初始位置（单位：米）和速度（单位：米/秒）
initial_pos = np.array([
    [57.9e9, 0, 0],  # 水星
    [108.2e9, 0, 0], # 金星
    [149.6e9, 0, 0], # 地球
    [228e9, 0, 0],   # 火星
    [778.5e9, 0, 0], # 木星
    [1432e9, 0, 0],  # 土星
    [2867e9, 0, 0],  # 天王星
    [4515e9, 0, 0]   # 海王星
])

initial_vel = np.array([
    [0, 47400, 0],
    [0, 35000, 0],
    [0, 29800, 0],
    [0, 24100, 0],
    [0, 13100, 0],
    [0, 9700, 0],
    [0, 6800, 0],
    [0, 5400, 0]
])

# 模拟参数
t_end = 0.004 * yr  # 总模拟时间
dt_constant = 0.1
save_interval = 10000  # 每10000步保存一次数据

# 仅保留上一步和当前步的pos/vel，无需全量数组
pos_prev = initial_pos.copy()
vel_prev = initial_vel.copy()

# 存储需要保存的数据
saved_pos = [pos_prev.copy()]  # 先保存初始状态
saved_vel = [vel_prev.copy()]
step_counter = 0
current_time = 0.0

while current_time < t_end:
    # 计算加速度
    r = np.linalg.norm(pos_prev, axis=1)
    acc = -G * M_Sun / r[:, np.newaxis]**3 * pos_prev

    # 计算自适应时间步长
    current_dt = dt_constant * np.sqrt(np.linalg.norm(pos_prev, axis=1)**3 / (G * M_Sun))
    min_dt = np.min(current_dt)

    # Leapfrog积分步骤
    half_vel = vel_prev + 0.5 * acc * min_dt
    pos_current = pos_prev + half_vel * min_dt

    # 用新位置重新计算加速度
    r_current = np.linalg.norm(pos_current, axis=1)
    acc_current = -G * M_Sun / r_current[:, np.newaxis]**3 * pos_current
    vel_current = half_vel + 0.5 * acc_current * min_dt

    # 更新时间和计数器
    current_time += min_dt
    step_counter += 1

    # 每save_interval步保存数据
    if step_counter % save_interval == 0:
        saved_pos.append(pos_current.copy())
        saved_vel.append(vel_current.copy())

    # 更新上一步数据为当前步
    pos_prev, vel_prev = pos_current, vel_current

# 转换为numpy数组方便后续处理
saved_pos = np.array(saved_pos)
saved_vel = np.array(saved_vel)

# 轨道绘图（仅用保存的数据绘图）
fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(111, projection='3d')

ax.scatter(0, 0, 0, color='yellow', marker='o', s=50, label='太阳')

for particle in range(num_particles):
    # 提取该行星所有保存的位置点，转换为千米
    x = saved_pos[:, particle, 0] / 1e3
    y = saved_pos[:, particle, 1] / 1e3
    z = saved_pos[:, particle, 2] / 1e3
    ax.plot(x, y, z, label=f'行星 {particle + 1} 轨道')

ax.set_xlabel('X (km)')
ax.set_ylabel('Y (km)')
ax.set_zlabel('Z (km)')
ax.legend(loc='upper right', bbox_to_anchor=(1.1, 1.1))
ax.set_title('行星绕日轨道模拟')
plt.show()

方案2：直接写入文件（无内存压力） #

如果后续不需要在内存中保留数据，直接将每10000步的pos/vel写入文件，完全不占用额外内存，适合超长时间模拟。

修改后代码 #

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 常量定义
M_Sun = 1.989e30  # 太阳质量
G = 6.67430e-11   # 引力常数 m^3 kg^(-1) s^(-2)
yr = 365 * 24 * 60 * 60  # 一年的秒数

# 粒子数量
num_particles = 8

# 初始位置（单位：米）和速度（单位：米/秒）
initial_pos = np.array([
    [57.9e9, 0, 0],  # 水星
    [108.2e9, 0, 0], # 金星
    [149.6e9, 0, 0], # 地球
    [228e9, 0, 0],   # 火星
    [778.5e9, 0, 0], # 木星
    [1432e9, 0, 0],  # 土星
    [2867e9, 0, 0],  # 天王星
    [4515e9, 0, 0]   # 海王星
])

initial_vel = np.array([
    [0, 47400, 0],
    [0, 35000, 0],
    [0, 29800, 0],
    [0, 24100, 0],
    [0, 13100, 0],
    [0, 9700, 0],
    [0, 6800, 0],
    [0, 5400, 0]
])

# 模拟参数
t_end = 0.004 * yr  # 总模拟时间
dt_constant = 0.1
save_interval = 10000  # 每10000步保存一次数据

# 初始化文件，写入表头
with open('orbit_data.txt', 'w') as f:
    # 表头格式：时间步, 行星ID, x_pos, y_pos, z_pos, x_vel, y_vel, z_vel
    f.write('step,planet_id,x_pos,y_pos,z_pos,x_vel,y_vel,z_vel\n')

# 仅保留上一步和当前步的pos/vel
pos_prev = initial_pos.copy()
vel_prev = initial_vel.copy()

step_counter = 0
current_time = 0.0

# 先写入初始状态
with open('orbit_data.txt', 'a') as f:
    for planet_id in range(num_particles):
        x, y, z = pos_prev[planet_id]
        vx, vy, vz = vel_prev[planet_id]
        f.write(f'{step_counter},{planet_id+1},{x},{y},{z},{vx},{vy},{vz}\n')

while current_time < t_end:
    # 计算加速度
    r = np.linalg.norm(pos_prev, axis=1)
    acc = -G * M_Sun / r[:, np.newaxis]**3 * pos_prev

    # 计算自适应时间步长
    current_dt = dt_constant * np.sqrt(np.linalg.norm(pos_prev, axis=1)**3 / (G * M_Sun))
    min_dt = np.min(current_dt)

    # Leapfrog积分步骤
    half_vel = vel_prev + 0.5 * acc * min_dt
    pos_current = pos_prev + half_vel * min_dt

    # 用新位置重新计算加速度
    r_current = np.linalg.norm(pos_current, axis=1)
    acc_current = -G * M_Sun / r_current[:, np.newaxis]**3 * pos_current
    vel_current = half_vel + 0.5 * acc_current * min_dt

    # 更新时间和计数器
    current_time += min_dt
    step_counter += 1

    # 每save_interval步写入文件
    if step_counter % save_interval == 0:
        with open('orbit_data.txt', 'a') as f:
            for planet_id in range(num_particles):
                x, y, z = pos_current[planet_id]
                vx, vy, vz = vel_current[planet_id]
                f.write(f'{step_counter},{planet_id+1},{x},{y},{z},{vx},{vy},{vz}\n')

    # 更新上一步数据为当前步
    pos_prev, vel_prev = pos_current, vel_current

# 可选：从文件读取数据绘图
saved_data = np.loadtxt('orbit_data.txt', delimiter=',', skiprows=1)
fig = plt.figure(figsize=(8, 8))
ax = fig.add_subplot(111, projection='3d')
ax.scatter(0, 0, 0, color='yellow', marker='o', s=50, label='太阳')

for planet_id in range(1, num_particles+1):
    planet_data = saved_data[saved_data[:,1]==planet_id]
    x = planet_data[:,2]/1e3
    y = planet_data[:,3]/1e3
    z = planet_data[:,4]/1e3
    ax.plot(x, y, z, label=f'行星 {planet_id} 轨道')

ax.set_xlabel('X (km)')
ax.set_ylabel('Y (km)')
ax.set_zlabel('Z (km)')
ax.legend(loc='upper right', bbox_to_anchor=(1.1, 1.1))
ax.set_title('行星绕日轨道模拟')
plt.show()

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