Skip to content
Snippets Groups Projects

BHT algorithm and Galaxy Collision Sim improvements

Closed BHT algorithm and Galaxy Collision Sim improvements
jt/bhtfin into main
Closed jakut77 requested to merge jt/bhtfin into main
Compare
and
Show latest version
1 file
+ 58
− 17
Compare changes
  • Side-by-side
  • Inline
tasks/src/ff_simulation_3D.py
+ 133
− 46
@@ -2,13 +2,16 @@
This code contains a class to simulate the collision of a Galaoxy with a given
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 the initial coditions desired, please
sure to type the correct keyword for the variable "initial" in the command
self.particles = self.initialize_particles(initial='random_2') in line 28.
When choosing "two", "four" or "four_2", make sure to adjust the number of
particles in "simulate.py" correspondingly.
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", 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
@@ -22,10 +25,12 @@ class GalaxyCollisionSimulation:
"""
def __init__(self, num_particles):
def __init__(self, initial, num_particles):
# Initialize the simulation with a given number of particles
self.num_particles = num_particles
self.particles = self.initialize_particles(initial='random_2') # Generate particles
self.initial = initial
self.particles = self.initialize_particles(initial=initial) # Generate particles
self.num_particles = len(self.particles)
self.barnes_hut = MainApp() # Initialize Barnes-Hut algorithm instance
self.barnes_hut.BuildTree(self.particles) # Build the Barnes-Hut tree with particles
self.barnes_hut.rootNode.ComputeMassDistribution() #Compute the center of mass of the tree nodes
@@ -42,8 +47,8 @@ class GalaxyCollisionSimulation:
for i, particle in enumerate(self.particles)
} # Store particle velocities for each time step
self.system_center_of_mass = {
0: (self.barnes_hut.rootNode.center_of_mass[0], self.barnes_hut.rootNode.center_of_mass[1],
self.barnes_hut.rootNode.center_of_mass[2])
0: [(self.barnes_hut.rootNode.center_of_mass[0], self.barnes_hut.rootNode.center_of_mass[1],
self.barnes_hut.rootNode.center_of_mass[2])]
}
def initialize_particles(self, initial='random'):
@@ -59,7 +64,58 @@ class GalaxyCollisionSimulation:
#particles=particles[:2]
return particles
elif initial == 'random':
elif initial == 'solar':
# Generate random particles and a big mass at the center of mass
particles = []
center_of_mass = np.zeros(3)
for _ in range(self.num_particles - 1):
x = np.random.uniform(0, 100)
y = np.random.uniform(0, 100)
z = np.random.uniform(48, 52)
vx = np.random.uniform(70, 100)
vy = np.random.uniform(70, 100)
vz = 0
mass = 1000
particle = Particle(x, y, z, mass)
particle.vx = vx
particle.vy = vy
particle.vz = vz
particles.append(particle)
center_of_mass += np.array([x, y, z])
center_of_mass /= self.num_particles - 1
sun = Particle(center_of_mass[0], center_of_mass[1], center_of_mass[2], mass**2)
particles.append(sun)
return particles
elif initial == 'galactic' and self.num_particles >= 3:
particles = []
for i in range(2):
center_of_mass = np.zeros(3)
for _ in range(self.num_particles - 1):
x = np.random.uniform(50 * i, 50 * (i + 1))
y = np.random.uniform(50 * i, 50 * (i + 1))
z = np.random.uniform(48, 52)
pm = (-1)**i
vx = np.random.uniform(40 * pm, 70 * pm)
vy = np.random.uniform(40 * pm, 70 * pm)
vz = 0
mass = 1000
particle = Particle(x, y, z, mass)
particle.vx = vx
particle.vy = vy
particle.vz = vz
particles.append(particle)
center_of_mass += np.array([x, y, z])
center_of_mass /= self.num_particles - 1
sun = Particle(center_of_mass[0], center_of_mass[1], center_of_mass[2], mass**2)
sun.vx = 50 * pm
sun.vy = -20 * pm
particles.append(sun)
return particles
elif initial == 'random_v':
# Generate random particles within a specified range with random initial velocities
particles = []
for _ in range(self.num_particles):
@@ -69,7 +125,6 @@ class GalaxyCollisionSimulation:
vx = np.random.uniform(-5, 5)
vy = np.random.uniform(-5, 5)
vz = np.random.uniform(-5, 5)
#mass = np.random.uniform(10, 100)
mass = 1000
particle = Particle(x, y, z, mass)
particle.vx = vx
@@ -79,27 +134,21 @@ class GalaxyCollisionSimulation:
return particles
elif initial == 'random_2':
elif initial == 'random':
# Generate random particles within a specified range without random initial velocities
particles = []
for _ in range(self.num_particles):
x = np.random.uniform(0, 100)
y = np.random.uniform(0, 100)
z = np.random.uniform(0, 100)
vx = 0
vy = 0
vz = 0
#mass = np.random.uniform(10, 100)
mass = 1000
particle = Particle(x, y, z, mass)
particle.vx = vx
particle.vy = vy
particle.vz = vz
particles.append(particle)
return particles
elif initial == 'four':
elif initial == 'four_v':
# Generate four particles with initial velocities
particles = [
Particle(60, 50, 40, 1000),
@@ -111,10 +160,9 @@ class GalaxyCollisionSimulation:
particles[1].vx, particles[1].vy, particles[1].vz = 0, -3, 3
particles[2].vx, particles[2].vy, particles[2].vz = -3, 0, -3
particles[3].vx, particles[3].vy, particles[3].vz = 3, 0, -3
#particles=particles[:2]
return particles
elif initial == 'four_2':
elif initial == 'four':
# Generate four particles without initial velocities
particles = [
Particle(60, 50, 40, 1000),
@@ -122,14 +170,11 @@ class GalaxyCollisionSimulation:
Particle(50, 60, 70, 1000),
Particle(50, 40, 30, 1000)
]
particles[0].vx, particles[0].vy, particles[0].vz = 0, 0, 0
particles[1].vx, particles[1].vy, particles[1].vz = 0, 0, 0
particles[2].vx, particles[2].vy, particles[2].vz = 0, 0, 0
particles[3].vx, particles[3].vy, particles[3].vz = 0, 0, 0
#particles=particles[:2]
return particles
else:
raise Exception("Initial condition not supported")
def simulate(self, num_steps, time_step):
def simulate(self, num_steps, time_step, print_tree):
"""
Function to simulate the galaxy collision for a certain number of step
"""
@@ -137,16 +182,15 @@ class GalaxyCollisionSimulation:
# For each step, build the tree
self.barnes_hut.BuildTree(self.particles)
self.barnes_hut.rootNode.ComputeMassDistribution()
if step in np.arange(0, num_steps, 1000):
n_particles, particles = self.barnes_hut.rootNode.particles_in_subtree()
print(f'===Octtree at step {step}, n_particles:{n_particles}===')
if print_tree and step in np.arange(0, num_steps, 1000):
print(f'===Octtree at step {step}===')
self.barnes_hut.rootNode.print_tree()
self.evaluate_forces() # Evaluate forces acting on each particle
self.integrate(time_step, 'velocity_verlet') # Integrate to update particle positions
self.system_center_of_mass[step] = (self.barnes_hut.rootNode.center_of_mass[0],
self.barnes_hut.rootNode.center_of_mass[1],
self.barnes_hut.rootNode.center_of_mass[2])
self.system_center_of_mass[0].append(
(self.barnes_hut.rootNode.center_of_mass[0], self.barnes_hut.rootNode.center_of_mass[1],
self.barnes_hut.rootNode.center_of_mass[2]))
# Store particle positions, velocities and forces for each time step
for i, particle in enumerate(self.particles):
self.particle_positions[i].append((particle.x, particle.y, particle.z))
@@ -229,24 +273,67 @@ class GalaxyCollisionSimulation:
ax.set_title(f'Day {step}')
ax.grid()
#print(f'Step {step}')
for particle_index, positions in self.particle_positions.items():
x, y, z = positions[step]
vx, vy, vz = self.particle_velocities[particle_index][step]
fx, fy, fz = self.particle_forces[particle_index][step]
#mass = self.particles[particle_index].mass
ax.scatter(x, y, z, c='blue', s=20, alpha=0.5) # Plot particle position for the current time step
c_x, c_y, c_z = self.system_center_of_mass[step]
ax.scatter(c_x, c_y, c_z, c='orange', s=300)
#print(
#f'i={particle_index}: x={round(x,2)}, y={round(y,2)}, z={round(z,2)}, vx={round(vx,2)}, vy={round(vy,2)}, vz={round(vz,2)}, fx={round(fx,2)}, fy={round(fy,2)}, fz={round(fz,2)}'
#)
#print(y)
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)
plt.show()
def plot_trajectory(self, update_interval=100):
"""
Method to visualize particles trajectories as a movie
"""
n_time_step = len(self.particle_positions[0])
n_bodies = len(self.particles)
fig = plt.figure()
ax = fig.add_subplot(projection='3d')
for i in range(0, n_time_step, update_interval):
ax.clear()
for j in range(n_bodies):
body_traj = self.particle_positions[j][i]
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)
ax.set_xlim(0, 100)
ax.set_ylim(0, 100)
ax.set_zlim(0, 100)
ax.set_xlabel("X")
ax.set_ylabel("Y")
ax.set_zlabel("Z")
ax.set_title(f"Day: {i}")
ax.set_xticks([])
ax.set_yticks([])
ax.set_zticks([])
ax.grid()
plt.pause(0.01)
if __name__ == "__main__":
sim = GalaxyCollisionSimulation(num_particles=10)
sim.simulate(num_steps=10000, time_step=0.001)
sim.display_snapshots(200, fix_axes=True)
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)