diff --git a/.gitignore b/.gitignore index 22c3d7f..bce39da 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,6 @@ +# Data temp directory +tmp/ + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/lib/integrator.py b/lib/integrator.py index 948ddfc..bd32ce2 100755 --- a/lib/integrator.py +++ b/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 diff --git a/lib/objects.py b/lib/objects.py index 0c007bb..ee535ec 100755 --- a/lib/objects.py +++ b/lib/objects.py @@ -77,6 +77,12 @@ class System: rij = np.linalg.norm(body.q-otherbody.q) 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__": diff --git a/lib/plots.py b/lib/plots.py index a1e99d5..dfcca90 100755 --- a/lib/plots.py +++ b/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(): @@ -16,24 +17,29 @@ class DynamicUpdate(): def on_launch(self): #Set up plot self.figure, self.ax = plt.subplots() - self.lines, = self.ax.plot([],[], 'o') + 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): diff --git a/main.py b/main.py index 80f490b..5ce4c87 100755 --- a/main.py +++ b/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__':