# Centrale 2 MP

import numpy as np
import matplotlib.pyplot as plt
import scipy.integrate as spi

## Données numériques

m = 1.0
U0 = 10.0
r0 = 1.0
d = .1
F0 = U0/d

## Composante radiale de la force

def F(r):
    return -F0*np.exp(-((r-r0)/d)**2)

valR = np.linspace(0, 3, 1000)
valF = F(valR)

plt.figure()
plt.grid()
plt.plot(valR, valF)
plt.xlabel(r'$r$')
plt.ylabel(r'$F_r(r)$')
plt.show()

## Énergie potentielle

def RHSgrad(U, r):
    return - F(r)

U0 = 0

res = spi.odeint(RHSgrad, U0, valR)

plt.figure()
plt.grid()
plt.plot(valR, res)
plt.xlabel(r'$r$')
plt.ylabel(r'$U(r)$')
plt.show()

## Équations du mouvement

def RHSpol(V, t): # Commentaire? Utilité?
    R, theta, Rp, thetap = V
    Rpp = R*thetap*thetap + F(R)/m
    thetapp = -2*Rp*thetap/R
    Vp = Rp, thetap, Rpp, thetapp
    return Vp

def RHScar(V, t): # Commentaire? Justification?
    X, Y, Xp, Yp = V
    R = (X*X+Y*Y)**.5
    a = F(R)/m/R
    Xpp = a*X
    Ypp = a*Y
    Vp = Xp, Yp, Xpp, Ypp
    return Vp

## Conditions initiales

plt.figure()
plt.grid()

theta = np.linspace(0, 2*np.pi, 1000)
xN = r0*np.cos(theta)
yN = r0*np.sin(theta)
plt.plot(xN, yN, 'k--')

x0 = -5
xp0 = 5.0
yp0 = 0

for y0 in (1.2, 1.1, 1.0, 0.9, 0.8, 0.7):
    V0 = x0, y0, xp0, yp0
    valT = np.linspace(0, 2.0, 1000)
    res = spi.odeint(RHScar, V0, valT)
    plt.plot(res[:,0], res[:,1])

plt.xlabel(r'$x(t)$')
plt.ylabel(r'$y(t)$')
plt.axis([-3, 3, -2, 2])
plt.axis('equal')
plt.show()