diff --git a/jobs/src/Galaxy_Collision_Simulation/3D/galaxysimulation.py b/jobs/src/Galaxy_Collision_Simulation/3D/galaxysimulation.py new file mode 100644 index 0000000000000000000000000000000000000000..670845486d02a2d29e9f43ce754b3b858c8bae22 --- /dev/null +++ b/jobs/src/Galaxy_Collision_Simulation/3D/galaxysimulation.py @@ -0,0 +1,236 @@ +""" +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