Schroedinger Equation for Hydrogen Atom¶

The Schroedinger equation is:

\begin{eqnarray} (-\frac{\hbar^2}{2m}\nabla^2-\frac{Z e^2}{4\pi\varepsilon_0 r})\psi(\vec{r})=E \psi(\vec{r}) \end{eqnarray}

using ansatz:

$\psi(\vec{r}) = Y_{lm}(\hat{r})\; u(r)/r$

and introducing dimensionless variables:

\begin{eqnarray} x = \frac{r}{r_B}\\ \varepsilon = \frac{E}{E_0} \end{eqnarray}

where \begin{eqnarray} && r_B = \frac{4\pi\varepsilon_0 \hbar^2}{m e^2} \approx 0.529 A\\ && E_0 = \frac{\hbar^2}{2 m r_B^2} == Ry \approx 13.6 eV \end{eqnarray}

we get the differential equation

\begin{eqnarray} u''(x)- \left(\frac{l(l+1)}{x^2}-\frac{2Z}{x}-\varepsilon\right)u(x)=0 \end{eqnarray}

Next we rewrite into the system of first order equations:

\begin{eqnarray} y = \left(u(x),u'(x)\right)\\ \frac{dy}{dx} = \left(u'(x),u''(x)\right) \end{eqnarray}

with boundary conditions \begin{eqnarray} &&u(0) = 0 \rightarrow \psi(0)<\infty\\ &&u(\infty)=0 \rightarrow \int |\psi(r)|^2 r^2 dr \propto \int u^2(r)dr < \infty \end{eqnarray}

Because boundary conditions are given at the two ends, we need so-called shooting method

Shooting algorithm:

Suppose the two boundary condistions are given at $a$ and $b$, i.e., $u(a)=u(b)=0$. Then

  • Choose $u(a)=0$ and $u'(a)=c$, with $c$ some constant.
  • Solve for $u(x)$ to the other end, and check if $u(b)=0$.
  • Using root finding routine find energy $\varepsilon$ for which u(b)=0. This is the bound state.
  • Continue with increasing energy $\varepsilon$ until sufficient number of bound states is found

Some remarks

  • It turns out that forward integration of the radial Sch. Eq. is unstable. It is better to start integrating from infinity, and then continue down to zero.
  • It is better to use logarithmic mesh for radial variable rather than linear. Radial functions need smaller number of points in logarithmic mesh

The implementation will follow these steps

  1. call SciPy routine
    integrate.odeint
    to integrate the one-electron Schroedinger equation. Note that the distance is measured in units of bohr radius and energy units is Rydberg ($1 Ry = 13.6058...eV$)
  2. The boundary conditions are $u(0)=0$ and $u(\infty)=0$.

    Use shooting method to obtain wave functions:

    1. Use logarithmic mesh of radial points for integration. Start integrating from a large distance ($R_{max} \sim 100$). At $R_{max}$ choose $u=0$ and some nonzero (not too large) derivative.
    2. Integrate the Schroedinger equation down to $r=0$. If your choice for the energy $\varepsilon$ corresponds to the bound state, the wave function at $u(r=0)$ will be zero.

  3. Start searching for the first bound state at sufficiently negative energy (for example $\sim -1.2 Z^2$) and increase energy in sufficiently small steps to bracket all necessary bound states. Ones the wave function at $r=0$ changes sign, use root finding routine, for example
    optimize.brentq,
    to compute zero to very high precision. Store the index and the energy of the bound state for further processing.
  4. Ones bound state energies are found, recompute $u(r)$ for all bound states. Normalize $u(r)$ and plot them.
  5. Compute electron density for various atoms (for example He, Li, ..) neglecting Coulomb repulsion:

    Populate first $Z$ lowest laying electron states and compute $\rho = \sum_{lm\in occupied} u_{lm}^2(r)/(4\pi r^2)$. Each state with quantum number $l$ can take $2(2l+1)$ electrons. Be carefull, if atom is not one of the Nobel gases (He, Ne, ...) the last orbital is only partially filled.

Recall:

\begin{eqnarray} && y = (u(r), u'(r) )\\ && dy/dr = (u'(r), u''(r)) \end{eqnarray}\begin{eqnarray} u''(r)= \left(\frac{l(l+1)}{r^2}-\frac{2Z}{r}-\varepsilon\right)u(r) \end{eqnarray}
In [1]:
from numpy import *
from scipy import integrate
from scipy import optimize
from numba import jit  # This is the new line with numba

@jit(nopython=True)
def Schroed_deriv(y,r,l,En):
    "Given y=[u,u'] returns dy/dr=[u',u''] "
    (u,up) = y
    return array([up, (l*(l+1)/r**2-2/r-En)*u])

First we try linear mesh and forward integration. It is supposed to be unstable. We know the ground state has energy $E_0=-1 Ry$ and we should get $1s$ state with integrating Scroedinger equation.

In [2]:
R = linspace(1e-10,20,500)
l=0
E0=-1.0 # the ground state, should give the first bound state.

ur = integrate.odeint(Schroed_deriv, [0.0, 1.0], R, args=(l,E0)) 
# choice of derivative determines normalization, but not the shape of the wave functions.
In [3]:
from pylab import *
%matplotlib inline

plot(R,ur[:,0],label='odeint')
plot(R,R*exp(-R),label='exact') # expected solution u = r * exp(-r)
legend(loc='best')
show()
No description has been provided for this image

