numpy 将多个列表以可读格式保存到文件中

eyh26e7m  于 9个月前  发布在  其他
关注(0)|答案(1)|浏览(87)

在我的测试粒子模拟太阳周围的8个行星中,我有pos_output和vel_output python列表。我想把它们保存到一个文件中,使其可读。这些是定期保存我的位置和速度。
我想在3列x, y, z的位置坐标和Vx, Vy, Vz的速度分量在其他3列。还有一个时间列,所以我知道在什么时候,是我的x,y,z, Vx, Vy, Vz保存。
列应为:t, x, y, z, Vx, Vy, Vz
我尝试了一些东西,但我只得到8个输出值的位置和速度。但有8个行星,它应该有(t_end/interval)输出。所以我应该有很多的位置和速度输出。请帮助我保存输出到一个可读的txt文件或类似的。
我的代码:

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 = 100000 #interval between time steps after which the position and velocities are 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
pos_output = []
vel_output = []

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:
        pos_output.append(pos[:,counter].copy())
        vel_output.append(vel[:,counter].copy())

# Convert lists to arrays
pos_output = np.concatenate(pos_output)
vel_output = np.concatenate(vel_output)

#save to file
np.savetxt('pos_output.txt', pos_output.reshape(-1, 3), fmt='%.6e', delimiter='\t', header='X\tY\tZ')
np.savetxt('vel_output.txt', vel_output.reshape(-1, 3), fmt='%.6e', delimiter='\t', header='VX\tVY\tVZ')

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

字符串

wb1gzix0

wb1gzix01#

我花了很多时间在你的代码上,所以我希望你能欣赏我的输入。也许我也误解了一切,请让我知道。

编辑:我改变了我原来的答案,以更好地适应你的问题,在评论中提供的信息之后。

提交您的问题

首先回答你的问题:我认为主要问题是你混淆了时间步和时间:t是当前时间步长的索引。但是,通过提示t < int(t_end),您假定每个时间步长的长度为1秒(显然它没有,因为你的自适应时间步进)。因此,你需要将while条件更改为实际时间(不是index)小于t_end。我将引入一个时间数组来跟踪时间,这意味着我们可以将其更改为:

while time[-1] <= t_end:

字符串

**注意:**请记住,这将使一个步骤后,时间通过t_end

由于自适应时间步进,我们事先不知道最终会有多少个时间步长。这意味着我们不能像您那样初始化posvel数组(我在以前的回答中),而是需要将每个新结果添加到数组中。这意味着我们只需从两个形状为(8, 1, 3)的数组开始(包含初始条件)。有趣的是,你使用一个自适应时间步长和一个辛积分器,比如蛙跳。因为你的时间步长并不总是相同的。通过初始化一个时间数组来跟踪当前时间步长当前所处的时间可能是一个很好的做法,也就是1d,当然从0开始

# Arrays to store pos and vel
pos = np.zeros((num_particles, 1, 3)) # shape: (#particles, timesteps taken + 1, #coordinates)
vel = np.zeros((num_particles, 1, 3)) # shape: (#particles, timesteps taken + 1, #coordinates)

# Initial conditions
time = np.zeros(1) # time array
pos[:, 0] = initial_pos
vel[:, 0] = initial_vel


然后在循环内部为每个时间步添加一个相同形状的数组。这也意味着我们不再需要变量t。相反,前一个时间步的数量可以通过轴1上的索引-1来访问。我将为每个时间步引入新的内部变量,用于当前位置,速度和时间:

_pos_t = pos[:, -1] + half_vel * min_dt # shape: (#particles, #coordinates)
_vel_t = half_vel + 0.5 * acc * min_dt # shape: (#particles, #coordinates)
_time_t = time[-1] + min_dt # shape: timesteps taken + 1


并在从(8,3)变为(8,1,3)后将它们连接为完整的数组:

pos = np.concatenate((pos, _pos_t[:, np.newaxis]), axis=1)
vel = np.concatenate((vel, _vel_t[:, np.newaxis]), axis=1)
time = np.append(time, _time_t)

每隔一段时间存储到新列表

现在我明白你背后的目标if counter % intervals == 0:.但你再次混淆了时间步和时间. counter只是第二个索引,这是过时的,在我看来.但intervals你如何描述它是一个时间间隔,所以在单位秒.然而,它也不正确替换counter与实际时间,因为我们不知道时间是否真的达到了间隔的倍数。我认为一个简单的解决方案是引入一个新的变量next_interval_time(如果你愿意,你也可以改变intervals),它记录了我们达到了intervals的多少倍。如果时间超过了一个间隔,将当前位置和速度附加到输出列表中,并记录时间:

