如何在每隔 N 个时间步长后保存数组状态？-Python教程-PHP中文网

如何在每隔 N 个时间步长后保存数组状态？

霞舞

发布： 2025-08-13 16:18:12

原创

937人浏览过

如何在每隔 n 个时间步长后保存数组状态？

本文旨在解决在模拟过程中，如何高效地保存数组状态，尤其是在需要控制内存使用，避免存储所有时间步数据的情况下。通过修改代码结构，实现在每隔 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] 还没有被更新，因此保存的是未更新的数据（初始值）。

解决方案：

Hugging Face

Hugging Face AI开源社区

135

查看详情

将 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成二维数组，方便写入

登录后复制

注意事项：

使用调试器可以帮助你更好地理解代码的执行过程，并找到潜在的问题。
在选择保存数据的方式时，需要根据具体的应用场景进行权衡。如果需要保留所有时间步的数据，并且内存空间足够，那么可以使用数组存储。如果内存空间有限，或者只需要部分时间步的数据，那么可以使用文件存储。
在写入文件时，需要注意文件的格式和编码方式，以便后续读取和处理。

总结：

本文介绍了两种在模拟过程中保存数组状态的方法：使用数组覆盖和写入文件。通过选择合适的方法，可以有效地减少内存的使用，并提高模拟的效率。同时，也强调了调试的重要性，并提供了一些注意事项。希望这些内容能够帮助你更好地进行数值模拟。

以上就是如何在每隔 N 个时间步长后保存数组状态？的详细内容，更多请关注php中文网其它相关文章！