Recal Euler's method and Runge Kutta method

In [4]:
from IPython.display import Image
Image('https://upload.wikimedia.org/wikipedia/commons/thumb/1/10/Euler_method.svg/2560px-Euler_method.svg.png')
Out[4]:
No description has been provided for this image
In [5]:
Image('https://upload.wikimedia.org/wikipedia/commons/thumb/7/7e/Runge-Kutta_slopes.svg/1920px-Runge-Kutta_slopes.svg.png')
Out[5]:
No description has been provided for this image

Indeed the integration is unstable, and needs to be done in opposite direction. Let's try from large R.

In [6]:
R = linspace(1e-10,20,500)
l=0
E0=-1.0
Rb=R[::-1] # invert the mesh

urb = integrate.odeint(Schroed_deriv, [0.0, -1e-7], Rb, args=(l,E0))
ur = urb[:,0][::-1] # we take u(r) and invert it in R.

norm=integrate.simps(ur**2,x=R)
ur *= 1./sqrt(norm)
In [7]:
plot(R,ur, label='numerical')
plot(R, R*exp(-R)*2, '.', label='exact') # with proper normalization, exact solution is 2*r*exp(-r)
legend(loc='best')
show()

plot(R,ur, '-')
plot(R, R*exp(-R)*2, 'o')
xlim(0,0.05)
ylim(0,0.10)
show()
No description has been provided for this image
No description has been provided for this image

Clearly the integration from infinity is stable, and we will use it here.

Logarithmic mesh is better suited for higher excited states, as they extend far away.

Lets create a subroutine of what we learned up to now:

In [8]:
def SolveSchroedinger(En,l,R):
    ur = integrate.odeint(Schroed_deriv, [0.0,-1e-7], R[::-1], args=(l,En))[:,0][::-1]
    ur *= 1./sqrt(integrate.simps(ur**2,x=R))
    return ur
In [9]:
l=1
n=3
En=-1./(n**2) # 3p orbital


#Ri = linspace(1e-6,20,500)   # linear mesh already fails for this case
Ri = linspace(1e-6,100,1000)
ui = SolveSchroedinger(En,l,Ri)
plot(Ri,ui,'o-', label='linear');
No description has been provided for this image
In [10]:
l=1
n=3
En=-1./(n**2)  # 3p orbital


#Ri = linspace(1e-6,20,500)   # linear mesh already fails for this case
Ri = linspace(1e-6,100,1000)
ui = SolveSchroedinger(En,l,Ri)


R = logspace(-6,2.,100)
ur = SolveSchroedinger(En,l,R)


#ylim([0,0.5])
plot(Ri,ui/Ri,'s-', label='u(r)/r linear')
plot(R,ur/R,'o-', label='u(r)/r logarithmic')
xlim([0,2])
ylim([-0.1,0.4])
legend(loc='best');
No description has been provided for this image

Shooting algorithm:

The boundary condistions are given at two points $a$ and $b$, i.e., $u(a)=u(b)=0$.

  • Choose $u(a)=0$ and $u'(a)=c$, with $c$ some constant.

  • Solve for $u(x)$ to the other end, and evaluate $u(b)$.

  • Using root finding routine find energy $\varepsilon$ for which u(b)=0. This is the bound state.

  • Continue with increasing energy $\varepsilon$ until sufficient number of bound states is found

In [11]:
def Shoot(En,R,l):
    ur = integrate.odeint(Schroed_deriv, [0.0,-1e-7], R[::-1], args=(l,En))[:,0][::-1]
    norm=integrate.simps(ur**2,x=R)# normalization is not essential here
    ur *= 1./sqrt(norm)            # once we normalize, the functions are all of the order of unity
    # we know that u(r)~ A r^(l+1), hence, u(r)/r^l ~ A r
    ur = ur/R**l
    # extrapolate to r=0
    f0 = ur[0]
    f1 = ur[1]
    f_at_0 = f0 + (f1-f0)*(0.0-R[0])/(R[1]-R[0])
    return f_at_0
In [12]:
R = logspace(-6,2.2,500)

Shoot(-1.,R,0), Shoot(-1/3**2, R, l=1)
Out[12]:
(-9.516227259926638e-09, 4076.972245144093)

Shooting algorithm:

The boundary condistions are given at two points $a$ and $b$, i.e., $u(a)=u(b)=0$.

  • Choose $u(a)=0$ and $u'(a)=c$, with $c$ some constant.

  • Solve for $u(x)$ to the other end, and evaluate $u(b)$.

  • Using root finding routine find energy $\varepsilon$ for which u(b)=0. This is the bound state.

  • Continue with increasing energy $\varepsilon$ until sufficient number of bound states is found

In [13]:
def FindBoundStates(R,l,nmax,Esearch):
    """ R       -- real space mesh
        l       -- orbital quantum number
        nmax    -- maximum number of bounds states we require
        Esearch -- energy mesh, which brackets all bound-states, i.e., every sign change of the wave function at u(0).
    """
    n=0
    Ebnd=[]                     # save all bound states
    u0 = Shoot(Esearch[0],R,l)  # u(r=0) for the first energy Esearch[0]
    for i in range(1,len(Esearch)):
        u1 = Shoot(Esearch[i],R,l) # evaluate u(r=0) and all Esearch points
        if u0*u1<0:
            Ebound = optimize.brentq(Shoot,Esearch[i-1],Esearch[i],xtol=1e-16,args=(R,l)) # root finding routine
            Ebnd.append((l,Ebound))
            if len(Ebnd)>nmax: break
            n+=1
            print(f"Found bound state at E={Ebound:14.9f} E_exact={-1.0/(n+l)**2:14.9f} l={l}")
        u0=u1
    
    return Ebnd
