如何在数值模拟中按固定时间间隔保存数组状态

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

# 常量定义
M_Sun = 1.989e30  # 太阳质量
G = 6.67430e-11  # 万有引力常量
yr = 365 * 24 * 60 * 60  # 一年秒数

# 粒子数量
num_particles = 8

# 初始条件
initial_pos = np.array([
    [57.9e9, 0, 0],  # Mercury
    [108.2e9, 0, 0],  # Venus
    [149.6e9, 0, 0],  # Earth
    [228e9, 0, 0],  # Mars
    [778.5e9, 0, 0],  # Jupiter
    [1432e9, 0, 0],  # Saturn
    [2867e9, 0, 0],  # Uranus
    [4515e9, 0, 0]   # Neptune
])

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
intervals = 10000  # 保存间隔

# 初始化位置和速度数组
pos = np.zeros((num_particles, int(t_end), 3))
vel = np.zeros((num_particles, int(t_end), 3))

# Leapfrog 积分
pos[:, 0] = initial_pos
vel[:, 0] = initial_vel
saved_pos = []
saved_vel = []

t = 1
while t < int(t_end):
    r = np.linalg.norm(pos[:, t - 1], axis=1)
    acc = -G * M_Sun / r[:, np.newaxis]**3 * pos[:, t - 1]

    # 计算当前时间步长
    current_dt = dt_constant * np.sqrt(np.linalg.norm(pos[:, t - 1], axis=1)**3 / (G * M_Sun))
    min_dt = np.min(current_dt)

    half_vel = vel[:, t - 1] + 0.5 * acc * min_dt
    pos[:, t] = pos[:, t - 1] + half_vel * min_dt

    # 重新计算加速度
    r = np.linalg.norm(pos[:, t], axis=1)
    acc = -G * M_Sun / r[:, np.newaxis]**3 * pos[:, t]
    vel[:, t] = half_vel + 0.5 * acc * min_dt

    # 保存位置和速度信息
    if t % intervals == 0:
        saved_pos.append(pos[:, t].copy())
        saved_vel.append(vel[:, t].copy())

    t += 1

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='Sun')

for particle in range(num_particles):
    x_particle = pos[particle, :, 0]
    y_particle = pos[particle, :, 1]
    z_particle = pos[particle, :, 2]
    ax.plot(x_particle, y_particle, z_particle, label=f'Particle {particle + 1} Orbit (km)')

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('Orbits of Planets around Sun (km)')
plt.show()