#!/usr/bin/env python
# coding: utf-8

# # ODE with periodic solution

# Consider the ODE system
# $$
# x' = -y, \qquad y' = x
# $$
# with initial condition
# $$
# x(0) = 1, \qquad y(0) = 0
# $$
# The exact solution is
# $$
# x(t) = \cos(t), \qquad y(t) = \sin(t)
# $$
# This solution is periodic. It also has a quadratic invariant
# $$
# x^2(t) + y^2(t) = 1, \qquad \forall t
# $$

# In[15]:


get_ipython().run_line_magic('matplotlib', 'inline')
get_ipython().run_line_magic('config', "InlineBackend.figure_format = 'svg'")
import numpy as np
from matplotlib import pyplot as plt


# In[16]:


theta = np.linspace(0.0, 2*np.pi, 500)
xe = np.cos(theta)
ye = np.sin(theta)


# ## Forward Euler scheme

# In[17]:


def ForwardEuler(h,T):
    N = int(T/h)
    x,y = np.zeros(N),np.zeros(N)
    x[0] = 1.0
    y[0] = 0.0
    for n in range(1,N):
        x[n] = x[n-1] - h*y[n-1]
        y[n] = y[n-1] + h*x[n-1]
    return x,y


# In[18]:


h = 0.02
T = 4.0*np.pi
x,y = ForwardEuler(h,T)

plt.plot(xe,ye,'r--',label='Exact')
plt.plot(x,y,label='FE')
plt.legend()
plt.grid(True)
plt.gca().set_aspect('equal')


# The phase space trajectory is spiralling outward.

# ## Backward Euler scheme

# $$
# x_n = x_{n-1} - h y_n, \qquad y_n = y_{n-1} + h x_n
# $$
# Eliminate $y_n$ from first equation to get
# $$
# x_n = \frac{x_{n-1} - h y_{n-1}}{1 + h^2}
# $$

# In[19]:


def BackwardEuler(h,T):
    N = int(T/h)
    x,y = np.zeros(N),np.zeros(N)
    x[0] = 1.0
    y[0] = 0.0
    for n in range(1,N):
        x[n] = (x[n-1] - h*y[n-1])/(1.0 + h**2)
        y[n] = y[n-1] + h*x[n]
    return x,y


# In[28]:


h = 0.02
T = 4.0*np.pi
x,y = BackwardEuler(h,T)

plt.plot(xe,ye,'r--',label='Exact')
plt.plot(x,y,label='BE')
plt.legend()
plt.grid(True)
plt.gca().set_aspect('equal')


# The phase space trajectory is spiralling inward.

# ## Trapezoidal scheme

# $$
# x_n = x_{n-1} - \frac{h}{2}(y_{n-1} + y_n), \qquad y_n = y_{n-1} + \frac{h}{2}(x_{n-1} + x_n)
# $$
# Eliminate $y_n$ from first equation
# $$
# x_n = \frac{ (1-\frac{1}{4}h^2) x_{n-1} - h y_{n-1} }{1 + \frac{1}{4}h^2}
# $$

# In[21]:


def Trapezoid(h,T):
    N = int(T/h)
    x,y = np.zeros(N),np.zeros(N)
    x[0] = 1.0
    y[0] = 0.0
    for n in range(1,N):
        x[n] = ((1.0-0.25*h**2)*x[n-1] - h*y[n-1])/(1.0 + 0.25*h**2)
        y[n] = y[n-1] + 0.5*h*(x[n-1] + x[n])
    return x,y


# In[27]:


h = 0.02
T = 4.0*np.pi
x,y = Trapezoid(h,T)

plt.plot(xe,ye,'r--',label='Exact')
plt.plot(x,y,label='Trap')
plt.legend()
plt.grid(True)
plt.gca().set_aspect('equal')


# The phase space trajectory is exactly the unit circle. 
# 
# Multiply first equation by $x_n + x_{n-1}$ and second equation by $y_n + y_{n-1}$
# $$
# (x_n + x_{n-1})(x_n - x_{n-1}) = - \frac{h}{2}(x_n + x_{n-1})(y_n + y_{n-1})
# $$
# $$
# (y_n + y_{n-1})(y_n - y_{n-1}) = + \frac{h}{2}(x_n + x_{n-1})(y_n + y_{n-1})
# $$
# Adding the two equations we get
# $$
# x_n^2 + y_n^2 = x_{n-1}^2 + y_{n-1}^2
# $$
# Thus the Trapezoidal method is able to preserve the invariant.
# 
# **Excercise**: Show that the energy increases for forward Euler and decreases for backward Euler, for any step size $h$.

# ## Partitioned Euler
# This is a symplectic Runge-Kutta method. $x$ is updated explicitly and $y$ is updated implicitly.
# $$
# x_n = x_{n-1} - h y_{n-1}, \qquad y_n = y_{n-1} + h x_{n}
# $$
# The update of $y$ uses the latest updated value of $x$.

# In[23]:


def PartEuler(h,T):
    N = int(T/h)
    x,y = np.zeros(N),np.zeros(N)
    x[0] = 1.0
    y[0] = 0.0
    for n in range(1,N):
        x[n] = x[n-1] - h * y[n-1]
        y[n] = y[n-1] + h * x[n]
    return x,y


# In[24]:


h = 0.02
T = 4.0*np.pi
x,y = PartEuler(h,T)

plt.plot(xe,ye,'r--',label='Exact')
plt.plot(x,y,label='PE')
plt.legend()
plt.grid(True)
plt.gca().set_aspect('equal')


# We get very good results, even though the method is only first order. We can also switch the implicit/explicit parts
# 
# $$
# y_n = y_{n-1} + h x_{n-1}, \qquad x_n = x_{n-1} - h y_n
# $$

# ## Classical RK4 scheme

# In[25]:


def rhs(u):
    return np.array([u[1], -u[0]])

def RK4(h,T):
    N = int(T/h)
    u = np.zeros((N,2))
    u[0,0] = 1.0 # x
    u[0,1] = 0.0 # y
    for n in range(0,N-1):
        w = u[n,:]
        k0 = rhs(w)
        k1 = rhs(w + 0.5*h*k0)
        k2 = rhs(w + 0.5*h*k1)
        k3 = rhs(w + h*k2)
        u[n+1,:] = w + (h/6)*(k0 + k1 + k2 + k3)
    return u[:,0], u[:,1]


# We test this method for a longer time.

# In[26]:


h = 0.02
T = 10.0*np.pi
x,y = RK4(h,T)

plt.plot(xe,ye,'r--',label='Exact')
plt.plot(x,y,label='RK4')
plt.legend()
plt.grid(True)
plt.gca().set_aspect('equal')


# RK4 is a more accurate method compared to forward/backward Euler schemes, but it still loses total energy.