In [14]:
Esearch = -1.2/arange(1,20,0.2)**2
Esearch
Out[14]:
array([-1.2       , -0.83333333, -0.6122449 , -0.46875   , -0.37037037,
       -0.3       , -0.24793388, -0.20833333, -0.17751479, -0.15306122,
       -0.13333333, -0.1171875 , -0.10380623, -0.09259259, -0.08310249,
       -0.075     , -0.06802721, -0.06198347, -0.05671078, -0.05208333,
       -0.048     , -0.0443787 , -0.04115226, -0.03826531, -0.03567182,
       -0.03333333, -0.03121748, -0.02929688, -0.02754821, -0.02595156,
       -0.0244898 , -0.02314815, -0.02191381, -0.02077562, -0.01972387,
       -0.01875   , -0.01784652, -0.0170068 , -0.01622499, -0.01549587,
       -0.01481481, -0.01417769, -0.01358081, -0.01302083, -0.01249479,
       -0.012     , -0.01153403, -0.01109467, -0.01067996, -0.01028807,
       -0.00991736, -0.00956633, -0.00923361, -0.00891795, -0.00861821,
       -0.00833333, -0.00806235, -0.00780437, -0.00755858, -0.00732422,
       -0.00710059, -0.00688705, -0.006683  , -0.00648789, -0.0063012 ,
       -0.00612245, -0.0059512 , -0.00578704, -0.00562957, -0.00547845,
       -0.00533333, -0.00519391, -0.00505988, -0.00493097, -0.00480692,
       -0.0046875 , -0.00457247, -0.00446163, -0.00435477, -0.0042517 ,
       -0.00415225, -0.00405625, -0.00396354, -0.00387397, -0.0037874 ,
       -0.0037037 , -0.00362275, -0.00354442, -0.00346861, -0.0033952 ,
       -0.0033241 , -0.00325521, -0.00318844, -0.0031237 , -0.00306091])
In [15]:
Esearch = -1.2/arange(1,20,0.2)**2
R = logspace(-6,2.2,500)
nmax=7

Bnd=[]
for l in range(nmax-1):
    Bnd += FindBoundStates(R,l,nmax-l,Esearch)
Found bound state at E=  -1.000000018 E_exact=  -1.000000000 l=0
Found bound state at E=  -0.249999997 E_exact=  -0.250000000 l=0
Found bound state at E=  -0.111111111 E_exact=  -0.111111111 l=0
Found bound state at E=  -0.062500000 E_exact=  -0.062500000 l=0
Found bound state at E=  -0.040000002 E_exact=  -0.040000000 l=0
Found bound state at E=  -0.027777773 E_exact=  -0.027777778 l=0
Found bound state at E=  -0.020409865 E_exact=  -0.020408163 l=0
Found bound state at E=  -0.249999997 E_exact=  -0.250000000 l=1
Found bound state at E=  -0.111111112 E_exact=  -0.111111111 l=1
Found bound state at E=  -0.062500001 E_exact=  -0.062500000 l=1
Found bound state at E=  -0.040000003 E_exact=  -0.040000000 l=1
Found bound state at E=  -0.027777889 E_exact=  -0.027777778 l=1
Found bound state at E=  -0.020421586 E_exact=  -0.020408163 l=1
Found bound state at E=  -0.111111111 E_exact=  -0.111111111 l=2
Found bound state at E=  -0.062500001 E_exact=  -0.062500000 l=2
Found bound state at E=  -0.040000004 E_exact=  -0.040000000 l=2
Found bound state at E=  -0.027778516 E_exact=  -0.027777778 l=2
Found bound state at E=  -0.020414271 E_exact=  -0.020408163 l=2
Found bound state at E=  -0.062500000 E_exact=  -0.062500000 l=3
Found bound state at E=  -0.040000001 E_exact=  -0.040000000 l=3
Found bound state at E=  -0.027778126 E_exact=  -0.027777778 l=3
Found bound state at E=  -0.020405512 E_exact=  -0.020408163 l=3
Found bound state at E=  -0.040000000 E_exact=  -0.040000000 l=4
Found bound state at E=  -0.027777799 E_exact=  -0.027777778 l=4
Found bound state at E=  -0.020410583 E_exact=  -0.020408163 l=4
Found bound state at E=  -0.027777787 E_exact=  -0.027777778 l=5
Found bound state at E=  -0.020406833 E_exact=  -0.020408163 l=5
In [16]:
Bnd
Out[16]:
[(0, -1.0000000175237493),
 (0, -0.24999999672119474),
 (0, -0.11111111116473228),
 (0, -0.06250000042496895),
 (0, -0.040000002241628314),
 (0, -0.02777777274566952),
 (0, -0.020409864997276275),
 (0, -0.015266291726552637),
 (1, -0.24999999669248765),
 (1, -0.11111111161933523),
 (1, -0.062500001453899),
 (1, -0.040000002949991245),
 (1, -0.02777788871543995),
 (1, -0.020421586294474957),
 (1, -0.015413169413305942),
 (2, -0.11111111126358458),
 (2, -0.06250000053821912),
 (2, -0.040000004196094376),
 (2, -0.027778515786198273),
 (2, -0.020414270711027666),
 (2, -0.015484354553456571),
 (3, -0.062499999809775046),
 (3, -0.04000000123306971),
 (3, -0.027778125963281237),
 (3, -0.020405511668886427),
 (3, -0.015519806412158195),
 (4, -0.039999999954891764),
 (4, -0.02777779850099035),
 (4, -0.020410582926254875),
 (4, -0.015576802781241763),
 (5, -0.027777787425561036),
 (5, -0.020406832575417467),
 (5, -0.015601410734180693)]

