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.

Delete this entry (admin only).