Newer
Older
"""
This unittest tests implementation of integrators
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
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],
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 linear momentum conservation
P = p[:, :2] + p[:, 2:]
self.assertTrue(np.greater(1e-10 + np.zeros(P[0].shape), P - P[0]).all())
## checking total angular 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())