如何在数值模拟中按时间步长保存数组状态

第一段引用上面的摘要：本文介绍了在数值模拟中，如何以指定的时间步长间隔保存数组状态，从而有效降低内存占用。通过修改保存状态的逻辑，确保在正确的时间点记录数组数据，并提供代码示例，帮助读者理解和实现这一优化策略，适用于需要长时间模拟但又不想存储所有中间状态的场景。

在进行长时间的数值模拟时，例如行星轨道模拟、分子动力学模拟等，存储每个时间步长的状态数据（例如位置、速度）可能会导致内存占用过大，甚至超出可用内存。为了解决这个问题，可以只保存每隔一定时间步长的状态数据，从而在一定程度上降低内存需求，同时仍然能够获得足够的信息来分析模拟结果。

优化数组状态保存策略

问题的核心在于理解保存状态的时机。在原始代码中，counter 变量用于跟踪时间步数，并在 counter 达到 intervals 的倍数时保存状态。然而，counter 的更新发生在状态更新之前，导致保存的是下一个时间步的状态，而此时该状态尚未被计算出来，因此是零。

要解决这个问题，有以下几种方法：

1. 保存 pos[:, counter] 而不是 pos[:, t]: 因为 counter 实际反映的是上一个时间步的状态，所以应该保存 pos[:, counter]。

2. 调整 if 条件: 将 if counter % intervals == 0 改为 if t % intervals == 1，因为 t 从 1 开始计数。

3. 调整 t 的更新顺序: 在保存状态之后再更新 t 的值。

以下代码展示了第一种解决方案的实现：

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

    t += 1
    counter += 1

    # Save the pos and vel here
    if counter % intervals == 0:
        saved_pos.append(pos[:,counter].copy()) # 修改这里
        saved_vel.append(vel[:,counter].copy()) # 修改这里

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()

将 saved_pos.append(pos[:,t].copy()) 和 saved_vel.append(vel[:,t].copy()) 分别修改为 saved_pos.append(pos[:,counter].copy()) 和 saved_vel.append(vel[:,counter].copy())，确保保存的是正确时间步的状态。

使用文件存储代替数组

除了修改保存状态的逻辑，还可以考虑使用文件存储代替数组来保存状态数据。这种方法可以进一步降低内存占用，特别是当模拟时间很长时。以下代码展示了如何将位置和速度数据写入文件：

import numpy as np

# 假设 pos 和 vel 是当前时间步的位置和速度数组
# intervals 是保存间隔

def save_state_to_file(pos, vel, filename_pos, filename_vel):
    """将位置和速度数据追加到文件中。"""
    with open(filename_pos, 'a') as f_pos, open(filename_vel, 'a') as f_vel:
        np.savetxt(f_pos, pos.reshape(1, -1))  # 将 pos 展平为一行
        np.savetxt(f_vel, vel.reshape(1, -1))  # 将 vel 展平为一行

# 模拟循环内
if counter % intervals == 0:
    save_state_to_file(pos[:, t], vel[:, t], 'pos_data.txt', 'vel_data.txt')

在这个例子中，save_state_to_file 函数将位置和速度数据追加到指定的文件中。np.savetxt 函数用于将数组数据写入文件，pos.reshape(1, -1) 和 vel.reshape(1, -1) 将多维数组展平为一维数组，以便于写入文件。

注意：

• 使用文件存储时，需要注意文件的打开和关闭，以避免资源泄漏。可以使用 with open(...) as f: 语句来自动管理文件的打开和关闭。
• 使用文件存储时，读取数据时需要重新将数据reshape为正确的维度。

总结

通过优化数组状态保存策略，可以有效降低数值模拟的内存占用。选择哪种方法取决于具体的应用场景和需求。如果只需要保存少量状态数据，修改保存逻辑可能更简单。如果需要保存大量状态数据，使用文件存储可能更合适。在实际应用中，可以结合使用这些方法，以达到最佳的性能和内存利用率。