#!/usr/bin/env python
"""

Plotter forskjellige losninger av TUSL for HO.
Vi antar m\omega/\hbar = 1 nm^{-2} og m\omega^2 = 0.1 eV nm^{-2}

Hermite polynomene i numpy er definert som i Rottmann

@author Are Raklev
"""

from numpy import *
from matplotlib.pyplot import *
from numpy.polynomial.hermite_e import hermeval


# Definerer losninger av TUSL
def H_n( x, n ):
    coef = [0] * (n+1)
    coef[n] = 1
    return hermeval(x, coef)

def psi_n( x, n ):
    coef = [0] * (n+1)
    coef[n] = 1
    h = hermeval(sqrt(2.)*x, coef)
    return (1./pi)**0.25*sqrt(1./math.factorial(n))*h*exp( -0.5*x**2 )

# Definer HO potensialet
def HO( x ):
    return 0.5*0.1*x*x

# Lager x-verdier
x = linspace( -5, 5, 1e3 )

# Lager foerst en figur og plotter \psi_n i den
figure()
plot( x, HO( x ), label='$V(x)$')
plot( x, psi_n( x, 0 ), label='$\psi_0(x)$')
plot( x, psi_n( x, 1 ), label='$\psi_1(x)$')
plot( x, psi_n( x, 2 ), label='$\psi_2(x)$')
#plot( x, psi_n( x, 3 ), label='$\psi_3(x)$')
#plot( x, psi_n( x, 4 ), label='$\psi_4(x)$')

# Setter tittel
title('Harmonisk oscillator')

# Tekst langs x-aksen
xlabel('$x$ [nm]')

# Tekst langs y-aksen
ylabel('$\psi_0 (x)$ [nm$^{-1/2}$]')

#
legend(loc='best')

# Lagre som .eps fil
savefig('ho012.eps')

# Lager nok en figur for aa plotte Hermite polynomer
figure()
plot( x, H_n(x, 0), label='$H_0(x)$')
plot( x, H_n(x, 1), label='$H_1(x)$')
plot( x, H_n(x, 2), label='$H_2(x)$')

# Setter tittel
title('Hermite polynomer')

# Tekst langs x-aksen
xlabel('$x$')

# Tekst langs y-aksen
ylabel('$H_n (x)$')

#
legend(loc='best')

# Viser plottet
show()