Testing with NASA data

jobs/inputs/earth_ephemeris.txt

+!$$SOF
+MAKE_EPHEM=YES
+COMMAND=399
+EPHEM_TYPE=VECTORS
+CENTER='coord@10'
+COORD_TYPE=GEODETIC
+SITE_COORD='13.35075983141236,52.497297553950155,0.0354'
+START_TIME='2023-01-01'
+STOP_TIME='2023-05-01'
+STEP_SIZE='1 DAYS'
+VEC_TABLE='2x'
+REF_SYSTEM='ICRF'
+REF_PLANE='ECLIPTIC'
+VEC_CORR='NONE'
+CAL_TYPE='M'
+OUT_UNITS='KM-S'
+VEC_LABELS='YES'
+VEC_DELTA_T='NO'
+CSV_FORMAT='NO'
+OBJ_DATA='YES'
\ No newline at end of file

jobs/inputs/moon_ephemeris.txt

+!$$SOF
+MAKE_EPHEM=YES
+COMMAND=301
+EPHEM_TYPE=VECTORS
+CENTER='coord@10'
+COORD_TYPE=GEODETIC
+SITE_COORD='13.35075983141236,52.497297553950155,0.0354'
+START_TIME='2023-01-01'
+STOP_TIME='2023-05-01'
+STEP_SIZE='1 DAYS'
+VEC_TABLE='2x'
+REF_SYSTEM='ICRF'
+REF_PLANE='ECLIPTIC'
+VEC_CORR='NONE'
+CAL_TYPE='M'
+OUT_UNITS='KM-S'
+VEC_LABELS='YES'
+VEC_DELTA_T='NO'
+CSV_FORMAT='NO'
+OBJ_DATA='YES'
\ No newline at end of file

jobs/inputs/sun_ephemeris.txt

+!$$SOF
+MAKE_EPHEM='YES'
+COMMAND='10'
+EPHEM_TYPE='VECTORS'
+CENTER='coord@399'
+COORD_TYPE='GEODETIC'
+SITE_COORD='13.35075983141236,52.497297553950155,0.0354'
+START_TIME='2023-01-01'
+STOP_TIME='2023-05-01'
+STEP_SIZE='1 DAYS'
+VEC_TABLE='2x'
+REF_SYSTEM='ICRF'
+REF_PLANE='ECLIPTIC'
+VEC_CORR='NONE'
+CAL_TYPE='M'
+OUT_UNITS='KM-S'
+VEC_LABELS='YES'
+VEC_DELTA_T='NO'
+CSV_FORMAT='YES'
+OBJ_DATA='NO'
\ No newline at end of file

jobs/src/a_job.py

-"""
-Just some toy code for demonstration
-"""
-import numpy as np
-
-
-def sigma_x() -> np.ndarray:
-    """
-    Pauli gate X (or NOT gate) 
-    """
-
-    return np.array([[0, 1], [1, 0]])
-
-
-def sigma_y() -> np.ndarray:
-    """
-    Pauli gate Y
-    """
-
-    return np.array([[0, -1j], [1j, 0]])
-
-
-def sigma_z() -> np.ndarray:
-    """
-    Pauli gate Z
-    """
-
-    return np.array([[1, 0], [0, -1]])

jobs/src/integrators.py

@@ -3,6 +3,7 @@ This module contains 2 integrators, which is at least
 of second consistency order and approrimately preserves 
 the energy over long times.
 """
+
 from collections.abc import Callable
 
 import numpy as np

jobs/src/jpl_data_query.py

 """
