"""

Wave Function Animation of a Particle in a Box

"""

import numpy as np

from pylab import *

from matplotlib import pyplot as plt

from matplotlib import animation

#matplotlib.use('GTK3Agg')

# from matplotlib import interactive

# interactive(True)

########################################Variables######################################################################################################

N_Slices = 1000 #Number of slices in the box

Time_step = 1e-18 #Time step for each iteration

Mass = 9.109e-31 #mass of electron

plank = 1.0546e-36 #Planks constant

L_Box = 1e-9 #Length of the box

Grid = L_Box/N_Slices #Lenght of each slice

#####################################Si(0) using the given equation ###############################################################################

Si_0 = np.zeros(N_Slices+1,complex) #Initiating Si funtion at time step = 0

x = np.linspace(0,L_Box,N_Slices+1)

def G_Equation(x):

x_0 = L_Box/2

Sig = 1e-10

k = 5e10

result = exp(-(x-x_0)**2/2/Sig**2)*exp(1j*k*x) #Given Equation at t = 0

return result

Si_0[:] = G_Equation(x) #Si funtion at time step = 0

#######################################V = Bxsi(0)################################################################################

a_1 = 1 + Time_step*plank*1j/(2*Mass*(Grid**2)) #Diagonal of A Tridiagonal matrix

a_2 = -Time_step*plank*1j/(4*Mass*Grid**2) #Up and Down to A Tridiagonal matrix

b_1 = 1 - Time_step*plank*1j/(2*Mass*(Grid**2)) #Diagonal of B Tridiagonal matrix

b_2 = Time_step*plank*1j/(4*Mass*Grid**2) #Up and Down to B Tridiagonal matrix

BxSi_0 = [] #V = BxSi and si funtion at x = 0

for i in range(1000):

if i == 0:

BxSi_0.append(b_1*Si_0[0] + b_2*(Si_0[1])) #V can be maipulated by the equation in Text book

else:

BxSi_0.append(b_1*Si_0[i] + b_2*(Si_0[i+1] + Si_0[i-1]))

BxSi_0 = np.array(BxSi_0)

#####################################Tri Diagonal matrix algorithm#####################################################################################

def TDMAsolver(a, b, c, d): #Instead of solving using Numpy.linalg, it is prefered to Use

nf = len(d) #Tri Diagonal Matrix algorithm

ac, bc, cc, dc = map(np.array, (a, b, c, d)) # a,b,c's are up,dia,down element in tridiagonl matrix A

for it in range(1, nf): #AX = d

mc = ac[it-1]/bc[it-1]

bc[it] = bc[it] - mc*cc[it-1]

dc[it] = dc[it] - mc*dc[it-1]

xc = bc

xc[-1] = dc[-1]/bc[-1]

for il in range(nf-2, -1, -1):

xc[il] = (dc[il]-cc[il]*xc[il+1])/bc[il]

return xc

global a #A matrix is fixed through out the problem, so it is good to globalize the variables

global b

global c

b = N_Slices*[a_1] #In A matrix, Both Up,Down elements are a_2 and diag matrix is a_1

a = (N_Slices-1)*[a_2]

c = (N_Slices-1)*[a_2]

####################################Si 1st funtion solver####################################################################################

global Si_1 #First si_funtion usinf Axsi(0+h) = BxSi(0)

Si_1 = TDMAsolver(a, b, c, BxSi_0) #This can be solved by TDM(A,BxSi(0))

###################################A funtion which caliculates si at each step#####################################################################################

global Si_sd #AxSi_stepup = BxSi_stepdown

Si_sd = {} #At first Buckting Si_stepdown in to directry which we can using for finding Si_stepup

def sifuntion(i): #In next iteration, Last iteration Si_stepup will be this iteration's Si_stepdown

if i == 0:

Si_sd[0] = Si_1

return Si_1

else:

Si_stepdown = Si_sd[i-1]

V = np.zeros(N_Slices,complex)

