Skip to content
Snippets Groups Projects
Commit fb1db657 authored by nicoa96's avatar nicoa96
Browse files

galaxysimulation

parent e901728b
Branches
No related tags found
1 merge request!10Nicoa96 main patch 73354
Pipeline #59648 passed
Stage: test-jobs
Show whitespace changes
Inline Side-by-side
jobs/src/Galaxy_Collision_Simulation/3D/galaxysimulation.py 0 → 100644
+ 236
− 0
View file @ fb1db657
"""
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.
"""
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
from bht_algorithm_3D 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, particle.z)]
for i, particle in enumerate(self.particles)
} # Store particle positions for each time step
self.particle_velocities = {
i: [(particle.vx, particle.vy, particle.vz)]
for i, particle in enumerate(self.particles)
} # Store particle velocities for each time step
self.particle_forces = {
i: [(particle.fx, particle.fy, particle.fz)]
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])
}
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, 40, 1000), Particle(40, 50, 60, 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=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)
z = np.random.uniform(0, 100)
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
particle.vy = vy
particle.vz = vz
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)
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':
# Generate four particles with initial velocities
particles = [Particle(60, 50, 40, 1000), Particle(40, 50, 60, 1000), Particle(50,60,70,1000), Particle(50,40,30,1000)]
particles[0].vx, particles[0].vy, particles[0].vz= 0, 3, 3
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':
# Generate four particles without initial velocities
particles = [Particle(60, 50, 40, 1000), Particle(40, 50, 60, 1000), 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
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],
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))
self.particle_velocities[i].append((particle.vx, particle.vy, particle.vz))
self.particle_forces[i].append((particle.fx, particle.fy, particle.fz))
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, particle.fz = 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, particle.fz = 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.vz += particle.fz * time_step / particle.mass
particle.x += particle.vx * time_step # Update particle positions using velocity
particle.y += particle.vy * time_step
particle.z += particle.vz * 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.vz += particle.fz * time_step / (2 * particle.mass)
particle.x += particle.vx * time_step # Full step for position
particle.y += particle.vy * time_step
particle.z += particle.vz * 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)
particle.vz += particle.fz * 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]}, {pos[2]})")
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 = plt.figure()
ax = fig.add_subplot(111, projection='3d')
if fix_axes:
ax.set_xlim(0, 100)
ax.set_ylim(0, 100)
ax.set_zlim(0, 100)
ax.set_xlabel('X-axis')
ax.set_ylabel('Y-axis')
ax.set_zlabel('Z-axis')
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)
plt.show()
\ No newline at end of file
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment