diff --git a/tasks/src/ff_simulation_3D.py b/tasks/src/ff_simulation_3D.py index c6995cfdd48619a0c4f1039738789d84a0f3cd38..699a0f4fd95208e33bac6266836b1bb0b947bfb5 100644 --- a/tasks/src/ff_simulation_3D.py +++ b/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)