Double potential well with Python

I'm trying to understand the Schrödinger equation and solving it a bit better, and I'm running into some doubts while coding, even though I am adapting the code to this situation. Also I tried asking this on the physics stack exchange and I was directed here. this is the situation I am trying to solve: It's supposed to be an infinite well with a barrier in the middle. alpha b is just b multiplied by some constant. I have a running code, but since I am basing it off of someone else's code there's one setting I don't seem to understand (probably some basic concept). I'm having trouble defining the Wave function initial states so how can I define the initial state of psi?

I'm adding the code in case if anyone wants to play with this code.

from pylab import *
from scipy.integrate import odeint
from scipy.optimize import brentq

a=1
B=4
L= B+a
Vmax= 50
Vpot = False
N = 1000                  # number of points to take
psi = np.zeros([N,2])     # Wave function values and its derivative (psi and psi')
psi0 = array([0,1])   # Wave function initial states
Vo = 50
E = 0.0                   # global variable Energy  needed for Sch.Eq, changed in function "Wave function"
b = L                     # point outside of well where we need to check if the function diverges
x = linspace(-B-a, L, N)    # x-axis

def V(x):
'''
#Potential function in the finite square well.
'''
if -a <=x <=a:
val = Vo
elif x<=-a-B:
val = Vmax
elif x>=L:
val = Vmax
else:
val = 0
if Vpot==True:
if -a-B-(10/N) < x <= L+(1/N):
Ypotential.append(val)
Xpotential.append(x)
return val

def SE(psi, x):
"""
Returns derivatives for the 1D schrodinger eq.
Requires global value E to be set somewhere. State0 is first derivative of the
wave function psi, and state1 is its second derivative.
"""
state0 = psi[1]
state1 = 2.0*(V(x) - E)*psi[0]
return array([state0, state1])

def Wave_function(energy):
"""
Calculates wave function psi for the given value
of energy E and returns value at point b
"""
global psi
global E
E = energy
psi = odeint(SE, psi0, x)
return psi[-1,0]

def find_all_zeroes(x,y):
"""
Gives all zeroes in y = Psi(x)
"""
all_zeroes = []
s = sign(y)
for i in range(len(y)-1):
if s[i]+s[i+1] == 0:
zero = brentq(Wave_function, x[i], x[i+1])
all_zeroes.append(zero)
return all_zeroes

def find_analytic_energies(en):
"""
Calculates Energy values for the finite square well using analytical
model (Griffiths, Introduction to Quantum Mechanics, 1st edition, page 62.)
"""
z = sqrt(2*en)
z0 = sqrt(2*Vo)
z_zeroes = []
f_sym = lambda z: tan(z)-sqrt((z0/z)**2-1)      # Formula 2.138, symmetrical case
f_asym = lambda z: -1/tan(z)-sqrt((z0/z)**2-1)  # Formula 2.138, antisymmetrical case

# first find the zeroes for the symmetrical case
s = sign(f_sym(z))
for i in range(len(s)-1):   # find zeroes of this crazy function
if s[i]+s[i+1] == 0:
zero = brentq(f_sym, z[i], z[i+1])
z_zeroes.append(zero)
print ("Energies from the analyitical model are: ")
print ("Symmetrical case)")
for i in range(0, len(z_zeroes),2):   # discard z=(2n-1)pi/2 solutions cause that's where tan(z) is discontinous
print ("%.4f" %(z_zeroes[i]**2/2))
# Now for the asymmetrical
z_zeroes = []
s = sign(f_asym(z))
for i in range(len(s)-1):   # find zeroes of this crazy function
if s[i]+s[i+1] == 0:
zero = brentq(f_asym, z[i], z[i+1])
z_zeroes.append(zero)
print ("(Antisymmetrical case)")
for i in range(0, len(z_zeroes),2):   # discard z=npi solutions cause that's where ctg(z) is discontinous
print ("%.4f" %(z_zeroes[i]**2/2))

def main():
# main program

