# Demo entry 3295503

Sccrodinger

Submitted by anonymous on Dec 07, 2015 at 15:47
Language: Python. Code size: 8.4 kB.

```# -*- coding: utf-8 -*-
"""
Created on Thu Dec  3 09:51:27 2015

@author: barakovh
"""

#=============================================================================
#Projet d'informatique
#Equation de Schrodinger

import numpy as np
import pylab
import math as mt
# Set pylab to interactive mode so plots update when run outside ipython
pylab.ion()
#=============================================================================
def Gaussian(x,t,sigma):
return np.exp(-(x-t)**2/(2*sigma**2))
def free(npts):
return np.zeros(npts)
def step(npts,v0):
v = free(npts)
v[npts/2:] = v0
return v
def barrier(npts,v0,thickness):
v = free(npts)
v[npts/2:npts/2+thickness] = v0
return v
def harmonic(npts,v0):
v = free(npts)
x = dx*np.linspace(0,N,N)
center = 600
v = v0*0.001*(x-center)**2-0.5
## Multiplication by 0.001 makes the potential wider
## The factor -0.5 makes a translation of the potential
return v
def coulomb(npts,v0,thickness):

v = free(npts)

for i in xrange(npts/2+200,1100,1): ## Defines the witdth of the potential

v[i] = 0.1*v0*npts/(i-600) ## Defines the potential
for i in xrange(1100,1200,1):
v[i] =-10*v0              ## Defines the witdth of the potential
for i in xrange(100,npts/2-200,1): ## Defines the witdth of the potential

v[i] = 0.1*v0*npts/(600-i) ## Defines the potential
for i in xrange(0,100,1):
v[i] =-10*v0              ## Defines the witdth of the potential
return v
def double_barrier(npts,v0,thickness):
v = free(npts)
v[npts/4-1*thickness:npts/4] = 1.*v0                ## Defines the first barrier
v[3*npts/4:3*npts/4+1*thickness] = 1.*v0            ## Defines the second barrier

return v

def barriers(npts,v0,thickness):
v = free(npts)
v[npts/8-1*thickness:npts/8] = 1.*v0
v[3*npts/8-1*thickness:3*npts/8] = 1.*v0
v[npts/4-1*thickness:npts/4] = 1.*v0                ## Defines the first barrier
v[npts/2:npts/2+1*thickness] = 1.*v0                ## Defines the second barrier
v[5*npts/8:5*npts/8+1*thickness] = 1.*v0            ## Defines the second barrier
v[6*npts/8:6*npts/8+1*thickness] = 1.*v0            ## Defines the second barrier
v[7*npts/8:7*npts/8+1*thickness] = 1.*v0            ## Defines the second barrier
return v
#=============================================================================
N    = 1200     #  Number of spatial points.
T    = 1*N      #  Number of time steps.  5*N is a nice value for terminating
#  before anything reaches the boundaries.
Tp   = 20       #  Number of time steps to increment before updating the plot.
dx   = 1.0e0    #  Spatial resolution
X    = dx*np.linspace(0,N,N)        #  Spatial axis.
V0   = 1.0e-2
THCK = 15
POTENTIAL = 'coulomb'
#  Initial wave function constants
sigma = 40.0
x0 = round(N/2)-1*sigma       # Time shift ===> no time shift the gaussian is centered at half of the position interval
# to translate it add a multiple of sigma
k0 = np.pi/40                # Wavevector
#=============================================================================
# Code begins
if POTENTIAL=='free':
V = free(N)
elif POTENTIAL=='step':
V = step(N,V0)
elif POTENTIAL=='barrier':
V = barrier(N,V0,THCK)
elif POTENTIAL=='harmonic':
V = harmonic(N,V0)
elif POTENTIAL=='double_barrier':
V = double_barrier(N,V0,THCK)
elif POTENTIAL=='coulomb':
V = coulomb(N,V0,THCK)
elif POTENTIAL=='barriers':
V = barriers(N,V0,THCK)
#  The maximum stable time step is a function of the potential, V.
#   All the constants are vectors except Vmax
Vmax = V.max()                      #  Maximum potential of the domain.
dt   = 1/(2/(dx**2)+Vmax)           #  Critical time step.
c1   = dt/(dx**2)                   #  Constant coefficient 1.
c2   = 2*dt                         #  Constant coefficient 2.
c2V  = c2*V                         # pre-compute outside of update loop

#  Wave functions. Three states represent past(i-1), present(i), and future(i+1)
psi_r = np.zeros((3,N))  #  Real
psi_i = np.zeros((3,N))  #  Imaginary
psi_p = np.zeros(N,)     #  Probability

#  Temporal indexing constants, used for accessing rows of the wavefunctions.
PA = 0                 #  Past
PR = 1                 #  Present
FU = 2                 #  Future
#==============================================
#  Initialize wave function.
xn = range(1,N/2)
x = X[xn]/dx    #  Normalized position coordinate
gg = Gaussian(x,x0,sigma)
cx = np.cos(k0*x)
sx = np.sin(k0*x)
psi_r[PR,xn] = cx*gg
psi_i[PR,xn] = sx*gg
psi_r[PA,xn] = cx*gg
psi_i[PA,xn] = sx*gg
# Initial normalization of wavefunctions
psi_p = psi_r[PR]**2 + psi_i[PR]**2
#  Normalize the wave functions so that the total probability in the simulation is equal to 1.
P   = dx * psi_p.sum()      #  Total probability. Integral of the norm squared
nrm = np.sqrt(P)
psi_r /= nrm                #  x/=y means x=x/y
psi_i /= nrm
psi_p /= P
#  Initialize the figure and axes.
pylab.figure()
xmin = X.min()
xmax = X.max()
ymax = 1.5*(psi_r[PR]).max()   ## Maybe it is better to put  psi_p[PR] althought it is very small if not multiplied
pylab.axis([xmin,xmax,-ymax,ymax])
#  Initialize the plots
lineR, = pylab.plot(X,psi_r[PR],'b',alpha=0.7,label='Real')
lineI, = pylab.plot(X,psi_i[PR],'r',alpha=0.7,label='Imag')
lineP, = pylab.plot(X,6*psi_p,'k',label='Prob')

#  Plot non-zero potentials
if Vmax !=0 :
Efac = ymax/2.0/Vmax
V_plot = V*Efac
pylab.plot(X,V_plot,':k',zorder=0)   #  Potential line.
pylab.legend(loc='lower right')
pylab.draw()
pylab.xlim(xmin,xmax)
#  Direct index assignment is MUCH faster than using a spatial FOR loop, so
#  these constants are used in the update equations.
IDX1 = range(1,N-1)                            #  psi [ k ]
IDX2 = range(2,N)                              #  psi [ k + 1 ]
IDX3 = range(0,N-2)                            #  psi [ k - 1 ]
IDXT = np.zeros (N)
IDXT[0:N]=600
psi_pR=np.zeros(N)
psi_pL=np.zeros(N)
Somme=0
for t in range(T+1):
# Precompute a couple of indexing constants, this speeds up the computation
psi_rPR = psi_r[PR]
psi_iPR = psi_i[PR]
#  Apply the update equations.(Schrodinger equation for the real and the imaginary part)
psi_i[FU,IDX1] = psi_i[PA,IDX1] +c1*(psi_rPR[IDX2] - 2*psi_rPR[IDX1] + psi_rPR[IDX3])

psi_i[FU] -= c2V*psi_r[PR]

psi_r[FU,IDX1] = psi_r[PA,IDX1] - c1*(psi_iPR[IDX2] - 2*psi_iPR[IDX1] +psi_iPR[IDX3])

psi_r[FU] += c2V*psi_i[PR]

#  a += b is essentially the same as a = a + b

#  Increment the time steps.  PR -> PA and FU -> PR
psi_r[PA] = psi_rPR
psi_r[PR] = psi_r[FU]
psi_i[PA] = psi_iPR
psi_i[PR] = psi_i[FU]
#  Only plot after a few iterations to make the simulation run faster- Tp is fixed in the begining
if t % Tp == 0:
#        "Calculating the probability on the left of the barrier and on the right"
Right=0
Left=0
for i in range (2,N/2-2):

if  IDX2[i]<=IDXT[i]:

psi_pL[i] = psi_r[PR,i]**2 + psi_i[PR,i]**2
Left=Left+psi_pL[i]
##Left=mt.ceil(Left*1000)/1000    ## Round up the numbers
print(Left)

for i in range (N/2-1,N-2):

if  IDX2[i]>=IDXT[i]:

psi_pR[i] = psi_r[PR,i]**2 + psi_i[PR,i]**2
Right=Right+psi_pR[i]
##Right=mt.ceil(Right*1000)/1000   ## Round up the numbers
print(Right)

Somme=Right+Left
print(Somme)
##=======================================================================================
#  Compute observable probability for the plot.
psi_p = psi_r[PR]**2 + psi_i[PR]**2
#  Update the plots.
lineR.set_ydata(psi_r[PR])
lineI.set_ydata(psi_i[PR])
# the probability density is amplified by a factor so we can see it
# Multiplied by 6
lineP.set_ydata(5*psi_p)
pylab.pause(0.001)
pylab.draw()
##print(psi_pG[t])
## Calculate the norm squared of the function at each time and stock it in a vector
## TO DO
pylab.ioff()
pylab.show()
```

This snippet took 0.01 seconds to highlight.

Back to the Entry List or Home.