%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)$.

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[0] = psi0
    psi[1] = 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

Radial Schrödinger Solver (adapted from solution to exercise 2.1)

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

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[1], 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[0], num_nodes)
        
        # bisect       
        if num_nodes == 0 and absolute(u[0]) <= 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] > 0:
            E_min = E
        else:
            assert num_nodes > 0, "expect #nodes > 0 since u[0] < 0 while u[infty] > 0"
            E_max = E
    
    raise Exception("Not converged after %d iterations; u(0) = %f" % (maxiter, u[0]))

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:

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

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