next_interval_time = intervals
while ... :
    # Check if the elapsed time has surpassed the next interval
    if _time_t >= next_interval_time:
        pos_output.append(_pos_t.copy())
        vel_output.append(_vel_t.copy())
        t_output.append(_time_t.copy())
        next_interval_time += intervals


在循环之后,我会将它转换为一个numpy数组并转置以获得相同顺序的轴:

pos_output = np.array(pos_output).transpose((1, 0, 2)) # shape: (#objects, #stored timesteps, #coordinates)
vel_output = np.array(vel_output).transpose((1, 0, 2)) # shape: (#objects, #stored timesteps, #coordinates)

数组存储到文件

最好不要将numpy数组存储为txt(你会失去精度并且效率低下),但最好使用numpy method存储数组:
np.save('array.npy', array)
你可以装上
np.load('array.npy')
最后你得到了一个完整的数组。列表也是如此。在这里你可以使用Python pickle method。但是你把所有的列表都转换成了数组。

我的建议代码

我解决了你的问题,并实现了所有我提出的修改。我改变了t_end =165*yr,这是一个海王星年,所以应该给你一个粒子8的轨道。我也改变了intervals=1000000,因为我发现时间步长对于间隔来说太大了(这意味着我们存储了每个时间步长)。所以这里是完整的代码:

# -- setup as defined by you --
# Steps
t_end = 165*yr # Total time of integration
dt_constant = 0.1
intervals = 1000000 # seconds after which to store pos & vel
next_interval_time = intervals

# Arrays to store pos and vel
pos = np.zeros((num_particles, 1, 3)) # shape: (#particles, timesteps taken, #coordinates)
vel = np.zeros((num_particles, 1, 3)) # shape: (#particles, timesteps taken, #coordinates)

# Initial conditions
time = np.zeros(1) # time array
pos[:, 0] = initial_pos
vel[:, 0] = initial_vel
pos_output = []
vel_output = []
t_output = []

while time[-1] <= t_end: # NOTE: We end up doing one more timestep after t_end
    r = np.linalg.norm(pos[:, -1], axis=1)
    acc = -G * M_Sun / r[:, np.newaxis]**3 * pos[:, -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[:, -1], axis=1)**3 / (G * M_Sun))
    min_dt = np.min(current_dt)  # Use the minimum time step for all particles

    # leap frog integration
    # calculate velocity at half timestep
    half_vel = vel[:, -1] + 0.5 * acc * min_dt
    # current position only used in this timestep
    _pos_t = pos[:, -1] + half_vel * min_dt # shape: (#particles, #coordinates)

    # Recalculate acceleration with the new position
    r = np.linalg.norm(_pos_t, axis=1)
    # Acceleration at timestep
    acc = -G * M_Sun / r[:, np.newaxis]**3 * _pos_t #np.newaxis for broadcasting with pos[:, i-1]
    # current velocity only used in this timestep
    _vel_t = half_vel + 0.5 * acc * min_dt # shape: (#particles, #coordinates)
    
    # time at timestep t
    _time_t = time[-1] + min_dt

   # Check if the elapsed time has surpassed the next interval
    if _time_t >= next_interval_time:
        pos_output.append(_pos_t.copy())
        vel_output.append(_vel_t.copy())
        t_output.append(_time_t.copy())
        next_interval_time += intervals

    # add axis at position 1 to allow concatenation
    pos = np.concatenate((pos, _pos_t[:, np.newaxis]), axis=1)
    vel = np.concatenate((vel, _vel_t[:, np.newaxis]), axis=1)
    time = np.append(time, _time_t)

    # show current status by printing timestep number (-1 because initial conditions)
    print(f'timestep: {time.size -1} [progress: {_time_t/t_end*100:.3f}%]')

pos_output = np.array(pos_output).transpose((1, 0, 2)) # shape: (#objects, #stored timesteps, #coordinates)
np.save('pos_output.npy', pos_output)
vel_output = np.array(vel_output).transpose((1, 0, 2)) # shape: (#objects, #stored timesteps, #coordinates)
np.save('vel_output.npy', vel_output)
t_output = np.array(t_output)
np.save('t_output.npy', t_output)

# Orbit Plot
# --- your plotting routine ---

结果

最后得到预期的图:x1c 0d1x
输出数组的形状为(8, 5203, 3)。结束时间为165*yr = 5.203E9,这意味着我们期望5203个输出(end_time/interval = 5203.44)。我们可以稍后加载数组,如所述。

改进

也许有几个注意事项,来到我的脑海中的改进:

  • 考虑使用无量纲量进行模拟
  • 考虑使用极坐标来最小化数值误差

干杯

相关问题