Calculate energy level of H-like 1s orbital
using 1s radial function and Slator's X-alpha potential

H1s-HF-LDA.py

import os
import sys
from math import exp, sqrt
import numpy as np
from math import log, exp
from numpy import arange
from scipy import integrate # 数値積分関数 integrateを読み込む
from scipy.interpolate import interp1d
from matplotlib import pyplot as plt



"""
Calculate energy level of H-like 1s orbital
using 1s radial function and Slator's X-alpha potential
"""



# constants
pi = 3.14159265358979323846
h = 6.6260755e-34 # Js";
hbar = 1.05459e-34 # "Js";
c = 2.99792458e8 # m/s";
e = 1.60218e-19 # C";
kB = 1.380658e-23 # JK-1";
me = 9.1093897e-31 # kg";
e0 = 8.854418782e-12; # C2N-1m-2";
e2_4pie0 = 2.30711e-28 # Nm2";
a0 = 5.29177e-11 * 1.0e10 # A
HartreeToeV = 27.2116 # eV";
RyToeV = HartreeToeV / 2.0
pi2 = pi + pi
pi4 = pi2 + pi2

# Nuclear and orbital parameters
Z = 1.0
n = 1
l = 0
m = 0

# Number of electrons, Slater's alpha
Ne = 1.0
alpha = 2.0 / 3.0

# Radius range and integration (quad()) parameters
Rmin = 0.0
Rmax = 8.0
Rmaxdata = None
nR = 1001
Rstep = (Rmax - Rmin) / (nR - 1)
nmaxdiv = 40
eps = 1.0-12

# Radial distribution function of H-like 1s orbital
# Normalized by integ(4pi*r*r*Rr(r)*Rr(r))[r=0, inf] = 1
R1s0 = 2.0 / sqrt(4.0 * pi)
def Rr(Z, n, l, m, r):
    global R1s0

    return R1s0 * exp(-Z * r)

def rho(Z, n, l, m, r):
    psi = Rr(Z, n, l, m, r)
    return psi * psi


#===================================================
# Calculate related quantities for visualization
#===================================================
r = [Rmin + i * Rstep for i in range(nR+100)]
yRr = [] # Raidal distributuion function
yQr = [] # Integrated charge inside r: integ_rm(4pi * rm * rm * rho(rm))_[rm = 0, r]
for i in range(len(r)):
    yRr.append(Rr(Z, n, l, m, r[i]))
    if r[i] == 0.0:
        Q = 0.0
    else:
        Q, errQ = integrate.quad(lambda r: pi4 * r * r * rho(Z, n, l, m, r), 0, r[i], limit = nmaxdiv, epsrel = eps)
    yQr.append(Q)

Rmaxdata = max(r)

# not used. for faster integration in future
def integrate_trapezoidal(func, E0, E1, h):
    n = int((E1 - E0) / h + 1.000001)
    h = (E1 - E0) / n
    y = [func(E0 + i * h) for i in range(n)]

    S = 0.5 * (y[0] + y[n-1]) + sum(y[1:n-1])
    return [h * S, -1.0]

# Integrated charge inside r: integ_rm(4pi * rm * rm * rho(rm))_[rm = 0, r]
def Q(R):
    global Rmin, Rmaxdata, r, Qr

    if R < Rmin:
        print("Error in Q(r): Given r={} exceeds the R range [{}, {}]".format(R, Rmin, Rmax))
        exit()
    if R > Rmaxdata:
        return 0.0

    qfunc = interp1d(r, yQr, kind = 'cubic')
    return qfunc(R)

# Electrostatic potential from nuclear
def UZ(r):
    global Z
    if r < 1.0e-3:
        return -Z / 1.0e-3
    else:
        return -Z / r

# Electrostatic potential formed at r by distributed charge density rho(r)
def Urho(r):
# Uerho(r) = integ(rho(m) / |r-rm|)_[rm=0, inf] [r, rm are vectors]
# = Q(r) / r + integ[4pi * rm * rm * rho(rm) / rm]_[rm = r, inf]) [r is vector]
    if r == 0.0:
        Uee1 = 0.0
    else:
        Uee1 = Q(r) / r

    Uee2, errUee2 = integrate.quad(lambda rm: pi4 * rm * rho(Z, n, l, m, rm), r, Rmax, limit = nmaxdiv, epsrel = eps)
    return Uee1 + Uee2

def UXa(r):
# UXa = -3.0 * alpha * (3/4pi*rho(r))^(1/3)
    return -3.0 * alpha * pow(3.0/pi4 * rho(Z, n, l, m, r), 1.0/3.0)

#===================================================
# Calculate related quantities for visualization
#===================================================
ypotZ = [] # U(Z) = -Z/r
ypotZrho = [] # U(e-Z) for Z = 1
ypotXa = [] # Xa potential
ypotrho = [] # Electrostatic potential by rho(r)
for i in range(len(r)):
    if r[i] == 0.0:
        phi = 0.0
    else:
        phi, errpot = integrate.quad(lambda r: pi4 * r * rho(Z, n, l, m, r),
                                Rmin, r[i], limit = nmaxdiv, epsrel = eps)
        potZ = -Z / r[i]
    ypotZ.append(UZ(r[i]))
    ypotZrho.append(-Z * phi)
    ypotrho.append(Urho(r[i]))
    ypotXa.append(UXa(r[i]))


