import numpy as np
import matplotlib.pyplot as plt
from scipy.linalg import eigh_tridiagonal
from scipy.integrate import quad
import sympy as sym
sym.init_printing()

$$ \begin{bmatrix}\frac{1}{\Delta y^2}+mL^2V_1 & -\frac{1}{2 \Delta y^2} & 0 & 0...\\ -\frac{1}{2 \Delta y^2} & \frac{1}{\Delta y^2}+mL^2V_2 & -\frac{1}{2 \Delta y^2} & 0... \\ ...& ... & ... & -\frac{1}{2 \Delta y^2}\\...0 & 0 & -\frac{1}{2 \Delta y^2} & \frac{1}{\Delta y^2}+mL^2V_{N-1} \\ \end{bmatrix} \begin{bmatrix} \psi_1 \\ \psi_2 \\ ... \\ \psi_{N-1} \end{bmatrix} = mL^2 E \begin{bmatrix} \psi_1 \\ \psi_2 \\ ... \\ \psi_{N-1} \end{bmatrix} $$

$$ \psi_0 = \psi_N = 0$$

Define what $N$ and $dy$ is

N = 100
dy = 1/N
y = np.linspace(0, 1, N+1)

Define potential $mL^2 V$

#np.where(((y<0.25) | (y>0.75)), 1e8, 0)

def mL2V(y):
    #return 1000*np.sin(20*y) * y**4 
    return np.where(((y<0.25) | (y>0.75)), 500, 0)

d = 1/dy**2 + mL2V(y)[1:-1]
e = -1/(2*dy**2) * np.ones(len(d)-1)

w, v = eigh_tridiagonal(d, e)

plt.figure(figsize=(10,5))
plt.plot(y, mL2V(y), lw=3)
plt.title('Potential', fontsize=30)
plt.ylabel('$V/V_0$', fontsize=25)
plt.xlabel('$x/L$', fontsize=25)
plt.grid()
#plt.savefig('v3p1.png', dpi=200)

plt.figure(figsize=(10,5))
plt.plot(y[1:-1], v.T[0], label="1")
plt.plot(y[1:-1], v.T[1], label="2")
plt.plot(y[1:-1], v.T[2], label="3")
plt.plot(y[1:-1], v.T[3], label="4")
plt.title('Eigenstates', fontsize=20)
plt.ylabel('$\psi$', fontsize=15)
plt.xlabel('$x/L$', fontsize=15)
plt.grid()
plt.legend()
plt.savefig('v3p2.png', dpi=200)

plt.bar(np.arange(0, 10, 1), w[0:10])
plt.ylabel('$mL^2 E/\hbar^2$', fontsize=20)

Text(0, 0.5, '$mL^2 E/\\hbar^2$')

Solving Semi-analytically for Solutions

We are assuming the solutions to the finite square well look like sinusoids in the well and decaying exponentials outside of it. That means we are looking for solutions of the form $ e^{-\alpha x} $ outside the well and of the form $ sin(k x) $ or $ cos(k x) $, but not a superposition of both due to potential symmetry, within the well.

By setting $ y = \alpha L/2 $ and $ x = k L / 2 $ we arrive at two solutions which can be solved simultaenously:

$ y = x tan (x) $ and $ x^2 + y^2 = R^2_0 $ where $ R^2_0 = \frac{2 m L^2 V_0}{4 \hbar } $

from sympy import Matrix, exp, sin, cos, integrate, oo, Eq, Rational, lambdify

alpha, L, k, E, hbar, V0, V, m = sym.symbols("alpha L k E hbar V_0 V m", real=True, positive=True)
alpha, L, k, E, hbar, V0, V, m

$\displaystyle \left( \alpha, \ L, \ k, \ E, \ \hbar, \ V_{0}, \ V, \ m\right)$

fig = plt.figure(figsize=(10,6))

plt.title('Finite Square Well', fontsize=20)
plt.ylim(0, 10)


line_color = None
piranges = [(0, np.pi/2), (np.pi/2, 3*np.pi/2), (3*np.pi/2, 5*np.pi/2),  (5*np.pi/2, 7*np.pi/2), ]
for xlow, xhi in piranges:
    x = np.linspace(xlow+0.01, xhi-.01, 200)
    y = x*np.tan(x)
    if not line_color:
        lines = plt.plot(x,y, label='x tan(x)')
        line_color = lines[-1].get_color()
    else:
        plt.plot(x,y, color=line_color)


plt.ylabel('y = kL / 2', fontsize=15)
plt.xlabel('$ x = \\alpha L/2$', fontsize=15)
plt.grid()
#xmin, xmax, ymin, ymax = plt.axis()

# Plot the circles
theta = np.linspace(0, np.pi, 100)
for r in (2,4,6,8,9):
    xc = r * np.sin(theta)
    yc = r * np.cos(theta)
    plt.plot(xc,yc, label="$x^2 + y^2={}$".format(r*r), linestyle='dotted')

# Plot the asymptotes
ymin, ymax = plt.gca().get_ylim()
xvlines = ( [avline[1] for avline in piranges])
plt.vlines(xvlines, ymin, ymax, linestyle='dashed', color='orange', lw=3)
plt.legend()

<matplotlib.legend.Legend at 0x7fe586e60340>

It was kind of tricky to work out how many roots I would expect to find and whereabouts I would expect to find them. This is kind of important if we want to bracket the roots in a root finding algorithm.

def fn(r):
    n  = int(r/np.pi)+1
    print("Expect to find {} roots".format(n))
    for i in range(0,n):
        print ("Root {} between {} and {}".format(i+1, i*np.pi, (i+0.5)*np.pi))

fn(4)

Expect to find 2 roots
Root 1 between 0.0 and 1.5707963267948966
Root 2 between 3.141592653589793 and 4.71238898038469

We are looking for some piecewise solutions to Schrödinger's equation.

$$ \left [ \frac{- \hbar^2}{2m} \frac{\partial}{\partial x^2} + V(x,t) \right ] \Psi (x,t) = E \, \Psi (x,t) $$

If we multiply through by $ -\frac{2m}{\hbar^2} $ and do some rearranging we end up with:

$$ \frac{\partial \Psi (x,t) }{\partial x^2} + \frac{2m}{\hbar^2} (E - V ) \, \Psi (x,t) = 0 $$

If $ E \gt V $ then the solutions must be a decaying exponential. If $ E \lt V $ then the solutions will be sinusoids of various frequencies.

The intersectons of the circles and the graphs determine the fixed energies where:

$$ k^2 = \frac{2mE}{\hbar^2} \mbox{ and } \alpha^2 = \frac{2m(V_0 - E)}{\hbar^2} $$

and the square of the radius is proportional the to the depth of the finite potential:

$$ R^2 = k^2 + \alpha^2 = \left( \frac{2 m V_0}{\hbar^2} \right) $$

alpha, L, k, E, hbar, V0, x, A, m = sym.symbols("alpha L k E hbar V_0 x A m", real=True, positive=True)
alpha, L, k, E, hbar, V0, x, A, m

$\displaystyle \left( \alpha, \ L, \ k, \ E, \ \hbar, \ V_{0}, \ x, \ A, \ m\right)$

Now we want to set up a joining of the two wavefunctions $ e^{-\alpha x} $ and $ \cos(kx) $ at the point x = L/2 to be continuous and also we want its derivatives joined be continuous also.

Therefore:

$$ e^{-\alpha x} = \cos(kx) $$

and it's derivative is:

$$ -\alpha e^{-\alpha x} = -k \sin (kx) $$

If we substitute one of the join points x = L/2 we get

$$ e^{-\frac{\alpha L}{2}} = \cos \left( \frac{k L}{2} \right) $$

If we agree to set L = 2 then this simplifies to

$$ e^{- \alpha} = \cos \left( k \right) $$

and

$$ -\alpha e^{-\alpha } = -k \sin (k) $$

m1 = sym.Matrix([[exp(-alpha), -cos(k)],[-alpha*exp(-alpha), k*sin(k)]])
m1

$\displaystyle \left[\begin{matrix}e^{- \alpha} & - \cos{\left(k \right)}\\- \alpha e^{- \alpha} & k \sin{\left(k \right)}\end{matrix}\right]$

m1.det()

$\displaystyle - \alpha e^{- \alpha} \cos{\left(k \right)} + k e^{- \alpha} \sin{\left(k \right)}$

mdet = sym.simplify(m1.det()/(sym.exp(-alpha)))
sym.Eq(mdet, 0)

$\displaystyle - \alpha \cos{\left(k \right)} + k \sin{\left(k \right)} = 0$

mdet.subs(alpha, k*sym.tan(k))

$\displaystyle k \sin{\left(k \right)} - k \cos{\left(k \right)} \tan{\left(k \right)}$

sym.simplify(_)

$\displaystyle 0$

slice = m1[0,:]*alpha
m1[0,:] = slice
m1

$\displaystyle \left[\begin{matrix}\alpha e^{- \alpha} & - \alpha \cos{\left(k \right)}\\- \alpha e^{- \alpha} & k \sin{\left(k \right)}\end{matrix}\right]$

row = m1[1,:] + sym.Matrix(slice)
row

$\displaystyle \left[\begin{matrix}0 & - \alpha \cos{\left(k \right)} + k \sin{\left(k \right)}\end{matrix}\right]$

m1[1,:] = row
m1

$\displaystyle \left[\begin{matrix}\alpha e^{- \alpha} & - \alpha \cos{\left(k \right)}\\0 & - \alpha \cos{\left(k \right)} + k \sin{\left(k \right)}\end{matrix}\right]$

However $ \displaystyle \alpha \cos{\left(k \right)} - k \sin{\left(k \right )} = 0 $ so that value ought to be replaced by zero in the matrix also.

m1[1,1]=0
m1

$\displaystyle \left[\begin{matrix}\alpha e^{- \alpha} & - \alpha \cos{\left(k \right)}\\0 & 0\end{matrix}\right]$

This means $ \begin{pmatrix}e^{-\alpha} & -\cos(k)\\ 0 & 0 \end{pmatrix} \begin{pmatrix}A\\B\end{pmatrix} = \begin{pmatrix}0\\0\end{pmatrix} $ or that $ A e^{- \alpha} -B \cos(k) = 0 $

Therefore we can write $ A = B \cos(k) e^{ \alpha} $ or $ B = A \sec(k) e^{ -\alpha } $ and we can set $ A = 1 $ for the pre normalised value.

Now if assume by the same symmetry argument that half of the wave function probability density function integral must be 0.5 then we should be able to compute both A and B.

m2 = sym.Matrix([[exp(-alpha), -sin(k)],[-alpha*exp(-alpha), -k*cos(k)]])
m2

$\displaystyle \left[\begin{matrix}e^{- \alpha} & - \sin{\left(k \right)}\\- \alpha e^{- \alpha} & - k \cos{\left(k \right)}\end{matrix}\right]$

m2.det()

$\displaystyle - \alpha e^{- \alpha} \sin{\left(k \right)} - k e^{- \alpha} \cos{\left(k \right)}$

mdet = sym.simplify(m2.det()/(sym.exp(-alpha)))
sym.Eq(mdet, 0)

$\displaystyle - \alpha \sin{\left(k \right)} - k \cos{\left(k \right)} = 0$

Compute the normalisation constant for cos(kx)

f1 = exp(- alpha * x)
f1

$\displaystyle e^{- \alpha x}$

f2 = A * cos(k*x)
f2

$\displaystyle A \cos{\left(k x \right)}$

f3 = A*sin(k*x)
f3

$\displaystyle A \sin{\left(k x \right)}$

nf1 = integrate(f2**2, (x, 0, 1)) + integrate(f1**2, (x, 1, oo))
nf1

$\displaystyle \frac{A^{2} \left(\frac{k}{2} + \frac{\sin{\left(k \right)} \cos{\left(k \right)}}{2}\right)}{k} + \frac{e^{- 2 \alpha}}{2 \alpha}$

