diff --git a/lib/integrator.py b/lib/integrator.py index 097f592..948ddfc 100755 --- a/lib/integrator.py +++ b/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) diff --git a/lib/objects.py b/lib/objects.py index 4b262cc..0c007bb 100755 --- a/lib/objects.py +++ b/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 diff --git a/main.py b/main.py index 7f141e3..80f490b 100755 --- a/main.py +++ b/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__':