modify leapfrog to take velocities instead of momenta, add new initial conditions

lib/integrator.py

@@ -14,18 +14,18 @@ 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 dp_dt(m_array, q_array):
+def dv_dt(m_array, q_array):
     """
-    Time derivative of the momentum, given by the position derivative of the Hamiltonian.
-    dp/dt = -dH/dq
+    Time derivative of the velocity, given by the position derivative of the Hamiltonian.
+    dv/dt = -1/m*dH/dq
     """
-    dp_array = np.zeros(q_array.shape)
+    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)
-        dp_array[i] = -G*m_array[i]*np.sum(m_j/np.sum(np.sqrt(np.sum((q_j-q_array[i])**2, axis=0)))**3*(q_j-q_array[i]), axis=0)
-    dp_array[np.isnan(dp_array)] = 0.
-    return dp_array
+        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):
     """
@@ -34,10 +34,10 @@ def frogleap(duration, step, dyn_syst, recover_param=False, display=False):
     """
     N = np.ceil(duration/step).astype(int)
     q_array = dyn_syst.get_positions()
-    p_array = dyn_syst.get_momenta()
+    v_array = dyn_syst.get_velocities()
     masses = dyn_syst.get_masses()
-    m_array = np.ones(p_array.shape)
-    for i in range(p_array.shape[0]):
+    m_array = np.ones(q_array.shape)
+    for i in range(q_array.shape[0]):
         m_array[i,:] = masses[i]
     
     E = np.zeros(N)
@@ -49,32 +49,31 @@ def frogleap(duration, step, dyn_syst, recover_param=False, display=False):
         except IOError:
             system("rm tmp/*")
         d = DynamicUpdate()
-        d.min_x, d.max_x = -1.5*np.abs(q_array).max(), +1.5*np.abs(q_array).max()
         d.on_launch()
     for j in range(N):
         # half-step drift
-        q_array, p_array = q_array + step/2*p_array/m_array , p_array
+        q_array, v_array = q_array + step/2*v_array , v_array
         # full-step kick
-        q_array, p_array = q_array , p_array - step*dp_dt(m_array, q_array)
+        q_array, v_array = q_array , v_array - step*dv_dt(m_array, q_array)
         # half-step drift
-        q_array, p_array = q_array + step/2*p_array/m_array , p_array
+        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.p = p_array[i]
-            if body.m != 0.:
-                body.v = body.p/body.m
+            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:
-            # In center of mass frame
-            q_cm = np.array([0,0,0])#np.sum(m_array*q_array, axis=0)/masses.sum()
             # display progression
-            d.on_running(q_array[:,0]-q_cm[0], q_array[:,1]-q_cm[1], q_array[:,2]-q_cm[2], step=j, label="step {0:d}/{1:d}".format(j,N))
-            time.sleep(1e-4)
+            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))
+            time.sleep(1e-5)
     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")

lib/objects.py

@@ -34,6 +34,9 @@ class System:
     def get_positions(self): #return the positions of the bodies
         return np.array([body.q for body in self.bodylist])
     
+    def get_velocities(self): #return the positions of the bodies
+        return np.array([body.v for body in self.bodylist])
+    
     def get_momenta(self): #return the momenta of the bodies
         return np.array([body.p for body in self.bodylist])
 

lib/plots.py

@@ -14,6 +14,11 @@ class DynamicUpdate():
 
     plt.ion()
 
+    def set_lims(self, factor=1.5):
+        self.ax.set_xlim(factor*self.min_x, factor*self.max_x)
+        self.ax.set_ylim(factor*self.min_x, factor*self.max_x)
+        self.ax.set_zlim(factor*self.min_x, factor*self.max_x)
+
     def on_launch(self):
         #Set up plot
         self.fig = plt.figure()
@@ -21,27 +26,27 @@ class DynamicUpdate():
         self.lines, = self.ax.plot([],[],[],'o')
         #Autoscale on unknown axis and known lims on the other
         self.ax.set_autoscaley_on(True)