en = linspace(0, Vo, 1000)   # vector of energies where we look for the stable states

psi_b = []      # vector of wave function at x = b for all of the energies in en
for e1 in en:
psi_b.append(Wave_function(e1))     # for each energy e1 find the the psi(x) at x = b
E_zeroes = find_all_zeroes(en, psi_b)   # now find the energies where psi(b) = 0

# Print energies for the bound states
print ("Energies for the bound states are: ")
for E in E_zeroes:
print ("%.2f" %E)
# Print energies of each bound state from the analytical model
find_analytic_energies(en)

# Plot wave function values at b vs energy vector
figure()
plot(en/Vo,psi_b)
title('Values of the $$\Psi(b)$$ vs. Energy')
xlabel('Energy, $$E/V_0$$')
ylabel('$$\Psi(x = b)$$', rotation='horizontal')
for E in E_zeroes:
plot(E/Vo, [0], 'go')
annotate("E = %.2f"%E, xy = (E/Vo, 0), xytext=(E/Vo, 30))
grid()

# Plot the wavefunctions for first 4 eigenstates
figure(2)
for E in E_zeroes[0:4]:
Wave_function(E)
plot(x, psi[:,0], label="E = %.2f"%E)
legend(loc="upper right")
title('Wave function')
xlabel('x, $$x/L$$')
ylabel('$$\Psi(x)$$', rotation='horizontal', fontsize = 15)
grid()

figure(3)
pot =[]
for i in x:
pot.append(V(i))
plot(x,pot)
show()
if __name__ == "__main__":
main()

• I've solved the finite square well potential using scipy.integrate.solve_bvp here. – Russell Anderson Aug 13 '19 at 8:41
• From experience it’s much easier and faster to solve this equation with tridiagonal matrix approach. See here, docs.solcore.solar/en/master/Examples/example_QWs.html (I was one of the authors of this module) – boyfarrell Aug 16 '19 at 16:57

Although this site is not intended for questions of the "please debug my program" kind, I believe that what's needed here is a clarification of the scientific computing aspects.

The method you are adopting, from the link that you gave, is the shooting method. The program uses a numerical ordinary differential equation solver for the Schrodinger equation, which takes initial conditions $$\psi(x_0)$$ and $$\psi'(x_0)$$, for some chosen starting point $$x_0$$, and integrates forward in $$x$$. The output is a wavefunction at discrete values of $$x$$; the program carries out this integration for a range of energies $$E$$.

Although your question refers to an infinite potential well, your diagram, your program, and your original source all deal with a finite potential well. So I'm going to stick with that. (To be honest, in the original program given at https://helentronica.com/2014/09/04/quantum-mechanics-with-the-python/, I can't be sure that the sign of $$V_0$$ is correct! but I haven't looked closely). If you really mean an infinite well, then there are only slight differences, which I return to at the end. For simplicity I'm also going to stick with the symmetric case, $$\alpha=1$$.

The main part of the program is mainly interested in checking the value of $$\psi(b)$$ for a position $$x=b$$ lying far outside the potential well; this is why the Wavefunction function only returns the last element of the computed $$\psi$$ array. Only values of $$E$$ which give non-divergent wavefunctions are wanted; in fact, the boundary condition that we want to enforce is $$\psi(x)\rightarrow 0$$ as $$x\rightarrow\infty$$. The program uses the array of values of $$\psi(b)$$ for different energies to determine intervals of $$E$$ where $$\psi(b)$$ changes sign, and then uses Brent's method to refine the value of $$E$$ for which $$\psi(b)=0$$. To a good approximation, this will be one of the allowable eigenvalues and the corresponding $$\psi(x)$$ will be the eigenfunction.

