"""
The file solves the quantum oscillator with the pde
	i eps u_t = - eps^2 u_{xx} + (x^2) u
with a Gaussian Initial Condition and Periodic Boundary Conditions.
 
The PDE is solved using a spectral method which gave us the system of equations:
	U_t = - i eps q^2 U - i/eps F{ x^2 F^-1{U} }
with k=-N/2...N/2-1
and
	q = pi k/L
	F{} = Fourier Transform
	F^-1{} = Inverse Fourier Transform
 
"""
 
from numpy import *
import pylab as pl
from scipy import integrate
import include.CreateMovie as movie
 
# Constant Coefficients
L = 12.0
n = 256
epsilon = 1.0
sigma = 1.0/1
mu = L/2
 
# Create Spatial and WaveNumber Grid
x = arange(n)*2*L/n - L
k = arange(-float(n)/2, float(n)/2)
q = pi/L*fft.ifftshift(k)
t = list(linspace(0,10, 800))
 
# Find Nyquist Mode
nq = n/2
 
# Create initial conditions (off-centered Gaussian)
psi_0 = exp( -(x-mu)**2/(2*sigma**2) )
psi_hat_0 = list(fft.fft(psi_0))
 
# seperate real and imaginary components
# The Python ODE solver can't solve complex ODEs.
# This means we need to break up the real and imaginary
# numbers into a system of ODEs
decouple_psi = empty(2*len(psi_hat_0),dtype=float64)
complex_psi = empty(len(psi_hat_0),dtype=complex128)
decouple_psi[:n] = real(psi_hat_0)
decouple_psi[n:] = imag(psi_hat_0)
 
 
def ode(psi, t):
	# Convert vector psi into the complex components
	complex_psi.real = psi[:n]
	complex_psi.imag = psi[n:]
 
	# Zero the Nyquist Mode.
	complex_psi[nq] = 0
 
	result = -1.0j*epsilon*(q**2)*complex_psi- \
			1.0j/epsilon* fft.fft(x**2*fft.ifft(complex_psi))
 
	# seperate real and imaginary components
	decouple_psi[:n] = real(result)
	decouple_psi[n:] = imag(result)
 
	return list(decouple_psi)
 
# Solve ODE
soln = integrate.odeint(ode, decouple_psi, t)
 
# Function to plot results
def plotFunction( plot_t ):
	complex_psi.real = soln[plot_t][:n]
	complex_psi.imag = soln[plot_t][n:]
 
	pl.title('Quantum Oscillator')
	pl.plot(x,real(fft.ifft(complex_psi)),'b', label='Real Component')
	pl.plot(x,imag(fft.ifft(complex_psi)),'r', label='Imag Component')
	pl.legend()
	pl.axis([-L, L, -1, 1])
 
# Create Movie
movie.CreateMovie(plotFunction, len(t), 20)