general debug and display angular momentum and mechanical energy curves

lib/integrator.py

@@ -34,7 +34,7 @@ def frogleap(duration, step, dyn_syst, recover_param=False, display=False):
     q_array = dyn_syst.get_positions()
     p_array = dyn_syst.get_momenta()
     E = np.zeros(N)
-    L = np.zeros(N)
+    L = np.zeros((N,3))
     
     if display:
         d = DynamicUpdate()
@@ -53,13 +53,14 @@ def frogleap(duration, step, dyn_syst, recover_param=False, display=False):
             body.p = p_array[i]
             if body.m != 0.:
                 body.v = body.p/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.])#np.sum(m_array.reshape((q_array.shape[0],1))*q_array, axis=0)/m_array.sum()
+            q_cm = 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_cm[0], q_array[:,1]-q_cm[1])
             time.sleep(0.01)

lib/objects.py

@@ -55,9 +55,8 @@ class System:
         return coord
 
     def COMShift(self): #Shift coordinates of bodies in system to COM frame and set COM at rest
-        comcoord = self.COM
         for body in self.bodylist:
-            body.q = body.q-self.COM()
+            body.q = body.q - self.COM()
             body.p = body.p - self.COMV()
         return 0
 

main.py

@@ -2,20 +2,21 @@
 # -*- coding:utf-8 -*-
 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
 
 def main():
     #initialisation
-    m = np.array([1e10, 1, 0])
+    m = np.array([1e5, 1e5, 0.1])
 
-    x1 = np.array([0, 0, 0])
+    x1 = np.array([-1, 0, 0])
     x2 = np.array([1, 0, 0])
-    x3 = np.array([10, 0, 0])
+    x3 = np.array([100, 0, 0])
     q = np.array([x1, x2, x3])
 
     v1 = np.array([0, 0, 0])
-    v2 = np.array([0, 0, 0])
+    v2 = np.array([0, 1, 0])
     v3 = np.array([0, 0, 0])
     v = np.array([v1, v2, v3])
 
@@ -23,9 +24,18 @@ def main():
     for i in range(3):
         bodylist.append(Body(m[i], q[i], v[i]))
     dyn_syst = System(bodylist)
+    dyn_syst.COMShift()
 
-    new_dyn_syst = frogleap(10, 0.1, dyn_syst, display=False)
-    print(np.all(new_dyn_syst.get_positions() == dyn_syst.get_positions()))
+    E, L = frogleap(10, 0.01, dyn_syst, recover_param=True, display=True)
+    fig1 = plt.figure()
+    ax1 = fig1.add_subplot(111)
+    ax1.plot(np.arange(E.shape[0]), E, label=r"$E_m$")
+    ax1.legend()
+    fig2 = plt.figure()
+    ax2 = fig2.add_subplot(111)
+    ax2.plot(np.arange(L.shape[0]), np.sum(L**2,axis=1), label=r"$L^2$")
+    ax2.legend()
+    plt.show()
     return 0
 
 if __name__ == '__main__':