from graphenemodeling.graphene import monolayer as mlg
from scipy.constants import elementary_charge, hbar
eV = elementary_charge
gamma=0.012 * eV / hbar
eF = 0.4*eV
kF = mlg.FermiWavenumber(eF,model='LowEnergy')
q = np.linspace(1e-3,3,num=200) * kF
disp_intra = mlg.PlasmonDispersion(q,gamma,eF,eps1=1,eps2=1,T=0,model='intra')
disp_local = mlg.PlasmonDispersion(q,gamma,eF,eps1=1,eps2=1,T=0,model='local')
disp_nonlocal = mlg.PlasmonDispersion(q,gamma,eF,eps1=1,eps2=1,T=0,model='nonlocal')
fig, ax = plt.subplots(figsize=(6,6))
ax.plot(q/kF,hbar*disp_intra/eF,'--',label='Intraband')
# <...
ax.plot(q/kF,hbar*disp_local/eF,'r-.',label='Local')
# <...
ax.plot(q[:190]/kF,hbar*disp_nonlocal[:190]/eF,'g',label='Nonlocal')
# <...
ax.plot(q/kF,q/kF,color='gray',linestyle='--')
# <...
ax.set_xlabel('$q/k_F$')
ax.set_ylabel('$\hbar\omega/E_F$')
ax.set_xlim(0,3)
ax.set_ylim(0,3)
plt.legend()
plt.show()