nf2 = integrate(f3**2, (x, 0, 1)) + integrate(f1**2, (x, 1, oo))
nf2

$\displaystyle \frac{A^{2} \left(\frac{k}{2} - \frac{\sin{\left(k \right)} \cos{\left(k \right)}}{2}\right)}{k} + \frac{e^{- 2 \alpha}}{2 \alpha}$

e1 = Eq(nf1, Rational(1,2))
e1

$\displaystyle \frac{A^{2} \left(\frac{k}{2} + \frac{\sin{\left(k \right)} \cos{\left(k \right)}}{2}\right)}{k} + \frac{e^{- 2 \alpha}}{2 \alpha} = \frac{1}{2}$

e2 = Eq(nf2, Rational(1,2))
e2

$\displaystyle \frac{A^{2} \left(\frac{k}{2} - \frac{\sin{\left(k \right)} \cos{\left(k \right)}}{2}\right)}{k} + \frac{e^{- 2 \alpha}}{2 \alpha} = \frac{1}{2}$

# Compute the normalisation value
normalisation1 = sym.simplify(sym.solveset(e1, A).args[0])
Eq(A, normalisation1)

$\displaystyle A = \frac{\sqrt{2} \sqrt{k} \sqrt{\frac{\alpha e^{2 \alpha} - 1}{2 k + \sin{\left(2 k \right)}}} e^{- \alpha}}{\sqrt{\alpha}}$

# Compute the normalisation value
normalisation2 = sym.simplify(sym.solveset(e2, A).args[0])
Eq(A, normalisation2)

$\displaystyle A = \frac{\sqrt{2} \sqrt{k} \sqrt{\frac{\alpha e^{2 \alpha} - 1}{2 k - \sin{\left(2 k \right)}}} e^{- \alpha}}{\sqrt{\alpha}}$

Numerical Solution

Now to find the solutions for $ \alpha $ and k which we are going to have to solve for numerically.

from scipy.optimize import root_scalar

$$ \alpha^2 + k^2 = R_0^2 $$

and since we know

$$ \alpha = k \tan k $$

then we can substitute that into our first equation to get:

$$ \left( k^2 \tan^2(k) \right) + k^2 = R_0^2 \Rightarrow k^2( 1 + \tan^2(k)) = k^2 \sec^2(k) = R_0^2 $$

def rootfn(x, r0):
    return x**2/np.cos(x)**2-r0**2

x1 = np.linspace(0, np.pi/2, 10)
y1 = rootfn(x1, 1)
print (y1)

[-1.00000000e+00 -9.68591166e-01 -8.62011436e-01 -6.34459096e-01
 -1.69448083e-01  8.43146052e-01  3.38649084e+00  1.17599051e+01
  6.36538288e+01  6.58079015e+32]

rootval = root_scalar(rootfn, args=(2,), method='bisect', bracket=(0, np.pi/2), rtol=0.0001)
rootval

      converged: True
           flag: 'converged'
 function_calls: 16
     iterations: 14
           root: 1.0297804776674795

rootfn(rootval.root, 2)

$\displaystyle -0.00181398318733272$

allroots=[]
for r in range(1,10):
    # calculate the number of roots we expect to find for this radius
    n = int(r/np.pi)+1
    roots = []
    for ii in range(n):
        lowrange = ii*np.pi
        hirange = (ii+.5)*np.pi
        try:
            root = root_scalar(rootfn, args=(r,), method='bisect', bracket=(lowrange+0.01, hirange-0.01), x0=lowrange + 0.01, x1=hirange-0.01, rtol=0.0001)
            if not root.converged:
                print("Warning root {} for radius {} did not converge".format(ii+1, r))
            roots.append(root.root)
        except ValueError:
            print (roots)
            print(r, lowrange, hirange)
            raise
    print (r, roots)
    allroots.append((r, roots.copy()))

        

