Double potential well with Python

Question

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()

Answer 1

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.