[Python] Self consistent field approximation in Python to calculate the Ionization energy

I am trying to compute the ionization energy of different elements given the electron structure using a self consistent field method where I'm incorporating electron-electron interactions with a mean-field potential. The approach is basically the following:

• 
  B-spline basis functions are used to represent wavefunctions and solve differential equations.
  
  
• 
  Poisson equation is used for computing the direct Coulomb potential $V_{\text{dir}}(r)$ from the electron density.
  
  
• 
  The Slater exchange potential is used as an approximation for exchange effects.
  

And a self-consistent loop where I:

1. Solve the Schrödinger equation for V=0.
2. Calculate the density.
3. Calculate V from the new density.
4. Solve the Schrödinger equation for new V.
5. Repeat until convergence.

And my goal is to run this for different atoms and ions to get the ionization energy .

Currently, I'm debugging the Poisson solver since the computed direct potential seems to blow up after the first iteration, I have been debugging for hours (Bsplines construction, BC implementation, the size of the grid etc), but I cant find where I went wrong! I thought that a second eye could help! any tips/suggestion?

Here is my code(for He and He+):

import numpy as np
from scipy.linalg import eigh, solve
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt
from numpy.linalg import matrix_rank

k = 4
r_min = 0.0
delta = 1e-6
r_max = 20.0
N_pts = 40
N_quad = 10
eta = 0.4
max_iter = 40
conv_tol = 1e-3
N_space = 500


def bsplgen(x_vals, tknot, k):
N_basis = len(tknot) - k
n_x = len(x_vals)
Bval = np.zeros((n_x, N_basis))
dBval = np.zeros((n_x, N_basis))
d2Bval = np.zeros((n_x, N_basis))

for xi, x in enumerate(x_vals):
B = np.zeros((len(tknot), k))
for i in range(N_basis + k - 1):
if tknot <= x < tknot[i + 1] or (x == tknot[-1] and tknot[i + 1] == tknot[-1]):
B[i, 0] = 1.0
for j in range(1, k):
for i in range(N_basis + k - j - 1):
t1, t2 = tknot, tknot[i + j]
t3, t4 = tknot[i + 1], tknot[i + j + 1]
left = (x - t1) / (t2 - t1) * B[i, j - 1] if t2 != t1 else 0.0
right = (t4 - x) / (t4 - t3) * B[i + 1, j - 1] if t4 != t3 else 0.0
B[i, j] = left + right

for i in range(N_basis):
Bval[xi, i] = B[i, k - 1]
if tknot[i + k - 1] != tknot:
dBval[xi, i] += (k - 1) * B[i, k - 2] / (tknot[i + k - 1] - tknot)
if tknot[i + k] != tknot[i + 1]:
dBval[xi, i] -= (k - 1) * B[i + 1, k - 2] / (tknot[i + k] - tknot[i + 1])

for i in range(N_basis):
term1 = term2 = term3 = term4 = 0.0
if k > 2:
if tknot[i + k - 1] != tknot and tknot[i + k - 2] != tknot:
term1 = B[i, k - 3] / ((tknot[i + k - 1] - tknot) * (tknot[i + k - 2] - tknot))
if tknot[i + k - 1] != tknot and tknot[i + k - 1] != tknot[i + 1]:
term2 = B[i + 1, k - 3] / ((tknot[i + k - 1] - tknot) * (tknot[i + k - 1] - tknot[i + 1]))
if tknot[i + k] != tknot[i + 1] and tknot[i + k - 1] != tknot[i + 1]:
term3 = B[i + 1, k - 3] / ((tknot[i + k] - tknot[i + 1]) * (tknot[i + k - 1] - tknot[i + 1]))
if tknot[i + k] != tknot[i + 1] and tknot[i + k] != tknot[i + 2]:
term4 = B[i + 2, k - 3] / ((tknot[i + k] - tknot[i + 1]) * (tknot[i + k] - tknot[i + 2]))
d2Bval[xi, i] = (k - 1) * (k - 2) * (term1 - term2 - term3 + term4)