V[0] = b_1*Si_stepdown[0] + b_2*(Si_stepdown[1])

V[1:N_Slices-1] = b_1*Si_stepdown[1:N_Slices-1] + b_2*(Si_stepdown[2:N_Slices] + Si_stepdown[0:N_Slices-2])

V[N_Slices-1] = b_1*Si_stepdown[N_Slices-1]+ b_2*(Si_stepdown[N_Slices-2])

Si_stepup = TDMAsolver(a, b, c, V)

Si_sd[i] = Si_stepup

x = Si_stepup.real

return x

####################################Animating#######################################################################################

fig = plt.figure()

ax = plt.axes(xlim=(0, 1000), ylim=(-1.5, 1.5))

line, = ax.plot([], [], lw=5)

ax.legend(prop=dict(size=20))

ax.set_facecolor('black')

ax.patch.set_alpha(0.8)

ax.set_xlabel('$x$',fontsize = 15,color = 'blue')

ax.set_ylabel(r'$|\psi(x)|$',fontsize = 15,color = 'blue')

ax.grid(color = 'blue')

ax.set_title(r'$|\psi(x)|$ vs $x$', color='blue',fontsize = 15 )

line, = ax.step([], [])

def init():

line.set_data([], [])

return line,

def animate(i):

x = np.linspace(0, 1000, num=1000)

y = sifuntion(i)

line.set_data(x, y)

line.set_color('red')

return line,

anim = animation.FuncAnimation(fig, animate, init_func=init,

frames=10**5, interval=1, blit=True)#5*10**5

plt.vlines(1, -5, 5, linestyles = 'solid', color= 'green',lw=5)

plt.vlines(999, -5, 5, linestyles = 'solid', color= 'green',lw=5)

plt.text(1,1,'Boundary',rotation=90,color= 'green' )

plt.text(975,1,'Boundary',rotation=90,color= 'green' )

plt.figure(figsize=(10,10))

plt.show()

# Writer = animation.writers['ffmpeg']

# writer = Writer()

# anim.save('im.mp4', writer=writer)

"""

General Numerical Solver for the 1D Time-Dependent Schrodinger's equation.

author: Jake Vanderplas

email: vanderplas@astro.washington.edu

website: http://jakevdp.github.com

license: BSD

Please feel free to use and modify this, but keep the above information. Thanks!

"""

import numpy as np

from matplotlib import pyplot as pl

from matplotlib import animation

from scipy.fftpack import fft,ifft

class Schrodinger(object):

"""

Class which implements a numerical solution of the time-dependent

Schrodinger equation for an arbitrary potential

"""

def __init__(self, x, psi_x0, V_x,

k0 = None, hbar=1, m=1, t0=0.0):

"""

Parameters

----------

x : array_like, float

length-N array of evenly spaced spatial coordinates

psi_x0 : array_like, complex

length-N array of the initial wave function at time t0

V_x : array_like, float

length-N array giving the potential at each x

k0 : float

the minimum value of k. Note that, because of the workings of the

fast fourier transform, the momentum wave-number will be defined

in the range

k0 < k < 2*pi / dx

where dx = x[1]-x[0]. If you expect nonzero momentum outside this

range, you must modify the inputs accordingly. If not specified,

k0 will be calculated such that the range is [-k0,k0]

hbar : float

value of planck's constant (default = 1)

m : float

particle mass (default = 1)

t0 : float

initial tile (default = 0)

"""

# Validation of array inputs

self.x, psi_x0, self.V_x = map(np.asarray, (x, psi_x0, V_x))

N = self.x.size

assert self.x.shape == (N,)

assert psi_x0.shape == (N,)

assert self.V_x.shape == (N,)

# Set internal parameters

self.hbar = hbar

self.m = m

self.t = t0

self.dt_ = None

self.N = len(x)

self.dx = self.x[1] - self.x[0]

self.dk = 2 * np.pi / (self.N * self.dx)

