make leapfrog integrator as System method

lib/integrator.py

@@ -1,82 +0,0 @@
-#!/usr/bin/python
-# -*- coding:utf-8 -*-
-"""
-Implementation of the various integrators for numerical integration.
-
-Comes from the assumption that the problem is analytically defined in position-momentum (q-p) space for a given hamiltonian H.
-"""
-from os import system
-import numpy as np
-from lib.plots import DynamicUpdate
-
-globals()['G'] = 6.67e-11 #Gravitational constant in SI units
-globals()['Ms'] = 2e30 #Solar mass in kg
-globals()['au'] = 1.5e11 #Astronomical unit in m
-
-def dv_dt(m_array, q_array):
-    """
-    Time derivative of the velocity, given by the position derivative of the Hamiltonian.
-    dv/dt = -1/m*dH/dq
-    """
-    dv_array = np.zeros(q_array.shape)
-    for i in range(q_array.shape[0]):
-        q_j = np.delete(q_array, i, 0)
-        m_j = np.delete(m_array, i, 0)
-        dv_array[i] = -G*np.sum((m_j*(q_j-q_array[i])).T/np.sqrt(np.sum((q_j-q_array[i])**2, axis=1))**3, axis=1).T
-    dv_array[np.isnan(dv_array)] = 0.
-    return dv_array
-
-
-def frogleap(duration, step, dyn_syst, recover_param=False, display=False):
-    """
-    Leapfrog integrator for first order partial differential equations.
-    iteration : half-step drift -> full-step kick -> half-step drift
-    """
-    N = np.ceil(duration/step).astype(int)
-    q_array = dyn_syst.get_positions()
-    v_array = dyn_syst.get_velocities()
-    masses = dyn_syst.get_masses()
-    m_array = np.ones(q_array.shape)
-    for i in range(q_array.shape[0]):
-        m_array[i,:] = masses[i]
-    
-    E = np.zeros(N)
-    L = np.zeros((N,3))
-    
-    if display:
-        try:
-            system("mkdir tmp")
-        except IOError:
-            system("rm tmp/*")
-        d = DynamicUpdate()
-        d.on_launch()
-    for j in range(N):
-        # half-step drift
-        q_array, v_array = q_array + step/2*v_array , v_array
-        # full-step kick
-        q_array, v_array = q_array , v_array - step*dv_dt(m_array, q_array)
-        # half-step drift
-        q_array, v_array = q_array + step/2*v_array , v_array
-        
-        for i, body in enumerate(dyn_syst.bodylist):
-            body.q = q_array[i]
-            body.v = v_array[i]
-            body.p = body.v*body.m
-        dyn_syst.COMShift()
-        
-        E[j] = dyn_syst.Eval()
-        L[j] = dyn_syst.Lval()
-
-        if display:
-            # display progression
-            if len(dyn_syst.bodylist) == 1:
-                d.on_running(q_array[0], q_array[1], q_array[2], step=j, label="step {0:d}/{1:d}".format(j,N))
-            else:
-                d.on_running(q_array[:,0], q_array[:,1], q_array[:,2], step=j, label="step {0:d}/{1:d}".format(j,N))
-
-    if display:
-        system("convert -delay 5 -loop 0 tmp/?????.png tmp/temp.gif && rm tmp/?????.png")
-        system("convert tmp/temp.gif -fuzz 30% -layers Optimize plots/dynsyst.gif && rm tmp/temp.gif")
-        
-    if recover_param:
-        return E, L

lib/objects.py

@@ -91,12 +91,61 @@ class System:
                     rij = np.linalg.norm(body.q-otherbody.q)
                     W = W - G*body.m*otherbody.m/rij
         return T + W
-    
-    def __repr__(self): # Called upon "print(system)"
-        return str([print(body) for body in self.bodylist])
 
-    def __str__(self): # Called upon "str(system)"
-        return str([str(body) for body in self.bodylist])
+    def Drift(self, dt):
+        for body in self.bodylist:
+            body.q = body.q + dt*body.v
+    
+    def Kick(self, dt):
+        for body in self.bodylist:
+            body.a = np.zeros(3)
+            for otherbody in self.bodylist:
+                if body != otherbody:
+                    rij = np.linalg.norm(body.q-otherbody.q)
+                    body.a = body.a - (body.q-otherbody.q)*G*otherbody.m/(rij**3)
+            body.v = body.v + dt*body.a
+        
+    def LP(self, dt):
+        self.COMShift()
+        self.Drift(dt/2)
+        self.Kick(dt)
+        self.Drift(dt/2)
+        self.time = self.time + dt
+        for body in self.bodylist:
+            body.p = body.v*body.m
+    
+    def leapfrog(self, duration, dt, recover_param=False, display=False):
+        if display:
+            try:
+                system("mkdir tmp")
+            except IOError:
+                system("rm tmp/*")
+            d = DynamicUpdate()
+            d.on_launch()
+
+        N = np.ceil(duration/dt).astype(int)
+        E = np.zeros(N)
+        L = np.zeros((N,3))
+        for j in range(N):
+            self.LP(dt)
+
+            E[j] = self.Eval()
+            L[j] = self.Lval()
+
+            if display and j%100==0:
+                # display progression
+                q_array = self.get_positions()
+                if len(self.bodylist) == 1:
+                    d.on_running(q_array[0], q_array[1], q_array[2], step=j, label="step {0:d}/{1:d}".format(j,N))
+                else:
+                    d.on_running(q_array[:,0], q_array[:,1], q_array[:,2], step=j, label="step {0:d}/{1:d}".format(j,N))
+        if display:
+            system("convert -delay 5 -loop 0 tmp/??????.png tmp/temp.gif && rm tmp/??????.png")
+            system("convert tmp/temp.gif -fuzz 30% -layers Optimize plots/dynsyst.gif && rm tmp/temp.gif")
+     
+        if recover_param:
+            return E, L
+         
 
     def Update_a(self): #update acceleration of bodies in system
         for body in self.bodylist:
