[FIXED] How to animate the particles at each timestep instead of their orbits in my simulation?

Issue

I want to animate my simulation in a better way:

a) Current code is simulating their orbits as a function of time. It is really some zigzag lines at the moment which is "ugly" to look at.

b) I want the num_particles = 100 to be animated at each instant of time. Like a ball moving in space. I don't want to trace it's trajectory/orbit.

c) So at each timestep, I should only have my num_particles as a cloud.

d) Save the animation as a mp4 or similar file.

My simulation code :

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation

# Constants
M_Sun = 1.989e30  # Solar Mass
M_bh = 4.3e6 *M_Sun #Mass of Sgr A*
G = 6.67430e-11  # m^3 kg^(-1) s^(-2)
yr = 365 * 24 * 60 * 60  # 1 year in seconds
R = 6.371e3 #Radius of Sphere

# Number of particles
num_particles = 100

#Uniformly distributed particles in Sphere of radius R
phi = np.random.uniform(0, 2*np.pi, size=num_particles)
costheta = np.random.uniform(-1, 1, size=num_particles)
u = np.random.uniform(0, 1, size=num_particles)

theta = np.arccos(costheta)
r = R * (u**(1/3))

x = r * np.sin(theta) * np.cos(phi)
y = r * np.sin(theta) * np.sin(phi)
z = r * np.cos(theta)

# Initial conditions for the particles (m and m/s)
initial_pos = np.random.uniform(0.8e13, 1.1e14, (num_particles, 3))
initial_vel = np.random.uniform(550e3, 730e3, (num_particles, 3))

# Steps
t_end = 1*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_bh / 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_bh))
    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_bh / 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)

print(pos_output.shape)
print(vel_output.shape)
print(t_output.shape)

#ANIMATION
from orbit_animation import animate_orbits
animate_orbits(pos_output)

My animation code:

# orbit_animation.py

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.animation import FuncAnimation

def animate_orbits(pos, intervals=1000000, interval=500):

    fig = plt.figure(figsize=(8, 8))
    ax = fig.add_subplot(111, projection='3d')

    # Scatter plot for Sgr A*
    sgr_a_plot = ax.scatter([0], [0], [0], color='black', marker='o', s=50, label='Sgr A*')

    # Initialize an empty line for the cloud particles
    cloud_plot, = ax.plot([], [], [], label='Cloud Particles')

    # Set plot labels and title
    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 cloud particles around Sgr A*')

    # Initialize axis limits
    x_min, x_max = np.min(pos[:, :, 0]), np.max(pos[:, :, 0])
    y_min, y_max = np.min(pos[:, :, 1]), np.max(pos[:, :, 1])
    z_min, z_max = np.min(pos[:, :, 2]), np.max(pos[:, :, 2])

    # Animation update function
    def update(frame):
        # Update Sgr A* position
        sgr_a_plot._offsets3d = ([0], [0], [0])

        # Update cloud particles
        cloud_plot.set_data(pos[:, frame, 0], pos[:, frame, 1])
        cloud_plot.set_3d_properties(pos[:, frame, 2])

        # Update axis limits dynamically
        x_min, x_max = np.min(pos[:, :, 0]), np.max(pos[:, :, 0])
        y_min, y_max = np.min(pos[:, :, 1]), np.max(pos[:, :, 1])
        z_min, z_max = np.min(pos[:, :, 2]), np.max(pos[:, :, 2])

        ax.set_xlim(x_min, x_max)
        ax.set_ylim(y_min, y_max)
        ax.set_zlim(z_min, z_max)

        return sgr_a_plot, cloud_plot

    # Create the animation
    animation = FuncAnimation(fig, update, frames=pos.shape[1], interval=interval, blit=True)
    plt.show()

I hope I could explain what I am trying to achieve here! Thanks in advance for your help!

Solution

You are very close.

To remove the line connecting your particles, set the linestyle to "none". To plot circle markers instead, set the marker to o.

cloud_plot, = ax.plot([], [], [], linestyle="none", marker='o', label='Cloud Particles')

The interval=500 (i.e. 0.5 fps) makes the animation very choppy, I would reduce that number to 50 (i.e. 20 fps).

To save the animation as an mp4, just give the filepath a .mp4 extension. It makes sense to match the fps of the save file to the animation.

animation.save("test.mp4", fps=20) 

Answered By - Paul Brodersen