Search code examples
pythonindexingsimulationphysics

Problem in calculating length to breadth ratio for test particle simulation in Python

I keep getting this error and I am unable to solve for this. I want length and width (averaged for all particles) for each time step. So the length and width array should have (19,) shaped instead of the current (1000,) shape which means that it is calculating for each particle.

The error:

‘plt.plot(t_output, ratio)

ValueError: x and y must have same first dimension, but have shapes (19,) and (1000,)’

I have tried to the best of my "coding" knowledge and this should be pretty simple for experienced and expert coders, but sadly, I am not one. In the last few lines of the script, I have commented what I was supposed to do and what I tried starting from "# Extract velocities for all particles and all timesteps"

Please help me fix this errors and achieve what I am trying to do!

import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial.transform import Rotation
from mpl_toolkits.mplot3d import Axes3D


# Constants
M_Sun = 1.989e33  # Solar Mass in g
M_bh = 4.3e6 *M_Sun # Mass of BH
G = 6.67430e-8  # cm^3 g^(-1) s^(-2)
yr = 365 * 24 * 60 * 60  # 1 year in seconds
pc = 3.086e18 # Parsec in cm
R = .225 * 0.04 * pc / 2 # Radius of cloud

# Number of particles
num_particles = 1000

# Uniformly distributed particles in Sphere of radius R
phi = np.random.uniform(0, 2*np.pi, size=num_particles)
costheta = np.random.uniform(0, 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

RA = .425 # Right Ascension
Decl = .523 # Declination

vx = 550e5 # x-component of velocity
vy = 730e5 # y-component of velocity
vz = 0 # z-component of velocity

# Transform RA and Decl to x,y,z components (1 arcsec = 0.04 pc in the galactic center)
x_0 = RA * 0.04 * pc
y_0 = Decl * 0.04 * pc
z_0 = 0

# Empty lists for position and velocity
initial_pos = []
initial_vel = []

for i in range(num_particles):
    initial_position = (x_0 + x[i], y_0 + y[i], z_0 + z[i]) # x_0 + x (x is already multiplied by R), same for y and z axes
    initial_velocity = (vx, vy, vz) # Fixed velocity for all particles
    initial_pos.append(initial_position)
    initial_vel.append(initial_velocity)

initial_pos = np.array(initial_pos) #Shape (num_particles, coordinates)
initial_vel = np.array(initial_vel) #Shape (num_particles, coordinates)

# Time
t_end = 100*yr # Total time of integration
dt_constant = 0.1 # Magic number for time stepping scheme
intervals = 1000000 # seconds after which to store pos & vel
next_interval_time = intervals

# Initial conditions
time = np.zeros(1)
_pos_t = initial_pos.copy()
_vel_t = initial_vel.copy()

# Lists to store outputs
pos_output = []
vel_output = []
t_output = []

# Number of Stored outputs
output = 20

while time[-1] <= t_end: #We end up doing one more timestep after t_end
    r = np.linalg.norm(_pos_t, axis=1)
    acc = -G * M_bh / r[:, np.newaxis]**3 * _pos_t #np.newaxis for broadcasting with _pos_t

    # Calculate the time step for the current particle
    current_dt = dt_constant * np.sqrt(np.linalg.norm(_pos_t, 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_t + 0.5 * acc * min_dt
    # current position only used in this timestep
    _pos_t = _pos_t + 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_t
    # 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: # >= because time steps aren't constant and it may not be an exact multiple
        # Store data at intervals
        pos_output.append(_pos_t.copy())
        vel_output.append(_vel_t.copy())
        t_output.append(_time_t.copy())
        next_interval_time += intervals

    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)
vel_output = np.array(vel_output).transpose((1, 0, 2))  # shape: (#objects, #stored timesteps, #coordinates)
t_output = np.array(t_output)

# Plotting the ratio of length to breadth as a function of time for the entire cloud 
# find the velocity vector at each timestep 
# move to COM frame
# find the angle of x or y position to the vel vector #dot product
# rotate one axis so it aligns with vel vector # use rotation matrix
# then x and y are perpendicular and one is aligned with vel vector
# then do the ratio calculation

# Extract velocities for all particles and all timesteps
velocities = vel_output[:, :, :]

# Move to the Center of Mass (COM) frame
com_velocity = np.mean(velocities, axis=0)  # Calculate the COM velocity for all timesteps
velocities_rel_com = velocities - com_velocity  # Relative velocities to the COM frame

# Calculate the angle between x axis and the velocity vector
dot_products = np.sum(pos_output[:, :, 0] * velocities_rel_com[:, :, 0], axis=-1) # dot product for each particle and each timestep (x and velocity vector)

# Find the angles between the position and velocity vectors
magnitude_pos = np.linalg.norm(pos_output[:, :, 0]) # magnitude of the position vector
magnitudes_vel = np.linalg.norm(velocities_rel_com[:, :, 0]) # magnitude of the velocity vector
angles = np.arccos(dot_products / (magnitude_pos * magnitudes_vel)) # angle between the position and velocity vectors

# Rotate the axes so that one axis aligns with the velocity vector
rotation = Rotation.from_euler('y', angles)
rotation_matrix = rotation.as_matrix()

positions_ = pos_output.copy()
# Rotate the position vectors
pos_rotated = np.einsum('ijk,ikl->ijl', positions_, rotation_matrix)

# Calculate the length and width for each timestep and all particles
length = np.max(pos_rotated[:, :, 0], axis=-1) - np.min(pos_rotated[:, :, 0], axis=-1)
width = np.max(pos_rotated[:, :, 1], axis=-1) - np.min(pos_rotated[:, :, 1], axis=-1)

# Calculate the ratio for each timestep
ratio = length / width

# Plot the ratio as a function of time
plt.plot(t_output, ratio)
plt.xlabel('Time (s)')
plt.ylabel('Ratio')
plt.title('Ratio of Length to Width of Cloud')
plt.show()
Solution

  • The error occurs in plt.plot(t_output, ratio). The size of the vectors needs to be the same for plotting, but the error is saying that t_output is length 19, whereas ratio is length 1000. I made the following modification

    length = np.max(pos_rotated[:, :, 0], axis=0) - np.min(pos_rotated[:, :, 0], axis=0)
width = np.max(pos_rotated[:, :, 1], axis=0) - np.min(pos_rotated[:, :, 1], axis=0)

    enter image description here

    Before, it was axis=-1. I changed it to axis=0. Suppose you take pos_rotated[:, :, 0] - this array has shape (1000, 19) - not (1000, 19, 3) which is what pos_rotated is. Given that it's shaped (1000, 19), if we want to find the maximum value across all 1000 particles for each of the 19 time steps, we need to specify axis=0 i.e. perform np.max over the particles dimension. You'll end up with an array shaped (19) that, for each timepoint, is the maximum over 1000 particles.