docstring fix

tasks/src/ff_simulation_3D.py

@@ -4,10 +4,13 @@ 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
@@ -274,11 +277,14 @@ class GalaxyCollisionSimulation:
                 x, y, z = positions[step]
                 vx, vy, vz = self.particle_velocities[particle_index][step]
                 fx, fy, fz = self.particle_forces[particle_index][step]
-                color = 'blue'
                 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)
@@ -300,9 +306,12 @@ class GalaxyCollisionSimulation:
             ax.clear()
             for j in range(n_bodies):
                 body_traj = self.particle_positions[j][i]
-                color = 'blue'
                 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)
@@ -324,7 +333,7 @@ class GalaxyCollisionSimulation:
 
 
 if __name__ == "__main__":
-    sim = GalaxyCollisionSimulation(initial='galactic', num_particles=3)
+    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)