modify initialisation parameters and add gif animation saving

.gitignore

@@ -1,3 +1,6 @@
+# Data temp directory
+tmp/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

lib/integrator.py

@@ -5,8 +5,9 @@ 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.
 """
-import numpy as np
+from os import system
 import time
+import numpy as np
 from lib.plots import DynamicUpdate
 
 globals()["G"] = 1. #Gravitationnal constant
@@ -19,7 +20,7 @@ def dp_dt(m_array, q_array):
     dp_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).reshape((q_j.shape[0],1))
+        m_j = np.delete(m_array, i, 0)#.reshape((q_j.shape[0],1))
         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
@@ -30,9 +31,13 @@ def frogleap(duration, step, dyn_syst, recover_param=False, display=False):
     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()
+    masses = dyn_syst.get_masses()
+    m_array = np.ones(p_array.shape)
+    for i in range(p_array.shape[0]):
+        m_array[i,:] = masses[i]
+    
     E = np.zeros(N)
     L = np.zeros((N,3))
     
@@ -60,10 +65,13 @@ def frogleap(duration, step, dyn_syst, recover_param=False, display=False):
 
         if display:
             # 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_cm = np.array([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])
-            time.sleep(0.01)
+            d.on_running(q_array[:,0]-q_cm[0], q_array[:,1]-q_cm[1], step=j, label="step {0:d}/{1:d}".format(j,N))
+            time.sleep(1e-4)
+    if display:
+        system("convert -delay 5 -loop 0 tmp/????.png tmp/temp.gif && rm tmp/?????.png")
+        system("convert tmp/temp.gif -fuzz 10% -layers Optimize dynsyst.gif && rm tmp/temp.gif")
         
     if recover_param:
         return E, L

lib/objects.py

@@ -78,6 +78,12 @@ class System:
                     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])
+
 
 if __name__ == "__main__":
     # initialisation mass

lib/plots.py

@@ -4,6 +4,7 @@
 Implementation of the plotting and visualization functions.
 """
 import numpy as np
+import time
 import matplotlib.pyplot as plt
 
 class DynamicUpdate():
@@ -19,21 +20,26 @@ 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_xlim(self.min_x, self.max_x)
+        self.ax.set_ylim(self.min_x, self.max_x)
         #Other stuff
         self.ax.grid()
+        self.ax.set_aspect('equal')
 
-    def on_running(self, xdata, ydata):
+    def on_running(self, xdata, ydata, step=None, label=None):
         #Update data (with the new _and_ the old points)
         self.lines.set_xdata(xdata)
         self.lines.set_ydata(ydata)
+        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.figure.canvas.draw()
         self.figure.canvas.flush_events()
+        if not step is None and step%10==0:
+            self.figure.savefig("tmp/{0:05d}.png".format(step),bbox_inches="tight")
 
     #Example
     def __call__(self):

main.py

@@ -8,34 +8,35 @@ from lib.objects import Body, System
 
 def main():
     #initialisation
-    m = np.array([1e5, 1e5, 0.1])
+    m = np.array([1, 1, 1e-5])
 
     x1 = np.array([-1, 0, 0])
     x2 = np.array([1, 0, 0])
     x3 = np.array([100, 0, 0])
     q = np.array([x1, x2, x3])
 
-    v1 = np.array([0, 0, 0])
-    v2 = np.array([0, 1, 0])
+    v1 = np.array([0, -0.35, 0])
+    v2 = np.array([0, 0.35, 0])
     v3 = np.array([0, 0, 0])
     v = np.array([v1, v2, v3])
 
     bodylist = []
-    for i in range(3):
+    for i in range(2):
         bodylist.append(Body(m[i], q[i], v[i]))
     dyn_syst = System(bodylist)
     dyn_syst.COMShift()
 
-    E, L = frogleap(10, 0.01, dyn_syst, recover_param=True, display=True)
+    duration, step = 100, 0.01
+    E, L = frogleap(duration, step, 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.plot(np.arange(E.shape[0])/duration, 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.plot(np.arange(L.shape[0])/duration, np.sum(L**2,axis=1), label=r"$L^2$")
     ax2.legend()
-    plt.show()
+    plt.show(block=True)
     return 0
 
 if __name__ == '__main__':