diff --git a/lib/integrator.py b/lib/integrator.py index a54ff7e..20c80e1 100644 --- a/lib/integrator.py +++ b/lib/integrator.py @@ -2,7 +2,35 @@ # -*- coding:utf-8 -*- """ Implementation of the various integrators for numerical integration. -""" -def leapfrog(): - return 0 +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 + +def dp_dt(m_array, q_array): + """ + Time derivative of the momentum, given by the position derivative of the Hamiltonian. + dp/dt = -dH/dq + """ + dp_array = np.zeros((q_array.shape[0],3)) + for i in range(q_array.shape[0]): + m_j = np.delete(m_array) + q_j = np.delete(q_array) + dp_array[i] = m_array[i]*np.sum(m_j*/(q_j-q_array[i])**2, axis=0) + return dp_array + +def leapfrog(duration, step, m_array, q_array, p_array): + """ + 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) + for _ in range(N): + # half-step drift + q_array, p_array = q_array + step/2*p_array/m_array , p_array + # full-step kick + q_array, p_array = q_array , p_array - step*dp_dt(m_array, q_array) + # half-step drift + q_array, p_array = q_array + step/2*p_array/m_array , p_array + + return q_array, p_array