Python 2.7 : Energy eigenvalues and wavefunctions of Simple Harmonic Oscillator (SHO) using Finite Difference Method

Source code:

import numpy as np

import matplotlib.pyplot as plt

def Vpot(x):

    return 0.5*x**2

a = float(input('Provide lower limit of the domain: '))

b = float(input('Provide upper limit of the domain: '))

N = int(input('Provide number of grid points: '))

x = np.linspace(a,b,N)

h = x[1]-x[0]

T = np.zeros((N-2)**2).reshape(N-2,N-2)

for i in range(N-2):

    for j in range(N-2):

        if i==j:

            T[i,j]= -2

        elif np.abs(i-j)==1:

            T[i,j]=1

        else:

            T[i,j]=0

            V = np.zeros((N-2)**2).reshape(N-2,N-2)

for i in range(N-2):

    for j in range(N-2):

        if i==j:

            V[i,j]= Vpot(x[i+1])

        else:

            V[i,j]=0

H = -T/(2*h**2) + V

val,vec=np.linalg.eig(H)

z = np.argsort(val)

z = z[0:4]

energies=0.5*(val[z]/val[z][0])

print(energies)

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

for i in range(len(z)):

    y = []

    y = np.append(y,vec[:,z[i]])

    y = np.append(y,0)

    y = np.insert(y,0,0)

    plt.plot(x,y,lw=3, label="{} ".format(i))

    plt.xlabel('x', size=14)

    plt.ylabel('$\psi$(x)',size=14)

plt.legend()

plt.title('Wavefunctions for a harmonic oscillator using finite difference method',size=14)

plt.show()

Result:

Wave functions

Executing the code