[FIXED] How to save the state of an array every n time steps?

Issue

Reposting here since physics stack exchange said it belongs here. I am doing a simple test particle simulation of the orbits of 8 planets around the sun. I need to optimize the code a bit. My stored position and velocity arrays are very large (if I put a large t_end) and I would like to save the pos and vel every n=10000 (for example) time steps instead of storing all the positions and velocities. I tried different methods but couldn't really solve the problem. I want to do it such that the pos and vel arrays are overwritten and not stored for every time step, but only stored for every n=10000 time steps. So, maybe instead of storing them in arrays, I want to write the pos and vel to a file and update the file after every n steps so as to add the new values in a new row/line.

My current code is just me creating a new array where I store the values every n steps.

I know this might affect my current orbit plot, any inputs are appreciated but this is just a structure for a much more complex simulation later, where I won't need every value from my simulation.

Here is the 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 = 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[:,t].copy())
        saved_vel.append(vel[:,t].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()

Solution

It is easy to understand why this happens if you step through your program in a debugger. If you aren't familiar with a debugger, I highly recommend reading How to debug small programs and What is a debugger and how can it help me diagnose problems?

Add a breakpoint at the saved_pos.append(pos[:,t].copy()) line, and inspect what's being appended to the saved_pos array. Since counter lags t by 1, you have:

• counter is 10000, which activates the if condition
• t is 10001

However, you haven't yet done anything to pos[:, t], which is why it is still zeros and that gets appended to the saved_pos.

You want to either append pos[:, counter] to the list, or check if t % intervals == 1 (since t starts at 1)

Or, increment t after you've appended that array to the list

Answered By - pho