Skip to content
Snippets Groups Projects
Commit 4f7d9f22 authored by nguyed99's avatar nguyed99
Browse files

Naming changes in BHT

parent 476083b0
No related branches found
No related tags found
No related merge requests found
Pipeline #59659 passed
"""
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)
...@@ -139,7 +139,7 @@ class GalaxyCollisionSimulation: ...@@ -139,7 +139,7 @@ class GalaxyCollisionSimulation:
self.barnes_hut.rootNode.ComputeMassDistribution() self.barnes_hut.rootNode.ComputeMassDistribution()
if step in np.arange(0, num_steps, 1000): if step in np.arange(0, num_steps, 1000):
n_particles, particles = self.barnes_hut.rootNode.particles_in_subtree() n_particles, particles = self.barnes_hut.rootNode.particles_in_subtree()
print(f'===Quadtree at step {step}, n_particles:{n_particles}===') print(f'===Octtree at step {step}, n_particles:{n_particles}===')
self.barnes_hut.rootNode.print_tree() self.barnes_hut.rootNode.print_tree()
self.evaluate_forces() # Evaluate forces acting on each particle self.evaluate_forces() # Evaluate forces acting on each particle
self.integrate(time_step, 'velocity_verlet') # Integrate to update particle positions self.integrate(time_step, 'velocity_verlet') # Integrate to update particle positions
...@@ -211,7 +211,7 @@ class GalaxyCollisionSimulation: ...@@ -211,7 +211,7 @@ class GalaxyCollisionSimulation:
Function to display pictures for each time step Function to display pictures for each time step
""" """
num_time_steps = len(next(iter(self.particle_positions.values()))) num_time_steps = len(next(iter(self.particle_positions.values())))
print('===Quadtree at the end===') print('===Octtree at the end===')
self.barnes_hut.rootNode.print_tree() # print the tree self.barnes_hut.rootNode.print_tree() # print the tree
axes_type = "fixed" if fix_axes else "unfixed" axes_type = "fixed" if fix_axes else "unfixed"
#Fixed axis limits #Fixed axis limits
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment