#!/usr/bin/python3
# Le pendule simple avec l'algorithme RK4

#!/usr/bin/python3
# Simple gravity pendulum, using RK4 method

import numpy as np
import matplotlib.pyplot as plt

g = 9.81  # gravitational acceleration
L = 0.1   # pendulum length = 10 cm
theta0 = 1.5 # starting angle at t=0
omega0 = 0.  # angular speed at t=0
a, b = 0, 3  # time interval
N = 300      # number of steps
delta = 1.0E-7    # target local accuracy

# Vector-valued function, return RHS = a NumPy array [ftheta, fomega] 
def f(r, t):
    theta = r[0]
    omega = r[1]
    ftheta = omega
    fomega = -g / L * np.sin(theta)
    return np.array([ftheta, fomega])
    
# Take one RK4 step of width h 
def rk4step(r, t, h):
    k1 = h * f(r, t)
    k2 = h * f(r + k1/2, t + h/2)
    k3 = h * f(r + k2/2, t + h/2)
    k4 = h * f(r + k3, t + h)
    return r + (k1 + 2*k2 + 2*k3 + k4)/6
        
h = (b - a) * delta**(1/4)       # initial increment
thetapoints = []                 # store theta(t) here
omegapoints = []                 # store omega(t) here
r = np.array([theta0, omega0])   # starting r from initial conditions
t = a                            # starting t
tpoints = []                     # intermediate times

while t < b:
    tpoints += [t]
    thetapoints += [r[0]]     
    omegapoints += [r[1]]
    r0 = rk4step(r, t, h)
    r1 = rk4step(r0, t + h, h)
    r2 = rk4step(r, t, 2*h)
    # estimate the error w.r.t. theta = r[0]
    rho = 30 * h * delta / abs(r2[0] - r1[0])
    if rho < 1:   # target precision not reached: repeat with smaller h
        h *= rho**(1/4)
        r0 = rk4step(r, t, h)
        r1 = rk4step(r0, t + h, h)
        t += 2*h
    else:         # target precision surpassed: increase h for next iteration
        t += 2*h
        h *= min(rho**(1/4), 3.0)  # increase at most by a factor of 3
    r = r1
    
plt.plot(tpoints, L * np.sin(thetapoints))
plt.plot(tpoints[::5], L*np.sin(thetapoints[::5]), 'bo')
for t in tpoints[::5]: 
    plt.axvline(t) 
plt.xlabel("t [s]")
plt.ylabel("x(t) [m]")
plt.ylim(-0.12, 0.12)

plt.show()