Skip to content
Snippets Groups Projects

BHT algorithm and Galaxy Collision Sim improvements

Closed jakut77 requested to merge jt/bhtfin into main
@@ -4,10 +4,13 @@ number of particles/planets. To choose the initial conditions regadring initial
positions, velocities and sometimes the number of particles, please check the
function "initialize_particles". To choose initial state of the system,
specify the "initial" variable and number of particles. Possible initializations:
"random", "random_v", "two", "four", "four_v", "solar". The initializations marked with v
generate particles with initial velocities. The "solar" intialization generates random planets
and a heavy mass at the center of mass of the system. The "num_particles" parameter is only
relevant when choosing the "random" and "solar" options.
"random", "random_v", "two", "four", "four_v", "solar", and "galactic".
The initializations marked with v generate particles with initial velocities.
The "solar" intialization generates random planets and a heavy mass at the center of mass
of the system. The "galactic" option creates two large masses orbiting each other and
small masses orbiting each of them - the total number of particles will be double the input.
The "num_particles" parameter is only relevant when choosing the "random", "solar", and
"galactic" options.
"""
import matplotlib.pyplot as plt
import numpy as np
@@ -274,11 +277,14 @@ class GalaxyCollisionSimulation:
x, y, z = positions[step]
vx, vy, vz = self.particle_velocities[particle_index][step]
fx, fy, fz = self.particle_forces[particle_index][step]
color = 'blue'
if self.initial == 'galactic' and particle_index in [
len(self.particles) / 2 - 1, len(self.particles) - 1
]:
color = 'red'
elif self.initial == 'solar' and particle_index == len(self.particles) - 1:
color = 'orange'
else:
color = 'blue'
ax.scatter(x, y, z, c=color, s=20, alpha=0.5) # Plot particle position for the current time step
c_x, c_y, c_z = self.system_center_of_mass[0][step]
ax.scatter(c_x, c_y, c_z, c='orange', s=40)
@@ -300,9 +306,12 @@ class GalaxyCollisionSimulation:
ax.clear()
for j in range(n_bodies):
body_traj = self.particle_positions[j][i]
color = 'blue'
if self.initial == 'galactic' and j in [n_bodies / 2 - 1, n_bodies - 1]:
color = 'red'
elif self.initial == 'solar' and j == n_bodies - 1:
color = 'orange'
else:
color = 'blue'
ax.scatter(body_traj[0], body_traj[1], body_traj[2], c=color, alpha=0.5)
c_x, c_y, c_z = self.system_center_of_mass[0][i]
ax.scatter(c_x, c_y, c_z, c='orange', s=100)
@@ -324,7 +333,7 @@ class GalaxyCollisionSimulation:
if __name__ == "__main__":
sim = GalaxyCollisionSimulation(initial='galactic', num_particles=3)
sim = GalaxyCollisionSimulation(initial='solar', num_particles=5)
sim.simulate(num_steps=10000, time_step=0.001, print_tree=False)
#sim.display_snapshots(50, fix_axes=True)
sim.plot_trajectory(update_interval=10)
Loading