# set momentum scale

if k0 == None:

self.k0 = -0.5 * self.N * self.dk

else:

self.k0 = k0

self.k = self.k0 + self.dk * np.arange(self.N)

self.psi_x = psi_x0

self.compute_k_from_x()

# variables which hold steps in evolution of the

self.x_evolve_half = None

self.x_evolve = None

self.k_evolve = None

# attributes used for dynamic plotting

self.psi_x_line = None

self.psi_k_line = None

self.V_x_line = None

def _set_psi_x(self, psi_x):

self.psi_mod_x = (psi_x * np.exp(-1j * self.k[0] * self.x)

* self.dx / np.sqrt(2 * np.pi))

def _get_psi_x(self):

return (self.psi_mod_x * np.exp(1j * self.k[0] * self.x)

* np.sqrt(2 * np.pi) / self.dx)

def _set_psi_k(self, psi_k):

self.psi_mod_k = psi_k * np.exp(1j * self.x[0]

* self.dk * np.arange(self.N))

def _get_psi_k(self):

return self.psi_mod_k * np.exp(-1j * self.x[0] *

self.dk * np.arange(self.N))

def _get_dt(self):

return self.dt_

def _set_dt(self, dt):

if dt != self.dt_:

self.dt_ = dt

self.x_evolve_half = np.exp(-0.5 * 1j * self.V_x

/ self.hbar * dt )

self.x_evolve = self.x_evolve_half * self.x_evolve_half

self.k_evolve = np.exp(-0.5 * 1j * self.hbar /

self.m * (self.k * self.k) * dt)

psi_x = property(_get_psi_x, _set_psi_x)

psi_k = property(_get_psi_k, _set_psi_k)

dt = property(_get_dt, _set_dt)

def compute_k_from_x(self):

self.psi_mod_k = fft(self.psi_mod_x)

def compute_x_from_k(self):

self.psi_mod_x = ifft(self.psi_mod_k)

def time_step(self, dt, Nsteps = 1):

"""

Perform a series of time-steps via the time-dependent

Schrodinger Equation.

Parameters

----------

dt : float

the small time interval over which to integrate

Nsteps : float, optional

the number of intervals to compute. The total change

in time at the end of this method will be dt * Nsteps.

default is N = 1

"""

self.dt = dt

if Nsteps > 0:

self.psi_mod_x *= self.x_evolve_half

for i in range(Nsteps - 1):

self.compute_k_from_x()

self.psi_mod_k *= self.k_evolve

self.compute_x_from_k()

self.psi_mod_x *= self.x_evolve

self.compute_k_from_x()

self.psi_mod_k *= self.k_evolve

self.compute_x_from_k()

self.psi_mod_x *= self.x_evolve_half

self.compute_k_from_x()

self.t += dt * Nsteps

######################################################################

# Helper functions for gaussian wave-packets

def gauss_x(x, a, x0, k0):

"""

a gaussian wave packet of width a, centered at x0, with momentum k0

"""

return ((a * np.sqrt(np.pi)) ** (-0.5)

* np.exp(-0.5 * ((x - x0) * 1. / a) ** 2 + 1j * x * k0))

def gauss_k(k,a,x0,k0):

"""

analytical fourier transform of gauss_x(x), above

"""

return ((a / np.sqrt(np.pi))**0.5

* np.exp(-0.5 * (a * (k - k0)) ** 2 - 1j * (k - k0) * x0))

######################################################################

# Utility functions for running the animation

def theta(x):

"""

theta function :

returns 0 if x<=0, and 1 if x>0

"""

x = np.asarray(x)

y = np.zeros(x.shape)

y[x > 0] = 1.0

return y

def square_barrier(x, width, height):

return height * (theta(x) - theta(x - width))

######################################################################

# Create the animation

# specify time steps and duration

dt = 0.01

N_steps = 50

t_max = 120

frames = int(t_max / float(N_steps * dt))

