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

Solves Schrodinger's equation for the harmonic oscillator:

psi'' = (x^2-K)*psi

The equation is rewritten as a set of two coupled first order differential
equations with the functions psi[0] and psi[1]:

psi'[0] = psi[1]  and psi'[1] = (x^2-K)*psi[0]

Then a numeric integration routine odeint is used to find psi[0] and psi[1] 
at some given value of x and given initial conditions for psi[0] and psi[1]
at the first value of x

For details on how odeint works, see
http://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.odeint.html#scipy.integrate.odeint

@author Are Raklev
Date: 20.02.13
Email: ahye@fys.uio.no
"""

#  Numerical and plotting libraries
from scipy.integrate import odeint
from numpy import *
from pylab import *

# Function to find derivatives of the array psi containing psi[0] and psi[1]
def deriv( psi, x ):
    K = 1.00000
    return array([ psi[1], (x**2-K)*psi[0] ])

# Values of x we want to plot for (split in two)
xpos = linspace(0.0,5.0,1000)
xneg = linspace(0.0,-5.0,1000)

# Initial values (at first value of x used, here x=0) for psi[0] and psi[1]
psi_init = array([1.0,0.0])

# Solve for psi
psi = odeint( deriv, psi_init, xpos)

# Plot figure
figure()
plot( xpos, psi[:,0], color='b')  # psi[:,0] is the first column of psi
plot( xneg, psi[:,0], color='b')  # Mirror
#plot( xpos, psi[:,1], color='r') # Plot the derivative

# Set title
title('Harmonisk oscillator')

# Axis lables and range
xlabel('$x$')
ylabel('$\psi(x)$')
ylim([-1.1,1.2])

# Show plot
show()