1 [0.7390654728477594]
2 [1.0298871106698615]
3 [1.170162999116519]
4 [1.252321886546815, 3.59513706199805]
5 [1.3064634575262268, 3.837259681867517]
6 [1.344703308357839, 3.9856757365803097]
7 [1.373288543385431, 4.088658305156534, 6.616141671133699]
8 [1.3954373678770082, 4.165138006819758, 6.830436280744518]
9 [1.413042843754929, 4.224201538797298, 6.96825118869211]

fig = plt.figure(figsize=(10,6))

plt.title('Finite Square Well With Discrete Energies', fontsize=15)
plt.ylim(0, 10)


line_color = None
piranges = [(0, np.pi/2), (np.pi/2, 3*np.pi/2), (3*np.pi/2, 5*np.pi/2),  (5*np.pi/2, 7*np.pi/2), ]
for xlow, xhi in piranges:
    x1 = np.linspace(xlow+0.01, xhi-.01, 200)
    y1 = x1*np.tan(x1)
    if not line_color:
        lines = plt.plot(x1,y1, label='k tan(k)')
        line_color = lines[-1].get_color()
    else:
        plt.plot(x1,y1, color=line_color)


plt.xlabel('k', fontsize=15)
plt.ylabel(r'$ \alpha $', fontsize=15)
plt.grid()
#xmin, xmax, ymin, ymax = plt.axis()

# Plot the circles
theta = np.linspace(0, np.pi, 100)
for r in range(1,10):
    xc = r * np.sin(theta)
    yc = r * np.cos(theta)
    plt.plot(xc,yc, color='lightgreen', linestyle='dotted', lw=3)
    
# Plot the intersections of the circles and the x tan x lines.
scatterx = []
scattery = []
energies = []
for r, roots in allroots:
    for root in roots:
        k1 = root
        alpha1 = k1*np.tan(k1)
        scatterx.append(k1)
        scattery.append(alpha1)
        
        abfactor = np.exp(-alpha1)/np.cos(k1)
        print ("r={}, k={:.4f}, α={:.4f}, AB factor={:7.4f} E={:.4f} hbar^2/2m"\
               .format(r, k1,alpha1, abfactor, k1**2))
        energies.append((r, k1,alpha1, abfactor))
plt.scatter(scatterx, scattery, c='red')

# Plot the asymptotes
ymin, ymax = plt.gca().get_ylim()
xvlines = ( [avline[1] for avline in piranges])
plt.vlines(xvlines, ymin, ymax, linestyle='dashed', color='orange', lw=3)
plt.legend()

r=1, k=0.7391, α=0.6736, AB factor= 0.6899 E=0.5462 hbar^2/2m
r=2, k=1.0299, α=1.7146, AB factor= 0.3497 E=1.0607 hbar^2/2m
r=3, k=1.1702, α=2.7628, AB factor= 0.1618 E=1.3693 hbar^2/2m
r=4, k=1.2523, α=3.7984, AB factor= 0.0716 E=1.5683 hbar^2/2m
r=4, k=3.5951, α=1.7524, AB factor=-0.1929 E=12.9250 hbar^2/2m
r=5, k=1.3065, α=4.8268, AB factor= 0.0307 E=1.7068 hbar^2/2m
r=5, k=3.8373, α=3.2038, AB factor=-0.0529 E=14.7246 hbar^2/2m
r=6, k=1.3447, α=5.8459, AB factor= 0.0129 E=1.8082 hbar^2/2m
r=6, k=3.9857, α=4.4832, AB factor=-0.0170 E=15.8856 hbar^2/2m
r=7, k=1.3733, α=6.8624, AB factor= 0.0053 E=1.8859 hbar^2/2m
r=7, k=4.0887, α=5.6822, AB factor=-0.0058 E=16.7171 hbar^2/2m
r=7, k=6.6161, α=2.2881, AB factor= 0.1074 E=43.7733 hbar^2/2m
r=8, k=1.3954, α=7.8759, AB factor= 0.0022 E=1.9472 hbar^2/2m
r=8, k=4.1651, α=6.8356, AB factor=-0.0021 E=17.3484 hbar^2/2m
r=8, k=6.8304, α=4.1620, AB factor= 0.0182 E=46.6549 hbar^2/2m
r=9, k=1.4130, α=8.8829, AB factor= 0.0009 E=1.9967 hbar^2/2m
r=9, k=4.2242, α=7.9543, AB factor=-0.0007 E=17.8439 hbar^2/2m
r=9, k=6.9683, α=5.6936, AB factor= 0.0043 E=48.5565 hbar^2/2m

