from hf.basis import (
AT_CHG, AT_X, AT_Y, AT_Z, bas_contract, bas_nao_cart, bas_nc, bas_ncart, bas_powers, PR_EXP,
SH_ATOM, SH_CART0, SH_L, SH_NCART, SH_NPRIM, SH_P0,
)
import hf.boys
from hf.hermite import herm_boys_fill, herm_efill, herm_esize, herm_grid, HERM_GRID_XMAX, herm_rfill
MD1_PI = 3.141592653589793
MD1_S = 0
MD1_T = 1
MD1_V = 2
# the three matrices a pair block holds: S, T and V
MD1_NMAT = 3
# rows of lib/hfref/prims' shell table per case (one per function of an ERI quartet)
MD1_PS_ROWS = 4
# a prims case's atoms: the two shell centers and the point charge
MD1_T_NATOM = 3
# a row of md1_prims_fix: (case, pick, kind, value)
MD1_FIX_ROW = 4
#| The Boys grid a one-electron build needs: orders up to `2 lmax`.
def md1_grid(shells: Arr) {
lmax = 0
for s to len(shells) {
lmax = max(lmax, shells[s, SH_L])
}
return herm_grid(2 * lmax, HERM_GRID_XMAX)
}
#| S, T and V for shell pair `(sa, sb)`, as `[3, ncart_a * ncart_b]` with the
#| component pair `(ia, ib)` at column `ia * ncart_b + ib`. `cn` comes from
#| `bas_contract` and `G` from `md1_grid`.
def md1_pair(atoms: Arr, shells: Arr, prims: Arr, cn: Arr, G: Arr, sa, sb) {
la = shells[sa, SH_L]
lb = shells[sb, SH_L]
na = bas_ncart(la)
nb = bas_ncart(lb)
PA = bas_powers(la)
PB = bas_powers(lb)
ca = round(shells[sa, SH_ATOM])
cb = round(shells[sb, SH_ATOM])
Ax = atoms[ca, AT_X]
Ay = atoms[ca, AT_Y]
Az = atoms[ca, AT_Z]
Bx = atoms[cb, AT_X]
By = atoms[cb, AT_Y]
Bz = atoms[cb, AT_Z]
# component factors, one per row and column
NA = zeros(na)
NB = zeros(nb)
for i to na {
NA[i] = bas_nc(PA[i, 0], PA[i, 1], PA[i, 2])
}
for j to nb {
NB[j] = bas_nc(PB[j, 0], PB[j, 1], PB[j, 2])
}
#: two extra powers on b, for the kinetic integral
L = la + lb
lb2 = lb + 2
ld = lb2 + 1
nt = la + lb2 + 1
EX = zeros(herm_esize(la, lb2))
EY = zeros(herm_esize(la, lb2))
EZ = zeros(herm_esize(la, lb2))
D = L + 1
RS = zeros(2 * D * D * D)
F = zeros(L + 1)
out = zeros(MD1_NMAT, na * nb)
pa0 = round(shells[sa, SH_P0])
pb0 = round(shells[sb, SH_P0])
for ip to shells[sa, SH_NPRIM] {
a = prims[pa0 + ip, PR_EXP]
for jp to shells[sb, SH_NPRIM] {
b = prims[pb0 + jp, PR_EXP]
p = a + b
c = cn[pa0 + ip] * cn[pb0 + jp]
herm_efill(EX, 0, la, lb2, a, b, Ax - Bx)
herm_efill(EY, 0, la, lb2, a, b, Ay - By)
herm_efill(EZ, 0, la, lb2, a, b, Az - Bz)
sq = sqrt(MD1_PI / p)
Px = (a * Ax + b * Bx) / p
Py = (a * Ay + b * By) / p
Pz = (a * Az + b * Bz) / p
#: T per axis: -2b^2 S(j+2) + b(2j+1) S(j) - j(j-1)/2 S(j-2)
for i to na {
ax = PA[i, 0]
ay = PA[i, 1]
az = PA[i, 2]
for j to nb {
bx = PB[j, 0]
by = PB[j, 1]
bz = PB[j, 2]
sx = EX[(ax * ld + bx) * nt] * sq
sy = EY[(ay * ld + by) * nt] * sq
sz = EZ[(az * ld + bz) * nt] * sq
tx = -2.0 * b * b * EX[(ax * ld + bx + 2) * nt] * sq + b * (2 * bx + 1) * sx
ty = -2.0 * b * b * EY[(ay * ld + by + 2) * nt] * sq + b * (2 * by + 1) * sy
tz = -2.0 * b * b * EZ[(az * ld + bz + 2) * nt] * sq + b * (2 * bz + 1) * sz
if bx > 1 {
tx = tx - 0.5 * bx * (bx - 1) * EX[(ax * ld + bx - 2) * nt] * sq
}
if by > 1 {
ty = ty - 0.5 * by * (by - 1) * EY[(ay * ld + by - 2) * nt] * sq
}
if bz > 1 {
tz = tz - 0.5 * bz * (bz - 1) * EZ[(az * ld + bz - 2) * nt] * sq
}
w = c * NA[i] * NB[j]
k = i * nb + j
out[MD1_S, k] = out[MD1_S, k] + w * sx * sy * sz
out[MD1_T, k] = out[MD1_T, k] + w * (tx * sy * sz + sx * ty * sz + sx * sy * tz)
}
}
#: nuclear attraction: an R table at P - C per charged atom
for at to len(atoms) {
Zc = atoms[at, AT_CHG]
if Zc == 0.0 {
continue
}
X = Px - atoms[at, AT_X]
Y = Py - atoms[at, AT_Y]
Z = Pz - atoms[at, AT_Z]
herm_boys_fill(G, L, p * (X * X + Y * Y + Z * Z), F)
herm_rfill(RS, L, p, X, Y, Z, F, -Zc * 2.0 * MD1_PI / p)
for i to na {
ax = PA[i, 0]
ay = PA[i, 1]
az = PA[i, 2]
for j to nb {
bx = PB[j, 0]
by = PB[j, 1]
bz = PB[j, 2]
ox = (ax * ld + bx) * nt
oy = (ay * ld + by) * nt
oz = (az * ld + bz) * nt
acc = 0.0
for t to ax + bx + 1 {
for u to ay + by + 1 {
eu = EX[ox + t] * EY[oy + u]
for v to az + bz + 1 {
acc = acc + eu * EZ[oz + v] * RS[(t * D + u) * D + v]
}
}
}
k = i * nb + j
out[MD1_V, k] = out[MD1_V, k] + c * NA[i] * NB[j] * acc
}
}
}
}
}
return out
}
#| S, T and V for the whole basis, `[3, n, n]` over Cartesian AOs. Each shell
#| pair is computed once, for `sa >= sb`, and mirrored, so every matrix is
#| exactly symmetric.
def md1_all(atoms: Arr, shells: Arr, prims: Arr) {
n = bas_nao_cart(shells)
cn = bas_contract(shells, prims)
G = md1_grid(shells)
M = zeros(MD1_NMAT, n, n)
for sa to len(shells) {
for sb to sa + 1 {
md1_scatter(M, md1_pair(atoms, shells, prims, cn, G, sa, sb), shells, sa, sb)
}
}
return M
}
#| Write a `md1_pair` block into `M` at both `(sa, sb)` and `(sb, sa)`.
def md1_scatter(M: Arr, B: Arr, shells: Arr, sa, sb) {
na = shells[sa, SH_NCART]
nb = shells[sb, SH_NCART]
a0 = shells[sa, SH_CART0]
b0 = shells[sb, SH_CART0]
for k to MD1_NMAT {
for i to na {
for j to nb {
M[k, a0 + i, b0 + j] = B[k, i * nb + j]
M[k, b0 + j, a0 + i] = B[k, i * nb + j]
}
}
}
return 0
}
#| One matrix out of `md1_all`.
def md1_pick(atoms: Arr, shells: Arr, prims: Arr, which) {
return slab(md1_all(atoms, shells, prims), which)
}
#| Integral-stage members `(atoms, shells, prims) -> [n, n]`.
def md1_s(atoms: Arr, shells: Arr, prims: Arr) -> Arr {
return md1_pick(atoms, shells, prims, MD1_S)
}
def md1_t(atoms: Arr, shells: Arr, prims: Arr) -> Arr {
return md1_pick(atoms, shells, prims, MD1_T)
}
def md1_v(atoms: Arr, shells: Arr, prims: Arr) -> Arr {
return md1_pick(atoms, shells, prims, MD1_V)
}
#| The core Hamiltonian `T + V`.
def md1_h(atoms: Arr, shells: Arr, prims: Arr) -> Arr {
M = md1_all(atoms, shells, prims)
return slab(M, MD1_T) + slab(M, MD1_V)
}
#| Test helper: the largest absolute gap between matrix `k` of `M` and `R`.
def md1_gap(M: Arr, k, R: Arr) {
return max(abs(slab(M, k) - R))
}
# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors