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.

"""
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. Also, use eigenvalue 
functions from scipy (instead of those from numpy).

History:
This file begun: Friday 2020 May 22 19:50:26 IST 
This version: Saturday 2020 May 23 16:07:52 IST 
"""

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( (self.nInteriorNodes, self.nDomainNodes) )
        for c in range( nCntVecs ):
            aPsi = aaEigVecs[ :, c ]
            # Find the area under the prob. curve
            aPsiSq = aPsi * aPsi
            dArea = simps( aPsiSq, aX[ 1 : self.nDomainNodes-1 ] )
            # 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 = aPsi[ 0 ]
                d1 = aPsi[ 1 ]
                if d1 < d0:
                    aPsi *= -1
            aaPsi[ c, 1 : self.nDomainNodes-1 ] = 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

Equations using the matrix notation for a simple artificial neural network

Update on 2020.01.30 16:58 IST:

I have taken the document offline. Yes, it was too incomplete, tentative, and in fact also had errors. The biggest and most obvious error was about the error vector. 🙂

No, no one pointed any of the errors or flaws to me.

Yes, I will post an expanded and revised version later, hopefully in the first week of February. (I started work on this document on 24th Jan., but also was looking into other issues.) When I am done, I will delete this entire post, and make a new entry to announce the availability of the corrected, expanded and revised document.

The original post appears below.


Go, read [^].

Let me know about the typo’s. Also, errors. [Though I don’t expect you to do that. [I will eat my estimate of your moral character, on this count, at least for the time being.]]

I can, and might, take it out of the published domain, any time I want. [Yes, I am irresponsible, careless, unreliable, etc. [Also, “imposing” type. Without “team-spirit”. One who looks down on his colleagues as being way, way, inferior to me.]]

I will also improve on it—I mean the document. In fact, I even intend to expand it, with some brief notes to be added on various activation- and loss-functions.

Eventually, I may even publish it at GitHub. Or, at arXiv. [If they let me do that. But then, another consideration: there are physicists there!]

 


A song I like:

(Hindi) “कभी तो मिलेगी” (kabhee to milegi)
Singer: Lata
Music: Roshan
Lyrics: Majrooh Sultanpuri

[Credits happily listed in a mostly random order. [The issue was only with the last two; it was clear which one had to appear first.]

 

 

Work Is Punishment

Work is not worship—they said.

It’s a punishment, full stop!—they said.

One that is to be smilingly borne.

Else you lose your job.

And so lose everything else too. …


Hmmm… I said. … I was confused.

Work is enjoyment, actually. … I then discovered.

I told them.


They didn’t believe.

Not when I said it.

Not because they ceased believing in me.

It’s just that. They. Simply. Didn’t. Believe. In. It.

And they professed to believe in

a lot of things that never did make

any sense to themselves.

They said so.

And it was so.


A long many years have passed by, since then.


Now, whether they believe in it or not,

I have come to believe in this gem:

Work is punishment—full stop.


That’s the principle on the basis of which I am henceforth going to operate.

And yes! This still is a poem alright?

[What do you think most poems written these days are like?]

It remains a poem.


And I am going to make money. A handsome amount of money.

For once in my life-time.

After all, one can make money and still also write poems.

That’s what they say.

Or do science. Real science. Physics. Even physics for that matter.

Or, work. Real work, too.


It’s better than having no money and…

.