From 76f6245986c7e6cc78386adfc4f39b66bf76f248 Mon Sep 17 00:00:00 2001
From: jakut77 <jakut77@zedat.fu-berlin.de>
Date: Fri, 23 Feb 2024 00:16:10 +0100
Subject: [PATCH] BHT unit test
---
jobs/src/bht_algorithm.py | 3 -
.../bht.py => jobs/src/bht_algorithm_2d.py | 9 +-
jobs/{unittests => tests}/__init__.py | 0
jobs/{unittests => tests}/test_a_job.py | 0
jobs/tests/test_bht_algorithm.py | 150 ++++++++++++++++++
jobs/tests/test_integrator.py | 139 ++++++++++++++++
.../test_jpl_data_query.py | 0
jobs/unittests/test_bht_algorithm.py | 0
jobs/unittests/test_integrator.py | 5 -
tasks/bht/2D/{bht.py => bht_quad.py} | 0
tasks/bht/{ => 3D}/bht_oct.py | 0
tasks/bht/{ => 3D}/galaxysimulation.py | 0
tasks/bht/{ => 3D}/simulate.py | 0
13 files changed, 294 insertions(+), 12 deletions(-)
delete mode 100644 jobs/src/bht_algorithm.py
rename tasks/bht/bht.py => jobs/src/bht_algorithm_2d.py (98%)
rename jobs/{unittests => tests}/__init__.py (100%)
rename jobs/{unittests => tests}/test_a_job.py (100%)
create mode 100644 jobs/tests/test_bht_algorithm.py
create mode 100644 jobs/tests/test_integrator.py
rename jobs/{unittests => tests}/test_jpl_data_query.py (100%)
delete mode 100644 jobs/unittests/test_bht_algorithm.py
delete mode 100644 jobs/unittests/test_integrator.py
rename tasks/bht/2D/{bht.py => bht_quad.py} (100%)
rename tasks/bht/{ => 3D}/bht_oct.py (100%)
rename tasks/bht/{ => 3D}/galaxysimulation.py (100%)
rename tasks/bht/{ => 3D}/simulate.py (100%)
diff --git a/jobs/src/bht_algorithm.py b/jobs/src/bht_algorithm.py
deleted file mode 100644
index c6630ca..0000000
--- a/jobs/src/bht_algorithm.py
+++ /dev/null
@@ -1,3 +0,0 @@
-"""
-This module implements the Barnes Hut Tree Algorithm
-"""
\ No newline at end of file
diff --git a/tasks/bht/bht.py b/jobs/src/bht_algorithm_2d.py
similarity index 98%
rename from tasks/bht/bht.py
rename to jobs/src/bht_algorithm_2d.py
index 30f345a..071da04 100644
--- a/tasks/bht/bht.py
+++ b/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)
diff --git a/jobs/unittests/__init__.py b/jobs/tests/__init__.py
similarity index 100%
rename from jobs/unittests/__init__.py
rename to jobs/tests/__init__.py
diff --git a/jobs/unittests/test_a_job.py b/jobs/tests/test_a_job.py
similarity index 100%
rename from jobs/unittests/test_a_job.py
rename to jobs/tests/test_a_job.py
diff --git a/jobs/tests/test_bht_algorithm.py b/jobs/tests/test_bht_algorithm.py
new file mode 100644
index 0000000..c25c8ad
--- /dev/null
+++ b/jobs/tests/test_bht_algorithm.py
@@ -0,0 +1,150 @@
+"""
+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()
diff --git a/jobs/tests/test_integrator.py b/jobs/tests/test_integrator.py
new file mode 100644
index 0000000..f772a13
--- /dev/null
+++ b/jobs/tests/test_integrator.py
@@ -0,0 +1,139 @@
+"""
+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()
diff --git a/jobs/unittests/test_jpl_data_query.py b/jobs/tests/test_jpl_data_query.py
similarity index 100%
rename from jobs/unittests/test_jpl_data_query.py
rename to jobs/tests/test_jpl_data_query.py
diff --git a/jobs/unittests/test_bht_algorithm.py b/jobs/unittests/test_bht_algorithm.py
deleted file mode 100644
index e69de29..0000000
diff --git a/jobs/unittests/test_integrator.py b/jobs/unittests/test_integrator.py
deleted file mode 100644
index 6d1ad12..0000000
--- a/jobs/unittests/test_integrator.py
+++ /dev/null
@@ -1,5 +0,0 @@
-"""
-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
diff --git a/tasks/bht/2D/bht.py b/tasks/bht/2D/bht_quad.py
similarity index 100%
rename from tasks/bht/2D/bht.py
rename to tasks/bht/2D/bht_quad.py
diff --git a/tasks/bht/bht_oct.py b/tasks/bht/3D/bht_oct.py
similarity index 100%
rename from tasks/bht/bht_oct.py
rename to tasks/bht/3D/bht_oct.py
diff --git a/tasks/bht/galaxysimulation.py b/tasks/bht/3D/galaxysimulation.py
similarity index 100%
rename from tasks/bht/galaxysimulation.py
rename to tasks/bht/3D/galaxysimulation.py
diff --git a/tasks/bht/simulate.py b/tasks/bht/3D/simulate.py
similarity index 100%
rename from tasks/bht/simulate.py
rename to tasks/bht/3D/simulate.py
--
GitLab