# -*- coding: utf-8 -*-
"""
OSCILLATIONS DANS LE POTENTIEL DE MORSE
Comparaison de la méthode d'Euler ordre 2 et de la fonction odeint de scipy

Code proposé par Étienne Thibierge
https://www.etienne-thibierge.fr/
"""

import numpy as np
import matplotlib.pyplot as plt
plt.close('all')

### PARAMÈTRES DE LA SIMULATION =================================================
dt = 0.001      # pas de temps adimensionné
t_max = 70      # durée de la simulation adimensionnée
N = int(t_max/dt)  # nombre de points de la simulation

### Potentiel de Morse adimensionné
def f(x):
    return - np.exp(-x) * (1-np.exp(-x))

### Conditions initiales :
E = .5 # on fixe Em = Ec(0)
x0 = 0  # départ du fond du puits (pas d'énergie potentielle)
v0 = np.sqrt(E)


### RÉSOLUTION PAR LA MÉTHODE D'EULER ORDRE 2 ===================================

### Initialisation des listes :
t = [n*dt for n in range(N)]  # liste des temps
x = [None for n in range(N)]  # liste des positions
v = [None for n in range(N)]  # liste des vitesses

x[0] = x0
v[0] = v0

### Relations de récurrence :
for n in range(N-1):
    x[n+1] = x[n] + v[n]*dt
    v[n+1] = v[n] + f(x[n])*dt

### Tracés :
plt.figure()
plt.plot(t,x)
plt.xlabel('$t$ (adimensionné)')
plt.ylabel('$x$ (adimensionné)')


### RÉSOLUTION AVEC ODEINT ===================================================
from scipy.integrate import odeint
t = np.linspace(0, t_max, N)  # liste des temps

### Fonction d'évolution
def f_odeint(X,t):
    x, v = X
    return [v, f(x)]

X0 = [x0, v0]   # condition initiale
X = odeint(f_odeint, X0, t)

### Tracés :
plt.figure()
plt.plot(t,X[:,0])
plt.xlabel('$t$ (adimensionné)')
plt.ylabel('$x$ (adimensionné)')