#infinite squre well potential

import numpy as np

import matplotlib.pyplot as plt

from scipy.integrate import odeint,simps

from scipy.optimize import bisect

from scipy.optimize import newton

def f(U,X,E):

y,z=U

f1,f2=z,(V-E)*y

return[f1,f2]

def shoot(E):

sol=odeint(f,U,X,args =(E,))

return sol[:,0][-1]

#parameters and variables

V=0

energies=np.arange(0,200,0.2)

X=np.linspace(0,1,100)

U=[0,0.001]

#shooting for a range of energy values

hits=[shoot(E) for E in energies]

#plot to locate eigen energies

plt.plot(energies,hits)

plt.axhline(linestyle="--")

plt.show()

# energy by root finding method

En=bisect(shoot,0,25)# ground state by bisection method

#eigen function(solution corresponding to eigen energy)

sol=odeint(f,U,X,args=(En,))

psi=sol[:,0]

#normalised eigen function

N=1/np.sqrt(simps(psi*psi,X))
psi_normal=N*psi

#plot of normalised wave function

plt.plot(X,psi_normal)

plt.show()

def psi_normal(En):

sol=odeint(f,U,X,args=(En,))

psi=sol[:,0]

N=1/np.sqrt(simps(psi*psi,X))

return N*psi

for n in range (1,5):

En=newton(shoot,9*n**2)#newton method

plt.plot(X,psi_normal(En))

plt.axhline(linestyle="--")

plt.xlim(0,1)

plt.ylim(-1.5,1.5)

plt.show()