The "state" which is required by the ODE solver, calculated by the SE function, will always be an array having two elements: the first derivative and the second derivative of $$\psi$$. This is because the Schrodinger equation is a second-order differential equation, which is being tackled here as two first order equations: \begin{align*} \frac{d\psi}{dx} &= \psi'(x) \\ \frac{d\psi'}{dx} &= 2[V(x)-E]\psi(x) \end{align*} in your units. The psi array stores the values of $$\psi(x)$$ and $$\psi'(x)$$; the SE function returns the two quantities on the right hand side of these equations. So, the changes you have made to the SE function, returning 6 values, are on the wrong track. The value of the second derivative in the different regions of the potential well will depend on $$x$$ through the function $$V(x)$$, which you set up at the start of the program. So nothing in SE needs changing from the original code which you downloaded.

The other aspects needing attention are the ranges of $$x$$. In the original code you started working with, the width of the well is determined by $$L$$ (which was set to 1 in the sample program) and the check on $$\psi(b)$$ was made at $$b\gg L$$ ($$b$$ was set to 2 in the example). You need to do something similar. Various starting points $$x_0$$ (for the shooting integration) were considered in the original example; for a symmetric potential well, there is something to be said for starting at the midpoint, $$x=0$$, with either $$\psi(0)=0$$ or $$\psi'(0)=0$$ to produce, respectively, antisymmetric or symmetric solutions. But you can also start at $$-b$$, where $$-b\ll -L$$, and this might be better in your case. Finally, as someone noted in this answer https://physics.stackexchange.com/a/438299/197851 to your question on Physics SE, the locations of the edges of the well in your diagram do not match the same locations in your code.

So, to summarize: there are really no fundamental changes needed to the original program, just changing $$V(x)$$ should work. The SE routine does not need changing from the originally downloaded version. You need to pay more attention to the limits of the $$x$$ range used in the shooting method: you should end at a position which is outside the potential well on the right, and you might start from a position outside the well, on the left. I recommend starting again with the original program, where the well extends from $$-L$$ to $$+L$$ where $$L=1$$, and adding a symmetrical bump around the origin from $$-a$$ to $$+a$$, where $$a<1$$, keeping $$b=2$$.

If you really want to deal with an infinite potential well, then you should set $$b=L$$ and enforce the boundary condition $$\psi(b)=0$$. In this case it also makes sense to start shooting at $$x=-b$$, with $$\psi(-b)=0$$ and $$\psi'(b)$$ nonzero.

Hopefully this helps. As I indicated, this site isn't for program debugging, but if you don't follow the idea behind the method you are using, I'm happy to try and clarify.

• what I was hoping for was an answer like yours, sorry for not expressing myself correctly – user169808 Nov 5 '18 at 16:47
• No worries. The intention of my first paragraph was to make clear that (IMHO) this question falls within the scope of the site, even though one needs to look at the code. Anyway, feel free to consider the answer, decide whether it is helpful, and whether to accept it or not (see scicomp.stackexchange.com/help/accepted-answer). There is no hurry: better answers from others might come along. Also, as I said, if it needs more clarification, just ask. – user28077 Nov 5 '18 at 16:56
• thanks, but you clarified some things for me. I'm trying to learn how to do these sorts of things on python, on my own, so I'm bound to be a bit lost. Thanks for your help. – user169808 Nov 5 '18 at 17:01
• hi, I'm trying to normalize the wave function but when I integrate the result I get is not equal to 1. I am integrating abs(psi)^2 from one side of the well to the other. Is there something I am missing to normalize these wavefunctions? In the end I want to calculate $\Delta x \Delta p$ – user169808 Nov 12 '18 at 12:21
• The Schrödinger equation is a linear equation, so if $\psi(x)$ satisfies it, so will $c\psi(x)$ where $c$ is any constant. This is equally true for your solutions: if you doubled the initial values of $\psi$ and $\psi'$, and redid the "shooting", you would just multiply your whole solution by $2$. It is quite usual to normalize the wavefunction "by hand" afterwards: calculate $C=\int |\psi(x)|^2 dx$ across the whole range of $x$, and then replace $\psi(x)\rightarrow \psi(x)/\sqrt{C}$. – user28077 Nov 12 '18 at 12:42