from scipy.integrate import odeint
import matplotlib.pyplot as plt
import numpy as np

# Constantes du système
g = 9.81  # m/s2
beta = 1e-3  # kg/m
m = 1       # kg


# Définition de l'équation différentielle
def equation(v,t):  # Attention à l'inversion de t et v
    return g-(beta/m)*v**2
t0 = 0
tf = 50  # seconde
v0 = 0  # Condition initiale
t = np.linspace(t0,tf,500) # Création des instants de calculs

# Résolution
v = odeint(equation,v0,t)
print(t)  # Affichage des instants t
print(v[:,0])  # Affichage des résultats

# Tracé de la vitesse au cours du temps
plt.plot(t,v[:,0])
plt.ylabel('v(t) en m/s')
plt.xlabel("Temps (s)")
plt.title("Vitesse lors d'une chute libre soumise à une force de frottement en v carré")
plt.grid(which='both')
plt.show()

# Comparaison à la solution théorique
x = np.linspace(0,50,500)
y = np.sqrt(m*g/beta)*(np.exp(2*g*x*np.sqrt(beta/(m*g)))-1)/(np.exp(2*g*x*np.sqrt(beta/(m*g)))+1)
plt.figure('théorie')
plt.title("Vitesse théorique lors d'une chute libre soumise à une force de frottement en v carré")
plt.plot(x,y)
plt.xlabel("temps en secondes")
plt.ylabel('v(t) en m/s')
plt.grid()
plt.show()