# -*- coding: utf-8 -*-
"""
Created on Thu Jan 22 11:32:51 2015

Listing 12.2 : résolution de l'équation de Schrödinger 
par discrétisation

@author: grivet
"""

import numpy as np
import  matplotlib.pyplot as plt
import scipy.linalg as lg

N = 30; L = 4
Ni = 2*N; Nx = 2*N-1
#le résultat doit être un nombre décimal !
h = 2.0*L/Ni
x = np.linspace(-L+h,L-h,Nx)
xxhh = h*h*x*x + 2
sousdiag = np.ones(Nx-1)
M = np.diag(xxhh) - np.diag(sousdiag,1) - np.diag(sousdiag,-1)
En,Psi = lg.eig(M)
xx = np.hstack([-L,x,L])
psi1 = np.hstack([0,Psi[:,0],0])
psi2 = np.hstack([0,Psi[:,1],0])
psi3 = np.hstack([0,Psi[:,2],0])
psi4 = np.hstack([0,Psi[:,3],0])
plt.plot(xx,psi1,xx,psi2,xx,psi3,xx,psi4)
plt.title("Fonctions d'onde de l'oscillateur harmonique")
plt.ylabel("$\psi_0, \psi_1,\psi_2, \psi_3$")
ee = np.real(0.5*En/(h*h))
print('valeurs propres : %8.4f %8.4f %8.4f %8.4f' %(ee[0],ee[1],ee[2],ee[3]))