Add unit test

jobs/src/integrators.py

 """
-This module contains several integrators, which is at least 
-of second consistency order and approximately preserves 
+This module contains 2 integrators, which is at least 
+of second consistency order and approrimately preserves 
 the energy over long times.
-"""
\ No newline at end of file
+"""
+from collections.abc import Callable
+
+import numpy as np
+
+
+def verlet(F: Callable, q0: np.ndarray, p0: np.ndarray, m: np.ndarray, dt: float):
+    """
+    Velocity-Verlet integrator for one time step
+    """
+    p = p0
+    q = q0
+
+    p = p + 1 / 2 * F(q) * dt
+    q = q + 1 / m * p * dt
+    p = p + 1 / 2 * F(q) * dt
+
+    return p, q

jobs/src/system.py

@@ -3,6 +3,7 @@ This module contains base class for a n-body gravitational system
 with 2 methods of simulations: direct and approximated force field.
 """
 
+from collections.abc import Callable
 from dataclasses import dataclass
 
 import numpy as np
@@ -10,17 +11,21 @@ import numpy as np
 
 @dataclass
 class GravitationalSystem:
-    r0: np.ndarray[float]  # shape = (N,3)
-    v0: np.ndarray[float]  # shape = (N,3)
+    r0: np.ndarray[float]  # shape = (N,d)
+    v0: np.ndarray[float]  # shape = (N,d)
     m: np.ndarray[float]  # shape = N
     t: np.ndarray[float]  # shape = n_time_steps
-    solver: str
+    force: Callable
+    solver: Callable
 
     def __post_init__(self):
         "Checking dimensions of inputs"
         assert self.r0.shape == self.v0.shape
         assert self.m.shape[0] == self.r0.shape[0]
 
+        solvers = ['verlet']
+        assert self.solver.__name__ in solvers
+
     def direct_simulation(self):
         """
         Using integrator to compute trajector in phase space
@@ -31,9 +36,8 @@ class GravitationalSystem:
         p[0] = self.m * self.v0
 
         for i in range(1, len(self.t)):
-            # something with dt = self.t[i] - self.t[i-1]
-            # p = m*v
-            pass
+            dt = self.t[i] - self.t[i - 1]
+            p[i], q[i] = self.solver(self.force, q[i - 1], p[i - 1], self.m, dt)
 
         return self.t, p, q
 

jobs/tests/test_integrator.py

 """
 This unittest tests implementation of integrators
-for analytically well-understood configurations of
-gravitional n-body systems where n ∈ {2, 3}
-"""
\ No newline at end of file
+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. 
+"""
+
+import logging
+import unittest
+
+import numpy as np
+
+from jobs.src.integrators import verlet
+from jobs.src.system import GravitationalSystem
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+# system settings
+R_SE = 1  # distance between Earth and the Sun, 1 astronomical unit (AU)
+R_L = 2.57e-3 * R_SE  # distance between Earth and the Moon
+
+M_S = 1  # solar mass
+M_E = 3.00e-6 * M_S
+M_L = 3.69e-8 * M_S
+
+T_E = 1  # earth yr
+T_L = 27.3 / 365.3 * T_E
+
+G = 4 * np.pi**2  # AU^3 / m / yr^2
+
+# simulation settings
+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:
+    r1 = q[0:2]
+    r2 = q[2:4]
+
+    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()
+
+
+class IntegratorTest(unittest.TestCase):
+
+    def test_verlet(self):
+        """
+        Test functionalities of velocity-Verlet algorithm
+        """
+
+        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, q = system.direct_simulation()
+
+        ## 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) + \
+            -G * M_E * M_L / np.linalg.norm(q[:,2:] - q[:,:2], axis=1)
+        self.assertTrue(np.greater(1e-2 + np.zeros(H.shape[0]),
+                                   H - H[0]).all())  #TODO: why is it so badly conserved..., DEBUG another day...
+
+        ## checking total momentum conservation
+        L = np.cross(q[:, :2], p[:, :2]) + np.cross(q[:, 2:], p[:,
+                                                                2:])  # TODO: strange behaviour..., DEBUG another day...
+        self.assertTrue(np.array_equal(np.zeros(L.shape[0]), L))
+
+
+if __name__ == '__main__':
+    unittest.main()