BHT algorithm and Galaxy Collision Sim improvements

tasks/src/ff_simulation_3D.py

@@ -4,9 +4,10 @@ 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". The initializations marked with v
-generate particles with initial velocities. The "num_particles" parameter is only
-relevant when choosing the "random" option.
+"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.
 """
 
 import matplotlib.pyplot as plt
@@ -59,6 +60,30 @@ class GalaxyCollisionSimulation:
             #particles=particles[:2]
             return particles
 
+        elif initial == 'solar':
+            # Generate random particles and a big mass at the center of mass
+            particles = []
+            center_of_mass = np.zeros(3)
+            for _ in range(self.num_particles - 1):
+                x = np.random.uniform(0, 100)
+                y = np.random.uniform(0, 100)
+                z = np.random.uniform(48, 52)
+                vx = np.random.uniform(70, 100)
+                vy = np.random.uniform(70, 100)
+                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
+            sun = Particle(center_of_mass[0], center_of_mass[1], center_of_mass[2], mass**2)
+            particles.append(sun)
+
+            return particles
+
         elif initial == 'random_v':
             # Generate random particles within a specified range with random initial velocities
             particles = []
@@ -248,7 +273,7 @@ class GalaxyCollisionSimulation:
                 body_traj = self.particle_positions[j][i]
                 ax.scatter(body_traj[0], body_traj[1], body_traj[2], c='blue', 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=40)
+                ax.scatter(c_x, c_y, c_z, c='orange', s=100)
 
             ax.set_xlim(0, 100)
             ax.set_ylim(0, 100)
@@ -257,7 +282,7 @@ class GalaxyCollisionSimulation:
             ax.set_xlabel("X")
             ax.set_ylabel("Y")
             ax.set_zlabel("Z")
-            ax.set_title(f"Time step: {i}")
+            ax.set_title(f"Day: {i}")
             ax.set_xticks([])
             ax.set_yticks([])
             ax.set_zticks([])
@@ -267,7 +292,7 @@ class GalaxyCollisionSimulation:
 
 
 if __name__ == "__main__":
-    sim = GalaxyCollisionSimulation(initial='random_v', num_particles=10)
+    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.plot_trajectory()
+    sim.plot_trajectory(update_interval=10)