-This module contains functions to query JPL Horizons 
-on-line solar system data (see https://ssd.jpl.nasa.gov/horizons).
-"""
+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).
+"""
\ No newline at end of file
+
+import re
+import requests
+
+import numpy as np
+
+
+def query_data(input_file: str) -> str:
+    """
+    Send request to the Horizons API with settings for the desired ephemeris results saved in an input file.
+    """
+    with open(input_file, 'r') as f:
+        file = f.read()
+
+    url = "https://ssd.jpl.nasa.gov/api/horizons_file.api"
+    r = requests.post(url, data={'format': 'json'}, files={'input': file})
+
+    return r.json()['result']
+
+
+def parse_output(output: str) -> np.ndarray:
+    """
+    Parse output of API calls for positions and velocities
+    """
+
+    output = list(filter(bool, output.split("\n")))
+    vector_list = [line.split(",") for line in output if line[1:4] == "XYZ"]
+    vector_list = [list(map(float, re.findall(r"[-+]?\d*\.?\d+(?:[Ee][-+]?\d+)?", i[0]))) for i in vector_list]
+
+    r = np.zeros((len(vector_list), 3), dtype=float)
+    v = np.zeros((len(vector_list), 3), dtype=float)
+
+    for i, vector in enumerate(vector_list):
+        r[i][0] = float(vector[0])
+        r[i][1] = float(vector[1])
+        r[i][2] = float(vector[2])
+        v[i][0] = float(vector[3])
+        v[i][1] = float(vector[4])
+        v[i][2] = float(vector[5])
+
+    return r, v

jobs/src/visualization.py

+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)

jobs/tests/test_a_job.py

-import logging
-import unittest
-
-import numpy as np
-
-from jobs.src.a_job import sigma_x, sigma_y, sigma_z
-
-logger = logging.getLogger(__name__)
-logging.basicConfig(level=logging.INFO)
-
-
-class AJobTest(unittest.TestCase):
-
-    def test_a_job(self):
-        """
-        Test functionalities of Pauli gates on single qubits
-        """
-        qubit_0 = np.array([1, 0])
-        qubit_1 = np.array([0, 1])
-
-        # ensure input is a single qubit
-        self.assertTrue(qubit_0.shape[0] == 2)
-        self.assertTrue(qubit_1.shape[0] == 2)
-
-        # test Pauli gate X
-        self.assertTrue(np.array_equal(np.dot(sigma_x(), qubit_0), qubit_1))
-        self.assertTrue(np.array_equal(np.dot(sigma_x(), qubit_1), qubit_0))
-
-        # test Pauli gate Y
-        self.assertTrue(np.array_equal(np.dot(sigma_y(), qubit_0), 1j * qubit_1))
-        self.assertTrue(np.array_equal(np.dot(sigma_y(), qubit_1), -1j * qubit_0))
-
-        # test Pauli gate Z
-        self.assertTrue(np.array_equal(np.dot(sigma_z(), qubit_0), qubit_0))
-        self.assertTrue(np.array_equal(np.dot(sigma_z(), qubit_1), -qubit_1))
-
-        logger.info("Logging for fun!")
-
-
-if __name__ == '__main__':
-    unittest.main()

tasks/src/direct_simulation.py

@@ -2,4 +2,56 @@
 This module contains functions for a direct simulation
 of gravitational n-body system using realistic data for various
 initial configurations
-"""
+"""
\ No newline at end of file
+
+import logging
+
+import numpy as np
+
+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"))
+
+# 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
+
+# simulation settings
+T = 0.3  # yr
+n_order = 6
+dt = 4**(-n_order)  # yr
+
+
+# force computation
+def force(q: np.ndarray) -> np.ndarray:
+    r1 = q[0:3]
+    r2 = q[3:6]
+
+    return np.array([
+        -G * M_E * (M_S * r1 / np.linalg.norm(r1)**3 + M_L * (r1 - r2) / np.linalg.norm(r1 - r2)**3),
+        -G * M_L * (M_S * r2 / np.linalg.norm(r2)**3 + M_E * (r2 - r1) / np.linalg.norm(r2 - r1)**3),
+    ]).flatten()
+
+
+system = GravitationalSystem(r0=np.hstack((r_E[0], r_M[0])),
+                             v0=np.hstack((v_E[0], v_M[0])),
+                             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,
+                             solver=verlet)
+
+t, p, q = system.simulation()
+
+plot_trajectory(q, mass=[M_E, M_L])