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