Python scripts for simulating QM, part 1: Time evolution of a particle in the infinite potential box

A Special note for the Potential Employers from the Data Science field:

Recently, in April 2020, I achieved a World Rank # 5 on the MNIST problem. The initial announcement can be found here [^], and a further status update, here [^].

All my data science-related posts can always be found here [^].


What’s been happening?

OK, with that special note done, let me now turn my attention to the two of you who regularly read this blog.

… After the MNIST classification problem, I turned my attention to using LSTM’s for time-series predictions, simply because I hadn’t tried much on this topic any time earlier. … So, I implemented a few models. I even seemed to get good accuracy levels.

However, after having continuously worked on Data Science straight for about 2.5 months, I began feeling like taking a little break from it.

I had grown tired a bit though I had not realized it while actually going through all those tedious trials. (The fact of my not earning any money might have added to some draining of the energy too.) In fact, I didn’t have the energy to pick up something on the reading side either.

Basically, after a lot of intense effort, I wanted something that was light but engaging.

In the end, I decided to look into Python code for QM.

Earlier, in late 2018, I had written a few scripts on QM. I had also blogged about it; see the “part 0” of this series [^]. (Somehow, it has received unusually greater number of hits after I announced my MNIST result.) However, after a gap of 1.5 years, I could not easily locate those scripts. … So, like any self-respecting programmer, I decided to code them again!

Below is the first result, a movie. … Though a movie, it should be boring to any one who is not interested in QM.


Movie of the wavefunction of an electron inside an infinite potential box:

 

An electron in an infinite potential box.

An electron inside a 3 cm long 1D box with infinite potentials on the sides. Time evolution of the second excited state (n = 3). In the standard QM, such states are supposed to be“stationary”.

 


The Python code: Main file:

Now, the code. It should be boring to any one who is not a Python programmer.

Update on 2020.05.30 16:00 IST: The code below has been updated to fix a bug in the square-normalization procedure. Earlier, the code was performing the Simpson integration over
the interior lattice points alone, which was incorrect. Now, the Simpson integration is performed over the entire domain, which is the correct procedure. Now, numerical and analytical solutions
are truly close.

"""
01.AJ_PIB_1D_Class.py

Written by and copyright (c) Ajit R. Jadhav. All rights reserved.

Particle in a Box.

Solves the Time-Independent Schrodinger Equation in 1D,
using the Finite Difference Method. Eigenvalues are found using
the direct matrix method.

Also shows the changes in the total wavefunction with time.
[The stationarity in the TISE is not static. In the mainstream 
QM, the stationarity is kinematical. In my new approach, it 
has been proposed to be kinetic. However, this simulation concerns
itself only with the standard, mainstream, QM.]

Environment: Developed and tested using:
Python 3.7.6 64-bit. All packages as in Anaconda 4.8.3 
on Ubuntu-MATE 20.04 (focal fossa) LTS of the date (below).

TBD: It would be nice to use sparse matrices, and to use 
eigenvalue functions from scipy (instead of those from numpy).

History:
0. This file begun: Friday 2020 May 22 19:50:26 IST 
1. Initially posted on the blog: Saturday 2020 May 23 16:07:52 IST 
2. This version: Saturday 2020 May 30 15:53:13 IST 
    -- Corrected a bug in ComputeNormalizedStationaryStates() 
    Earlier, the code was performing the Simpson integration over 
    the *interior* lattice points alone, which was incorrect. Now, the 
    Simpson integration is performed over the entire domain, which is 
    the correct procedure. Now, numerical and analytical solutions 
    are truly close.
"""

import numpy as np 
from scipy.integrate import simps
import matplotlib.pyplot as plt 
from matplotlib.animation import ImageMagickFileWriter

# SEE THE ACCOMPANYING FILE. THE NUMERICAL VALUES OF CONSTANTS 
# ARE DEFINED IN IT.
from FundaConstants import h, hbar, me, mP, eV2J

################################################################################
# THE MAIN CLASS 

