"""
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 27.
When choosing "two", "four" or "four_2", make sure to adjust the number of
particles in "simulate.py" correspondingly.
"""
import matplotlib.pyplot as plt
import numpy as np
from jobs.src.bht_algorithm_2D import MainApp, Particle
class GalaxyCollisionSimulation:
"""
Simulation class for galaxy collision using Barnes-Hut and Euler method
or Velocity Verlet
"""
def __init__(self, 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.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
self.particle_positions = {i: [(particle.x, particle.y)]
for i, particle in enumerate(self.particles)
} # Store particle positions for each time step
self.particle_velocities = {i: [(particle.vx, particle.vy)]
for i, particle in enumerate(self.particles)
} # Store particle velocities for each time step
self.particle_forces = {i: [(particle.fx, particle.fy)]
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])
}
def initialize_particles(self, initial='random'):
"""
Function to initialize particles with different initial conditions
"""
if initial == 'two':
# Generate two particles without initial velocities
particles = [Particle(60, 50, 1000), Particle(40, 50, 1000)]
particles[0].vx, particles[0].vy = 0, 0
particles[1].vx, particles[1].vy = 0, 0
#particles=particles[:2]
return particles
elif initial == 'random':
# Generate random particles within a specified range with random initial velocities
particles = []
for _ in range(self.num_particles):
x = np.random.uniform(0, 100)
y = np.random.uniform(0, 100)
vx = np.random.uniform(-5, 5)
vy = np.random.uniform(-5, 5)
#mass = np.random.uniform(10, 100)
mass = 1000
particle = Particle(x, y, mass)
particle.vx = vx
particle.vy = vy
particles.append(particle)
return particles
elif initial == 'random_2':
# 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)
vx = 0
vy = 0
#mass = np.random.uniform(10, 100)
mass = 1000
particle = Particle(x, y, mass)
particle.vx = vx
particle.vy = vy
particles.append(particle)
return particles
elif initial == 'four':
# Generate four particles with initial velocities
particles = [Particle(60, 50, 1000), Particle(40, 50, 1000), Particle(50, 60, 1000), Particle(50, 40, 1000)]
particles[0].vx, particles[0].vy = 0, 3
particles[1].vx, particles[1].vy = 0, -3
particles[2].vx, particles[2].vy = -3, 0
particles[3].vx, particles[3].vy = 3, 0
#particles=particles[:2]
return particles
elif initial == 'four_2':
# Generate four particles without initial velocities
particles = [Particle(60, 50, 1000), Particle(40, 50, 1000), Particle(50, 60, 1000), Particle(50, 40, 1000)]
particles[0].vx, particles[0].vy = 0, 0
particles[1].vx, particles[1].vy = 0, 0
particles[2].vx, particles[2].vy = 0, 0
particles[3].vx, particles[3].vy = 0, 0
#particles=particles[:2]
return particles
def simulate(self, num_steps, time_step):
"""
Function to simulate the galaxy collision for a certain number of step
"""
for step in range(num_steps):
# 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'===Quadtree at step {step}, n_particles:{n_particles}===')
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])
# 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))
self.particle_velocities[i].append((particle.vx, particle.vy))
self.particle_forces[i].append((particle.fx, particle.fy))
def evaluate_forces(self):
"""
Function to evaluate forces acting on each particle using Barnes-Hut algorithm
"""
for particle in self.particles:
force = self.barnes_hut.rootNode.CalculateForceFromTree(particle)
particle.fx, particle.fy = force
def calculate_force(self, particle):
"""
Function to calculate force acting on a single particle
"""
force = self.barnes_hut.rootNode.CalculateForceFromTree(particle)
particle.fx, particle.fy = force
def integrate(self, time_step, method):
"""
Function to integrate to update particle positions based on calculated forces
"""
if method == 'explicit_euler':
for particle in self.particles:
particle.vx += particle.fx * time_step / particle.mass # Update velocity components
particle.vy += particle.fy * time_step / particle.mass
particle.x += particle.vx * time_step # Update particle positions using velocity
particle.y += particle.vy * time_step
elif method == 'velocity_verlet':
for particle in self.particles:
particle.vx += particle.fx * time_step / (2 * particle.mass) # First half-step
particle.vy += particle.fy * time_step / (2 * particle.mass)
particle.x += particle.vx * time_step # Full step for position
particle.y += particle.vy * time_step
self.calculate_force(particle) # Recalculate force for the particle
particle.vx += particle.fx * time_step / (2 * particle.mass) # Second half-step
particle.vy += particle.fy * time_step / (2 * particle.mass)
def print_particle_positions(self):
"""
Function to print the positions of all particles for each time step
"""
for i, positions in self.particle_positions.items():
print(f"Particle {i + 1} Positions:")
for step, pos in enumerate(positions):
print(f" Time Step {step}: Position ({pos[0]}, {pos[1]})")
def display_snapshots(self, num_shots, fix_axes=True):
"""
Function to display pictures for each time step
"""
num_time_steps = len(next(iter(self.particle_positions.values())))
print('===Quadtree at the end===')
self.barnes_hut.rootNode.print_tree() # print the tree
axes_type = "fixed" if fix_axes else "unfixed"
#Fixed axis limits
print(f"For {axes_type} axis limits, num_particles={self.num_particles}")
for step in np.arange(0, num_time_steps, int(num_time_steps / num_shots)):
fig, ax = plt.subplots()
if fix_axes:
ax.set_xlim(0, 100)
ax.set_ylim(0, 100)
ax.set_xlabel('X-axis')
ax.set_ylabel('Y-axis')
ax.set_title(f'Day {step}')
ax.grid()
#print(f'Step {step}')
for particle_index, positions in self.particle_positions.items():
x, y = positions[step]
vx, vy = self.particle_velocities[particle_index][step]
fx, fy = self.particle_forces[particle_index][step]
#mass = self.particles[particle_index].mass
ax.scatter(x, y, c='blue', s=20, alpha=0.5) # Plot particle position for the current time step
c_x, c_y = self.system_center_of_mass[step]
ax.scatter(c_x, c_y, c='orange', s=300)
#print(
#f'i={particle_index}: x={round(x,2)}, y={round(y,2)}, vx={round(vx,2)}, vy={round(vy,2)}, fx={round(fx,2)}, fy={round(fy,2)}'
#)
#print(y)
plt.show()
if __name__ == "__main__":
sim = GalaxyCollisionSimulation(num_particles=3)
sim.simulate(num_steps=10000, time_step=0.001)
sim.display_snapshots(200, fix_axes=True)