-        self.ax.set_xlim(self.min_x, self.max_x)
-        self.ax.set_ylim(self.min_x, self.max_x)
-        self.ax.set_zlim(self.min_x, self.max_x)
+        self.set_lims()
         #Other stuff
         self.ax.grid()
         #self.ax.set_aspect('equal')
 
     def on_running(self, xdata, ydata, zdata, step=None, label=None):
+        values = np.sqrt(np.sum((np.array((xdata,ydata,zdata))**2).T,axis=1))
+        self.min_x, self.max_x = -np.abs(values).max(), np.abs(values).max()
+        self.set_lims()
         #Update data (with the new _and_ the old points)
+        self.lines.set_data_3d(xdata, ydata, zdata)
+        if not label is None:
+            self.ax.set_title(label)
+        #Need both of these in order to rescale
+        self.ax.relim()
+        self.ax.autoscale_view()
+        #We need to draw *and* flush
+        self.fig.canvas.draw()
+        self.fig.canvas.flush_events()
         if not step is None and step%100==0:
-            self.lines.set_data_3d(xdata, ydata, zdata)
-            if not label is None:
-                self.ax.set_title(label)
-            #Need both of these in order to rescale
-            self.ax.relim()
-            self.ax.autoscale_view()
-            #We need to draw *and* flush
-            self.fig.canvas.draw()
-            self.fig.canvas.flush_events()
-            if not step is None and step%100==0:
-                self.fig.savefig("tmp/{0:05d}.png".format(step),bbox_inches="tight")
+            self.fig.savefig("tmp/{0:05d}.png".format(step),bbox_inches="tight")
 
     #Example
     def __call__(self):

main.py

@@ -12,7 +12,8 @@ globals()['au'] = 1.5e11 #Astronomical unit in m
 
 def main():
     #initialisation
-    m = np.array([1., 1., 0.1])*Ms  # Masses in Solar mass
+    m = np.array([1, 1, 0.1])*Ms  # Masses in Solar mass
+    mu = m[0]*m[1]/(m[0]+m[1])
     a = np.array([1., 1., 5.])*au   # Semi-major axis in astronomical units
     psi = np.array([0., 0., 80.])*np.pi/180.    # Inclination of the orbital plane in degrees
 
@@ -21,9 +22,9 @@ def main():
     x3 = np.array([np.cos(psi[2]), 0., np.sin(psi[2])])*a[2]
     q = np.array([x1, x2, x3])
 
-    v1 = np.array([0, -np.sqrt(G*Ms/np.sqrt(np.sum(x1**2))), 0])
-    v2 = np.array([0, np.sqrt(G*Ms/np.sqrt(np.sum(x2**2))), 0])
-    v3 = np.array([0, np.sqrt(G*Ms*(2./np.sqrt(np.sum(x3**2))-1./a[2])), 0])
+    v1 = np.array([0., -np.sqrt(G*mu/np.sqrt(np.sum(x1**2))), 0])
+    v2 = np.array([0., np.sqrt(G*mu/np.sqrt(np.sum(x2**2))), 0.])
+    v3 = np.array([0., np.sqrt(G*(m[0]+m[1])*(2./np.sqrt(np.sum(x3**2))-1./a[2])), 0.])
     v = np.array([v1, v2, v3])
 
     bodylist = []
@@ -32,8 +33,9 @@ def main():
     dyn_syst = System(bodylist)
     dyn_syst.COMShift()
 
-    duration, step = 100, 0.01
-    E, L = frogleap(duration, step, dyn_syst, recover_param=True, display=True)
+    duration, step = 0.5*3e7, 1e1
+    E, L = frogleap(duration, step, dyn_syst, recover_param=True)#, display=True)
+    
     fig1 = plt.figure(figsize=(30,15))
     ax1 = fig1.add_subplot(111)
     ax1.plot(np.arange(E.shape[0])/duration, E, label=r"$E_m$")
@@ -45,6 +47,7 @@ def main():
     ax2.legend()
     fig2.savefig("plots/L2.png",bbox_inches="tight")
     plt.show(block=True)
+    
     return 0
 
 if __name__ == '__main__':