@@ -105,7 +154,6 @@ class System:
                 if body != otherbody:
                     rij = np.linalg.norm(body.q-otherbody.q)
                     body.a = body.a - (body.q-otherbody.q)*G*otherbody.m/(rij**3)
-        return 1
 
     def Update_j(self): #update jerk of bodies in system
         for body in self.bodylist:
@@ -117,14 +165,11 @@ class System:
                     deltar = (body.q-otherbody.q)
                     vr = deltav + 3.*deltar*np.inner(deltav,deltar)/(rij**2)
                     body.j = body.j - G*otherbody.m/(rij**3)*vr
-        return 1
 
     def Predict(self,dt):  # update predicted position and velocities of bodies in system
         for body in self.bodylist:
             body.qp = body.q +dt*body.v+((dt**2)*body.a/2.)+((dt**3)*body.j/6.)
             body.vp = body.v + dt*body.a + ((dt**2)*body.j/2.)
-            #print("v=",body.v," vp=" ,body.vp)
-        return 1
 
     def Update_ap(self): #update acceleration of bodies in system
         for body in self.bodylist:
@@ -133,7 +178,6 @@ class System:
                 if body != otherbody:
                     rij = np.linalg.norm(body.qp-otherbody.qp)
                     body.ap = body.ap - (body.qp-otherbody.qp)*G*otherbody.m/(rij**3)
-        return 1
 
     def Update_jp(self): #update jerk of bodies in system
         for body in self.bodylist:
@@ -145,7 +189,6 @@ class System:
                     deltar = (body.qp-otherbody.qp)
                     vr = deltav + 3.*deltar*np.inner(deltav,deltar)/(rij**2)
                     body.jp = body.jp - G*otherbody.m/(rij**3)*vr
-        return 1
 
     def Correct(self,dt):  # correct position and velocities of bodies in system
         for body in self.bodylist:
@@ -154,7 +197,6 @@ class System:
 
             body.q = body.qp +((dt**4)*a2/24.) + ((dt**5)*a3/120.)
             body.v = body.vp +((dt**3)*a2/6.) + ((dt**4)*a3/24.)
-        return 1
 
     def HPC(self, dt):  # update position and velocities of bodies in system with hermite predictor corrector
         self.COMShift()
@@ -186,14 +228,23 @@ class System:
             E[j] = self.Eval()
             L[j] = self.Lval()
 
-            if display:
+            if display and j%100==0:
                 # display progression
                 q_array = self.get_positions()
                 if len(self.bodylist) == 1:
                     d.on_running(q_array[0], q_array[1], q_array[2], step=j, label="step {0:d}/{1:d}".format(j,N))
                 else:
                     d.on_running(q_array[:,0], q_array[:,1], q_array[:,2], step=j, label="step {0:d}/{1:d}".format(j,N))
-        
+        if display:
+            system("convert -delay 5 -loop 0 tmp/??????.png tmp/temp.gif && rm tmp/??????.png")
+            system("convert tmp/temp.gif -fuzz 30% -layers Optimize plots/dynsyst.gif && rm tmp/temp.gif")
+     
         if recover_param:
             return E, L
+ 
+    def __repr__(self): # Called upon "print(system)"
+        return str([print(body) for body in self.bodylist])
+
+    def __str__(self): # Called upon "str(system)"
+        return str([str(body) for body in self.bodylist])
 

lib/plots.py

@@ -21,7 +21,7 @@ class DynamicUpdate():
 
     def on_launch(self):
         #Set up plot
-        self.fig = plt.figure()
+        self.fig = plt.figure(figsize=(10,10))
         self.ax = self.fig.add_subplot(projection='3d')
         self.lines, = self.ax.plot([],[],[],'o')
         #Autoscale on unknown axis and known lims on the other
@@ -45,8 +45,8 @@ class DynamicUpdate():
         #We need to draw *and* flush
         self.fig.canvas.draw()
         self.fig.canvas.flush_events()
-        if not step is None and step%10==0:
-            self.fig.savefig("tmp/{0:05d}.png".format(step),bbox_inches="tight")
+        if not step is None and step%1000==0:
+            self.fig.savefig("tmp/{0:06d}.png".format(step),bbox_inches="tight")
 
     #Example
     def __call__(self):

main.py

@@ -3,7 +3,6 @@
 from sys import exit as sysexit
 import numpy as np
 import matplotlib.pyplot as plt
-from lib.integrator import frogleap
 from lib.objects import Body, System
 
 globals()['G'] = 6.67e-11 #Gravitational constant in SI units
@@ -28,15 +27,15 @@ def main():
     v = np.array([v1, v2, v3])
 
     bodylist = []
-    for i in range(2):
+    for i in range(3):
         bodylist.append(Body(m[i], q[i], v[i]))
     dyn_syst = System(bodylist)
     dyn_syst.COMShift()
 
     duration, step = 100*3e7, 1e4
-    #E, L = frogleap(duration, step, dyn_syst, recover_param=True)#, display=True)
-    E, L = dyn_syst.hermite(duration,step, recover_param=True)#, display=True)
-    
+    E, L = dyn_syst.leapfrog(duration, step, recover_param=True, display=True)
+    #E, L = dyn_syst.hermite(duration,step, recover_param=True, display=True)
+    plt.close()
     fig1 = plt.figure(figsize=(30,15))
     ax1 = fig1.add_subplot(111)
     ax1.plot(np.arange(E.shape[0])/duration, E, label=r"$E_m$")