BHT unit test

jobs/src/bht_algorithm.py

-"""
-This module implements the Barnes Hut Tree Algorithm
-"""
\ No newline at end of file

tasks/bht/bht.py → jobs/src/bht_algorithm_2d.py

@@ -5,7 +5,7 @@ class MainApp:
 
     def __init__(self):
         """Initialize the MainApp with a root node."""
-        self.rootNode = TreeNode(x=0, y=0, width=100, height=100)
+        self.rootNode = TreeNode(x=-1, y=-1, width=2, height=2)
 
     def BuildTree(self, particles):
         """Build the Quadtree by inserting particles.
@@ -19,7 +19,7 @@ class MainApp:
 
     def ResetTree(self):
         """Reset the Quadtree by reinitializing the root node."""
-        self.rootNode = TreeNode(x=0, y=0, width=100, height=100)
+        self.rootNode = TreeNode(x=-1, y=-1, width=2, height=2)
 
 
 class Particle:
@@ -263,11 +263,12 @@ class TreeNode:
         np.ndarray: The gravitational force vector between particle1 and particle2.
         """
         #G = 6.674 * (10 ** -11)  # Gravitational constant
-        G = 1
+        #G = 1
+        G = 4 * np.pi**2  # AU^3 / m / yr^2
 
         dx = particle2.x - particle1.x
         dy = particle2.y - particle1.y
-        cutoff_radius = 5
+        cutoff_radius = 0
         r = max(np.sqrt(dx**2 + dy**2), cutoff_radius)
 
         force_magnitude = G * particle1.mass * particle2.mass / (r**2)

jobs/tests/test_bht_algorithm.py

+"""
+This unittest tests implementation of the BHT algorithm
+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 matplotlib.pyplot as plt
+import numpy as np
+
+from jobs.src.integrators import verlet
+from jobs.src.system import GravitationalSystem
+from jobs.src.bht_algorithm_2d import MainApp, Particle
+
+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
+
+
+# force computation
+def force(q: np.ndarray) -> np.ndarray:
+    r1 = q[0:2]  # Earth coordinates
+    r2 = q[2:4]  # Moon coordinates
+
+    sun = Particle(0, 0, M_S)
+    earth = Particle(r1[0], r1[1], M_E)
+    moon = Particle(r2[0], r2[1], M_L)
+
+    particles = [sun, earth, moon]
+
+    barnes_hut = MainApp()  # Initialize Barnes-Hut algorithm instance
+    barnes_hut.BuildTree(particles)  # Build the Barnes-Hut tree with particles
+    barnes_hut.rootNode.ComputeMassDistribution()  #Compute the center of mass of the tree nodes
+
+    f_earth = barnes_hut.rootNode.CalculateForceFromTree(earth)
+    f_moon = barnes_hut.rootNode.CalculateForceFromTree(moon)
+
+    return np.concatenate((f_earth, f_moon), axis=0)
+
+
+class IntegratorTest(unittest.TestCase):
+
+    def test_verlet(self):
+        """
+        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:],
+                                     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())
+
+        ## 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()

jobs/tests/test_integrator.py

+"""
+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 matplotlib.pyplot as plt
+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
+
+
+# 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
+        """
+        # 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:],
+                                     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())
+
+        ## 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()

jobs/unittests/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