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

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation

g = 9.81  # accélération gravitationnelle
L = 0.1   # longueur du pendule
theta0 = 1.5 # angle à t=0
omega0 = 0.  # vitesse angulaire à t=0
a, b = 0, 30  # intervalle de temps
N = 3000      # nombre de pas

# Fonction vectorielle, renvoie le vecteur [ftheta, fomega] des membres de droite
def f(r, t):
    theta = r[0]
    omega = r[1]
    ftheta = omega
    fomega = -g / L * np.sin(theta)
    return np.array([ftheta, fomega])
    
h = (b - a) / N                  # pas d'incrément
tpoints = np.arange(a, b, h)     # temps intermédiaires
thetapoints = []                 # on va mettre les theta(t) ici
omegapoints = []                 # on va mettre les omega(t) ici
r = np.array([theta0, omega0])   # conditions aux limites
    
for t in tpoints:
    thetapoints.append(r[0])     # ajouter nouvelle valuer à thetapoints...
    omegapoints.append(r[1])     # ... et à omegapoints. Ensuite: RK4.
    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 += (k1 + 2*k2 + 2*k3 + k4)/6

fig = plt.figure()
ax = plt.axes(xlim=(-0.12, 0.12), ylim=(-0.12, 0.02), aspect=1)
line, = ax.plot([], [], "o-", lw=2)

def init():
    line.set_data([], [])
    return line,

def animate(i):
    x = L*np.sin(thetapoints[i])
    y = -L*np.cos(thetapoints[i])
    line.set_data([0,x], [0,y])
    return line,

anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=3000, interval=20)
                           
plt.show()