We have many bound states, but they are not properly sorted in order in which to fill atomic shells.

In [17]:
sorted(Bnd, key= lambda x: x[1])
Out[17]:
[(0, -1.0000000175237493),
 (0, -0.24999999672119474),
 (1, -0.24999999669248765),
 (1, -0.11111111161933523),
 (2, -0.11111111126358458),
 (0, -0.11111111116473228),
 (1, -0.062500001453899),
 (2, -0.06250000053821912),
 (0, -0.06250000042496895),
 (3, -0.062499999809775046),
 (2, -0.040000004196094376),
 (1, -0.040000002949991245),
 (0, -0.040000002241628314),
 (3, -0.04000000123306971),
 (4, -0.039999999954891764),
 (2, -0.027778515786198273),
 (3, -0.027778125963281237),
 (1, -0.02777788871543995),
 (4, -0.02777779850099035),
 (5, -0.027777787425561036),
 (0, -0.02777777274566952),
 (1, -0.020421586294474957),
 (2, -0.020414270711027666),
 (4, -0.020410582926254875),
 (0, -0.020409864997276275),
 (5, -0.020406832575417467),
 (3, -0.020405511668886427),
 (5, -0.015601410734180693),
 (4, -0.015576802781241763),
 (3, -0.015519806412158195),
 (2, -0.015484354553456571),
 (1, -0.015413169413305942),
 (0, -0.015266291726552637)]

If energies are not degenerate, this would work. But for degenerate case, we need more careful sorting!

In [18]:
def cmpKey(x):
    return x[1] + x[0]/10000.  # energy has large wait, but degenerate energy states are sorted by l

Bnd = sorted(Bnd, key=cmpKey)
In [19]:
Bnd
Out[19]:
[(0, -1.0000000175237493),
 (0, -0.24999999672119474),
 (1, -0.24999999669248765),
 (0, -0.11111111116473228),
 (1, -0.11111111161933523),
 (2, -0.11111111126358458),
 (0, -0.06250000042496895),
 (1, -0.062500001453899),
 (2, -0.06250000053821912),
 (3, -0.062499999809775046),
 (0, -0.040000002241628314),
 (1, -0.040000002949991245),
 (2, -0.040000004196094376),
 (3, -0.04000000123306971),
 (4, -0.039999999954891764),
 (0, -0.02777777274566952),
 (1, -0.02777788871543995),
 (2, -0.027778515786198273),
 (3, -0.027778125963281237),
 (4, -0.02777779850099035),
 (5, -0.027777787425561036),
 (0, -0.020409864997276275),
 (1, -0.020421586294474957),
 (2, -0.020414270711027666),
 (3, -0.020405511668886427),
 (4, -0.020410582926254875),
 (5, -0.020406832575417467),
 (1, -0.015413169413305942),
 (2, -0.015484354553456571),
 (0, -0.015266291726552637),
 (3, -0.015519806412158195),
 (4, -0.015576802781241763),
 (5, -0.015601410734180693)]

Now the states are sorted as expected in periodic table.

Now let's try to compute the charge density for a heavier atom, example of Ni.

In [20]:
Z=46 # like Ni
N=0
rho=zeros(len(R))
for (l,En) in Bnd:
    ur = SolveSchroedinger(En,l,R)
    dN = 2*(2*l+1) # each radial function can store these many electrons
    if N+dN<=Z:
        ferm = 1.  # no fractional occupancy needed
    else:
        ferm = (Z-N)/float(dN)  # fractional occupancy, because the orbital is not fully filled
    drho = ur**2 * ferm * dN/(4*pi*R**2) # charge density per solid angle per radius: drho/(dOmega*dr)
    rho += drho
    N += dN
    print(f'adding state ({l:2d},{En:14.9f}) with fermi={ferm:4.2f} and current N={N:5.1f}')
    if N>=Z: break
adding state ( 0,  -1.000000018) with fermi=1.00 and current N=  2.0
adding state ( 0,  -0.249999997) with fermi=1.00 and current N=  4.0
adding state ( 1,  -0.249999997) with fermi=1.00 and current N= 10.0
adding state ( 0,  -0.111111111) with fermi=1.00 and current N= 12.0
adding state ( 1,  -0.111111112) with fermi=1.00 and current N= 18.0
adding state ( 2,  -0.111111111) with fermi=1.00 and current N= 28.0
adding state ( 0,  -0.062500000) with fermi=1.00 and current N= 30.0
adding state ( 1,  -0.062500001) with fermi=1.00 and current N= 36.0
adding state ( 2,  -0.062500001) with fermi=1.00 and current N= 46.0