#==========================
# Main start
#==========================
print("")
print("Analytical solution")
# T = e^2 / 8pi / e0 / a0
Tanal = e / 8.0 / pi / e0 / (a0*1.0e-10)
# U = -e^2 / 4pi / e0 / a0
Uanal = -e / 4.0 / pi / e0 / (a0*1.0e-10)
Eanal = Tanal + Uanal
print("T(analytical) = {} eV".format(Tanal))
print("U(analytical) = {} eV".format(Uanal))
print(" Etotl(analytical) = {} eV".format(Eanal))

print("")
print("Numerical integration")
Rr2tot, err = integrate.quad(lambda r: r * r * rho(Z, n, l, m, r), Rmin, Rmax, limit = nmaxdiv, epsrel = eps)
Rr2tot *= pi4
print("R(r) normalization check: 2pi * integ(r*r * Rr*2)dr = ", Rr2tot)

# Kinetic energy
# -1 / 2 * integ(4pi*r*r*R(r) * laplasian(R(r)))
# laplasian(Rr(r)) / R1s0 = 1/r^2 * d/dr[r^2dRr/dr] = 1/r^2 *d/dr[-Z * r^2*exp(-Zr)]
# = -Z / r^2 * [2r - Z * r*r]*exp(-Zr)
# r*r*R(r) * laplasian(R(r)) = -R1s0 * R1s0 * Z * r * [2 - Z * r] * exp(-2Zr)
T, errT = integrate.quad(lambda r: r * (2.0 - Z * r) * exp(-2.0 * Z * r), Rmin, Rmax, limit = nmaxdiv, epsrel = eps)
T *= 1.0 / 2.0 * 4.0 * pi * R1s0 * R1s0 * Z
print("T = {} eV (err = {})".format(T * HartreeToeV, errT))

# UeZ = integ(4pi * r*r * (-(Z / r) * rho(r))
# = -Z * 4pi * integ(r * rho(r))
UeZ, errUez = integrate.quad(lambda r: pi4 * r * rho(Z, n, l, m, r), Rmin, Rmax, limit = nmaxdiv, epsrel = eps)
UeZ *= -Z
print("U(e-Z) = {} eV (err = {})".format(UeZ * HartreeToeV, errUez))

# Uee = integ_r( integ_rm(rho(rm)/|rm-r|) * rho(r) ) [r and rm are vectors]
# = integ_r(4pi * r * r * Urho(r) * rho(r))_[r=0, inf] [r is scalar]
# Urho(r) = Q(r) / r + integ[4pi * rm * rm * rho(rm) / rm]_[rm = r, inf])
Uee, errUee = integrate.quad(lambda r: pi4 * r * r * Urho(r) * rho(Z, n, l, m, r), Rmin, Rmax, limit = nmaxdiv, epsrel = eps)
Uee *= Ne
print("U(e-e) = {} eV (err = {})".format(Uee * HartreeToeV, errUee))

# UXa = -3.0 * alpha * integ(4pi * r * r * (3/4pi*rho(r))^(1/3) * rho(r)))
UXa, errUez = integrate.quad(lambda r: pi4 * r * r * UXa(r) * rho(Z, n, l, m, r), Rmin, Rmax, limit = nmaxdiv, epsrel = eps)
UXa *= pow(Ne, 1.0/3.0)
print("U(Xa) = {} eV (err = {})".format(UXa * HartreeToeV, errUez))

Etot = T + UeZ + Uee + UXa
print(" E(tot) = {} eV".format(Etot * HartreeToeV))
print("")

Nstart = 0.0
Nend = 1.0
nN = 1001
Nstep = (Nend - Nstart) / (nN - 1)
xN = [Nstart + i * Nstep for i in range(nN)]
yE1s = []
for i in range(nN):
    yE1s.append(HartreeToeV * (T + UeZ + xN[i] * Uee + pow(xN[i], 1.0/3.0) * UXa))

#=============================
# Plot graphs
#=============================
fig = plt.figure(figsize = (12, 8))

ax1 = fig.add_subplot(2, 3, 1)
ax2 = fig.add_subplot(2, 3, 2)
ax3 = fig.add_subplot(2, 3, 3)
ax4 = fig.add_subplot(2, 3, 4)
ax5 = fig.add_subplot(2, 3, 5)
ax1.plot(r, yRr, label = 'Rr(r)', color = 'black')
ax1.set_xlim([Rmin, Rmax])
ax1.set_ylim([0, max(yRr)*1.1])
ax1.set_xlabel("r (bohr)")
ax1.set_ylabel("Rr(r)")
ax1.legend()
ax2.plot(r, yQr, label = 'Q(r)', color = 'black')
ax2.set_xlabel("r (bohr)")
ax2.set_ylabel("Q(r)")
ax2.legend()
ax3.plot(r, ypotZ, label = 'U(Z)', color = 'black')
ax3.plot(r, ypotZrho, label = 'U(Z-rho)', color = 'red')
ax3.plot(r, -np.array(ypotrho), label = '-U(rho)', color = 'blue')
ax3.plot(r, ypotXa, label = 'U(Xa)', color = 'green')
ax3.set_ylim([min(ypotXa) * 2.0, 0.0])
ax3.set_xlabel("r (bohr)")
ax3.set_ylabel("Potential / Energy (Hartree)")
ax3.legend()
ax4.plot(xN, yE1s, label = '', color = 'black')
ax4.plot([min(xN), max(xN)], [Eanal, Eanal], color = 'red', linestyle='dashed', linewidth = 0.5)
ax4.set_xlabel("r (bohr)")
ax4.set_ylabel(" (eV)")
ax4.legend()

plt.tight_layout()
plt.pause(0.1)

print("Press ENTER to exit>>", end = '')
input()

exit()