Modify DynamicUpdate for colorful display of N-body problem

lib/objects.py

@@ -26,10 +26,10 @@ class Body:
         self.vp = np.zeros(3)
 
     def __repr__(self): # Called upon "print(body)"
-        return "Body of mass: {0:.2f}kg, position: {1}, velocity: {2}".format(self.m, self.p, self.v)
+        return "Body of mass: {0:.2e}kg, position: {1}, velocity: {2}".format(self.m, self.q, self.v)
 
     def __str__(self): # Called upon "str(body)"
-        return "Body of mass: {0:.2f}kg, position: {1}, velocity: {2}".format(self.m, self.p, self.v)
+        return "Body of mass: {0:.2e}kg".format(self.m)
 
 class System:
 
@@ -120,7 +120,7 @@ class System:
                 system("mkdir tmp")
             except IOError:
                 system("rm tmp/*")
-            d = DynamicUpdate()
+            d = DynamicUpdate(self)
             d.on_launch()
 
         N = np.ceil(duration/dt).astype(int)
@@ -134,14 +134,13 @@ class System:
 
             if display and j%100==0:
                 # display progression
-                q_array = self.get_positions()
                 if len(self.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))
+                    d.on_running(self, 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))
+                    d.on_running(self, step=j, label="step {0:d}/{1:d}".format(j,N))
         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")
+            system("convert tmp/temp.gif -fuzz 10% -layers Optimize plots/dynsyst.gif")# && rm tmp/temp.gif")
      
         if recover_param:
             return E, L
@@ -216,7 +215,7 @@ class System:
                 system("mkdir tmp")
             except IOError:
                 system("rm tmp/*")
-            d = DynamicUpdate()
+            d = DynamicUpdate(self)
             d.on_launch()
 
         N = np.ceil(duration/dt).astype(int)
@@ -230,14 +229,13 @@ class System:
 
             if display and j%100==0:
                 # display progression
-                q_array = self.get_positions()
                 if len(self.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))
+                    d.on_running(self, 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))
+                    d.on_running(self, step=j, label="step {0:d}/{1:d}".format(j,N))
         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")
+            system("convert tmp/temp.gif -fuzz 10% -layers Optimize plots/dynsyst.gif")# && rm tmp/temp.gif")
      
         if recover_param:
             return E, L

lib/plots.py

@@ -14,6 +14,9 @@ class DynamicUpdate():
 
     plt.ion()
 
+    def __init__(self, dyn_syst):
+        self.dyn_syst = dyn_syst
+
     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)
@@ -23,20 +26,28 @@ class DynamicUpdate():
         #Set up plot
         self.fig = plt.figure(figsize=(10,10))
         self.ax = self.fig.add_subplot(projection='3d')
-        self.lines, = self.ax.plot([],[],[],'o')
+        self.lines = []
+        for i,body in enumerate(self.dyn_syst.bodylist):
+            x, y, z = body.q
+            lines, = self.ax.plot([x],[y],[z],'o',color="C{0:d}".format(i),label="{0:s}".format(str(body)))
+            self.lines.append(lines)
+        self.lines = np.array(self.lines)
         #Autoscale on unknown axis and known lims on the other
         self.ax.set_autoscaley_on(True)
         self.set_lims()
         #Other stuff
         self.ax.grid()
-        #self.ax.set_aspect('equal')
+        self.ax.legend()
 
-    def on_running(self, xdata, ydata, zdata, step=None, label=None):
+    def on_running(self, dyn_syst, step=None, label=None):
+        xdata, ydata, zdata = dyn_syst.get_positions()
         values = np.sqrt(np.sum((np.array((xdata,ydata,zdata))**2).T,axis=1))
         self.min_x, self.max_x = -np.max([np.abs(values).max(),self.max_x]), np.max([np.abs(values).max(),self.max_x])
         self.set_lims()
         #Update data (with the new _and_ the old points)
-        self.lines.set_data_3d(xdata, ydata, zdata)
+        for i,body in enumerate(dyn_syst.bodylist):
+            x, y, z = body.q
+            self.lines[i].set_data_3d([x], [y], [z])
         if not label is None:
             self.ax.set_title(label)
         #Need both of these in order to rescale