Resulting charge density for a Ni-like Hydrogen atom

In [21]:
from pylab import *
%matplotlib inline

plot(R,rho*4*pi*R**2,label='charge density')
xlim([0,60])
show()
No description has been provided for this image
In [22]:
integrate.simpson(rho*R**2 * 4*pi,x=R)  # total density
Out[22]:
45.99999999999999
In [23]:
l=3
En=-1/4**2
#l,En = Bnd[9]
#R = logspace(-6,2,1000)
R = logspace(-2,2.5,1000)
ur = SolveSchroedinger(En,l,R)

plot(R,ur,'-');
No description has been provided for this image

It seems that error accumulation still prevents one to calculate f-states near the nucleous. Only if we start sufficiently away from nucleout (0.1$R_b$) numerics works.

Can we do better?

Numerov algorithm¶

The general purpose integration routine is not the best method for solving the Schroedinger equation, which does not have first derivative terms.

Numerov algorithm is better fit for such equations, and its algorithm is summarized below.

The second order linear differential equation (DE) of the form

\begin{eqnarray} x''(t) = f(t) x(t) + u(t) \end{eqnarray}

is a target of Numerov algorithm.

Due to a special structure of the DE, the fourth order error cancels and leads to sixth order algorithm using second order integration scheme.

If we expand x(t) to some higher power and take into account the time reversal symmetry of the equation, all odd term cancel

\begin{eqnarray} x(h) = x(0)+h x'(0)+\frac{1}{2}h^2 x''(0)+\frac{1}{3!}h^3 x^{(3)}(0)+\frac{1}{4!}h^4 x^{(4)}(0)+\frac{1}{5!}h^5 x^{(5)}(0)+...\\ x(-h) = x(0)-h x'(0)+\frac{1}{2}h^2 x''(0)-\frac{1}{3!}h^3 x^{(3)}(0)+\frac{1}{4!}h^4 x^{(4)}(0)-\frac{1}{5!}h^5 x^{(5)}(0)+... \end{eqnarray}

hence

\begin{eqnarray} x(h)+x(-h) = 2x(0)+h^2 (f(0)x(0)+u(0))+\frac{2}{4!}h^4 x^{(4)}(0)+O(h^6) \end{eqnarray}

If we are happy with $O(h^4)$ algorithm, we can neglect $x^{(4)}$ term and get the following recursion relation

\begin{equation} x_{i+1} = 2 x_i - x_{i-1} + h^2 (f_i x_i+u_i) +O(h^4). \end{equation}

where we renaimed \begin{eqnarray} &x_{i-1}&= x(-h)\\ &x_i &= x(0)\\ &x_{i+1} &= x(h) \end{eqnarray}

But we know from the differential equation that

\begin{equation} x^{(4)} = \frac{d^2 x''(t)}{dt^2} = \frac{d^2}{dt^2}(f(t) x(t)+u(t)) \end{equation}

and we will use the well known descrete expression for the second order derivative \begin{eqnarray} g''(t) = \frac{g(t+h) - 2 g(t) + g(t-h)}{h^2} + O(h^2) \end{eqnarray} which can be approximated by

\begin{equation} x^{(4)}= \frac{f_{i+1}x_{i+1}+u_{i+1}-2 f_i x_i -2 u_i+ f_{i-1}x_{i-1}+u_{i-1}}{h^2} + O(h^2) \end{equation}

Inserting the fourth order derivative into the above recursive equation (forth equation in his chapter), we get

\begin{equation} x_{i+1}-2 x_i+x_{i-1}=h^2(f_i x_i+u_i)+\frac{h^2}{12}(f_{i+1}x_{i+1}+u_{i+1}-2 f_i x_i -2 u_i+ f_{i-1}x_{i-1}+u_{i-1}) + O(h^6) \end{equation}

If we switch to a new variable $w_i=x_i(1-\frac{h^2}{12} f_i)-\frac{h^2}{12}u_i$ we are left with the following equation

\begin{equation} w_{i+1} = 2 w_i - w_{i-1} + h^2 (f_i x_i + u_i)+O(h^6) \end{equation}

The variable $x$ needs to be recomputed at each step with \begin{equation} x_i=\frac{w_i+\frac{h^2}{12}u_i}{1-\frac{h^2}{12}f_i}. \end{equation}

In [24]:
@jit(nopython=True)
def Numerovc(f, x0, dx, dh):
    """Given precomputed function f(x), solves for x(t), which satisfies:
          x''(t) = f(t) x(t)
          dx = (dx(t)/dt)_{t=0}
          x0 = x(t=0)
    """
    x = zeros(len(f))
    x[0] = x0
    x[1] = x0+dh*dx
    h2 = dh**2
    h12 = h2/12.
    w0=x0*(1-h12*f[0])
    w1=x[1]*(1-h12*f[1])
    xi = x[1]
    fi = f[1]
    for i in range(2,f.size):
        w2 = 2*w1-w0+h2*fi*xi  # here fi,xi=f1,x1 at the first step
        fi = f[i]              # at this point fi=f2 in the first step
        xi = w2/(1-h12*fi)     # xi is not x2 in the first step
        x[i]=xi                # save x2 into x[2]
        w0,w1 = w1,w2
    return x

For Numerov algorithm we can evaluate derivative part $f(r)$ for all points at once:

