# Restricted Hartree-Fock for a linear chain of N hydrogen atoms in STO-3G.
#
# A scaled-down but genuine quantum-chemistry workload: the same machinery a
# package like Psi4 runs, restricted to s-type Gaussians. The cost is dominated
# by the O(N^4) two-electron integral build, which is parallelized with
# `parallel for` (OpenMP on the native build).
#
# Change `n_atoms` (keep it even for closed-shell RHF) to scale the work.
pi = 3.141592653589793
# STO-3G hydrogen: one contracted s-shell, three primitives.
alpha_h = [3.42525091, 0.62391373, 0.16885540]
d_h = [0.15432897, 0.53532814, 0.44463454]
n_prim = 3
# ---- Molecule: N hydrogens evenly spaced along z. ----
n_atoms = 8
R = 1.4
n_basis = n_atoms
n_occ = n_atoms / 2
centers = zeros(n_atoms, 3)
for i to n_atoms {
centers[i, 2] = i * R
}
Z = zeros(n_atoms)
for i to n_atoms {
Z[i] = 1.0
}
basis_atom = zeros(n_basis)
for i to n_basis {
basis_atom[i] = i
}
# ---- Boys function F_0(T) by direct Simpson integration. ----
def boys_F0(T) {
if T < 1.0e-8 {
return 1.0 - T / 3.0 + T * T / 10.0
}
return simpson(fn(t) -> exp(-T * t * t), 0.0, 1.0, 40)
}
def dist_sq(A, B) {
d0 = A[0] - B[0]
d1 = A[1] - B[1]
d2 = A[2] - B[2]
return d0 * d0 + d1 * d1 + d2 * d2
}
def gauss_product_center(a, A, b, B) {
p = a + b
P = zeros(3)
P[0] = (a * A[0] + b * B[0]) / p
P[1] = (a * A[1] + b * B[1]) / p
P[2] = (a * A[2] + b * B[2]) / p
return P
}
def prim_overlap(a, A, b, B) {
p = a + b
r2 = dist_sq(A, B)
K = exp(-a * b / p * r2)
Na = pow(2.0 * a / pi, 0.75)
Nb = pow(2.0 * b / pi, 0.75)
return Na * Nb * pow(pi / p, 1.5) * K
}
def prim_kinetic(a, A, b, B) {
p = a + b
r2 = dist_sq(A, B)
S = prim_overlap(a, A, b, B)
return a * b / p * (3.0 - 2.0 * a * b / p * r2) * S
}
def prim_nuclear(a, A, b, B, C, Zc) {
p = a + b
Na = pow(2.0 * a / pi, 0.75)
Nb = pow(2.0 * b / pi, 0.75)
P = gauss_product_center(a, A, b, B)
rAB = dist_sq(A, B)
rPC = dist_sq(P, C)
K = exp(-a * b / p * rAB)
return -2.0 * pi / p * Zc * Na * Nb * K * boys_F0(p * rPC)
}
def prim_eri(a, A, b, B, c, C, d, D) {
p = a + b
q = c + d
Na = pow(2.0 * a / pi, 0.75)
Nb = pow(2.0 * b / pi, 0.75)
Nc = pow(2.0 * c / pi, 0.75)
Nd = pow(2.0 * d / pi, 0.75)
rAB = dist_sq(A, B)
rCD = dist_sq(C, D)
P = gauss_product_center(a, A, b, B)
Q = gauss_product_center(c, C, d, D)
rPQ = dist_sq(P, Q)
K1 = exp(-a * b / p * rAB)
K2 = exp(-c * d / q * rCD)
pref = 2.0 * pow(pi, 2.5) / (p * q * sqrt(p + q))
return Na * Nb * Nc * Nd * pref * K1 * K2 * boys_F0(p * q / (p + q) * rPQ)
}
def center_of(atom_idx) {
R = zeros(3)
R[0] = centers[atom_idx, 0]
R[1] = centers[atom_idx, 1]
R[2] = centers[atom_idx, 2]
return R
}
def contracted_one_elec(mu, nu, kind, C, Zc) {
A = center_of(basis_atom[mu])
B = center_of(basis_atom[nu])
total = 0.0
for p to n_prim {
for q to n_prim {
a = alpha_h[p]
b = alpha_h[q]
dp = d_h[p]
dq = d_h[q]
v = 0.0
if kind == 0 { v = prim_overlap(a, A, b, B) }
if kind == 1 { v = prim_kinetic(a, A, b, B) }
if kind == 2 { v = prim_nuclear(a, A, b, B, C, Zc) }
total += dp * dq * v
}
}
return total
}
def contracted_eri(mu, nu, lam, sig) {
A = center_of(basis_atom[mu])
B = center_of(basis_atom[nu])
C = center_of(basis_atom[lam])
D = center_of(basis_atom[sig])
total = 0.0
for p to n_prim {
for q to n_prim {
for r to n_prim {
for s to n_prim {
a = alpha_h[p]
b = alpha_h[q]
c = alpha_h[r]
dd = alpha_h[s]
coef4 = d_h[p] * d_h[q] * d_h[r] * d_h[s]
total += coef4 * prim_eri(a, A, b, B, c, C, dd, D)
}
}
}
}
return total
}
# ---- One-electron matrices S, T, V and core Hamiltonian. ----
S = zeros(n_basis, n_basis)
T = zeros(n_basis, n_basis)
V = zeros(n_basis, n_basis)
H_core = zeros(n_basis, n_basis)
z3 = zeros(3)
for mu to n_basis {
for nu to n_basis {
S[mu, nu] = contracted_one_elec(mu, nu, 0, z3, 0.0)
T[mu, nu] = contracted_one_elec(mu, nu, 1, z3, 0.0)
v_total = 0.0
for atom to n_atoms {
Rc = center_of(atom)
v_total += contracted_one_elec(mu, nu, 2, Rc, Z[atom])
}
V[mu, nu] = v_total
H_core[mu, nu] = T[mu, nu] + V[mu, nu]
}
}
# ---- Two-electron integrals (the O(N^4) hot spot), parallelized. ----
narrate "Building two-electron integrals (parallel)..."
eri = zeros(n_basis, n_basis, n_basis, n_basis)
parallel for mu to n_basis {
for nu to n_basis {
for lam to n_basis {
for sig to n_basis {
eri[mu, nu, lam, sig] = contracted_eri(mu, nu, lam, sig)
}
}
}
}
# ---- Symmetric orthogonalization X = S^{-1/2}. ----
s_vals = zeros(n_basis)
s_vecs = zeros(n_basis, n_basis)
eig_sym(S, s_vals, s_vecs)
D_inv_sqrt = zeros(n_basis, n_basis)
for i to n_basis {
D_inv_sqrt[i, i] = 1.0 / sqrt(s_vals[i])
}
X = s_vecs @ D_inv_sqrt @ transpose(s_vecs)
# ---- Nuclear repulsion. ----
E_nuc = 0.0
for i to n_atoms {
for j to n_atoms {
if j > i {
E_nuc += Z[i] * Z[j] / sqrt(dist_sq(center_of(i), center_of(j)))
}
}
}
# ---- SCF loop. ----
P = zeros(n_basis, n_basis)
E_total = 0.0
tol = 1.0e-8
max_iter = 100
for iter to max_iter {
G = zeros(n_basis, n_basis)
for mu to n_basis {
for nu to n_basis {
g = 0.0
for lam to n_basis {
for sig to n_basis {
j_int = eri[mu, nu, lam, sig]
k_int = eri[mu, lam, nu, sig]
g += P[lam, sig] * (j_int - 0.5 * k_int)
}
}
G[mu, nu] = g
}
}
F = H_core + G
Fp = transpose(X) @ F @ X
eps_vals = zeros(n_basis)
Cp = zeros(n_basis, n_basis)
eig_sym(Fp, eps_vals, Cp)
C = X @ Cp
P_new = zeros(n_basis, n_basis)
for mu to n_basis {
for nu to n_basis {
psum = 0.0
for a to n_occ {
psum += C[mu, a] * C[nu, a]
}
P_new[mu, nu] = 2.0 * psum
}
}
E_elec = 0.0
for mu to n_basis {
for nu to n_basis {
E_elec += 0.5 * P_new[mu, nu] * (H_core[mu, nu] + F[mu, nu])
}
}
E_total = E_elec + E_nuc
delta = 0.0
for mu to n_basis {
for nu to n_basis {
d = P_new[mu, nu] - P[mu, nu]
if d < 0.0 { d = -d }
if d > delta { delta = d }
}
}
P = P_new
if delta < tol { break }
}
show n_atoms
show E_total