[FIXED] Saving multiple lists to a file in readable format

Issue

In my test particle simulation of 8 planets around the sun, I have pos_output and vel_output python lists. I want to save these to a file and make it readable. These are saving my positions and velocities at regular intervals.

I want the x, y, z position coordinates in 3 columns and Vx, Vy, Vz velocity components in other 3 columns. Also a time column so I know at what time, is my x,y,z, Vx, Vy, Vz saved.

The columns should be: t, x, y, z, Vx, Vy, Vz

I tried something but I am only getting 8 output values for position and velocity. But there are 8 planets and it should have (t_end/interval) outputs. So I should have a lot of position and velocity outputs. Please help me save the outputs to a readable txt file or similar.

My code:

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

Solution

I spent a lot of time with your code now, so I hope you appreciate my input. Maybe I also misunderstood everything, please let me know.

EDIT: I changed my original answer to fit better to your question, after the info provided in the comments.

Answering your Question

So first to answer your question: I think the main problem is you are mixing up timesteps and time: t is your index of the current timestep. However by prompting t < int(t_end) you assume every timestep has a length of one second (which it clearly doesn't have, because of your adaptive time stepping). Therfore you need to change your while condition to the actual time (not index) being smaller than t_end. I will introduce an time array to keep track to the time, which means we can change it to:

while time[-1] <= t_end:

Note: Keep in mind though that this will make one step after time passed t_end.

We don't know a priori how many timesteps we will be ending up with, because of the adaptive time stepping. This means we can not initialize the pos and vel arrays as you did (and I in my previous answer), but rather need to add every new result to the arrray. This means we just start with the two arrays being shape (8, 1, 3) (containing initial conditions). It is interesting that you use a adaptive time step with an symplectic integrator like leap frog. Because of that your timesteps are not always the same. It is probably good practise to keep track of the time that the current time step is currently at by initializing a time array, which will be 1d and of course starts with 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

Then inside the loop add an array of the same shape for every time step. Also means we don't need the variable t anymore. Instead quantities of the previous timestep can be accesed by index -1 on axis 1. I would introduce new internal variables for current position, velocity and time for every time step:

_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

and concatenate them to the complete array after changing from (8, 3) to (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)

Storing to a new list after every interval

Now I understand your goal behind if counter % intervals == 0:. But again you are mixing up timesteps and time. counter is just a second index, which is obsolute in my opinion. But intervals how you described it is a time interval so in unit seconds. However its also not correct replacing counter with the actual time, because we don't know if the time actually exactly reaches multiples of intervals. I think an easy solution for that is introducing a new variable next_interval_time (you can also change intervals if you prefer) that keeps track which multiple of intervals we have reached. If the time passed one interval, append the current position and velocity to the output lists and also keep track of the times:

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

After the loop, I would convert it to a numpy array and transpose to get the same order of axis:

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)

Storing arrays to file

It is preferable to store the numpy arrays not as txt (you loose precision and is inefficient), but better use the numpy method of storing arrays:

np.save('array.npy', array)

Which you can load with

np.load('array.npy')

and you end up with the complete array. The same can be said for lists. There you can use the Python pickle method. However you convert all your lists to arrays.

My proposed code

I solved your issue and implemented all my proposed changes. I changed t_end =165*yr, which is one Neptun year, so should show you one orbit of particle 8. Also I changed to intervals=1000000, because I found out the time step were too large for intervals (which meant we stored every timestep). So here goes the complete code:

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

Result

I end up with the expected plot:

The output arrays have a shape (8, 5203, 3). The end time is 165*yr = 5.203E9, which means we expect 5203 outputs (end_time/interval = 5203.44). We can load the arrays later like described.

Improvements

Maybe a few notes that came to my mind for improvements:

• consider using dimensionless quantities for simulation
• consider using polar coordinates to minimize numerical error

Cheers!

Answered By - mcjeb