lib/integrator.py

@@ -6,6 +6,8 @@ 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
+import time
+from lib.plots import DynamicUpdate
 
 def dp_dt(m_array, q_array):
     """
@@ -16,15 +18,19 @@ def dp_dt(m_array, q_array):
     for i in range(q_array.shape[0]):
         m_j = np.delete(m_array, i)
         q_j = np.delete(q_array, i, 0)
-        dp_array = m_array[i]*np.sum((m_j/np.sum((q_j-q_array[i])**3, axis=1)).reshape((q_j.shape[0],1))*(q_j-q_array[i]), axis=0)
+        dp_array = -m_array[i]*np.sum((m_j/np.sum((q_j-q_array[i])**3, axis=1)).reshape((q_j.shape[0],1))*(q_j-q_array[i]), axis=0)
     return dp_array
 
-def frogleap(duration, step, m_array, q_array, p_array):
+def frogleap(duration, step, m_array, q_array, p_array, 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)
+    if display:
+        d = DynamicUpdate()
+        d.min_x, d.max_x = -1.5*np.abs(q_array).max(), +1.5*np.abs(q_array).max()
+        d.on_launch()
     for _ in range(N):
         # half-step drift
         q_array, p_array = q_array + step/2*p_array/m_array , p_array
@@ -33,4 +39,10 @@ def frogleap(duration, step, m_array, q_array, p_array):
         # half-step drift
         q_array, p_array = q_array + step/2*p_array/m_array , p_array
 
+        if display:
+            # display progression
+            q_cm = np.sum(m_array.reshape((q_array.shape[0],1))*q_array, axis=0)/m_array.sum()
+            d.on_running(q_array[:,0]-q_cm[0], q_array[:,1]-q_cm[1])
+            time.sleep(0.1)
+
     return q_array, p_array