Merge branch 'main' into jima

build/pyproject.toml

@@ -8,6 +8,7 @@ readme = "README.md"
 [tool.poetry.dependencies]
 python = ">=3.10"
 numpy = "~1.24"
+matplotlib = "3.8.0"
 pytest = "^7.4.4"
 
 [build-system]

jobs/tests/test_integrator.py

@@ -9,6 +9,7 @@ so that vectors can be treated as two-dimensional.
 import logging
 import unittest
 
+import matplotlib.pyplot as plt
 import numpy as np
 
 from jobs.src.integrators import verlet
@@ -35,18 +36,6 @@ T = 0.3  # yr
 n_order = 6
 dt = 4**(-n_order)  # yr
 
-# vector of r0 and v0
-x0 = np.array([
-    R_SE,
-    0,
-    R_SE + R_L,
-    0,
-    0,
-    R_SE * 2 * np.pi / T_E,
-    0,
-    1 / M_E * M_E * R_SE * 2 * np.pi / T_E + 1 * R_L * 2 * np.pi / T_L,
-])
-
 
 # force computation
 def force(q: np.ndarray) -> np.ndarray:
@@ -65,6 +54,17 @@ class IntegratorTest(unittest.TestCase):
         """
         Test functionalities of velocity-Verlet algorithm
         """
+        # vector of r0 and v0
+        x0 = np.array([
+            R_SE,
+            0,
+            R_SE + R_L,
+            0,
+            0,
+            R_SE * 2 * np.pi / T_E,
+            0,
+            1 / M_E * M_E * R_SE * 2 * np.pi / T_E + 1 * R_L * 2 * np.pi / T_L,
+        ])
 
         system = GravitationalSystem(r0=x0[:4],
                                      v0=x0[4:],
@@ -90,6 +90,50 @@ class IntegratorTest(unittest.TestCase):
         L = np.cross(q[:, :2], p[:, :2]) + np.cross(q[:, 2:], p[:, 2:])
         self.assertTrue(np.greater(1e-10 + np.zeros(L.shape[0]), L - L[0]).all())
 
+        ## checking error
+        dts = [dt, 2 * dt, 4 * dt]
+        errors = []
+        ts = []
+        for i in dts:
+            system = GravitationalSystem(r0=x0[:4],
+                                         v0=x0[4:],
+                                         m=np.array([M_E, M_E, M_L, M_L]),
+                                         t=np.linspace(0, T, int(T // i)),
+                                         force=force,
+                                         solver=verlet)
+
+            t, p_t, q_t = system.direct_simulation()
+
+
+            H = np.linalg.norm(p_t[:,:2], axis=1)**2 / (2 * M_E) + np.linalg.norm(p_t[:,2:], axis=1)**2 / (2 * M_L) + \
+            -G * M_S * M_E / np.linalg.norm(q_t[:,:2], axis=1) - G * M_S * M_L / np.linalg.norm(q_t[:,2:], axis=1) + \
+            -G * M_E * M_L / np.linalg.norm(q_t[:,2:] - q_t[:,:2], axis=1)
+
+            errors.append((H - H[0]) / i**2)
+            ts.append(t)
+
+        plt.figure()
+        plt.plot(ts[0], errors[0], label="dt")
+        plt.plot(ts[1], errors[1], linestyle='--', label="2*dt")
+        plt.plot(ts[2], errors[2], linestyle=':', label="4*dt")
+        plt.xlabel("$t$")
+        plt.ylabel("$\delta E(t)/(\Delta t)^2$")
+        plt.legend()
+        plt.show()
+
+        ## checking time reversal: p -> -p
+        x0 = np.concatenate((q[-1, :], -1 * p[-1, :] / np.array([M_E, M_E, M_L, M_L])), axis=0)
+        system = GravitationalSystem(r0=x0[:4],
+                                     v0=x0[4:],
+                                     m=np.array([M_E, M_E, M_L, M_L]),
+                                     t=np.linspace(0, T, int(T // dt)),
+                                     force=force,
+                                     solver=verlet)
+
+        t, p_reverse, q_reverse = system.direct_simulation()
+        self.assertTrue(np.greater(1e-10 + np.zeros(4), q_reverse[-1] - q[0]).all())
+        self.assertTrue(np.greater(1e-10 + np.zeros(4), p_reverse[-1] - p[0]).all())
+
 
 if __name__ == '__main__':
     unittest.main()