Merge remote-tracking branch 'origin/main' into main

LeapFrog_1.0.py

@@ -1,20 +0,0 @@
-from math import *
-import numpy as np
-
-def main():
-
-    #initialisation
-    x1 = np.array([0, 0, 0])
-    x2 = np.array([1, 0, 0])
-    x3 = np.array([1, 0, 0])
-
-    v1 = np.array([0, 0, 0])
-    v2 = np.array([1, 0, 0])
-    v3 = np.array([1, 0, 0])
-
-    t = 0
-    dt = 0.1
-
-
-if __name__ == '__main__':
-    main()

lib/integrator.py

@@ -9,6 +9,8 @@ import numpy as np
 import time
 from lib.plots import DynamicUpdate
 
+globals()["G"] = 1. #Gravitationnal constant
+
 def dp_dt(m_array, q_array):
     """
     Time derivative of the momentum, given by the position derivative of the Hamiltonian.
@@ -18,17 +20,20 @@ def dp_dt(m_array, q_array):
     for i in range(q_array.shape[0]):
         q_j = np.delete(q_array, i, 0)
         m_j = np.delete(m_array, i).reshape((q_j.shape[0],1))
-        dp_array[i] = -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[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.
-    print(dp_array)
     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()
@@ -41,14 +46,17 @@ def frogleap(duration, step, m_array, q_array, p_array, display=False):
         # half-step drift
         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)
+        
+        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 q_array, p_array
+    return dyn_syst

lib/objects.py

@@ -24,6 +24,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.bodylist])
+    
+    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
@@ -65,6 +74,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
+from lib.objects import Body, System
 
 def main():
     #initialisation
-    m = np.array([1e5, 1, 1])
+    m = np.array([1e10, 1, 0])
 
     x1 = np.array([0, 0, 0])
     x2 = np.array([1, 0, 0])
@@ -16,9 +17,14 @@ def main():
     v1 = np.array([0, 0, 0])
     v2 = np.array([0, 0, 0])
     v3 = np.array([0, 0, 0])
-    p = m*np.array([v1, v2, v3])
+    v = np.array([v1, v2, v3])
 
-    q, p = frogleap(10, 0.01, m, q, p, display=True)
+    bodylist = []
+    for i in range(3):
+        bodylist.append(Body(m[i], q[i], v[i]))
+    dyn_syst = System(bodylist)
+
+    new_dyn_syst = frogleap(10, 0.01, dyn_syst, display=True)
     return 0
 
 if __name__ == '__main__':