'''
Solve the Kepler problem $H = p^2/2 - 4 \pi / | x |$.

Code unit:
   coordinate: initial semi-major axis $a$
   time: orbital period $P$

Output file:
    "./symplectic-method-for-keplerian-motion.svg"
'''

import numpy as np
import matplotlib.pyplot as plt
plt.rcParams["text.usetex"] = True
plt.rcParams["font.family"] = "DejaVu Sans"
plt.rcParams["font.size"] = 14

fig = plt.figure()
plt.subplots_adjust(left=0.15, right=0.75)
ax = fig.add_subplot(111)

### parameters ####################
e = 0.5      # eccentricity
dt = 1./128. # time step
t_end = 4    # calculation end time
###################################

def force(x):
    r = np.linalg.norm(x)
    return -4.*np.pi**2*x/r**3

def solve(method, param):
    ts, xs, vs = [0.], [np.array([1.+e, 0])], [np.array([0., 2.*np.pi*np.sqrt((1-e)/(1+e))])]
    while ts[-1] < t_end:
        x, v = method(xs[-1], vs[-1])
        ts.append(ts[-1] + dt)
        xs.append(x)
        vs.append(v)
    ts, xs, vs = np.array(ts), np.array(xs), np.array(vs)
    energy = np.sum(vs**2, axis=1)/2. - 4.*np.pi**2/np.linalg.norm(xs, axis=1)
    ax.plot(ts, np.fabs((energy - energy[0])/energy[0]), **param)
    return 0

# Euler
def euler(x, v):
    return x + v*dt, v + force(x)*dt
solve(euler, {"label":"Euler"})

# RK4
def rk4(x, v):
    dx1, dv1 = v,           force(x)
    dx2, dv2 = v+dv1*dt/2., force(x+dx1*dt/2.)
    dx3, dv3 = v+dv2*dt/2., force(x+dx2*dt/2.)
    dx4, dv4 = v+dv3*dt,    force(x+dx3*dt   )
    return x + (dx1 + 2.*(dx2 + dx3) + dx4)*dt/6., v + (dv1 + 2.*(dv2 + dv3) + dv4)*dt/6., 
solve(rk4, {"label": "RK4"})

# symp1
def symp1(x, v):
    x = x + v*dt
    return x, v + force(x)*dt
solve(symp1, {"label":"Symp1", "ls":"-."})

# symp2
def symp2(x, v, dt=dt):
    x = x + v*dt/2.
    v = v + force(x)*dt
    return x + v*dt/2., v
solve(symp2, {"label":"Symp2", "ls":":"})

# symp4
def symp4(x, v):
    d1 = 1./(2. - np.cbrt(2.))
    d2 = 1. - 2.*d1
    x, v = symp2(x, v, dt=dt*d1)
    x, v = symp2(x, v, dt=dt*d2)
    return symp2(x, v, dt=dt*d1)
solve(symp4, {"label":"Symp4", "ls":"-."})

ax.set_title("$e={}$".format(e))
ax.set_xlabel(r"time $t/P$")
ax.set_ylabel(r"energy error $\left| \Delta E / E_0 \right|$")
ax.set_yscale("log")
ax.set_xlim([0, t_end])
ax.set_ylim([1e-7, 2])
ax.grid()
ax.legend(bbox_to_anchor=(1.05, 1), loc='upper left', borderaxespad=0, fontsize=14)

plt.savefig("symplectic-method-for-keplerian-motion.svg")
plt.show()
plt.close()