#! /usr/bin/env python
# Author: Are Raklev
# Date: 20.02.13
# Email: ahye@fys.uio.no

# Solves Schrodinger's second order differential equation for the harmonic oscillator:
# psi'' = (xsi^2-K)*psi
# This is done by rewriting the equation 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] = (xsi^2-K)*psi[0]  
# Then a numeric integration routine odeint is used to find psi[0] and psi[1] at some given
# value of xsi and given initial conditions for psi[0] and psi[1] at the first value of xsi
# For details on how odeint works, see
# http://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.odeint.html#scipy.integrate.odeint

#  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,xsi):
    K = 2.000
    return array([ psi[1], (xsi**2-K)*psi[0] ])

# Values of xsi we want to plot for
xsi = linspace(-5.0,5.0,1000)

# Initial values (at first value of xsi used)
psiinit = array([0.01,0.01])

# Solve for psi
psi = odeint(deriv,psiinit,xsi)

# Plot figure
figure()
plot(xsi,psi[:,0])  # psi[:,0] is the first column of psi
xlabel('xsi')
ylabel('psi')
show()