在数值模拟中，经常需要保存数组（例如位置和速度）的状态，以便后续分析或可视化。然而，如果模拟时间较长，或者时间步长较小，存储所有时间步的状态可能会占用大量内存。本文将介绍如何在 Python 中，使用 NumPy 数组，以固定的时间步长间隔保存数组状态，从而有效减少内存占用。
核心思路
核心思路是利用取模运算 (%)，仅在特定的时间步保存数据。同时，需要注意循环中索引的更新顺序，以确保在保存数据时，数组已经更新到目标时间步的状态。
代码示例与解析
以下代码基于问题中的行星轨道模拟示例，展示了如何修改代码以实现按固定时间间隔保存位置和速度信息。
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()代码解释:
- if t % intervals == 0:: 这是核心的条件判断语句。t % intervals 计算当前时间步 t 除以保存间隔 intervals 的余数。当余数为 0 时，表示当前时间步是需要保存的时间步。
- saved_pos.append(pos[:, t].copy()): 将当前时间步的位置信息 pos[:, t] 添加到 saved_pos 列表中。.copy() 方法创建数组的副本，防止后续修改 pos 数组影响已保存的数据。速度信息的保存方式类似。
- 索引更新顺序 确保在保存 pos[:, t] 之前，已经完成了对 pos[:, t] 的更新。
注意事项
- 索引问题: 循环中的索引 t 从 1 开始，因此需要使用 t % intervals == 0 来判断是否需要保存数据。 如果 t 从 0 开始，则应该使用 (t + 1) % intervals == 0。
- 内存占用: 即使按固定间隔保存，如果模拟时间非常长，saved_pos 和 saved_vel 仍然可能占用大量内存。可以考虑将数据写入文件，而不是全部保存在内存中。
- 数据分析: 保存间隔的选择取决于具体的应用场景。需要根据模拟的精度要求和数据分析的需求，合理选择保存间隔。
- 时间步长: 确保时间步长足够小，以保证模拟的准确性。
总结
通过使用取模运算和注意索引更新顺序，可以在数值模拟中按固定时间间隔保存数组状态，有效减少内存占用。 这种方法适用于各种需要保存中间状态的数值模拟，例如物理模拟、机器学习模型训练等。 记住，合理选择保存间隔，并根据具体需求调整代码，才能获得最佳的性能和效果。