class AJ_PIB_1D( object ):

    def __init__( self, nInteriorNodes, dMass, dh ):
        self.nInteriorNodes = nInteriorNodes
        self.nDomainNodes = nInteriorNodes + 2
        self.dMass = dMass # Mass associated with the QM particle
        self.dh = dh # cell-size ( \Delta x ).

        # The following numpy ndarray's get allocated 
        # during computations.
        self.aaT = None 
        self.aaV = None 
        self.aaH = None 

        self.aaPsi = None 
        self.aE_n = None 

        self.A_ana = None 
        self.ak_ana = None 
        self.aE_ana = None 
        self.aPsi_ana = None 
        return 

    # Creates the kinetic energy matrix for the interior points of the
    # domain.
    def ComputeKE( self ):
        self.aaT = np.zeros( (self.nInteriorNodes, self.nInteriorNodes) )
        for i in range( self.nInteriorNodes ):
            self.aaT[ i, i ] = -2.0
        for i in range( self.nInteriorNodes-1 ):
            self.aaT[ i, i+1 ] = 1.0
        for i in range( 1, self.nInteriorNodes ):
            self.aaT[ i, i-1 ] = 1.0 
        dFactorKE = - hbar**2 / ( 2.0 * self.dMass * self.dh**2 )
        self.aaT *= dFactorKE
        return

    # Creates the potential energy matrix for the interior points of the
    # domain. You can supply an arbitrary potential function via the array
    # aV of size = interior points count, and values in joule.
    def ComputePE( self, aV= None ):
        self.aaV = np.zeros( (self.nInteriorNodes, self.nInteriorNodes) )
        if None != aV:
            for i in range( self.nInteriorNodes ):
                self.aaV[ i, i ] = aV[ i ]
        return

    def ComputeHamiltonian( self, aV= None ):
        self.ComputeKE() 
        self.ComputePE( aV )
        self.aaH = self.aaT + self.aaV 
        return 

    # Note, the argument aX has the size = the count of domain points, not 
    # the count of interior points.
    # QM operators are Hermitian. We exploit this fact by using the 
    # numpy.linalg.eigh function here. It is faster than numpy.linalg.eig, 
    # and, unlike the latter, also returns results sorted in the ascending 
    # order. 
    # HOWEVER, NOTE, the eigenvectors returned can have signs opposite 
    # of what the analytial solution gives. The eigh (or eig)-returned 
    # vectors still *are* *eigen* vectors. However, for easier comparison 
    # with the analytical solution, we here provide a quick fix. 
    # See below in this function.
    def ComputeNormalizedStationaryStates( self, aX, bQuickFixForSigns= False ):
        assert( self.nDomainNodes == len( aX ) )
        
        # Use the LAPACK library to compute the eigenvalues
        aEigVals, aaEigVecs = np.linalg.eigh( self.aaH )
        
        # SQUARE-NORMALIZE THE EIGENVECTORS

        # Note:
        # The eigenvectors were found on the interior part of the domain, 
        # i.e., after dropping the boundary points at extreme ends. But the 
        # wavefunctions are defined over the entire domain (with the 
        # Dirichlet condition of 0.0 specified at the boundary points).
        
        nCntVecs = aaEigVecs.shape[ 1 ]
        assert( nCntVecs == self.nInteriorNodes )

        # eigh returns vectors in *columns*. We prefer to store the 
        # normalized vectors in *rows*.
        aaPsi = np.zeros( (nCntVecs, self.nDomainNodes) )
        for c in range( nCntVecs ):
            aPsi = aaPsi[ c ] # For ease of readability
            aPsi[ 1 : self.nDomainNodes-1 ] = aaEigVecs[ :, c ]
            # Find the area under the prob. curve
            aPsiSq = aPsi * aPsi
            dArea = simps( aPsiSq, aX )
            # Use it to normalize the wavefunction
            aPsi /= np.sqrt( dArea )
            # The analytical solution always has the curve going up 
            # (with a +ve gradient) at the left end of the domain. 
            # We exploit this fact to have a quick fix for the signs.
            if bQuickFixForSigns is True:
                d0, d1 = aPsi[ 0 ], aPsi[ 1 ]
                if d1 < d0:
                    aPsi *= -1
            aaPsi[ c ] = aPsi
        self.aaPsi = aaPsi
        self.aE_n = aEigVals
        return 

    # Standard analytical solution. See, e.g., the Wiki article: 
    # "Particle in a box"
    def ComputeAnalyticalSolutions( self, aX ):

        xl = aX[ 0 ]
        xr = aX[ self.nDomainNodes-1 ]
        L = xr - xl 
        A = np.sqrt( 2.0 / L )
        self.A_ana = A 

        # There are as many eigenvalues as there are interior points
        self.ak_ana = np.zeros( self.nInteriorNodes )
        self.aE_ana = np.zeros( self.nInteriorNodes )
        self.aPsi_ana = np.zeros( (self.nInteriorNodes, self.nDomainNodes) )
        for n in range( self.nInteriorNodes ):
            # The wavevector. (Notice the absence of the factor of '2'. 
            # Realize, the 'L' here refers to half of the wavelength of 
            # the two travelling waves which make up the standing wave. 
            # That's why.)
            k_n = n * np.pi / L 
            # The energy.
            E_n = n**2 * h**2 / (8.0 * self.dMass * L**2)

            # A simplest coordinate transformation:
            # For x in [0,L], phase angle is given as
            # Phase angle = n \pi x / L = k_n x. 
            # We have arbitrary coordinates for the left- and 
            # right-boundary point. So, 
            # Phase angle = k_n (x - xl)
            ap = k_n * (aX - xl)
            
            aPsiAna = A * np.sin( ap )
            self.ak_ana[ n ] = k_n 
            self.aE_ana[ n ] = E_n 
            # We prefer to store the normalized wavefunction 
            # in rows. (Contrast: linalg.eigh and eig store the 
            # eigen vectors in columns.)
            self.aPsi_ana[ n, : ] = aPsiAna
        return 
        
    # This function gets the value that is the numerical equivalent to the 
    # max wave amplitude 'A', i.e., sqrt(2/L) in the analytical solution. 
    def GetMaxAmplNum( self ):
        dMax = np.max( np.abs(self.aaPsi) )
        return dMax


