'''Solve the harmonic oscillator $H = p^2/2 + q^2/2$'''

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)

dt = 1./16.
t_end = 8
def force(x, v):
    return -x

def solve(method, param):
    ts, xs, vs = [0.], [1.], [0.]
    while ts[-1] < 2.*np.pi*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)
    ax.plot(ts/2./np.pi, np.fabs(xs**2 + vs**2 - 1.), **param)
    return 0

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

# RK4
def rk4(x, v):
    dx1, dv1 = v, -x
    dx2, dv2 = v+dv1*dt/2., force(x+dx1*dt/2., v+dv1*dt/2.)
    dx3, dv3 = v+dv2*dt/2., force(x+dx2*dt/2., v+dv2*dt/2.)
    dx4, dv4 = v+dv3*dt,    force(x+dx3*dt,    v+dv3*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_tmp = x + v*dt
    return x + v*dt, v + force(x_tmp, v)*dt
solve(symp1, {"label":"Symp1", "ls":"-."})

# symp2
def symp2(x, v, dt=dt):
    x_tmp = x + v*dt/2.
    v_tmp = v + force(x_tmp, v)*dt
    return x_tmp + v_tmp*dt/2., v_tmp
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_xlabel("time $t/2\pi$")
ax.set_ylabel(r"energy error $\left| \Delta E \right| / E_0$")
ax.set_yscale("log")
ax.set_xlim([0, t_end])
ax.set_ylim(bottom=1e-11)
ax.grid()
ax.legend(bbox_to_anchor=(1.05, 1), loc='upper left', borderaxespad=0, fontsize=14)

plt.savefig("symplectic-method-for-harmonic-oscillator.svg")
plt.show()
plt.close()