本文旨在解决在模拟过程中，如何高效地保存数组状态，尤其是在需要控制内存使用，避免存储所有时间步数据的情况下。通过修改代码结构，实现在每隔 N 个时间步长后，将位置和速度数据写入文件或覆盖数组，从而优化存储空间，并提供相应的代码示例和调试建议。
在进行数值模拟时，经常需要保存模拟过程中的数据，例如位置、速度等。如果模拟时间很长，或者时间步长很小，那么存储所有时间步的数据将会占用大量的内存空间。为了解决这个问题，我们可以只保存每隔 N 个时间步长的数据，从而有效地减少内存的使用。
以下提供两种实现方式：
1. 使用数组覆盖
这种方法适用于不需要保留所有时间步数据的场景。通过在循环中覆盖数组的方式，只保留最新的 N 个时间步的数据。
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
# Constants
M_Sun = 1.989e30 #Solar Mass
G = 6.67430e-11 # m^3 kg^(-1) s^(-2)
yr = 365 * 24 * 60 * 60 #1 year in seconds
# Number of particles
num_particles = 8
# Initial conditions for the particles (m and m/s)
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]
])
# Steps
t_end = 0.004 * yr #Total time of integration
dt_constant = 0.1
intervals = 10000 #Number of outputs of pos and vel to be saved
# Arrays to store pos and vel
pos = np.zeros((num_particles, int(t_end), 3))
vel = np.zeros((num_particles, int(t_end), 3))
# Leapfrog Integration (2nd Order)
pos[:, 0] = initial_pos
vel[:, 0] = initial_vel
saved_pos = []
saved_vel = []
t = 1
counter = 0
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] #np.newaxis for broadcasting with pos[:, i-1]
# Calculate the time step for the current particle
current_dt = dt_constant * np.sqrt(np.linalg.norm(pos[:, t - 1], axis=1)**3 / (G * M_Sun))
min_dt = np.min(current_dt) # Use the minimum time step for all particles
half_vel = vel[:, t - 1] + 0.5 * acc * min_dt
pos[:, t] = pos[:, t - 1] + half_vel * min_dt
# Recalculate acceleration with the new position
r = np.linalg.norm(pos[:, t], axis=1)
acc = -G * M_Sun / r[:, np.newaxis]**3 * pos[:, t] #np.newaxis for broadcasting with pos[:, i-1]
vel[:, t] = half_vel + 0.5 * acc * min_dt
# Save the pos and vel here
if counter % intervals == 0:
saved_pos.append(pos[:,t].copy())
saved_vel.append(vel[:,t].copy())
t += 1
counter += 1
saved_pos = np.array(saved_pos)
saved_vel = np.array(saved_vel)
# Orbit Plot
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()需要注意的是，原代码存在一个问题：
在保存 pos 和 vel 时，counter 的值比 t 小 1。这意味着当 counter % intervals == 0 时，pos[:, t] 还没有被更新，因此保存的是未更新的数据（初始值）。
解决方案：
- 将 pos[:,t].copy() 改为 pos[:, t-1].copy() 和 vel[:, t-1].copy()：保存上一个时间步的数据，即已经计算过的数据。
- 将 t += 1 和 counter += 1 的位置互换：保证在保存数据时，t 和 counter 的值是一致的。
- 将判断条件修改为 if t % intervals == 0：由于 t 从 1 开始，因此需要使用 t % intervals == 0 来判断是否需要保存数据。
以下是修改后的代码：
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
# Constants
M_Sun = 1.989e30 #Solar Mass
G = 6.67430e-11 # m^3 kg^(-1) s^(-2)
yr = 365 * 24 * 60 * 60 #1 year in seconds
# Number of particles
num_particles = 8
# Initial conditions for the particles (m and m/s)
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]
])
# Steps
t_end = 0.004 * yr #Total time of integration
dt_constant = 0.1
intervals = 10000 #Number of outputs of pos and vel to be saved
# Arrays to store pos and vel
pos = np.zeros((num_particles, int(t_end), 3))
vel = np.zeros((num_particles, int(t_end), 3))
# Leapfrog Integration (2nd Order)
pos[:, 0] = initial_pos
vel[:, 0] = initial_vel
saved_pos = []
saved_vel = []
t = 1
counter = 0
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] #np.newaxis for broadcasting with pos[:, i-1]
# Calculate the time step for the current particle
current_dt = dt_constant * np.sqrt(np.linalg.norm(pos[:, t - 1], axis=1)**3 / (G * M_Sun))
min_dt = np.min(current_dt) # Use the minimum time step for all particles
half_vel = vel[:, t - 1] + 0.5 * acc * min_dt
pos[:, t] = pos[:, t - 1] + half_vel * min_dt
# Recalculate acceleration with the new position
r = np.linalg.norm(pos[:, t], axis=1)
acc = -G * M_Sun / r[:, np.newaxis]**3 * pos[:, t] #np.newaxis for broadcasting with pos[:, i-1]
vel[:, t] = half_vel + 0.5 * acc * min_dt
# Save the pos and vel here
if t % intervals == 0:
saved_pos.append(pos[:,t].copy())
saved_vel.append(vel[:,t].copy())
t += 1
counter += 1
saved_pos = np.array(saved_pos)
saved_vel = np.array(saved_vel)
# Orbit Plot
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()2. 写入文件
这种方法适用于需要保留所有时间步数据，但又不想占用大量内存空间的场景。通过将数据写入文件，可以避免将所有数据都加载到内存中。
import numpy as np
# 假设 pos 和 vel 是你的位置和速度数组
# N 是你想要保存数据的间隔
N = 10000 # 例如，每 10000 步保存一次
with open('positions.txt', 'w') as f_pos, open('velocities.txt', 'w') as f_vel:
for t in range(int(t_end)): # 替换为你的时间步长循环
# 你的模拟代码...
if t % N == 0:
# 将当前位置和速度写入文件
np.savetxt(f_pos, pos[:,t].reshape(1,-1)) # 将数据reshape成二维数组，方便写入
np.savetxt(f_vel, vel[:,t].reshape(1,-1)) # 将数据reshape成二维数组，方便写入注意事项：
- 使用调试器可以帮助你更好地理解代码的执行过程，并找到潜在的问题。
- 在选择保存数据的方式时，需要根据具体的应用场景进行权衡。如果需要保留所有时间步的数据，并且内存空间足够，那么可以使用数组存储。如果内存空间有限，或者只需要部分时间步的数据，那么可以使用文件存储。
- 在写入文件时，需要注意文件的格式和编码方式，以便后续读取和处理。
总结：
本文介绍了两种在模拟过程中保存数组状态的方法：使用数组覆盖和写入文件。通过选择合适的方法，可以有效地减少内存的使用，并提高模拟的效率。同时，也强调了调试的重要性，并提供了一些注意事项。希望这些内容能够帮助你更好地进行数值模拟。