################################################################################
# Utility functions

# NOTE: SAVING MOVIES CAN TAKE A LOT MORE TIME (7--10 MINUTES).
def Plot( model, n, nTimeSteps, bSaveMovie= False, sMovieFileName= None ):
    # The class computes and stores only the space-dependent part.
    aPsiSpace = model.aaPsi[ n-1 ]

    # Period = 2 \pi / \omega = 1 / \nu
    # Since E = h \nu, \nu = E/h, and so, Period = h/E
    nu = model.aE_n[ n-1 ] / h 
    dPeriod = 1.0 / nu 

    dt = dPeriod / (nTimeSteps-1) 

    # Plotting...

    plt.style.use( 'ggplot' )
    # Plot size is 9 inches X 6 inches. Reduce if you have smaller 
    # screen size.
    fig = plt.figure( figsize=(9,6) ) 

    if bSaveMovie is True:
        movieWriter = ImageMagickFileWriter()
        movieWriter.setup( fig, sMovieFileName )

    dMaxAmpl = model.GetMaxAmplNum() # Required for setting the plot limits.
    dTime = 0.0 # How much time has elapsed in the model?
    for t in range( nTimeSteps ):
        # TIME-DEPENDENT PART: 
        # \psi_t = e^{-i E_n/\hbar t} = e^{-i \omega_n t} = e^{-i 2 \pi nu t}
        # Compute the phase factor (which appears in the exponent).
        dTheta = 2.0 * np.pi * nu * dTime 
        # The Euler identity. Compute the *complete* wavefunction (space and time)
        # at this instant.
        aPsi_R_t = aPsiSpace * np.cos( dTheta )
        aPsi_I_t = - aPsiSpace * np.sin( dTheta )
        
        plt.clf()
        sTitle = "Particle in an infinite-potential box (n = %d)\n" % (n)
        sTitle += "Domain size: %7.4lf m. Oscillation period: %7.4lf s.\n" % (L, dPeriod)
        sTitle += "Time step: %3d/%3d. Time elapsed in simulation: %7.4lf s." % (t+1, nTimeSteps, dTime) 
        plt.title( sTitle )

        plt.xlabel( "Distance, m" )
        plt.ylabel( "Wavefunction amplitude, $m^{-1/2}$" )

        plt.grid( True )

        plt.xlim( (xl - L/10), (xr + L/10) )
        plt.ylim( -1.1*dMaxAmpl, 1.1*dMaxAmpl )

        plt.plot( aX, aPsi_R_t , color= 'darkcyan', label= r'Re($\Psi$)' )
        plt.plot( aX, aPsi_I_t , color= 'purple', label= r'Im($\Psi$)' )

        plt.legend( loc= 'upper right', shadow= True, fontsize= 'small' ) 
        
        if bSaveMovie is True:
            movieWriter.grab_frame()
        else:
            plt.pause( 0.001 )

        dTime += dt

    if bSaveMovie is True:
        movieWriter.finish()
    else:
        plt.show()


################################################################################
# MAIN DRIVER CODE
# We use the SI system throughout. [This is a program. It runs on a computer.]

# DOMAIN GEOMETRY

xl = -1.0e-02 # Left end (min. x)
xr = 2.0e-02 # Right end (max. x)
L = xr - xl # Length of the domain
xc = (xl + xr )/ 2.0 # Center point

# MESH
# Count of cells = Count of nodes in the domain - 1. 
# It's best to take an odd number for the count of domain nodes. This way, 
# the peak(s) of the wavefunction will not be missed.
nDomainNodes = 101
aX, dh = np.linspace(   start= xl, stop= xr, num= nDomainNodes, 
                        endpoint= True, retstep= True, dtype= np.float )

