From 2e33077b93ec5a12c2859b2e56cd06f65378e984 Mon Sep 17 00:00:00 2001
From: nguyed99 <nguyed99@zedat.fu-berlin.de>
Date: Fri, 8 Mar 2024 22:51:12 +0100
Subject: [PATCH] Add bht+jpl, etc
---
.gitignore | 3 +-
README.md | 4 +-
build/pyproject.toml | 7 +-
jobs/src/jpl_data_query.py | 6 +
jobs/src/random_perturbations.py | 4 -
jobs/src/visualization.py | 35 -----
jobs/tests/test_bht_algorithm_2D.py | 11 ++
jobs/tests/test_integrator.py | 11 ++
jobs/tests/test_jpl_data_query.py | 42 ++++++
pyproject.toml | 2 +-
tasks/src/a_task.py | 14 --
tasks/src/direct_simulation.py | 94 ++++++++++---
tasks/src/ff_simulation.py | 189 +++++++++++++++++++++++++-
tasks/src/utils.py | 33 +++++
tasks/tests/test_a_task.py | 43 ------
tasks/tests/test_direct_simulation.py | 0
tasks/tests/test_ff_simulation.py | 0
17 files changed, 375 insertions(+), 123 deletions(-)
delete mode 100644 jobs/src/random_perturbations.py
delete mode 100644 jobs/src/visualization.py
delete mode 100644 tasks/src/a_task.py
create mode 100644 tasks/src/utils.py
delete mode 100644 tasks/tests/test_a_task.py
delete mode 100644 tasks/tests/test_direct_simulation.py
delete mode 100644 tasks/tests/test_ff_simulation.py
diff --git a/.gitignore b/.gitignore
index 14fdbc8..af0c7ad 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,2 +1,3 @@
/__pycache__
-*.pyc
\ No newline at end of file
+*.pyc
+/playground
\ No newline at end of file
diff --git a/README.md b/README.md
index 398e8f9..109c893 100644
--- a/README.md
+++ b/README.md
@@ -39,11 +39,11 @@ Jobs and tasks are run inside of docker containers. There is a docker build scri
The image is built now and ready to be used! The tests in the Docker container built from the image can be sanity-checked via :
- docker run -it --rm git.imp.fu-berlin.de:5000/comp-sci-project/jobs-n-tasks:local python jobs/tests/test_a_job.py
+ docker run -it --rm git.imp.fu-berlin.de:5000/comp-sci-project/jobs-n-tasks:local pytest jobs/tests/test_a_job.py
or for testing a task:
- docker run -it --rm git.imp.fu-berlin.de:5000/comp-sci-project/jobs-n-tasks:local python tasks/tests/test_a_task.py
+ docker run -it --rm git.imp.fu-berlin.de:5000/comp-sci-project/jobs-n-tasks:local pytest tasks/tests/test_a_task.py
For live experience (meaning you enter the container), run:
diff --git a/build/pyproject.toml b/build/pyproject.toml
index 31a7120..45a3e87 100644
--- a/build/pyproject.toml
+++ b/build/pyproject.toml
@@ -6,10 +6,13 @@ authors = []
readme = "README.md"
[tool.poetry.dependencies]
-python = ">=3.10"
-numpy = "~1.24"
+python = "3.11.5"
+numpy = "1.24.4"
matplotlib = "3.8.0"
+astropy = "6.0.0"
+sunpy = "5.1.1"
pytest = "^7.4.4"
+astroquery = "^0.4.6"
[build-system]
requires = ["poetry-core"]
diff --git a/jobs/src/jpl_data_query.py b/jobs/src/jpl_data_query.py
index 6a2713e..1924a4c 100644
--- a/jobs/src/jpl_data_query.py
+++ b/jobs/src/jpl_data_query.py
@@ -1,6 +1,12 @@
"""
This module contains functions to query and parse JPL Horizons
on-line solar system data (see https://ssd-api.jpl.nasa.gov/doc/horizons.html).
+
+To use, run the following command
+
+ INPUT_PATH = "jobs/inputs/"
+ r_E, v_E = parse_output(query_data(INPUT_PATH + "earth_ephemeris.txt"))
+ r_M, v_M = parse_output(query_data(INPUT_PATH + "moon_ephemeris.txt"))
"""
import re
diff --git a/jobs/src/random_perturbations.py b/jobs/src/random_perturbations.py
deleted file mode 100644
index c950790..0000000
--- a/jobs/src/random_perturbations.py
+++ /dev/null
@@ -1,4 +0,0 @@
-"""
-This module contains functions to add Gaussian (white) noises
-as random small perturbations in the initial configuration.
-"""
\ No newline at end of file
diff --git a/jobs/src/visualization.py b/jobs/src/visualization.py
deleted file mode 100644
index 584d31a..0000000
--- a/jobs/src/visualization.py
+++ /dev/null
@@ -1,35 +0,0 @@
-import numpy as np
-import matplotlib.pyplot as plt
-
-
-def plot_trajectory(traj: np.ndarray, mass: list, update_interval=10):
- """
- Args:
- - traj: trajectory with shape (no_time_steps, no_bodies * dim)
- - mass: masses of bodies
- """
-
- n_time_step, n_bodies = traj.shape #
- n_bodies //= 3
-
- fig = plt.figure()
- ax = fig.add_subplot(projection='3d')
-
- for i in range(0, n_time_step, update_interval):
- ax.clear()
- for j in range(n_bodies):
- start_idx = j * 3
- end_idx = (j + 1) * 3
-
- body_traj = traj[i, start_idx:end_idx]
- ax.scatter(body_traj[0], body_traj[1], body_traj[2], s=10 * mass[j] / min(mass))
-
- ax.set_xlabel("X")
- ax.set_ylabel("Y")
- ax.set_zlabel("Z")
- ax.set_title(f"Time step: {i}")
- ax.set_xticks([])
- ax.set_yticks([])
- ax.set_zticks([])
-
- plt.pause(0.01)
diff --git a/jobs/tests/test_bht_algorithm_2D.py b/jobs/tests/test_bht_algorithm_2D.py
index 0318849..3ea2ec1 100644
--- a/jobs/tests/test_bht_algorithm_2D.py
+++ b/jobs/tests/test_bht_algorithm_2D.py
@@ -86,6 +86,17 @@ class IntegratorTest(unittest.TestCase):
t, p, q = system.simulation()
+ ## plotting trajectories
+ plt.figure()
+ plt.plot(q[:, 0], q[:, 1], label="Earth")
+ plt.legend()
+ plt.show()
+
+ plt.figure()
+ plt.plot(q[:, 2] - q[:, 0], q[:, 3] - q[:, 1], label="Moon seen from Earth")
+ plt.legend()
+ plt.show()
+
## checking total energy conservation
H = np.linalg.norm(p[:,:2], axis=1)**2 / (2 * M_E) + np.linalg.norm(p[:,2:], axis=1)**2 / (2 * M_L) + \
-G * M_S * M_E / np.linalg.norm(q[:,:2], axis=1) - G * M_S * M_L / np.linalg.norm(q[:,2:], axis=1) + \
diff --git a/jobs/tests/test_integrator.py b/jobs/tests/test_integrator.py
index fbd7008..e9bb889 100644
--- a/jobs/tests/test_integrator.py
+++ b/jobs/tests/test_integrator.py
@@ -75,6 +75,17 @@ class IntegratorTest(unittest.TestCase):
t, p, q = system.simulation()
+ ## plotting trajectories
+ plt.figure()
+ plt.plot(q[:, 0], q[:, 1], label="Earth")
+ plt.legend()
+ plt.show()
+
+ plt.figure()
+ plt.plot(q[:, 2] - q[:, 0], q[:, 3] - q[:, 1], label="Moon seen from Earth")
+ plt.legend()
+ plt.show()
+
## checking total energy conservation
H = np.linalg.norm(p[:,:2], axis=1)**2 / (2 * M_E) + np.linalg.norm(p[:,2:], axis=1)**2 / (2 * M_L) + \
-G * M_S * M_E / np.linalg.norm(q[:,:2], axis=1) - G * M_S * M_L / np.linalg.norm(q[:,2:], axis=1) + \
diff --git a/jobs/tests/test_jpl_data_query.py b/jobs/tests/test_jpl_data_query.py
index e69de29..5f88417 100644
--- a/jobs/tests/test_jpl_data_query.py
+++ b/jobs/tests/test_jpl_data_query.py
@@ -0,0 +1,42 @@
+"""
+This unittest tests data query from JPL Horizons database
+for the restricted three-body (sun-earth-moon) problem.
+The sun is fixed at the origin (center of mass). For
+simplicity, the moon is assumed to be in the eliptic plane,
+so that vectors can be treated as two-dimensional.
+The trajectories are plotted in Heliographic Stonyhurst coordinates.
+"""
+
+import unittest
+
+import matplotlib.pyplot as plt
+
+import astropy.units as u
+from sunpy.coordinates import (get_body_heliographic_stonyhurst, get_horizons_coord)
+from sunpy.time import parse_time
+
+time = parse_time('2024-01-01 00:00')
+
+
+class IntegratorTest(unittest.TestCase):
+
+ def test_data_query(self):
+ moon = get_horizons_coord('301', {'start': time - 30 * u.day, 'stop': time + 30 * u.day, 'step': '1d'})
+
+ earth = get_horizons_coord('399', {'start': time - 30 * u.day, 'stop': time + 30 * u.day, 'step': '1d'})
+
+ sun = get_body_heliographic_stonyhurst('Sun', time) # (lon, lat, radius) in (deg, deg, AU)
+
+ def coord_to_polar(coord):
+ return coord.lon.to_value('rad'), coord.radius.to_value('AU')
+
+ fig = plt.figure()
+ ax = fig.add_subplot(projection='polar')
+ ax.plot(0, 0, 'o', label='Sun', color='orange', markersize=10)
+ ax.plot(*coord_to_polar(earth.transform_to(sun)), label='Earth', color='blue', linestyle='-.')
+
+ ax.plot(*coord_to_polar(moon.transform_to(sun)), label='Moon', color='purple', linestyle='dashed')
+ ax.set_title('Stonyhurst heliographic coordinates')
+ ax.legend(loc='upper left', bbox_to_anchor=(1.0, 0.5))
+
+ plt.show()
diff --git a/pyproject.toml b/pyproject.toml
index abf9553..d719911 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,7 +6,7 @@ authors = []
readme = "README.md"
[tool.poetry.dependencies]
-python = "~3.11"
+python = "3.11.5"
toml = "^0.10.2"
[tool.poetry.group.dev.dependencies]
diff --git a/tasks/src/a_task.py b/tasks/src/a_task.py
deleted file mode 100644
index fa6c160..0000000
--- a/tasks/src/a_task.py
+++ /dev/null
@@ -1,14 +0,0 @@
-"""
-This module implements the SWAP gate using the Pauli gates
-"""
-
-import numpy as np
-
-from jobs.src.a_job import sigma_x, sigma_y, sigma_z
-
-
-def swap_gate() -> np.ndarray:
-
- swap_matrix = 1 / 2 * (np.kron(np.identity(2), np.identity(2)) + np.kron(sigma_x(), sigma_x()) +
- np.kron(sigma_y(), sigma_y()) + np.kron(sigma_z(), sigma_z()))
- return swap_matrix
diff --git a/tasks/src/direct_simulation.py b/tasks/src/direct_simulation.py
index 2d54a92..88ff8ee 100644
--- a/tasks/src/direct_simulation.py
+++ b/tasks/src/direct_simulation.py
@@ -1,35 +1,48 @@
"""
This module contains functions for a direct simulation
-of gravitational n-body system using realistic data for various
+of gravitational 2-body system using realistic data for various
initial configurations
"""
-import logging
-
+import astropy.units as u
+import matplotlib.pyplot as plt
import numpy as np
+from sunpy.coordinates import (get_body_heliographic_stonyhurst, get_horizons_coord)
+from sunpy.time import parse_time
from jobs.src.integrators import verlet
-from jobs.src.jpl_data_query import (parse_output, query_data)
from jobs.src.system import GravitationalSystem
-from jobs.src.visualization import plot_trajectory
-
-logger = logging.getLogger(__name__)
-logging.basicConfig(level=logging.INFO)
-
-INPUT_PATH = "jobs/inputs/"
-
-r_E, v_E = parse_output(query_data(INPUT_PATH + "earth_ephemeris.txt"))
-r_M, v_M = parse_output(query_data(INPUT_PATH + "moon_ephemeris.txt"))
+from tasks.src.utils import (coord_to_polar, xyz_to_polar)
# system settings
M_S = 1 # solar mass
M_E = 3.00e-6 * M_S
M_L = 3.69e-8 * M_S
-
G = 4 * np.pi**2 # AU^3 / m / yr^2
+# JPL data
+time = parse_time('2024-01-01 00:00')
+
+moon = get_horizons_coord('301', {
+ 'start': time - 40 * u.day,
+ 'stop': time + 40 * u.day,
+ 'step': '1d'
+},
+ include_velocity=True)
+
+earth = get_horizons_coord('399', {
+ 'start': time - 40 * u.day,
+ 'stop': time + 40 * u.day,
+ 'step': '1d'
+},
+ include_velocity=True)
+
+sun = get_body_heliographic_stonyhurst('Sun', time) # (lon, lat, radius) in (deg, deg, AU)
+moon = moon.transform_to(sun)
+earth = earth.transform_to(sun)
+
# simulation settings
-T = 0.3 # yr
+T = 0.23 # yr
n_order = 6
dt = 4**(-n_order) # yr
@@ -45,8 +58,18 @@ def force(q: np.ndarray) -> np.ndarray:
]).flatten()
-system = GravitationalSystem(r0=np.hstack((r_E[0], r_M[0])),
- v0=np.hstack((v_E[0], v_M[0])),
+# run direct simulation
+moon_0_pos = moon[0].represent_as('cartesian')
+earth_0_pos = earth[0].represent_as('cartesian')
+
+x0 = np.array([
+ earth_0_pos.x.value, earth_0_pos.y.value, earth_0_pos.z.value, moon_0_pos.x.value, moon_0_pos.y.value,
+ moon_0_pos.z.value, earth[0].velocity.d_x.value, earth[0].velocity.d_y.value, earth[0].velocity.d_z.value,
+ moon[0].velocity.d_x.value, moon[0].velocity.d_y.value, moon[0].velocity.d_z.value
+])
+
+system = GravitationalSystem(r0=x0[:6],
+ v0=x0[6:],
m=np.array([M_E, M_E, M_E, M_L, M_L, M_L]),
t=np.linspace(0, T, int(T // dt)),
force=force,
@@ -54,4 +77,39 @@ system = GravitationalSystem(r0=np.hstack((r_E[0], r_M[0])),
t, p, q = system.simulation()
-plot_trajectory(q, mass=[M_E, M_L])
+# plot trajectory in polar coordinates
+q_E_polar = np.array([xyz_to_polar(pos[0], pos[1], pos[2]) for pos in q])
+q_M_polar = np.array([xyz_to_polar(pos[3], pos[4], pos[5]) for pos in q])
+
+fig = plt.figure()
+ax = fig.add_subplot(projection='polar')
+ax.plot(0, 0, 'o', label='Sun', color='orange', markersize=10)
+ax.plot(q_E_polar[:, 0], q_E_polar[:, 1], label='Earth', color='blue', linestyle='-.')
+ax.plot(q_M_polar[:, 0], q_M_polar[:, 1], label='Moon', color='purple', linestyle='dashed')
+ax.plot(*coord_to_polar(earth.transform_to(sun)), label='Earth(JPL)', color='green', linestyle='-.')
+ax.plot(*coord_to_polar(moon.transform_to(sun)), label='Moon(JPL)', color='red', linestyle='dashed')
+ax.set_title('Stonyhurst heliographic coordinates')
+ax.legend(loc='upper left', bbox_to_anchor=(1.0, 0.5))
+plt.show()
+
+# plot trajectory in Cartesian coordinates
+moon_xyz = moon.represent_as('cartesian')
+earth_xyz = earth.represent_as('cartesian')
+moon_JPL = np.array([(i.x.value, i.y.value, i.z.value) for i in moon_xyz])
+earth_JPL = np.array([(i.x.value, i.y.value, i.z.value) for i in earth_xyz])
+
+fig = plt.figure()
+ax = fig.add_subplot(projection='3d')
+ax.plot(0, 0, 0, 'o', label='Sun', color='orange', markersize=10)
+ax.plot(q[:, 0], q[:, 1], q[:, 2], label="Earth")
+ax.plot(q[:, 3], q[:, 4], q[:, 5], label="Moon")
+ax.plot(earth_JPL[:, 0], earth_JPL[:, 1], earth_JPL[:, 2], '--', label='Earth(JPL)')
+ax.plot(moon_JPL[:, 0], moon_JPL[:, 1], moon_JPL[:, 2], '-.', label='Moon(JPL)')
+ax.set_xlabel("X")
+ax.set_ylabel("Y")
+ax.set_zlabel("Z")
+ax.set_xticks([])
+ax.set_yticks([])
+ax.set_zticks([])
+ax.legend()
+plt.show()
diff --git a/tasks/src/ff_simulation.py b/tasks/src/ff_simulation.py
index 8da049a..26fba3b 100644
--- a/tasks/src/ff_simulation.py
+++ b/tasks/src/ff_simulation.py
@@ -1,5 +1,188 @@
"""
This module contains functions for a simulation with force field
-of gravitational n-body system using realistic data for various
-initial configurations
-"""
\ No newline at end of file
+of gravitational n-body system
+"""
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from jobs.src.bht_algorithm_3D import MainApp, Particle
+
+
+# Simulation class for galaxy collision using Barnes-Hut and Euler method
+class GalaxyCollisionSimulation:
+
+ 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') # Generate particles
+ self.barnes_hut = MainApp(self.particles) # 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'):
+
+ if initial == 'two':
+ particles = [Particle(60, 50, 40, 1000), Particle(40, 50, 60, 1000)]
+ particles[0].vx, particles[0].vy, particles[0].vz = 5, 3, 3
+ particles[1].vx, particles[1].vy, particles[1].vz = 5, -3, 3
+ #particles=particles[:2]
+ return particles
+
+ elif initial == 'random':
+ # Generate random particles within a specified range
+ particles = []
+ for _ in range(self.num_particles):
+ x = np.random.uniform(50, 150)
+ y = np.random.uniform(50, 150)
+ z = np.random.uniform(50, 150)
+ 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 == 'four':
+ 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
+
+ def simulate(self, num_steps, time_step):
+ # Simulate the galaxy collision for a certain number of steps
+ 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):
+ print(f'===Quadtree at step {step}===')
+ 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):
+ # 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):
+ # 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):
+ # 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):
+ # 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):
+ # 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, 200)
+ ax.set_ylim(0, 200)
+ ax.set_zlim(0, 200)
+ ax.set_xlabel('X-axis')
+ ax.set_ylabel('Y-axis')
+ ax.set_zlabel('Z-axis')
+ ax.set_title(f'Time Step {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=20)
+ #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()
+
+
+if __name__ == "__main__":
+ sim = GalaxyCollisionSimulation(num_particles=10)
+ sim.simulate(num_steps=10000, time_step=0.001)
+ sim.display_snapshots(200, fix_axes=True)
diff --git a/tasks/src/utils.py b/tasks/src/utils.py
new file mode 100644
index 0000000..6e3b611
--- /dev/null
+++ b/tasks/src/utils.py
@@ -0,0 +1,33 @@
+import numpy as np
+from astropy.coordinates import CartesianRepresentation, SphericalRepresentation
+
+
+def rotation_matrix(phi, theta) -> np.ndarray:
+ """
+ Rotate a vector represented in spherical coordinates.
+
+ Args:
+ phi (float): Azimuthal angle in radians.
+ theta (float): Polar angle in radians.
+ """
+ # Construct the rotation matrix directly in spherical coordinates
+ rot_phi = np.array([[np.cos(phi), -np.sin(phi), 0], [np.sin(phi), np.cos(phi), 0], [0, 0,
+ 1]]) # Rotation around z-axis
+
+ rot_theta = np.array([[np.sin(theta), 0, np.cos(theta)], [0, 1, 0], [-np.cos(theta), 0,
+ np.sin(theta)]]) # Rotation around y-axis
+
+ # Combine the rotation matrices
+ rotation_matrix = np.dot(rot_phi, rot_theta)
+
+ return rotation_matrix
+
+
+def coord_to_polar(coord):
+ return coord.lon.to_value('rad'), coord.radius.to_value('AU')
+
+
+def xyz_to_polar(x, y, z):
+ cart = CartesianRepresentation(x, y, z)
+ sph = cart.represent_as(SphericalRepresentation)
+ return sph.lon.value, sph.distance.value
diff --git a/tasks/tests/test_a_task.py b/tasks/tests/test_a_task.py
deleted file mode 100644
index 8ca451b..0000000
--- a/tasks/tests/test_a_task.py
+++ /dev/null
@@ -1,43 +0,0 @@
-import logging
-import unittest
-
-import numpy as np
-
-from tasks.src.a_task import swap_gate
-
-logger = logging.getLogger(__name__)
-logging.basicConfig(level=logging.INFO)
-
-
-class ATaskTest(unittest.TestCase):
-
- def test_a_task(self):
- """
- Test functionalities of swap gate on 2-qubits
- """
-
- qubit_0 = np.array([1, 0])
- qubit_1 = np.array([0, 1])
-
- qubit_00 = np.kron(qubit_0, qubit_0)
- qubit_11 = np.kron(qubit_1, qubit_1)
- qubit_01 = np.kron(qubit_0, qubit_1)
- qubit_10 = np.kron(qubit_1, qubit_0)
-
- # ensure input is a 2-qubit
- self.assertTrue(qubit_00.shape[0] == 4)
- self.assertTrue(qubit_01.shape[0] == 4)
- self.assertTrue(qubit_10.shape[0] == 4)
- self.assertTrue(qubit_11.shape[0] == 4)
-
- self.assertTrue(np.array_equal(np.dot(swap_gate(), qubit_00), qubit_00))
- self.assertTrue(np.array_equal(np.dot(swap_gate(), qubit_11), qubit_11))
-
- self.assertTrue(np.array_equal(np.dot(swap_gate(), qubit_01), qubit_10))
- self.assertTrue(np.array_equal(np.dot(swap_gate(), qubit_10), qubit_01))
-
- logger.info("Logging for fun!")
-
-
-if __name__ == '__main__':
- unittest.main()
diff --git a/tasks/tests/test_direct_simulation.py b/tasks/tests/test_direct_simulation.py
deleted file mode 100644
index e69de29..0000000
diff --git a/tasks/tests/test_ff_simulation.py b/tasks/tests/test_ff_simulation.py
deleted file mode 100644
index e69de29..0000000
--
GitLab