Because Schroedinger Eq is: \begin{eqnarray} u''(r)= \left(\frac{l(l+1)}{r^2}-\frac{2Z}{r}-\varepsilon\right)u(r) \end{eqnarray} the $f$ function is \begin{eqnarray} f(r)= \left(\frac{l(l+1)}{r^2}-\frac{2Z}{r}-\varepsilon\right) \end{eqnarray}

In [25]:
def fSchrod(En, l, R):
    return l*(l+1.)/R**2-2./R-En

The Numerov algorithm is much faster, but the price we pay is the mesh has to be linear. We can not use logarithmic mesh in combination with Numerov algorithm.

In [26]:
Rl = linspace(1e-6,100,1000)
l,En=0,-1
f = fSchrod(En,l,Rl[::-1])     # here we turn mesh R around, so that f is given from large r down to r=0.
ur = Numerovc(f,0.0,1e-7,Rl[1]-Rl[0])[::-1] # turn around the solution, so that it starts with r=0
norm = integrate.simps(ur**2,x=Rl)
ur *= 1/sqrt(abs(norm))
In [27]:
from pylab import *
%matplotlib inline

plot(Rl,ur)
plot(Rl,exp(-Rl)*Rl*2.);
xlim(0,10)
Out[27]:
(0.0, 10.0)
No description has been provided for this image

Numerov seems much more precise than odeint, and avoids numerical problems we had before

In [28]:
Rl = linspace(1e-6,100,1000)
l,En=3,-1/4**2
f = fSchrod(En,l,Rl[::-1])     # here we turn mesh R around, so that f is given from large r down to r=0.
ur = Numerovc(f,0.0,1e-7,Rl[1]-Rl[0])[::-1] # turn around the solution, so that it starts with r=0
norm = integrate.simps(ur**2,x=Rl)
ur *= 1/sqrt(abs(norm))

plot(Rl,ur,'-')
xlim(0,60)
grid()
No description has been provided for this image
In [32]:
Rl = linspace(1e-6,100,1000)
l,En=1,-1/3**2   # 3p orbital


f = fSchrod(En,l,Rl[::-1])
ur = Numerovc(f,0.0,-1e-7,-Rl[1]+Rl[0])[::-1]
norm = integrate.simps(ur**2,x=Rl)
ur *= 1/sqrt(abs(norm))

plot(Rl,ur, label='u(r)')
plot(Rl,ur/Rl, label='u(r)/r')
legend(loc='best')
show()
No description has been provided for this image
In [33]:
def fSchrod(En, l, R):
    return l*(l+1.)/R**2-2./R-En

def ComputeSchrod(En,R,l):
    "Computes Schrod Eq." 
    f = fSchrod(En,l,R[::-1])
    ur = Numerovc(f,0.0,-1e-7,-R[1]+R[0])[::-1]
    norm = integrate.simps(ur**2,x=R)
    return ur*1/sqrt(abs(norm))

def Shoot(En,R,l):
    ur = ComputeSchrod(En,R,l)
    ur = ur/R**l
    f0,f1 = ur[0],ur[1]
    f_at_0 = f0 + (f1-f0)*(0.0-R[0])/(R[1]-R[0])
    return f_at_0

Put it all together is for the homework (see below).

It seems with Numerov we are getting substantial error-bar for the energy of 1s state. We could increase the number of points in the mesh, but the error decreases only linearly with the number of points used.

Where is the problem? What should be done?

In [35]:
R = linspace(1e-6,100,2000)

optimize.brentq(Shoot,-1.1,-0.99,xtol=1e-16,args=(R,0))
Out[35]:
-0.9999219225299896

Check that approximate solution gives smaller wave function at zero than exact energy, which confirms that root finding routine works fine.

In [36]:
Shoot(-1.0,R,l=0), Shoot(-0.9999221089559636,R,l=0)
Out[36]:
(9.742774445580766e-08, 2.3262872423201642e-10)

Let's check how the function looks like near zero

In [37]:
ur = ComputeSchrod(-1.0,R,0)
plot(R,ur,'o-')
xlim(0,0.2)
ylim(0,0.4)
Out[37]:
(0.0, 0.4)
No description has been provided for this image

Idea: The mesh is very sparse near zero, and in the range of the first few points, the curve is not linear enough. Linear extrapolation gives the error.

Can we do better?

Let's use cubic extrapolation with first 4 points.

In [38]:
help(polyfit)
help(polyval)
Help on _ArrayFunctionDispatcher in module numpy:

