## Oscillateur Masse-Ressort 2D sans frottement (cf cours)
## Armel MARTIN - Septembre 2021
## distribué sous licence CC-by-nc-sa
import numpy as np # pour les tableaux, simulation num
import matplotlib.pyplot as plt # pour les graphes

def OH2D(Y,l0,C):
    """
    Définit le second membre du problème de Cauchy correspondant à un système masse-ressort
    2D oscillant sans frottements sur un plan horizontal, ressort de longueur à vide l0, fixé 
    par l'autre extrémité en l'origine du repère du plan du mouvement.
    C représente la constante des aires, introduite pour éliminer le d.l. theta (système à
    1 d.l. effectif).
    """
    F = np.zeros(len(Y)) # len(Y) = 2 ici
    omega0 = 2*np.pi # rapport k/m fixé, échelle de temps T0 = 1s.
    F[0] = Y[1]
    F[1] = C**2/Y[0]**3  - omega0**2*(Y[0] - l0)
    return F                

def euler(pb_cauchy,l0,C,Y0,t0,tN,N):
    """
    Intègre le pb de Cauchy vectoriel 2D dY/dt = F(Y) correspondant à une ED 1D d'ordre 2, 
    sur [t0,tN] avec N pas de temps, cond initiale Y0. 
    """
    instants = np.zeros(N+1)
    rr = np.zeros(N+1) 
    vvr = np.zeros(N+1)
    h = (tN-t0)/N  # pas de temps
    # Initialisations
    t = t0
    Y = Y0 # conditions initiales
    instants[0] = t0
    rr[0] = Y[0]
    vvr[0] = Y[1]
    # boucle temporelle
    for n in range(N):
        t = t + h
#        Y = Y + h * pb_cauchy(Y,l0,C) # pas d'Euler (vectoriel: dimension 2) - méthode d'ordre 1
        Y = Y + h * pb_cauchy(Y + h/2*pb_cauchy(Y,l0,C),l0,C) # pas du Point Milieu (vectoriel: dimension 2) - méthode d'ordre 2
        instants[n+1] = t
        rr[n+1] = Y[0]
        vvr[n+1] = Y[1]
    return instants, rr, vvr ## instants, positions, vitesses

def position_theta(theta0,r,C,t0,tN,N):
    """
    Calcule theta en intégrant la loi des aires, par la méthode des trapèzes.
    """
    h = (tN-t0)/N
    ttheta = np.zeros(N+1)
    ttheta[0] = theta0
    ttheta_point = C / r**2 # calcul vectorisé
    for n in range(N):
        ttheta[n+1] = ttheta[n] + h/2*(ttheta_point[n] + ttheta_point[n+1]) # le facteur h/2 se factorise sur la somme puisque le pas est constant. 
    return ttheta

def simulation(F,l0,r0,theta0,v0,t0,tN,N):
    """
    Reçoit en argument la longueur à vide l0, la position initiale (r0,theta0),
    la vitesse initiale (v0r,theta_point). C est calculée in situ.
    """
    Y0 = np.array([r0,v0[0]]) # conditions initiales pour la variable r(t)
    C = r0**2*v0[1] # constante des aires
    solution = euler(F,l0,C,Y0,t0,tN,N)
    instants = solution[0]
    rr = solution[1]
    vvr = solution[2]
    ttheta = position_theta(theta0,rr,C,t0,tN,N)
    xx = rr*np.cos(ttheta)
    yy = rr*np.sin(ttheta)
    # Graphes
    plt.subplot(1,3,1)
    plt.plot(instants,rr,'r-')
    plt.grid()
    plt.xlabel('t')
    plt.ylabel('r')
    plt.subplot(1,3,2)
    plt.plot(instants,ttheta,'r-')
    plt.grid()
    plt.xlabel('t')
    plt.ylabel(r'$\theta$')
    plt.subplot(1,3,3)
    plt.plot(xx,yy,'r-')
    plt.grid()
    plt.xlabel('x')
    plt.ylabel('y')
    plt.axis('equal')
    plt.show() # afficher le graphe
    return 


# Faisons un premier test:
l0 = 1.
r0,theta0= 3., 0.
v0 = np.array([0.,0.5])
t0,tN,N = 0,10.,10000
simulation(OH2D,l0,r0,theta0,v0,t0,tN,N)
## Autres essais intéressants
#simulation(OH2D,l0,1.5,theta0,np.array([0.,1.]),t0,tN,N)
#simulation(OH2D,l0,1.5,theta0,np.array([0.,2.]),t0,tN,N)