## Oscillateur anharmonique 1D (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 oa(Y,a,l0):
    """
    Définit le second membre du problème de Cauchy correspondant à un système masse-ressort
    1D oscillant sans frottements sur un axe horizontal, ressort de longueur à vide l0, fixé 
    par l'autre extrémité à une distance a de l'axe du mouvement.
    """
    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] = - omega0**2*(1 - l0/np.sqrt(Y[0]**2+a**2))*Y[0]
    return F                

def euler(pb_cauchy,a,l0,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)
    xx = np.zeros(N+1)
    vvx = np.zeros(N+1)
    h = (tN-t0)/N  # pas de temps
    # Initialisations
    t = t0
    Y = Y0 # conditions initiales
    instants[0] = t0
    xx[0] = Y[0]
    vvx[0] = Y[1]
    # boucle temporelle
    for n in range(N):
        t = t + h
#        Y = Y + h * pb_cauchy(Y,a,l0) # pas d'Euler (vectoriel: dimension 2)
        Y = Y + h * pb_cauchy(Y + h/2*pb_cauchy(Y,a,l0),a,l0) # pas du Point Milieu (vectoriel: dimension 2)
        instants[n+1] = t
        xx[n+1] = Y[0]
        vvx[n+1] = Y[1]
    return instants, xx, vvx ## instants, positions, vitesses

def loi_horaire(F,a,l0,Y0,t0,tN,N):
    solution = euler(F,a,l0,Y0,t0,tN,N)
    instants = solution[0]
    xx = solution[1]
    vvx = solution[2]
    plt.plot(instants,xx,'r-')
    plt.grid()
    plt.xlabel('t')
    plt.ylabel('x')
#    plt.show() # afficher le graphe
    return 

def portrait(F,a,l0,Y0,t0,tN,N):
    solution = euler(F,a,l0,Y0,t0,tN,N)
    instants = solution[0]
    xx = solution[1]
    vvx = solution[2]
    plt.plot(xx,vvx,'r-')
    plt.grid()
    plt.xlabel('x')
    plt.ylabel('y') 
#    plt.show() # afficher le graphe
    return 

def simulation(F,a,l0,Y0,t0,tN,N):
    solution = euler(F,a,l0,Y0,t0,tN,N)
    instants = solution[0]
    xx = solution[1]
    vvx = solution[2]
    plt.subplot(1,2,1)
    plt.plot(instants,xx,'r-')
    plt.grid()
    plt.xlabel('t')
    plt.ylabel('x')
    plt.subplot(1,2,2)
    plt.plot(xx,vvx,'r-')
    plt.grid()
    plt.xlabel('x')
    plt.ylabel('y')
#    plt.show() # afficher le graphe
    return 

# Faisons un premier test:
a,l0 = 1.,2.
x0,v0= 0., 1.
t0,tN,N = 0,10.,10000
#loi_horaire(oa,a,l0,np.array([x0,v0]),t0,tN,N)
## graphes simulations
#plt.show()
#portrait(oa,a,l0,np.array([x0,v0]),t0,tN,N)
#plt.show()
#simulation(oa,a,l0,np.array([x0,v0]),t0,tN,N)
#plt.show()

for k in range(10):
    simulation(oa,a,l0,np.array([1.,1.*k]),t0,tN,N)
plt.show()