<matplotlib.legend.Legend at 0x7fe586d632e0>

def fn(x, alpha, k, abfactor):
    if -1 <= x <= 1:
        return abfactor * np.cos(k*x)
    else:
        return np.exp(-alpha * np.abs(x))

def fn2(x, alpha, k, abfactor):
    return fn(x, alpha, k, abfactor)**2

fnv = np.vectorize(fn)

N = 200

xval = np.linspace(-5, 5, N+1)

plt.figure(figsize=(10,5))
r, k, alp, abfactor = energies[0]

# Normalise the area under the curve
area = quad(fn2, -np.inf, np.inf, args=(alp, k, abfactor) )
yval = fnv(xval, alp, k, abfactor)/np.sqrt(area[0])
#display(exp(-alpha))
#print (area)


plt.plot(xval,yval)
plt.title('Wavefunction', fontsize=20)
plt.ylabel('$\Psi$', fontsize=15)
plt.xlabel('$x$', fontsize=15)
plt.grid()

ymin, ymax = plt.gca().get_ylim()
xvlines = [-1, 1]
plt.vlines(xvlines, ymin, ymax, linestyle='dashed', color='orange', lw=3)
#plt.legend()

#plt.savefig('v3p1.png', dpi=200)

<matplotlib.collections.LineCollection at 0x7fe5865d1040>

from IPython.display import display_latex
display_latex("$ 3x $", raw=True)

$ 3x $

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

#s = r"""$ \begin{{ matrix }}{:.4f} e^{{- {:.4f} x}} & \text{{if }} x \leq 1 \\ {:.4f} \sin({:.4f} x) & \text{{if }} x \ge 1  \end{{matrix}} $
s = r"$ v_0 = %d \frac{\hbar^2}{2m} \Rightarrow \Psi = \left \{ \begin{matrix} %.3f \, e^{- %.3f \, x} & \text{if } |x| \lt 1 " + \
    r"\\ %.3f \cos (%.3f x) & \text{if } |x| \ge 1   \end{matrix} \right \} $"
#print (repr(s))
eqns = []

for ii in range(0,3):
    r, k, alp, abfactor = energies[ii]
    area = quad(fn2, -np.inf, np.inf, args=(alp, k, abfactor) )
    #print(r, k, alp)
    #print (area)
    s1 = s % (ii+1, 1/area[0], alp, abfactor/area[0], k)
    #print (alp, np.exp(-alp), abfactor*np.cos(k))
    eqns.append((exp(-alp*x)/np.sqrt(area[0]), abfactor*cos(k*x)/np.sqrt(area[0])))
    #print (fnv(1, alp, k , abfactor))
    display_latex(s1, raw=True)
    #display(abfactor/area[0]*sin(k*x))
    
    
    y = fnv(xval, alp, k, abfactor)/np.sqrt(area[0])
    #y = fnv(x, alp, k, abfactor)
    plt.plot(xval,y, label=r"$V_0 ={} \frac{{\hbar^2}}{{2m}}$".format(ii+1))
    #print (r)
plt.title('Finite Square Well - 1st Energy Level - Various Potentials', fontsize=15)
plt.ylabel('$\Psi$', fontsize=15)
plt.xlabel('$x$', fontsize=15)
plt.grid()

ymin, ymax = plt.gca().get_ylim()
xvlines = [-1, 1]
plt.vlines(xvlines, ymin, ymax, linestyle='dashed', color='orange', lw=3)
plt.legend()

#plt.savefig('v3p1.png', dpi=200)