polyfit(x, y, deg, rcond=None, full=False, w=None, cov=False)
    Least squares polynomial fit.
    
    .. note::
       This forms part of the old polynomial API. Since version 1.4, the
       new polynomial API defined in `numpy.polynomial` is preferred.
       A summary of the differences can be found in the
       :doc:`transition guide </reference/routines.polynomials>`.
    
    Fit a polynomial ``p(x) = p[0] * x**deg + ... + p[deg]`` of degree `deg`
    to points `(x, y)`. Returns a vector of coefficients `p` that minimises
    the squared error in the order `deg`, `deg-1`, ... `0`.
    
    The `Polynomial.fit <numpy.polynomial.polynomial.Polynomial.fit>` class
    method is recommended for new code as it is more stable numerically. See
    the documentation of the method for more information.
    
    Parameters
    ----------
    x : array_like, shape (M,)
        x-coordinates of the M sample points ``(x[i], y[i])``.
    y : array_like, shape (M,) or (M, K)
        y-coordinates of the sample points. Several data sets of sample
        points sharing the same x-coordinates can be fitted at once by
        passing in a 2D-array that contains one dataset per column.
    deg : int
        Degree of the fitting polynomial
    rcond : float, optional
        Relative condition number of the fit. Singular values smaller than
        this relative to the largest singular value will be ignored. The
        default value is len(x)*eps, where eps is the relative precision of
        the float type, about 2e-16 in most cases.
    full : bool, optional
        Switch determining nature of return value. When it is False (the
        default) just the coefficients are returned, when True diagnostic
        information from the singular value decomposition is also returned.
    w : array_like, shape (M,), optional
        Weights. If not None, the weight ``w[i]`` applies to the unsquared
        residual ``y[i] - y_hat[i]`` at ``x[i]``. Ideally the weights are
        chosen so that the errors of the products ``w[i]*y[i]`` all have the
        same variance.  When using inverse-variance weighting, use
        ``w[i] = 1/sigma(y[i])``.  The default value is None.
    cov : bool or str, optional
        If given and not `False`, return not just the estimate but also its
        covariance matrix. By default, the covariance are scaled by
        chi2/dof, where dof = M - (deg + 1), i.e., the weights are presumed
        to be unreliable except in a relative sense and everything is scaled
        such that the reduced chi2 is unity. This scaling is omitted if
        ``cov='unscaled'``, as is relevant for the case that the weights are
        w = 1/sigma, with sigma known to be a reliable estimate of the
        uncertainty.
    
    Returns
    -------
    p : ndarray, shape (deg + 1,) or (deg + 1, K)
        Polynomial coefficients, highest power first.  If `y` was 2-D, the
        coefficients for `k`-th data set are in ``p[:,k]``.
    
    residuals, rank, singular_values, rcond
        These values are only returned if ``full == True``
    
        - residuals -- sum of squared residuals of the least squares fit
        - rank -- the effective rank of the scaled Vandermonde
           coefficient matrix
        - singular_values -- singular values of the scaled Vandermonde
           coefficient matrix
        - rcond -- value of `rcond`.
    
        For more details, see `numpy.linalg.lstsq`.
    
    V : ndarray, shape (M,M) or (M,M,K)
        Present only if ``full == False`` and ``cov == True``.  The covariance
        matrix of the polynomial coefficient estimates.  The diagonal of
        this matrix are the variance estimates for each coefficient.  If y
        is a 2-D array, then the covariance matrix for the `k`-th data set
        are in ``V[:,:,k]``
    
    
    Warns
    -----
    RankWarning
        The rank of the coefficient matrix in the least-squares fit is
        deficient. The warning is only raised if ``full == False``.
    
        The warnings can be turned off by
    
        >>> import warnings
        >>> warnings.simplefilter('ignore', np.RankWarning)
    
    See Also
    --------
    polyval : Compute polynomial values.
    linalg.lstsq : Computes a least-squares fit.
    scipy.interpolate.UnivariateSpline : Computes spline fits.
    
    Notes
    -----
    The solution minimizes the squared error
    
    .. math::
        E = \sum_{j=0}^k |p(x_j) - y_j|^2
    
    in the equations::
    
        x[0]**n * p[0] + ... + x[0] * p[n-1] + p[n] = y[0]
        x[1]**n * p[0] + ... + x[1] * p[n-1] + p[n] = y[1]
        ...
        x[k]**n * p[0] + ... + x[k] * p[n-1] + p[n] = y[k]
    
    The coefficient matrix of the coefficients `p` is a Vandermonde matrix.
    
    `polyfit` issues a `RankWarning` when the least-squares fit is badly
    conditioned. This implies that the best fit is not well-defined due
    to numerical error. The results may be improved by lowering the polynomial
    degree or by replacing `x` by `x` - `x`.mean(). The `rcond` parameter
    can also be set to a value smaller than its default, but the resulting
    fit may be spurious: including contributions from the small singular
    values can add numerical noise to the result.
    
    Note that fitting polynomial coefficients is inherently badly conditioned
    when the degree of the polynomial is large or the interval of sample points
    is badly centered. The quality of the fit should always be checked in these
    cases. When polynomial fits are not satisfactory, splines may be a good
    alternative.
    
    References
    ----------
    .. [1] Wikipedia, "Curve fitting",
           https://en.wikipedia.org/wiki/Curve_fitting
    .. [2] Wikipedia, "Polynomial interpolation",
           https://en.wikipedia.org/wiki/Polynomial_interpolation
    
    Examples
    --------
    >>> import warnings
    >>> x = np.array([0.0, 1.0, 2.0, 3.0,  4.0,  5.0])
    >>> y = np.array([0.0, 0.8, 0.9, 0.1, -0.8, -1.0])
    >>> z = np.polyfit(x, y, 3)
    >>> z
    array([ 0.08703704, -0.81349206,  1.69312169, -0.03968254]) # may vary
    
    It is convenient to use `poly1d` objects for dealing with polynomials:
    
    >>> p = np.poly1d(z)
    >>> p(0.5)
    0.6143849206349179 # may vary
    >>> p(3.5)
    -0.34732142857143039 # may vary
    >>> p(10)
    22.579365079365115 # may vary
    
    High-order polynomials may oscillate wildly:
    
    >>> with warnings.catch_warnings():
    ...     warnings.simplefilter('ignore', np.RankWarning)
    ...     p30 = np.poly1d(np.polyfit(x, y, 30))
    ...
    >>> p30(4)
    -0.80000000000000204 # may vary
    >>> p30(5)
    -0.99999999999999445 # may vary
    >>> p30(4.5)
    -0.10547061179440398 # may vary
    
    Illustration:
    
    >>> import matplotlib.pyplot as plt
    >>> xp = np.linspace(-2, 6, 100)
    >>> _ = plt.plot(x, y, '.', xp, p(xp), '-', xp, p30(xp), '--')
    >>> plt.ylim(-2,2)
    (-2, 2)
    >>> plt.show()

