BHT algorithm and Galaxy Collision Sim improvements

tasks/src/ff_simulation_3D.py

@@ -4,12 +4,14 @@ 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 initial state of the system, 
 specify the "initial" variable and number of particles. Possible initializations:
-"random", "random_v", "two", "four", "four_v", "solar". 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 "num_particles" parameter is only
-relevant when choosing the "random" and "solar" options.
+"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
 
@@ -26,7 +28,9 @@ class GalaxyCollisionSimulation:
     def __init__(self, initial, num_particles):
         # Initialize the simulation with a given number of particles
         self.num_particles = num_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
@@ -78,12 +82,39 @@ class GalaxyCollisionSimulation:
                 particle.vz = vz
                 particles.append(particle)
                 center_of_mass += np.array([x, y, z])
-            center_of_mass /= self.num_particles
+            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 = []
@@ -140,6 +171,8 @@ class GalaxyCollisionSimulation:
                 Particle(50, 40, 30, 1000)
             ]
             return particles
+        else:
+            raise Exception("Initial condition not supported")
 
     def simulate(self, num_steps, time_step, print_tree):
         """
@@ -240,19 +273,21 @@ 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
+                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)
-                #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()
 
@@ -271,7 +306,13 @@ class GalaxyCollisionSimulation:
             ax.clear()
             for j in range(n_bodies):
                 body_traj = self.particle_positions[j][i]
-                ax.scatter(body_traj[0], body_traj[1], body_traj[2], c='blue', alpha=0.5)
+                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)
 
@@ -292,7 +333,7 @@ class GalaxyCollisionSimulation:
 
 
 if __name__ == "__main__":
-    sim = GalaxyCollisionSimulation(initial='solar', num_particles=6)
-    sim.simulate(num_steps=10000, time_step=0.001, print_tree=True)
-    #sim.display_snapshots(20, 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)