return Bval, dBval, d2Bval


def linear_knots(r_min, r_max, N_pts, k):
knots = np.linspace(r_min+delta, r_max, N_pts, endpoint=False)
knots = np.concatenate((np.full(k, r_min), knots, np.full(k, r_max)))
return knots


def slater_exchange_potential(rho):
return -3 * (3 * rho / (8 * np.pi)) ** (1 / 3)


def build_H_S(tknot, k, ell, Vtot_func):
N_basis = len(tknot) - k
H = np.zeros((N_basis, N_basis))
S = np.zeros((N_basis, N_basis))
xi, wi = np.polynomial.legendre.leggauss(N_quad)

for i in range(N_basis):
for j in range(N_basis):
left = max(tknot, tknot[j])
right = min(tknot[i + k], tknot[j + k])
if left >= right:
continue
m_start = np.searchsorted(tknot, left, side='left') - 1
m_end = np.searchsorted(tknot, right, side='right') - 1
for m in range(m_start, m_end):
a, b = tknot[m], tknot[m + 1]
if b <= a:
continue
xq = 0.5 * (b - a) * xi + 0.5 * (b + a)
wq = 0.5 * (b - a) * wi
Bval, dBval, _ = bsplgen(xq, tknot, k)
Veff = Vtot_func(xq)
integrand_H = (
0.5 * dBval[:, i] * dBval[:, j]
+ 0.5 * ell * (ell + 1) * Bval[:, i] * Bval[:, j] / (xq ** 2)
- Z / xq * Bval[:, i] * Bval[:, j]
+ Veff * Bval[:, i] * Bval[:, j]
)
integrand_S = Bval[:, i] * Bval[:, j]
H[i, j] += np.sum(wq * integrand_H)
S[i, j] += np.sum(wq * integrand_S)
return H[1:-1, 1:-1], S[1:-1, 1:-1]


def solve_radial_schrodinger(tknot, k, ell, Vtot_func):
H, S = build_H_S(tknot, k, ell, Vtot_func)
eigvals, eigvecs = eigh(H, S)
return eigvals, eigvecs


def compute_density(knots, eigvecs, occ_orbitals):
r_vals = np.linspace(r_min + delta, r_max, N_space)
Bval, _, _ = bsplgen(r_vals, knots, k)

rho = np.zeros_like(r_vals)

for idx, Nj in occ_orbitals:
Pnl = Bval[:, 1:-1] @ eigvecs[:, idx]

norm = np.sqrt(np.trapz(Pnl**2, r_vals))
Pnl /= norm

rho += Nj * (Pnl**2) / (4 * np.pi * (r_vals)**2)

return r_vals, rho

def solve_poisson(r_vals, rho, tknot):
N_basis = len(tknot) - k
N_coll = N_basis - 2
r_coll = np.linspace(r_vals[0], r_vals[-1], N_coll)

_, _, d2Bval = bsplgen(r_coll, tknot, k)
A = d2Bval[:, 1:]

rho_interp = interp1d(r_vals, rho, kind='cubic', bounds_error=False, fill_value=0.0)

B = -r_coll * rho_interp(r_coll) * 4 * np.pi

Bval_bc, _, _ = bsplgen([r_max - delta], tknot, k)
A = np.vstack([A, Bval_bc[0, 1:]])
B = np.append(B, np.sum(rho * r_vals**2) * 4 * np.pi)

coeffs = solve(A, B)
coeffs = np.insert(coeffs, 0, 0.0)
Bval_eval, _, _ = bsplgen(r_vals, tknot, k)
phi_r = Bval_eval @ coeffs

return phi_r / r_vals


def self_consistent_field(Z_in, config):
old_eigvals = None
global Z
Z = Z_in
tknot = linear_knots(r_min, r_max, N_pts, k)
r_vals = np.linspace(r_min+delta, r_max, N_space)
Vdir = np.zeros_like(r_vals)
Vexch = np.zeros_like(r_vals)
orbitals = [None] * len(config)
occupations = [occ for _, _, occ in config]
ell_values = [ell for _, ell, _ in config]

