In :
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt, mpld3
import matplotlib.cm as cm
from numpy import sqrt, exp, sinh, conj, absolute, pi
from scipy import constants, special
plt.rcParams['figure.figsize'] = 12, 8
mpld3.enable_notebook()

C:\Users\Michael\Anaconda\lib\site-packages\numpy\oldnumeric\__init__.py:11: ModuleDeprecationWarning: The oldnumeric module will be dropped in Numpy 1.9
warnings.warn(_msg, ModuleDeprecationWarning)


We use atomic units; all wave functions are represented such that $u[i]$ is the value at the $i$-th bin; normalization then corresponds to $\lVert u \rVert_{L^2}^2 = dx \sum_i u[i]^2 = 1$.

# (a) Schrödinger Solver¶

### Numerov Integration (from solution to exercise 3.2)¶

Integrate $-\psi''(x) = k(x) \psi(x)$.

In :
def numerov_step(k0, k1, k2, psi0, psi1, dx):
"""compute psi2 via a single numerov step"""
dx_sqr = dx**2
c0 = (1 + 1/12. * dx_sqr * k0)
c1 = 2 * (1 - 5/12. * dx_sqr * k1)
c2 = (1 + 1/12. * dx_sqr * k2)
return (c1 * psi1 - c0 * psi0) / c2

def numerov(ks, psi0, psi1, dx):
"""compute psi = [psi0, psi1, psi2, ...] for ks = [k0, k1, k2, ...] via iterated numerov steps"""
psi = np.zeros_like(ks)
psi = psi0
psi = psi1
for i in range(2, len(ks)):
psi[i] = numerov_step(ks[i-2], ks[i-1], ks[i], psi[i-2], psi[i-1], dx)
return psi


Find ground state for radial Schrödinger equation $-\frac 1 2 u''(r) + V(r) u(r) = E u(r)$.

In :
r_min = 0
r_max = 20
r_steps = 50000
rs, dr = np.linspace(r_min, r_max, r_steps + 1, retstep=True)
r_min, rs = rs, rs[1:]  # skip r=0

def solve_schrodinger(Vs, E_min=-10, E_max=0, maxiter=100, stoptol = 0.0001, verbose=False, retnumiter=False):
"""returns ground state energy, ground state wave function, and -- optionally -- the number of numerov iterations that were required"""
for i in range(maxiter):
E = (E_min + E_max) / 2.
if verbose:
print "%02d: trying E = %f in [%f, %f]" % (i + 1, E, E_min, E_max),

# numerov integration starting from r*exp(-r) at r_max to r_min
ks = 2 * (E - Vs)
psi_last = rs[-1] * exp(-rs[-1])
psi_2nd_last = rs[-2] * exp(-rs[-2])
u = numerov(ks[::-1], psi_last, psi_2nd_last, dr)
u = u[::-1]
num_nodes = np.sum(u[1:] * u[:-1] < 0)
if verbose:
print ', u(0) = %f, #nodes = %d' % (u, num_nodes)

# bisect
if num_nodes == 0 and absolute(u) <= stoptol:
if verbose:
print "  -> converged after", i + 1, "iterations: E_0 =", E
return (E, u, i+1) if retnumiter else (E, u)
if num_nodes == 0 and u > 0:
E_min = E
else:
assert num_nodes > 0, "expect #nodes > 0 since u < 0 while u[infty] > 0"
E_max = E

raise Exception("Not converged after %d iterations; u(0) = %f" % (maxiter, u))

E, u = solve_schrodinger(-1./rs, verbose=True)

01: trying E = -5.000000 in [-10.000000, 0.000000] , u(0) = 12884645614632673280.000000, #nodes = 0
02: trying E = -2.500000 in [-5.000000, 0.000000] , u(0) = 65058210171.040329, #nodes = 0
03: trying E = -1.250000 in [-2.500000, 0.000000] , u(0) = 54164.193034, #nodes = 0
04: trying E = -0.625000 in [-1.250000, 0.000000] , u(0) = 0.750139, #nodes = 0
05: trying E = -0.312500 in [-0.625000, 0.000000] , u(0) = -0.000956, #nodes = 1
06: trying E = -0.468750 in [-0.625000, -0.312500] , u(0) = -0.007805, #nodes = 1
07: trying E = -0.546875 in [-0.625000, -0.468750] , u(0) = 0.063054, #nodes = 0
08: trying E = -0.507812 in [-0.546875, -0.468750] , u(0) = 0.005109, #nodes = 0
09: trying E = -0.488281 in [-0.507812, -0.468750] , u(0) = -0.004283, #nodes = 1
10: trying E = -0.498047 in [-0.507812, -0.488281] , u(0) = -0.000560, #nodes = 1
11: trying E = -0.502930 in [-0.507812, -0.498047] , u(0) = 0.001995, #nodes = 0
12: trying E = -0.500488 in [-0.502930, -0.498047] , u(0) = 0.000652, #nodes = 0
13: trying E = -0.499268 in [-0.500488, -0.498047] , u(0) = 0.000031, #nodes = 0
-> converged after 13 iterations: E_0 = -0.499267578125


Compare with exact solution:

In :
E_exact = 0.5
u_exact = rs * exp(-rs)

In :
# normalize u & u_exact are the values of a wave function normalized to norm 1
u /= np.linalg.norm(u) * sqrt(dr)
u_exact /= np.linalg.norm(u_exact) * sqrt(dr)

# plot
plt.figure()
plt.title('Hydrogen wavefunction ($l=0$)')
plt.plot(rs, u, label='numerov')
plt.plot(rs, u_exact, label='exact')
plt.xlabel('r')
plt.ylabel('u(r)')
plt.legend(loc='best')
plt.show()