# specify constants

hbar = 1.0 # planck's constant

m = 1.9 # particle mass

# specify range in x coordinate

N = 2 ** 11

dx = 0.1

x = dx * (np.arange(N) - 0.5 * N)

# specify potential

V0 = 1.5

L = hbar / np.sqrt(2 * m * V0)

a = 3 * L

x0 = -60 * L

V_x = square_barrier(x, a, V0)

V_x[x < -98] = 1E6

V_x[x > 98] = 1E6

# specify initial momentum and quantities derived from it

p0 = np.sqrt(2 * m * 0.2 * V0)

dp2 = p0 * p0 * 1./80

d = hbar / np.sqrt(2 * dp2)

k0 = p0 / hbar

v0 = p0 / m

psi_x0 = gauss_x(x, d, x0, k0)

# define the Schrodinger object which performs the calculations

S = Schrodinger(x=x,

psi_x0=psi_x0,

V_x=V_x,

hbar=hbar,

m=m,

k0=-28)

######################################################################

# Set up plot

fig = pl.figure()

# plotting limits

xlim = (-100, 100)

klim = (-5, 5)

# top axes show the x-space data

ymin = 0

ymax = V0

ax1 = fig.add_subplot(211, xlim=xlim,

ylim=(ymin - 0.2 * (ymax - ymin),

ymax + 0.2 * (ymax - ymin)))

psi_x_line, = ax1.plot([], [], c='r', label=r'$|\psi(x)|$')

V_x_line, = ax1.plot([], [], c='k', label=r'$V(x)$')

center_line = ax1.axvline(0, c='k', ls=':',

label = r"$x_0 + v_0t$")

title = ax1.set_title("")

ax1.legend(prop=dict(size=12))

ax1.set_xlabel('$x$')

ax1.set_ylabel(r'$|\psi(x)|$')

# bottom axes show the k-space data

ymin = abs(S.psi_k).min()

ymax = abs(S.psi_k).max()

ax2 = fig.add_subplot(212, xlim=klim,

ylim=(ymin - 0.2 * (ymax - ymin),

ymax + 0.2 * (ymax - ymin)))

psi_k_line, = ax2.plot([], [], c='r', label=r'$|\psi(k)|$')

p0_line1 = ax2.axvline(-p0 / hbar, c='k', ls=':', label=r'$\pm p_0$')

p0_line2 = ax2.axvline(p0 / hbar, c='k', ls=':')

mV_line = ax2.axvline(np.sqrt(2 * V0) / hbar, c='k', ls='--',

label=r'$\sqrt{2mV_0}$')

ax2.legend(prop=dict(size=12))

ax2.set_xlabel('$k$')

ax2.set_ylabel(r'$|\psi(k)|$')

V_x_line.set_data(S.x, S.V_x)

######################################################################

# Animate plot

def init():

psi_x_line.set_data([], [])

V_x_line.set_data([], [])

center_line.set_data([], [])

psi_k_line.set_data([], [])

title.set_text("")

return (psi_x_line, V_x_line, center_line, psi_k_line, title)

def animate(i):

S.time_step(dt, N_steps)

psi_x_line.set_data(S.x, 4 * abs(S.psi_x))

V_x_line.set_data(S.x, S.V_x)

center_line.set_data(2 * [x0 + S.t * p0 / m], [0, 1])

psi_k_line.set_data(S.k, abs(S.psi_k))

title.set_text("t = %.2f" % S.t)

return (psi_x_line, V_x_line, center_line, psi_k_line, title)

# call the animator. blit=True means only re-draw the parts that have changed.

anim = animation.FuncAnimation(fig, animate, init_func=init,

frames=frames, interval=30, blit=True)

# uncomment the following line to save the video in mp4 format. This

# requires either mencoder or ffmpeg to be installed on your system

#anim.save('schrodinger_barrier.mp4', fps=15, extra_args=['-vcodec', 'libx264'])

pl.show()