import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
plt.close('all')
E1=1
E2=4*E1

def psi2(x,t):
    return np.sin(np.pi*x)**2 + np.sin(2*np.pi*x)**2+2*(np.sin(np.pi*x)*np.sin(2*np.pi*x))*np.cos((E2-E1)*t)


fig, ax = plt.subplots()
plt.title('Evolution de $|\psi(x,t)|^2$ pour un système à deux niveaux')
plt.xlabel('x')
plt.yticks([])


x = np.linspace(0,1,1000)
line, = ax.plot(x, psi2(x,0))
time_template = 't = %.1f T'
time_text = ax.text(0,3,'')




def animate(t):
    line.set_ydata(psi2(x,t/100))  # update the data.
    time_text.set_text(time_template % (t/200))
    return line, time_text


ani = animation.FuncAnimation(
    fig, animate, interval=20)


plt.show()