diff --git a/jobs/src/integrators.py b/jobs/src/integrators.py
index 22adaac6b21c8e2ea88272a37c5078b4d8c389a8..35d513a468772b3a936e7dff1a2f15a22e018f44 100644
--- a/jobs/src/integrators.py
+++ b/jobs/src/integrators.py
@@ -1,5 +1,22 @@
"""
-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
diff --git a/jobs/src/system.py b/jobs/src/system.py
index 3c123c9dfd4f0ea5db6a57c9edb86eb5f4bf30c1..264dcfdf22572b1208902fe034f99542928a374c 100644
--- a/jobs/src/system.py
+++ b/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
diff --git a/jobs/tests/test_integrator.py b/jobs/tests/test_integrator.py
index 6d1ad12fa4ac8fca09efaa6d638f6b65c763b8ba..3192c19b2163ad2e26d21807da12792c631605ac 100644
--- a/jobs/tests/test_integrator.py
+++ b/jobs/tests/test_integrator.py
@@ -1,5 +1,92 @@
"""
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()