Help on _ArrayFunctionDispatcher in module numpy:

polyval(p, x)
    Evaluate a polynomial at specific values.
    
    .. note::
       This forms part of the old polynomial API. Since version 1.4, the
       new polynomial API defined in `numpy.polynomial` is preferred.
       A summary of the differences can be found in the
       :doc:`transition guide </reference/routines.polynomials>`.
    
    If `p` is of length N, this function returns the value:
    
        ``p[0]*x**(N-1) + p[1]*x**(N-2) + ... + p[N-2]*x + p[N-1]``
    
    If `x` is a sequence, then ``p(x)`` is returned for each element of ``x``.
    If `x` is another polynomial then the composite polynomial ``p(x(t))``
    is returned.
    
    Parameters
    ----------
    p : array_like or poly1d object
       1D array of polynomial coefficients (including coefficients equal
       to zero) from highest degree to the constant term, or an
       instance of poly1d.
    x : array_like or poly1d object
       A number, an array of numbers, or an instance of poly1d, at
       which to evaluate `p`.
    
    Returns
    -------
    values : ndarray or poly1d
       If `x` is a poly1d instance, the result is the composition of the two
       polynomials, i.e., `x` is "substituted" in `p` and the simplified
       result is returned. In addition, the type of `x` - array_like or
       poly1d - governs the type of the output: `x` array_like => `values`
       array_like, `x` a poly1d object => `values` is also.
    
    See Also
    --------
    poly1d: A polynomial class.
    
    Notes
    -----
    Horner's scheme [1]_ is used to evaluate the polynomial. Even so,
    for polynomials of high degree the values may be inaccurate due to
    rounding errors. Use carefully.
    
    If `x` is a subtype of `ndarray` the return value will be of the same type.
    
    References
    ----------
    .. [1] I. N. Bronshtein, K. A. Semendyayev, and K. A. Hirsch (Eng.
       trans. Ed.), *Handbook of Mathematics*, New York, Van Nostrand
       Reinhold Co., 1985, pg. 720.
    
    Examples
    --------
    >>> np.polyval([3,0,1], 5)  # 3 * 5**2 + 0 * 5**1 + 1
    76
    >>> np.polyval([3,0,1], np.poly1d(5))
    poly1d([76])
    >>> np.polyval(np.poly1d([3,0,1]), 5)
    76
    >>> np.polyval(np.poly1d([3,0,1]), np.poly1d(5))
    poly1d([76])

In [39]:
def Shoot2(En,R,l):
    ur = ComputeSchrod(En,R,l)
    ur = ur/R**l
    poly = polyfit(R[:4], ur[:4], deg=3)
    return polyval(poly, 0.0)
In [40]:
Shoot2(-1,R,l=0), Shoot2(-0.9999221089559636,R,l=0)
Out[40]:
(7.081086849536259e-11, -9.638638706416136e-08)
In [41]:
optimize.brentq(Shoot2,-1.1,-0.9,xtol=1e-16,args=(R,0))
Out[41]:
-0.9999999428188899

Indeed we get $10^{-8}$ error as compared to $10^{-5}$ error before. So, the extrapolation must be improved to reduce the error.

Increasing the number of points for 10-times reduces the error factor of 1000, i.e., $O(1/N^3)$.

Better cubic extrapolation reduces the error for the same factor of 1000, equivalent to 10-times more points.

In [42]:
R = linspace(1e-6,100,2000)
print('Error with linear extrapolation:', optimize.brentq(Shoot,-1.1,-0.9,xtol=1e-16,args=(R,0)) + 1.0 )
print('Error with cubic extrapolation:',  optimize.brentq(Shoot2,-1.1,-0.9,xtol=1e-16,args=(R,0)) + 1.0 )
Error with linear extrapolation: 7.807747000609933e-05
Error with cubic extrapolation: 5.718111006913773e-08

It also helps to increase the number of points. 10-times denser grid gives roughly $10^{-3}$ smaller error.

In [43]:
R = linspace(1e-8,100,20000)
print('Error with linear extrapolation:', optimize.brentq(Shoot,-1.1,-0.9,xtol=1e-16,args=(R,0)) + 1.0 )
print('Error with cubic extrapolation:',  optimize.brentq(Shoot2,-1.1,-0.9,xtol=1e-16,args=(R,0)) + 1.0 )
Error with linear extrapolation: 8.297733611328795e-08
Error with cubic extrapolation: -1.089017764854816e-11

Homework¶

Compute charge density for atom with Z=28 (Ni), Z=92(U) and Z=94 (Pu).

In [ ]: