import numpy as np
import matplotlib.pyplot as plt

""" Paramètres du plasma dilué """

omegap=6e7  # pulsation plasma (en rad/s)
c=3e8   # célérité OEM dans le vide (m/s)
omega0=8e7  # pulsation centrale du paquet d'onde (en rad/s)
sigma=5e6   # écart-type du paquet d'onde gaussien (en rad/s)
N=100   # nombre de composantes dans le paquet d'onde
s0=1    # amplitude centrale du paquet d'onde gaussien (u.a.)

""" Construction du paquet d'onde """
# A est un tableau contenant en colonne 1 les pulsations (rad/s) et en colonne 2 les amplitudes (ua)
A=np.zeros((N,2))
    
omeg=np.linspace(omega0-2*sigma,omega0+2*sigma,N)
A[:,0]=omeg
    
A[:,1]=s0*np.exp(-(omeg-omega0)**2/(2*sigma**2))

""" Propagation du paquet d'onde """

def disp(omega):
    # Input : omega (rad/s)
    # Output : k (/m)
    
    # return omega/c
    return np.sqrt(omega**2-omegap**2)/c

def signal(x,t):
    # Output [float ou array] : s(x,t) pour tous les x donnés (en u.a.)
    s=0
    for i in range(0,N):
        omega=A[i,0]
        k=disp(omega)
        s+=A[i,1]*np.cos(omega*t-k*x)
    return s

""" Résolution et affichage graphique """

x=np.linspace(-1e3,4e3,10000)     # liste des positions (en m)
dates=[0,0.5e-5,1e-5]

for k in range(1,4):
    plt.subplot(3,1,k)
    s=signal(x,dates[k-1])
    plt.plot(x,s)
    plt.ylim(-60,60)
    plt.xlabel('Abscisse x (en m)')
    plt.ylabel('Amplitude (en u.a.)')
    plt.title('Paquet d\'ondes à t = '+str(dates[k-1])+' s')
    plt.grid()
plt.show()