# In the PIB model, infinite potential exists at either ends. So, we apply 
# the Dirichlet BC of \Psi(x,t) = 0 at all times. Even if the discretized 
# Laplacian were to be computed for the entire domain, in handling the 
# homogeneous BC, both the boundary-points would get dropped during the
# matrix partitioning. Similarly, V(x) would be infinite there. That's why,
# we allocate the Laplacian and Potential Energy matrices only for the
# interior points. 
nInteriorNodes = nDomainNodes - 2 

# We model the electron here. 
# Constants are defined in a separate file: 'FundaConstants.py'
# Suggestion: Try mP for the proton, and check the \Psi amplitudes.
dMass = me 

# Instantiate the main model class.
model = AJ_PIB_1D( nInteriorNodes, dMass, dh )

# Compute the system Hamiltonian.

# If you want, supply a custom-made potential function as an ndarray of 
# size nInteriorNodes, as an argument. Values should be in joules.
# 'None' means 0 joules everywhere inside the box.
model.ComputeHamiltonian( None )

# Compute the stationary states. For the second argument, see the 
# note in the function implementation.
model.ComputeNormalizedStationaryStates( aX, True )

# You can also have the analytical solution computed. Uncomment the 
# line below. The numerical and analytical solutions are kept in 
# completely different arrays inside the class. However, the plotting
# code has to be careful.
### model.ComputeAnalyticalSolutions( aX )


# PLOT THE STATIONARY STATES, AND SHOW THEIR OSCILLATIONS WITH TIME.
# (TISE *is* dynamic; the stationarity is dynamical.) 

# Note, here, we choose n to be a 1-based index, as is the practice
# in physics. Thus, the ground state is given by n = 1, and not n = 0.
# However, NOTE, the computed arrays of wavefunctions have 0-based 
# indices. If such dual-usage for the indices gets confusing, simple! 
# Just change the code!

n = 3 
# No. of frames to be plotted for a single period of oscillations
# The 0-th and the (nTimeSteps-1)th state is identical because the 
# Hamiltonian here is time-independent. 
nTimeSteps = 200 

# You can save a movie, but note, animated GIFs take a lot more time, even 
# ~10 minutes or more, depending on the screen-size and dpi.
# Note, ImageMagickFileWriter will write the temp .png files in the current 
# directory (i.e. the same directory where this Python file resides). 
# In case the program crashes (or you stop the program before it finishes), 
# you will have to manually delete the temporary .png files from the 
# program directory! (Even if you specify a separate directory for the 
# movie, the *temporary* files still get generated in the program directory.)
### Plot( model, n, nTimeSteps, True, './AJ_PIB_e_%d.gif' % (n) )

Plot( model, n, nTimeSteps, False, None )

################################################################################
# EOF
################################################################################


The ancillary file:

The main file imports the following file. It has nothing but the values of fundamental physical constants noted in it (together with the sources). Here it is:

"""
FundaConstants.py

Written by and copyright (c) Ajit R. Jadhav. All rights reserved.

Begun: Thursday 2020 May 21 20:55:37 IST 
This version: Saturday 2020 May 23 20:39:22 IST 
"""

import numpy as np 

"""
Planck's constant
https://en.wikipedia.org/wiki/Planck_constant 
``The Planck constant is defined to have the exact value h = 6.62607015×10−34 J⋅s in SI units.'' 
"""
h = 6.62607015e-34 # J⋅s. Exact value.
hbar = h / (2.0 * np.pi) # J⋅s. 

"""
Electron rest mass
https://en.wikipedia.org/wiki/Electron_rest_mass
``9.1093837015(28)×10−31'' 2018 CODATA value. NIST
"""
me = 9.1093837015e-31 # kg. 


"""
Proton rest mass
https://en.wikipedia.org/wiki/Proton 
1.67262192369(51)×10−27 kg
"""
mP = 1.67262192369e-27 # kg

"""
eV to Joule
https://en.wikipedia.org/wiki/Electronvolt 
1 eV = 1.602176634×10−19 J
"""
eV2J = 1.602176634e-19 # Conversion factor

And, that’s about it, folks! No documentation. But I have added a lot of (otherwise unnecessary) comments.

Take care, and bye for now.


A song I like:

(Marathi) सावली उन्हामध्ये (“saavali unaamadhe”)
Music: Chinar Mahesh
Lyrics: Chandrashekhar Sanekar
Singer: Swapnil Bandodkar