from hf.basis import bas_nao_cart, bas_nao_sph, SH_CART0, SH_L, SH_NCART, SH_NSPH, SH_SPH0
from hf.eri_store import es_index8, es_len8, es_pack8, es_pair, es_unpack8
import hf.cart2sph
from hf.md_eri import mderi_screened
# the fuzz test's seeds, 1 to 8
TR4_FUZZ_NSEED = 8
# tr4_rand_shells: 3 to 5 shells, each s to f (4 values of l), and a shell
# that would take the Cartesian count past 15 becomes an s
TR4_RAND_NSH_MIN = 3
TR4_RAND_NSH_SPREAD = 3.0
TR4_RAND_NL = 4.0
TR4_RAND_NCART_MAX = 15
# ---- one-electron matrices ----------------------------------------------------
#| `M^T X M` as two full matrix products, `[nao_sph, nao_sph]`. The shells
#| pack is used only for the two sizes.
def tr2_dense(X: Arr, M: Arr, shells: Arr) -> Arr {
return transpose(M) @ (X @ M)
}
#| `M^T X M` one shell pair at a time: the (a, b) spherical block is
#| `M_a^T X_ab M_b`, where `M_a` is shell a's diagonal block of M. Every pair
#| is done, so X need not be symmetric.
#@ 21711c66
def tr2_blocks(X: Arr, M: Arr, shells: Arr) -> Arr {
ns = bas_nao_sph(shells)
Y = zeros(ns, ns)
for a to len(shells) {
ca = round(shells[a, SH_CART0])
sa = round(shells[a, SH_SPH0])
na = round(shells[a, SH_NCART])
ma = round(shells[a, SH_NSPH])
for b to len(shells) {
cb = round(shells[b, SH_CART0])
sb = round(shells[b, SH_SPH0])
nb = round(shells[b, SH_NCART])
mb = round(shells[b, SH_NSPH])
# X_ab M_b, then M_a^T times that
H = X[ca:ca + na, cb:cb + nb] @ M[cb:cb + nb, sb:sb + mb]
Y[sa:sa + ma, sb:sb + mb] = transpose(M[ca:ca + na, sa:sa + ma]) @ H
}
}
return Y
}
#| `M D M^T`, `[nao_cart, nao_cart]`: a spherical density as the Cartesian
#| one with the same contractions, since Tr(D M^T X M) = Tr((M D M^T) X).
def tr2_back_dense(D: Arr, M: Arr, shells: Arr) -> Arr {
return M @ D @ transpose(M)
}
#| `M D M^T` one shell pair at a time.
def tr2_back_blocks(D: Arr, M: Arr, shells: Arr) -> Arr {
nc = bas_nao_cart(shells)
Y = zeros(nc, nc)
for a to len(shells) {
ca = round(shells[a, SH_CART0])
sa = round(shells[a, SH_SPH0])
na = round(shells[a, SH_NCART])
ma = round(shells[a, SH_NSPH])
for b to len(shells) {
cb = round(shells[b, SH_CART0])
sb = round(shells[b, SH_SPH0])
nb = round(shells[b, SH_NCART])
mb = round(shells[b, SH_NSPH])
H = M[ca:ca + na, sa:sa + ma] @ D[sa:sa + ma, sb:sb + mb]
Y[ca:ca + na, cb:cb + nb] = H @ transpose(M[cb:cb + nb, sb:sb + mb])
}
}
return Y
}
# ---- electron repulsion integrals ---------------------------------------------
#| The dense four-index transform of a dense `[nc, nc, nc, nc]` tensor, one
#| index per pass: l, then k, then j, then i. Returns `[ns, ns, ns, ns]`.
def tr4_dense(G: Arr, M: Arr, nc, ns) -> Arr {
#: one index per pass: nc^4 ns work each, not nc^4 ns^4
A = zeros(nc, nc, nc, ns)
for i to nc {
for j to nc {
for k to nc {
for s to ns {
t = 0.0
for l to nc {
t = t + G[i, j, k, l] * M[l, s]
}
A[i, j, k, s] = t
}
}
}
}
B = zeros(nc, nc, ns, ns)
for i to nc {
for j to nc {
for r to ns {
for s to ns {
t = 0.0
for k to nc {
t = t + A[i, j, k, s] * M[k, r]
}
B[i, j, r, s] = t
}
}
}
}
C = zeros(nc, ns, ns, ns)
for i to nc {
for q to ns {
for r to ns {
for s to ns {
t = 0.0
for j to nc {
t = t + B[i, j, r, s] * M[j, q]
}
C[i, q, r, s] = t
}
}
}
}
Y = zeros(ns, ns, ns, ns)
for p to ns {
for q to ns {
for r to ns {
for s to ns {
t = 0.0
for i to nc {
t = t + C[i, q, r, s] * M[i, p]
}
Y[p, q, r, s] = t
}
}
}
}
return Y
}
#| Packed Cartesian ERIs to packed spherical ones through the dense view and
#| four quarter transformations. Meant for n <= 30, like `es_unpack8`.
def tr4_quarter(Gp: Arr, M: Arr, shells: Arr) -> Arr {
nc = bas_nao_cart(shells)
ns = bas_nao_sph(shells)
G = es_unpack8(Gp, nc)
Y = tr4_dense(G, M, nc, ns)
return es_pack8(Y, ns)
}
#| The spherical block `[ma, mb, mc, md]` of shells (a, b, c, d), from the
#| Cartesian block read out of the packed vector and the four `M_l` blocks.
#| Any order of the four shells works, because `es_index8` accepts any order.
#| A shell up to p is the same in both sets, so its pass is skipped.
#@ df9ccdbc
def tr4_quartet(Gp: Arr, M: Arr, shells: Arr, a, b, c, d) -> Arr {
ca = round(shells[a, SH_CART0])
cb = round(shells[b, SH_CART0])
cc = round(shells[c, SH_CART0])
cd = round(shells[d, SH_CART0])
sa = round(shells[a, SH_SPH0])
sb = round(shells[b, SH_SPH0])
sc = round(shells[c, SH_SPH0])
sd = round(shells[d, SH_SPH0])
na = round(shells[a, SH_NCART])
nb = round(shells[b, SH_NCART])
nk = round(shells[c, SH_NCART])
nd = round(shells[d, SH_NCART])
ma = round(shells[a, SH_NSPH])
mb = round(shells[b, SH_NSPH])
mc = round(shells[c, SH_NSPH])
md = round(shells[d, SH_NSPH])
g = zeros(na, nb, nk, nd)
for i to na {
for j to nb {
for k to nk {
for l to nd {
g[i, j, k, l] = Gp[es_index8(ca + i, cb + j, cc + k, cd + l)]
}
}
}
}
# the last index first, then the third, second and first. A shell up to
# p has the same functions in both sets (its block of M is the identity,
# [test "a shell up to p is the same in both sets"]), so its pass is
# skipped and the block goes on as it is.
h = g
if shells[d, SH_L] > 1 {
h = zeros(na, nb, nk, md)
for i to na {
for j to nb {
for k to nk {
for s to md {
t = 0.0
for l to nd {
t = t + g[i, j, k, l] * M[cd + l, sd + s]
}
h[i, j, k, s] = t
}
}
}
}
}
e = h
if shells[c, SH_L] > 1 {
e = zeros(na, nb, mc, md)
for i to na {
for j to nb {
for r to mc {
for s to md {
t = 0.0
for k to nk {
t = t + h[i, j, k, s] * M[cc + k, sc + r]
}
e[i, j, r, s] = t
}
}
}
}
}
f = e
if shells[b, SH_L] > 1 {
f = zeros(na, mb, mc, md)
for i to na {
for q to mb {
for r to mc {
for s to md {
t = 0.0
for j to nb {
t = t + e[i, j, r, s] * M[cb + j, sb + q]
}
f[i, q, r, s] = t
}
}
}
}
}
y = f
if shells[a, SH_L] > 1 {
y = zeros(ma, mb, mc, md)
for p to ma {
for q to mb {
for r to mc {
for s to md {
t = 0.0
for i to na {
t = t + f[i, q, r, s] * M[ca + i, sa + p]
}
y[p, q, r, s] = t
}
}
}
}
}
return y
}
#| Packed Cartesian ERIs to packed spherical ones, one canonical shell
#| quartet at a time: a >= b, c >= d and pair(a, b) >= pair(c, d). Every
#| canonical spherical (pq|rs) lies in one of those quartets, since the
#| symmetry that maps a function quartet to its canonical order maps its
#| shell quartet the same way.
#@ 1d51440b
def tr4_blocks(Gp: Arr, M: Arr, shells: Arr) -> Arr {
ns = bas_nao_sph(shells)
out = zeros(es_len8(ns))
nsh = len(shells)
#: canonical shell quartets only (a >= b, c >= d, cd <= ab)
for a to nsh {
for b to a + 1 {
for c to a + 1 {
for d to c + 1 {
if es_pair(c, d) <= es_pair(a, b) {
y = tr4_quartet(Gp, M, shells, a, b, c, d)
sa = round(shells[a, SH_SPH0])
sb = round(shells[b, SH_SPH0])
sc = round(shells[c, SH_SPH0])
sd = round(shells[d, SH_SPH0])
for p to round(shells[a, SH_NSPH]) {
for q to round(shells[b, SH_NSPH]) {
for r to round(shells[c, SH_NSPH]) {
for s to round(shells[d, SH_NSPH]) {
out[es_index8(sa + p, sb + q, sc + r, sd + s)] = y[p, q, r, s]
}
}
}
}
}
}
}
}
}
return out
}
#| The ERIs over spherical functions in two passes: every Cartesian quartet
#| by `mderi_screened` (lib/hf/md_eri), then `tr4_blocks`. The production
#| build, `mderi_sph`, moves each quartet as it is computed; this is the path
#| it is held to, with the same signature, `(atoms, shells, prims, M)`.
#@ 0aa86860
def tr4_two_pass(atoms: Arr, shells: Arr, prims: Arr, M: Arr) -> Arr {
return tr4_blocks(mderi_screened(atoms, shells, prims), M, shells)
}
# ---- test helpers ---------------------------------------------------------------
#| The per-quartet form computed for a shell quartet and for three of its
#| symmetric images, (b a | d c), (c d | a b) and (d c | b a); the largest
#| gap between matching entries. `tr4_blocks` computes only the canonical
#| order, so this checks that the order it picks does not matter.
def tr4_block_sym_defect(Gp: Arr, M: Arr, shells: Arr) {
nsh = len(shells)
w = 0.0
for a to nsh {
for b to nsh {
for c to nsh {
for d to nsh {
y = tr4_quartet(Gp, M, shells, a, b, c, d)
y1 = tr4_quartet(Gp, M, shells, b, a, d, c)
y2 = tr4_quartet(Gp, M, shells, c, d, a, b)
y3 = tr4_quartet(Gp, M, shells, d, c, b, a)
for p to round(shells[a, SH_NSPH]) {
for q to round(shells[b, SH_NSPH]) {
for r to round(shells[c, SH_NSPH]) {
for s to round(shells[d, SH_NSPH]) {
v = y[p, q, r, s]
w = max(w, abs(v - y1[q, p, s, r]))
w = max(w, abs(v - y2[r, s, p, q]))
w = max(w, abs(v - y3[s, r, q, p]))
}
}
}
}
}
}
}
}
return w
}
# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors