# III-3 Evolution temporelle d'une particule confinée dans une superposition d'états
####################################################################################
import matplotlib.pyplot as plt
import numpy as np

plt.ion()

L = 1                                                  # Largeur du puits de potentiel infini
m = 1                                                  # Masse de la particule
hbarre = 1                                             # Constante de Planck réduite

def phin(x,n):                                         # Fonction d'onde spatiale de l'état stationnaire n dans un puits infini
    return np.sqrt(2/L)*np.sin(n*np.pi*x/L)

def En(n):                                             # Energie de l'état stationnaire n dans un puits infini
    return (n*np.pi*hbarre)**2/(2*m*L**2)

def proba(x,t,n1,n2):                                  # Densité de probabilités de présence de la particule dans un état de superposition de deux états stationnaires

    psi1 = phin(x,n1)*np.exp(-1j*t*En(n1)/hbarre)      # Fonction d'onde de l'état stationnaire 1
    psi2 = phin(x,n2)*np.exp(-1j*t*En(n2)/hbarre)      # Fonction d'onde de l'état stationnaire 2
    psi = 1/np.sqrt(2)*(psi1 + psi2)                   # Fonction d'onde normalisée de l'état de superposition

    return abs(psi)**2                                 # Densité linéique de probabilité de présence = module au carré de psi

x = np.linspace(0,L,1000)                              # Plage d'abscisse
temps = np.linspace(0,1,100)                           # Plage de temps

n1 = 1                                                 # Choix du mode de l'état stationnaire 1
n2 = 2                                                 # Choix du mode de l'état stationnaire 2

for t in temps:                                        # Animation de la densité linéique de probabilité de présence

    plt.clf()

    plt.vlines(0,0,10,linewidth=5)
    plt.vlines(1,0,10,linewidth=5)
    plt.hlines(0,0,1,linewidth=5)
    plt.axis([0,1,0,4])
    plt.xlabel('Abscisse x')
    plt.ylabel("Densité de probabilité de présence")

    plt.plot(x,proba(x,t,n1,n2),'r')

    plt.title("Oscillations quantiques avec $n_1$ = {}".format(n1) + " et $n_2$ = {}".format(n2), fontweight='bold')
    plt.grid()
    plt.show()
    plt.pause(0.02)
#############################################################################################