""" This unittest tests implementation of integrators 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-10 + np.zeros(H.shape[0]), H - H[0]).all()) ## checking total momentum conservation 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()) if __name__ == '__main__': unittest.main()