for iteration in range(max_iter):
Vtot = Vdir + Vexch
Vtot_func = interp1d(r_vals, Vtot, kind='cubic', fill_value='extrapolate')
eigvals_all = []
orbitals_new = []
for idx, (n, ell, occ) in enumerate(config):
eigvals, eigvecs = solve_radial_schrodinger(tknot, k, ell, Vtot_func)
Bval, _, _ = bsplgen(r_vals, tknot, k)
psi = Bval[:, 1:-1] @ eigvecs[:, n - 1]
norm = np.trapz(psi ** 2, r_vals)
orbitals_new.append(psi / np.sqrt(norm))
eigvals_all.append(eigvals[n - 1])

occ_orbitals = list(enumerate(occupations))
r_vals, rho = compute_density(tknot, eigvecs, occ_orbitals)

Vdir_new = solve_poisson(r_vals, rho, tknot)
Vexch_new = slater_exchange_potential(rho)

if np.sum(occupations) <= 1:
Vdir_new *= 0
Vexch_new *= 0

Vdir = (1 - eta) * Vdir_new + eta * Vdir
Vexch = (1 - eta) * Vexch_new + eta * Vexch

if iteration > 0 and np.allclose(eigvals_all, old_eigvals, atol=conv_tol):
print(f"\nConverged after {iteration} iterations.")
break
old_eigvals = eigvals_all
orbitals = orbitals_new

total_energy = 0
print("\n{:5} {:>20} {:>20} {:>20} {:>20}".format("", "Orbital energy", "< -Vee/2 >", "Occupation", "Total"))
print("=" * 90)
for i, (eps, psi, occ) in enumerate(zip(eigvals_all, orbitals, occupations)):
Vee = Vdir + Vexch
Vee_half = -0.5 * np.trapz(psi * Vee * psi, r_vals)
contrib = occ * (eps + Vee_half)
total_energy += contrib
orb_label = f"{i+1}s" if ell_values == 0 else f"{i+1}p" if ell_values == 1 else f"{i+1}d"
print(f"{orb_label:<5} {eps:20.5f} {Vee_half:20.5f} {occ:20} {contrib:20.2f}")
print("=" * 90)
print(f"{'Sum':<5} {'':20} {'':20} {'':20} {total_energy:20.2f}")
return eigvals_all, total_energy, r_vals, rho


def plot_density(r_vals, rho, label=""):
plt.figure(figsize=(6, 4))
plt.plot(r_vals, 4 * np.pi * r_vals**2 * rho, label=label)
plt.xlabel("r (a.u.)")
plt.ylabel(r"$4\pi r^2 \rho(r)$")
plt.title("Radial Electron Density")
plt.grid(True)
if label:
plt.legend()
plt.tight_layout()
plt.show()


def ionization_energy(Z, config_neutral, config_ion):
_, E_total_neutral, _, _ = self_consistent_field(Z, config_neutral)
_, E_total_ion, _, _ = self_consistent_field(Z, config_ion)
delta_E = E_total_ion - E_total_neutral
print(f"\nIonization energy (Hartree): {abs(delta_E):.5f}")
print(f"Ionization energy (eV): {abs(delta_E) * 27.2114:.2f} eV")
return delta_E


if __name__ == "__main__":
print("\n===== Helium (Z=2, 1s^2) =====")
eigvals, E_total, r_vals, rho = self_consistent_field(2, [(1, 0, 2)])
plot_density(r_vals, rho, label="He")

print("\n===== Helium+ (Z=2, 1s^1) =====")
eigvals_plus, E_total_plus, r_vals_plus, rho_plus = self_consistent_field(2, [(1, 0, 1)])
plot_density(r_vals_plus, rho_plus, label="He+")

ionization_energy(2, [(1, 0, 2)], [(1, 0, 1)])

Continue reading...