# wave_packet.py - propagates a wave packet with a dispersion relation omega(k)
# by Bjoern Malte Schaefer, GSFP+/Heidelberg, bjoern.malte.schaefer@uni-heidelberg.de

# import libraries numpy and matplotlib
import numpy as np
import pylab as plt

plt.close()

# set width sigma and the number of pixels ngrid
sigma = 0.1		# width of the wave packet
ngrid = 1024		# number of pixels
velocity = 0.01		# group velocity
dispersion = 0.0002	# dispersion
t = 10			# time 

# definition of the Gaussian envelope
def psi_gauss(x,sigma):
	aux = x / sigma
	result = np.exp(-aux**2 / 2.0)
	return(result)
	
# definition of the tophat envelope
def psi_tophat(x,sigma):
	return(np.asarray(np.abs(x/sigma)<1.0,dtype=float))
	
# dispersion relation omega(k)
def omega(k):
	result =  velocity * k + dispersion * k**2
	return(result)

# set up wave packet in real space and plot at t=0
x = np.linspace(-1.0,1.0,ngrid)
k = np.fft.fftfreq(ngrid,2.0/ngrid) * 2.0 * np.pi

psi_initial = psi_gauss(x,sigma)
plt.plot(x,psi_initial,'r--',label='initial wave packet at 0')

# Fourier transform
psi_fourier = np.fft.fft(psi_initial)

# multiply with k-dependent phase factor
phase = np.exp(-1j * omega(k) * t)
psi_fourier *= phase

# Fourier transform back to real space
psi_final = np.fft.ifft(psi_fourier)

# plot wavepacket at t
plt.plot(x,psi_final,'g-',label='final wave packet at t')

# add nice things
plt.legend(loc='upper left')
plt.xlabel('$x$-axis')
plt.ylabel('$\psi$-axis')
plt.ylim([-0.1,2.0])
plt.show()