Add get masses/positions/momenta methods and modify frogleap input/output

lib/integrator.py

@@ -24,12 +24,16 @@ def dp_dt(m_array, q_array):
     dp_array[np.isnan(dp_array)] = 0.
     return dp_array
 
-def frogleap(duration, step, m_array, q_array, p_array, display=False):
+def frogleap(duration, step, dyn_syst, 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)
+    m_array = dyn_syst.get_masses()
+    q_array = dyn_syst.get_positions()
+    p_array = dyn_syst.get_momenta()
+    
     if display:
         d = DynamicUpdate()
         d.min_x, d.max_x = -1.5*np.abs(q_array).max(), +1.5*np.abs(q_array).max()
@@ -43,13 +47,16 @@ def frogleap(duration, step, m_array, q_array, p_array, display=False):
         q_array, p_array = q_array + step/2*p_array/m_array , p_array
         #print(p_array)
 
-        # In center of mass frame
-        q_cm = np.sum(m_array.reshape((q_array.shape[0],1))*q_array, axis=0)/m_array.sum()
-        q_array -= q_cm
-
         if display:
+            # In center of mass frame
+            q_cm = np.array([0.,0.])#np.sum(m_array.reshape((q_array.shape[0],1))*q_array, axis=0)/m_array.sum()
             # display progression
-            d.on_running(q_array[:,0], q_array[:,1])
+            d.on_running(q_array[:,0]-q_cm[0], q_array[:,1]-q_cm[1])
             time.sleep(0.01)
         
-    return q_array, p_array
+        for i, body in enumerate(dyn_syst.bodylist):
+            body.q = q_array[i]
+            body.p = p_array[i]
+            body.v = body.p/body.m
+
+    return dyn_syst

lib/objects.py

@@ -25,6 +25,15 @@ class System:
     def __init__(self, bodylist):
         self.bodylist = bodylist
     
+    def get_masses(self): #return the masses of each object
+        return np.array([body.m for body in self.bodylit])
+    
+    def get_positions(self): #return the positions of the bodies
+        return np.array([body.q 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])
+
     def Mass(self): #return total system mass
         mass = 0
         for body in self.bodylist:
@@ -51,6 +60,7 @@ class System:
             L = L + np.cross(comq[i],body.p)
             i = i+1
         return L
+
     def Eval(self,Lbodylist): #return total energy of bodies in bodylist
         G = 1. #Gravitational constant (here normalized)
         T = 0

main.py

@@ -3,10 +3,11 @@
 from sys import exit as sysexit
 import numpy as np
 from lib.integrator import frogleap
+import lib.objects
 
 def main():
     #initialisation
-    m = np.array([1e10, 1, 1])
+    m = np.array([1e10, 1, 0])
 
     x1 = np.array([0, 0, 0])
     x2 = np.array([1, 0, 0])