"""
Ex E.21. Implement the RungeKutta4 method as a function, based
on the ForwardEuler function.
"""

import numpy as np
import matplotlib.pyplot as plt

def ForwardEuler(f, U0, T, n):
    """Solve u'=f(u,t), u(0)=U0, with n steps until t=T."""
    t = np.zeros(n+1)
    u = np.zeros(n+1)  # u[k] is the solution at time t[k]
    u[0] = U0
    t[0] = 0
    dt = T/float(n)
    for k in range(n):
        t[k+1] = t[k] + dt
        u[k+1] = u[k] + dt*f(u[k], t[k])
    return u, t

def RungeKutta4(f, U0, T, n):
    """Solve u'=f(u,t), u(0)=U0, with n steps until t=T."""
    t = np.zeros(n+1)
    u = np.zeros(n+1)  # u[k] is the solution at time t[k]
    u[0] = U0
    dt = T/float(n)
    for k in range(n):
        t[k+1] = t[k] + dt
        K1 = dt*f(u[k], t[k])
        K2 = dt*f(u[k]+0.5*K1,t[k]+0.5*dt)
        K3 = dt*f(u[k]+0.5*K2,t[k]+0.5*dt)
        K4 = dt*f(u[k]+K3,t[k]+dt)
        u[k+1] = u[k] + 1.0/6*(K1+2*K2+2*K3+K4)
    return u, t

"""
Example use:
solve the equation u' = u, u(0) = 1
compare RK4, ForwardEuler and exact solution u = exp(t)
"""
def f(u,t):
    return u

U0 = 1
T = 4
n = 5

u, t = ForwardEuler(f,U0,T,n)
u2, t2 = RungeKutta4(f,U0,T,n)

plt.plot(t,u,t2,u2,t,np.exp(t))
plt.legend(['Euler','RK4','Exact'])
plt.show()