implement leapfrog and stormer_verlet

jobs/src/integrators.py

@@ -20,3 +20,54 @@ def verlet(F: Callable, q0: np.ndarray, p0: np.ndarray, m: np.ndarray, dt: float
     p = p + 1 / 2 * F(q) * dt
 
     return p, q
+
+
+
+def leapfrog(F: Callable, q0: np.ndarray, p0: np.ndarray, m: np.ndarray, dt: float):
+    '''
+        # using leapfrog integration. https://en.wikipedia.org/wiki/Leapfrog_integration
+        The basic idea of the Leapfrog method is to update the position
+        and momentum of a particle or system at half-integer time steps.
+        The algorithm proceeds in discrete steps, and the update scheme looks like this:
+        1. Update Velocity at Half Step:
+        2. Update Position:
+        3. Update Velocity at Full Step:
+    '''
+    p = p0
+    q = q0
+
+    # 1. Update Momentum at Half Step:
+    p = p + dt / 2 * F(q)
+
+    # 2. Update Position:
+    q = q + p / m * dt
+
+    # 3. Update Momentum at Full Step:
+    p = p + dt / 2 * F(q)
+
+    # print("-------------------")
+    # print(f'accelerations={accelerations}')
+    # print(f'velocity={velocity}')
+    return p, q
+
+
+
+def stormer_verlet(F: Callable, q0: np.ndarray, p0: np.ndarray, m: np.ndarray, dt: float):
+    """
+      stormer_verlet integrator for one time step
+      # https://www.physics.udel.edu/~bnikolic/teaching/phys660/numerical_ode/node5.html
+      # https://en.wikipedia.org/wiki/Verlet_integration#Velocity_Verlet
+      # https://www.algorithm-archive.org/contents/verlet_integration/verlet_integration.html
+    """
+    p = p0
+    q = q0
+
+    # 1. Update Position at Half Step:
+    acc = F(q) / m
+    q = q + p / m * dt + acc * (dt ** 2 * 0.5)
+
+    # 2. Update Momentum:
+    p = p + acc * (dt * 0.5) * m + F(q) * (dt * 0.5)
+
+    return p, q
+

jobs/src/system.py

@@ -23,7 +23,7 @@ class GravitationalSystem:
         assert self.r0.shape == self.v0.shape
         assert self.m.shape[0] == self.r0.shape[0]
 
-        solvers = ['verlet']
+        solvers = ['verlet', "stormer_verlet", "leapfrog"]
         assert self.solver.__name__ in solvers
 
     def direct_simulation(self):

jobs/tests/test_integrator.py

@@ -12,7 +12,7 @@ import unittest
 import matplotlib.pyplot as plt
 import numpy as np
 
-from jobs.src.integrators import verlet
+from jobs.src.integrators import verlet, leapfrog, stormer_verlet
 from jobs.src.system import GravitationalSystem
 
 logger = logging.getLogger(__name__)
@@ -134,6 +134,168 @@ class IntegratorTest(unittest.TestCase):
         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())
 
+    def test_leapfrog(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=leapfrog)
+
+        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=leapfrog)
+
+            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=leapfrog)
+
+        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())
+
+    def test_stormer_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=stormer_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=stormer_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=stormer_verlet)
 if __name__ == '__main__':
     unittest.main()