$ v_0 = 1 \frac{\hbar^2}{2m} \Rightarrow \Psi = \left \{ \begin{matrix} 0.846 \, e^{- 0.674 \, x} & \text{if } |x| \lt 1 \\ 0.583 \cos (0.739 x) & \text{if } |x| \ge 1 \end{matrix} \right \} $

$ v_0 = 2 \frac{\hbar^2}{2m} \Rightarrow \Psi = \left \{ \begin{matrix} 5.166 \, e^{- 1.715 \, x} & \text{if } |x| \lt 1 \\ 1.806 \cos (1.030 x) & \text{if } |x| \ge 1 \end{matrix} \right \} $

$ v_0 = 3 \frac{\hbar^2}{2m} \Rightarrow \Psi = \left \{ \begin{matrix} 28.036 \, e^{- 2.763 \, x} & \text{if } |x| \lt 1 \\ 4.537 \cos (1.170 x) & \text{if } |x| \ge 1 \end{matrix} \right \} $

<matplotlib.legend.Legend at 0x7fe5855aca90>

Now recall:

$$ \frac{\partial \Psi (x,t) }{\partial x^2} + \frac{2m}{\hbar^2} (E - V ) \, \Psi (x,t) = 0 $$

Perform a sanity test on the derived equations

fig = plt.figure(figsize=(10,6))

plt.title('Finite Square Well', fontsize=15)

colors = ['blue','orange','green']

for clr, elmt in zip(colors, eqns):
    expx, cosx = elmt
    
    xseg1 = np.linspace(-5,-1, N)
    f1 = lambdify(x, expx)
    yval = f1(-xseg1)
    plt.plot(xseg1, yval, color=clr)
    
    xseg2 = np.linspace(1,5, N)
    yval = f1(xseg2)
    plt.plot(xseg2, yval, color=clr)
    
    f2 = lambdify(x, cosx)
    xseg3 = np.linspace(-1,1, N)
    yval = f2(xseg3)
    plt.plot(xseg3, yval, color=clr)
    
ymin, ymax = plt.gca().get_ylim()
xvlines = [-1, 1]
plt.vlines(xvlines, ymin, ymax, linestyle='dashed', color='orange', lw=2)
    
plt.grid()   

# The code under here is doubtful given the way it tries to calculate the energies and ends up
# with contradictory answers.
s = ""
for elmt in eqns:
    expx, cosx = elmt
    expxdiff = expx.diff(x,x)
    cosxdiff = cosx.diff(x,x)
    expr_exp = sym.solveset(sym.simplify((expx.diff(x,x)+2*m/hbar**2*(E-V0)*expx)/expx), E)
    cosx_exp = sym.solveset(sym.simplify((cosx.diff(x,x)+2*m/hbar**2*(E-V0)*cosx)/cosx), E)
    #display(Eq(E,expr_exp.args[0]), Eq(E, cosx_exp.args[0]))
    e1latex = (sym.latex(Eq(E,expr_exp.args[0])))
    e2latex = (sym.latex(Eq(E,cosx_exp.args[0])))
    s += "{} & {} ".format(e1latex, e2latex) 
    
    # use this to detect if at last index of for loop
    if elmt is not eqns[-1]:
        s += r" \\ "
    #display_latex ("$ "+e1latex+" $", raw=True)

    #display_latex('E = V_{0} - \\frac{0.226846595732537 \\hbar^{2}}{m}', raw=True)
s = r"$  \left \{ \begin{matrix} " + s + r"\end{matrix} \right \} $"

display_latex(s, raw=True)

$ \left \{ \begin{matrix} E = V_{0} - \frac{0.226846595732537 \hbar^{2}}{m} & E = V_{0} + \frac{0.273108886577841 \hbar^{2}}{m} \\ E = V_{0} - \frac{1.46988327214104 \hbar^{2}}{m} & E = V_{0} + \frac{0.530333730361958 \hbar^{2}}{m} \\ E = V_{0} - \frac{3.81657632689217 \hbar^{2}}{m} & E = V_{0} + \frac{0.684640722250683 \hbar^{2}}{m} \end{matrix} \right \} $