#!/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

# Vector-valued function, return RHS = a NumPy array [ftheta, fomega] 
def f(r, t):
    theta, omega = r  # r = array containing theta(t) and omega(t)
    ftheta = omega
    fomega = -g / L * np.sin(theta)
    return np.array([ftheta, fomega])
    
h = (b - a) / N                  # time increment
tpoints = np.linspace(a, b, N+1) # intermediate times
thetapoints = []                 # store theta(t) here
omegapoints = []                 # store omega(t) here
r = np.array([theta0, omega0])   # starting r from initial conditions
    
for t in tpoints:
    thetapoints += [r[0]]     # append new values to thetapoints...
    omegapoints += [r[1]]     # ... and to omegapoints. Then take an RK4 step.
    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)
    r = r + (k1 + 2*k2 + 2*k3 + k4)/6
    
plt.plot(tpoints, L * np.sin(thetapoints))

# For comparison: harmonic approximation
omegah = np.sqrt(g/L)    # frequency of a harmonic pendulum
plt.plot(tpoints, L*np.sin(theta0)*np.sin(omegah * tpoints + theta0))

plt.xlabel("t [s]")
plt.ylabel("x(t) [m]")
plt.ylim(-0.12, 0.12)

# Phase space plot:
plt.figure()
plt.xlim(-0.12, 0.12)
plt.plot(L * np.sin(thetapoints), L * (omegapoints * np.cos(thetapoints)))
plt.xlabel("x(t) [m]")
plt.ylabel("dx/dt [m/s]")

plt.show()