Notebooks · lib/hf/os_onebody

Overlap, kinetic and nuclear attraction by the Obara-Saika recurrences

This module computes the three one-electron matrices an SCF needs, S, T and V, over unit-normalized Cartesian Gaussians, by the recurrences of Obara and Saika (J. Chem. Phys. 84, 3963, 1986). It is the second of the library's one-electron families. It shares no code with the McMurchie-Davidson family (hermite.viv, md_*), so when the two agree to 10⁻¹² the agreement says something about both.

The idea. A Cartesian Gaussian on center A is x_A^i y_A^j z_A^k exp(-a r_A^2), with x_A = x - A_x. Raising the power i by one is multiplying by x_A, and the product of two Gaussians on A and B is one Gaussian on P = (aA + bB)/p, p = a + b. So x_A = (x - P_x) + (P_x - A_x), and one integration by parts turns the (x - P_x) piece into lower powers. For the overlap along one axis that gives

S(i+1, j) = PA S(i, j) + (i S(i-1, j) + j S(i, j-1)) / (2p)
S(i, j+1) = PB S(i, j) + (i S(i-1, j) + j S(i, j-1)) / (2p)

starting from S(0, 0) = sqrt(pi/p) exp(-mu AB^2) with mu = ab/p. The 3-D overlap is the product of the three axes.

Kinetic energy. T = -1/2 <a|d^2/dx^2 + d^2/dy^2 + d^2/dz^2|b>, which is Tx Sy Sz + Sx Ty Sz + Sx Sy Tz. Along one axis the same argument gives

T(0, 0)   = mu (1 - 2 mu AB^2) S(0, 0)
T(i+1, j) = PA T(i, j) + (i T(i-1, j) + j T(i, j-1)) / (2p)
            + (b/p) (2a S(i+1, j) - i S(i-1, j))
T(i, j+1) = PB T(i, j) + (i T(i-1, j) + j T(i, j-1)) / (2p)
            + (a/p) (2b S(i, j+1) - j S(i, j-1))

(Helgaker, Jorgensen and Olsen, section 9.3.4).

Nuclear attraction does not factor by axis, because 1/r_C couples them. Obara and Saika carry an auxiliary index m: the integral Theta(e)^(m) over the Gaussian product with e on A is built from

Theta(0)^(m)    = (2 pi / p) exp(-mu AB^2) F_m(p PC^2)
Theta(e+1_i)^(m) = PA_i Theta(e)^(m) - PC_i Theta(e)^(m+1)
                   + e_i/(2p) (Theta(e-1_i)^(m) - Theta(e-1_i)^(m+1))

where F_m is the Boys function and only m = 0 is kept at the end. Rather than recur on both centers, the recurrence builds every power on A up to la + lb and the horizontal step

(a, b+1_i) = (a+1_i, b) + AB_i (a, b)

moves powers from A to B. That step holds for any operator that does not depend on the basis centers, and it is applied once per shell pair after the primitives and nuclei are summed, since it is linear.

Components. Every Cartesian component up to total power L gets one index: l(l+1)(l+2)/6 for the components below l, then PySCF's order within the shell (lx counting down, then ly). os_ctab(L) tabulates the powers and the neighbor one step up and down along each axis, so the recurrences look up indices instead of computing them.

Normalization. bas_contract gives cn[p], which carries the radial norm of each primitive and the shell's contraction scale; each component then gets bas_nc(lx, ly, lz) (see lib/hf/basis.viv). The result is PySCF's int1e_*_cart divided by F F^T.

The stage members os_overlap, os_kinetic and os_nuclear take (atoms, shells, prims) and return [nao_cart, nao_cart]. The per-pair blocks (os_overlap_block and the others) take shell indices, for tests that look at one pair.

Used with import hf.os_onebody.

Look inside this program

Atlas
The program’s map: every function and which calls which.
There’s no recorded run of this program yet.

Shift+Enter in a cell runs it and moves to the next. Running a cell runs every cell above it too, as one program.

[1]lines 73–454

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_NPRIM, SH_P0,
)
from hf.boys import boys_values

OS_PI = 3.141592653589793

OS_CT_LX = 0
OS_CT_LY = 1
OS_CT_LZ = 2
OS_CT_L = 3
OS_CT_DN = 4
OS_CT_UP = 7
OS_CT_PIV = 10
OS_CT_LEN = 11
# rows of lib/hfref/prims' shell table per case: four single-primitive shells
OS_PS_ROWS = 4

#| Components with total power below `l`.
def os_ncum(l) {
    return l * (l + 1) * (l + 2) / 6
}

#| Index of component (lx, ly, lz) in the table of `os_ctab`.
def os_cidx(lx, ly, lz) {
    l = lx + ly + lz
    a = l - lx
    return os_ncum(l) + a * (a + 1) / 2 + lz
}

#| One row per Cartesian component with total power up to `L`: the powers,
#| `l`, the index one step down along x, y, z (-1 where that power is 0), one
#| step up (-1 at `l = L`), and the pivot, the first axis with a nonzero power
#| (-1 for s). The recurrences lower the pivot.
def os_ctab(L) {
    T = zeros(os_ncum(L + 1), OS_CT_LEN)
    for l to L + 1 {
        for a to l + 1 {
            lx = l - a
            for lz to a + 1 {
                ly = a - lz
                r = os_cidx(lx, ly, lz)
                T[r, OS_CT_LX] = lx
                T[r, OS_CT_LY] = ly
                T[r, OS_CT_LZ] = lz
                T[r, OS_CT_L] = l
                T[r, OS_CT_DN] = -1
                T[r, OS_CT_DN + 1] = -1
                T[r, OS_CT_DN + 2] = -1
                T[r, OS_CT_UP] = -1
                T[r, OS_CT_UP + 1] = -1
                T[r, OS_CT_UP + 2] = -1
                T[r, OS_CT_PIV] = -1
                if lx > 0 {
                    T[r, OS_CT_DN] = os_cidx(lx - 1, ly, lz)
                }
                if ly > 0 {
                    T[r, OS_CT_DN + 1] = os_cidx(lx, ly - 1, lz)
                }
                if lz > 0 {
                    T[r, OS_CT_DN + 2] = os_cidx(lx, ly, lz - 1)
                }
                if l < L {
                    T[r, OS_CT_UP] = os_cidx(lx + 1, ly, lz)
                    T[r, OS_CT_UP + 1] = os_cidx(lx, ly + 1, lz)
                    T[r, OS_CT_UP + 2] = os_cidx(lx, ly, lz + 1)
                }
                if lz > 0 {
                    T[r, OS_CT_PIV] = 2
                }
                if ly > 0 {
                    T[r, OS_CT_PIV] = 1
                }
                if lx > 0 {
                    T[r, OS_CT_PIV] = 0
                }
            }
        }
    }
    return T
}

#| The 1-D overlap table `S[i, j]`, `i <= la`, `j <= lb`, for exponents `a`
#| and `b` on centers `A` and `B` along one axis. The table is written into
#| `S` (at least `[la+1, lb+1]`).
def os_s1d(a, b, A, B, la, lb, S: Arr) {
    p = a + b
    P = (a * A + b * B) / p
    PA = P - A
    PB = P - B
    h = 0.5 / p
    S[0, 0] = sqrt(OS_PI / p) * exp(-a * b / p * (A - B) * (A - B))
    for i to la {
        v = PA * S[i, 0]
        if i > 0 {
            v = v + h * i * S[i - 1, 0]
        }
        S[i + 1, 0] = v
    }
    for j to lb {
        for i to la + 1 {
            v = PB * S[i, j]
            if i > 0 {
                v = v + h * i * S[i - 1, j]
            }
            if j > 0 {
                v = v + h * j * S[i, j - 1]
            }
            S[i, j + 1] = v
        }
    }
    return 0.0
}

#| The 1-D kinetic table `T[i, j]` from the overlap table `S` of the same
#| pair (filled by `os_s1d`). Written into `T`.
def os_t1d(a, b, A, B, la, lb, S: Arr, T: Arr) {
    p = a + b
    P = (a * A + b * B) / p
    PA = P - A
    PB = P - B
    h = 0.5 / p
    mu = a * b / p
    T[0, 0] = mu * (1.0 - 2.0 * mu * (A - B) * (A - B)) * S[0, 0]
    for i to la {
        v = PA * T[i, 0] + b / p * 2.0 * a * S[i + 1, 0]
        if i > 0 {
            v = v + h * i * T[i - 1, 0] - b / p * i * S[i - 1, 0]
        }
        T[i + 1, 0] = v
    }
    for j to lb {
        for i to la + 1 {
            v = PB * T[i, j] + a / p * 2.0 * b * S[i, j + 1]
            if i > 0 {
                v = v + h * i * T[i - 1, j]
            }
            if j > 0 {
                v = v + h * j * T[i, j - 1] - a / p * j * S[i, j - 1]
            }
            T[i, j + 1] = v
        }
    }
    return 0.0
}

#| Overlap (`kind` 0) or kinetic (`kind` 1) block of shells `sa`, `sb`,
#| `[ncart(la), ncart(lb)]`, unit-Cartesian.
def os_st_block(atoms: Arr, shells: Arr, prims: Arr, cn: Arr, sa, sb, kind) {
    la = shells[sa, SH_L]
    lb = shells[sb, SH_L]
    ia = shells[sa, SH_ATOM]
    ib = shells[sb, SH_ATOM]
    pa0 = shells[sa, SH_P0]
    pb0 = shells[sb, SH_P0]
    Pa = bas_powers(la)
    Pb = bas_powers(lb)
    na = len(Pa)
    nb = len(Pb)
    out = zeros(na, nb)
    Sx = zeros(la + 2, lb + 2)
    Sy = zeros(la + 2, lb + 2)
    Sz = zeros(la + 2, lb + 2)
    Tx = zeros(la + 2, lb + 2)
    Ty = zeros(la + 2, lb + 2)
    Tz = zeros(la + 2, lb + 2)
    for u to shells[sa, SH_NPRIM] {
        a = prims[pa0 + u, PR_EXP]
        for v to shells[sb, SH_NPRIM] {
            b = prims[pb0 + v, PR_EXP]
            cc = cn[pa0 + u] * cn[pb0 + v]
            os_s1d(a, b, atoms[ia, AT_X], atoms[ib, AT_X], la, lb, Sx)
            os_s1d(a, b, atoms[ia, AT_Y], atoms[ib, AT_Y], la, lb, Sy)
            os_s1d(a, b, atoms[ia, AT_Z], atoms[ib, AT_Z], la, lb, Sz)
            if kind == 1 {
                os_t1d(a, b, atoms[ia, AT_X], atoms[ib, AT_X], la, lb, Sx, Tx)
                os_t1d(a, b, atoms[ia, AT_Y], atoms[ib, AT_Y], la, lb, Sy, Ty)
                os_t1d(a, b, atoms[ia, AT_Z], atoms[ib, AT_Z], la, lb, Sz, Tz)
            }
            for r to na {
                x1 = Pa[r, 0]
                y1 = Pa[r, 1]
                z1 = Pa[r, 2]
                for c to nb {
                    x2 = Pb[c, 0]
                    y2 = Pb[c, 1]
                    z2 = Pb[c, 2]
                    if kind == 1 {
                        val = Tx[x1, x2] * Sy[y1, y2] * Sz[z1, z2] + Sx[x1, x2] * Ty[y1, y2] * Sz[z1, z2] + Sx[x1, x2] * Sy[y1, y2] * Tz[z1, z2]
                    } else {
                        val = Sx[x1, x2] * Sy[y1, y2] * Sz[z1, z2]
                    }
                    out[r, c] = out[r, c] + cc * val
                }
            }
        }
    }
    for r to na {
        for c to nb {
            out[r, c] = out[r, c] * bas_nc(Pa[r, 0], Pa[r, 1], Pa[r, 2]) * bas_nc(Pb[c, 0], Pb[c, 1], Pb[c, 2])
        }
    }
    return out
}

#| Overlap block of shells `sa`, `sb`.
def os_overlap_block(atoms: Arr, shells: Arr, prims: Arr, cn: Arr, sa, sb) {
    return os_st_block(atoms, shells, prims, cn, sa, sb, 0)
}

#| Kinetic block of shells `sa`, `sb`.
def os_kinetic_block(atoms: Arr, shells: Arr, prims: Arr, cn: Arr, sa, sb) {
    return os_st_block(atoms, shells, prims, cn, sa, sb, 1)
}

#| Nuclear attraction block of shells `sa`, `sb`: `-sum_C Z_C <a|1/r_C|b>`
#| over every row of `atoms`. The vertical recurrence builds `Theta(e)^(m)`
#| for every `e` up to `la + lb` on A; primitives and nuclei are summed into
#| `acc[e]` for `la <= l_e <= la + lb`, and the horizontal step then moves
#| `lb` units of power to B.
def os_nuclear_block(atoms: Arr, shells: Arr, prims: Arr, cn: Arr, sa, sb) {
    la = shells[sa, SH_L]
    lb = shells[sb, SH_L]
    L = la + lb
    CT = os_ctab(L)
    ne = os_ncum(L + 1)
    A = atoms[shells[sa, SH_ATOM], AT_X:AT_Z + 1]
    B = atoms[shells[sb, SH_ATOM], AT_X:AT_Z + 1]
    AB2 = sum((A - B) * (A - B))
    pa0 = shells[sa, SH_P0]
    pb0 = shells[sb, SH_P0]
    TH = zeros(ne, L + 1)
    acc = zeros(ne)
    for u to shells[sa, SH_NPRIM] {
        a = prims[pa0 + u, PR_EXP]
        for v to shells[sb, SH_NPRIM] {
            b = prims[pb0 + v, PR_EXP]
            p = a + b
            h = 0.5 / p
            P = (a * A + b * B) / p
            PA = P - A
            pref = 2.0 * OS_PI / p * exp(-a * b / p * AB2) * cn[pa0 + u] * cn[pb0 + v]
            for C to len(atoms) {
                Z = atoms[C, AT_CHG]
                if Z != 0.0 {
                    PC = P - atoms[C, AT_X:AT_Z + 1]
                    TH[0, 0:L + 1] = pref * boys_values(L, p * sum(PC * PC))
                    for q to ne - 1 {
                        e = q + 1
                        i = CT[e, OS_CT_PIV]
                        e1 = CT[e, OS_CT_DN + i]
                        e2 = CT[e1, OS_CT_DN + i]
                        ni = CT[e1, i]
                        for m to L + 1 - CT[e, OS_CT_L] {
                            w = PA[i] * TH[e1, m] - PC[i] * TH[e1, m + 1]
                            if ni > 0 {
                                w = w + h * ni * (TH[e2, m] - TH[e2, m + 1])
                            }
                            TH[e, m] = w
                        }
                    }
                    acc[os_ncum(la):ne] = acc[os_ncum(la):ne] - Z * TH[os_ncum(la):ne, 0]
                }
            }
        }
    }
    return os_hrr_pair(acc, CT, la, lb, A - B)
}

#| The horizontal step on a one-electron pair: from `acc[e]` (every `e` with
#| `la <= l_e <= la + lb`, on A) to the block `[ncart(la), ncart(lb)]`,
#| times the component factors. `(a, b+1_i) = (a+1_i, b) + AB_i (a, b)`.
def os_hrr_pair(acc: Arr, CT: Arr, la, lb, AB: Arr) {
    L = la + lb
    ne = os_ncum(L + 1)
    nbc = os_ncum(lb + 1)
    # H[e, f]: e on A, f on B (f indexed from 0 over all powers up to lb)
    H = zeros(ne, nbc)
    for q to ne - os_ncum(la) {
        e = os_ncum(la) + q
        H[e, 0] = acc[e]
    }
    for q to nbc - 1 {
        f = q + 1
        i = CT[f, OS_CT_PIV]
        f1 = CT[f, OS_CT_DN + i]
        lf = CT[f, OS_CT_L]
        for t to os_ncum(L - lf + 1) - os_ncum(la) {
            e = os_ncum(la) + t
            H[e, f] = H[CT[e, OS_CT_UP + i], f1] + AB[i] * H[e, f1]
        }
    }
    na = bas_ncart(la)
    nb = bas_ncart(lb)
    out = zeros(na, nb)
    for r to na {
        e = os_ncum(la) + r
        for c to nb {
            f = os_ncum(lb) + c
            out[r, c] = H[e, f] * bas_nc(CT[e, 0], CT[e, 1], CT[e, 2]) * bas_nc(CT[f, 0], CT[f, 1], CT[f, 2])
        }
    }
    return out
}

#| Fill a whole matrix from per-pair blocks: `kind` 0 overlap, 1 kinetic,
#| 2 nuclear attraction. Only `sa >= sb` is computed; the other triangle is
#| the transpose.
def os_matrix(atoms: Arr, shells: Arr, prims: Arr, kind) {
    cn = bas_contract(shells, prims)
    n = bas_nao_cart(shells)
    M = zeros(n, n)
    for sa to len(shells) {
        for sb to sa + 1 {
            if kind == 2 {
                blk = os_nuclear_block(atoms, shells, prims, cn, sa, sb)
            } else {
                blk = os_st_block(atoms, shells, prims, cn, sa, sb, kind)
            }
            a0 = shells[sa, SH_CART0]
            b0 = shells[sb, SH_CART0]
            for r to len(blk) {
                for c to bas_ncart(shells[sb, SH_L]) {
                    M[a0 + r, b0 + c] = blk[r, c]
                    M[b0 + c, a0 + r] = blk[r, c]
                }
            }
        }
    }
    return M
}

#| Stage member: the overlap matrix, `[nao_cart, nao_cart]`.
def os_overlap(atoms: Arr, shells: Arr, prims: Arr) -> Arr {
    return os_matrix(atoms, shells, prims, 0)
}

#| Stage member: the kinetic energy matrix.
def os_kinetic(atoms: Arr, shells: Arr, prims: Arr) -> Arr {
    return os_matrix(atoms, shells, prims, 1)
}

#| Stage member: the nuclear attraction matrix, over every atom's charge.
def os_nuclear(atoms: Arr, shells: Arr, prims: Arr) -> Arr {
    return os_matrix(atoms, shells, prims, 2)
}

# ---- tests ----------------------------------------------------------------------
# The reference modules used above exist for these tests. An import cannot sit
# inside a test block, so a program that uses this module also compiles their
# literal-returning defs; nothing calls them, and the measured cost of two of
# them (water_631gs_cart and water_tzvp) is 0.55 s of cc -O2.

test "the component table lists every power once, in PySCF's order within a shell" {
    # every shell up to l = 6 (i)
    os_lmax = 6
    os_ct = os_ctab(os_lmax)
    assert_near(len(os_ct), os_ncum(os_lmax + 1), 0.0)
    os_bad = 0
    for os_l to os_lmax + 1 {
        os_pw = bas_powers(os_l)
        for os_c to len(os_pw) {
            os_r = os_ncum(os_l) + os_c
            if os_ct[os_r, OS_CT_LX] != os_pw[os_c, 0] or os_ct[os_r, OS_CT_LY] != os_pw[os_c, 1] or os_ct[os_r, OS_CT_LZ] != os_pw[os_c, 2] {
                os_bad = os_bad + 1
            }
            # one step down and back up along the pivot returns to the same row
            if os_l > 0 {
                os_i = os_ct[os_r, OS_CT_PIV]
                os_dn = os_ct[os_r, OS_CT_DN + os_i]
                if os_ct[os_dn, OS_CT_UP + os_i] != os_r {
                    os_bad = os_bad + 1
                }
            }
        }
    }
    assert_near(os_bad, 0.0, 0.0)
}

# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors

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_NPRIM, SH_P0,
)
from hf.boys import boys_values

OS_PI = 3.141592653589793

OS_CT_LX = 0
OS_CT_LY = 1
OS_CT_LZ = 2
OS_CT_L = 3
OS_CT_DN = 4
OS_CT_UP = 7
OS_CT_PIV = 10
OS_CT_LEN = 11
# rows of lib/hfref/prims' shell table per case: four single-primitive shells
OS_PS_ROWS = 4

#| Components with total power below `l`.
def os_ncum(l) {
    return l * (l + 1) * (l + 2) / 6
}

#| Index of component (lx, ly, lz) in the table of `os_ctab`.
def os_cidx(lx, ly, lz) {
    l = lx + ly + lz
    a = l - lx
    return os_ncum(l) + a * (a + 1) / 2 + lz
}

#| One row per Cartesian component with total power up to `L`: the powers,
#| `l`, the index one step down along x, y, z (-1 where that power is 0), one
#| step up (-1 at `l = L`), and the pivot, the first axis with a nonzero power
#| (-1 for s). The recurrences lower the pivot.
def os_ctab(L) {
    T = zeros(os_ncum(L + 1), OS_CT_LEN)
    for l to L + 1 {
        for a to l + 1 {
            lx = l - a
            for lz to a + 1 {
                ly = a - lz
                r = os_cidx(lx, ly, lz)
                T[r, OS_CT_LX] = lx
                T[r, OS_CT_LY] = ly
                T[r, OS_CT_LZ] = lz
                T[r, OS_CT_L] = l
                T[r, OS_CT_DN] = -1
                T[r, OS_CT_DN + 1] = -1
                T[r, OS_CT_DN + 2] = -1
                T[r, OS_CT_UP] = -1
                T[r, OS_CT_UP + 1] = -1
                T[r, OS_CT_UP + 2] = -1
                T[r, OS_CT_PIV] = -1
                if lx > 0 {
                    T[r, OS_CT_DN] = os_cidx(lx - 1, ly, lz)
                }
                if ly > 0 {
                    T[r, OS_CT_DN + 1] = os_cidx(lx, ly - 1, lz)
                }
                if lz > 0 {
                    T[r, OS_CT_DN + 2] = os_cidx(lx, ly, lz - 1)
                }
                if l < L {
                    T[r, OS_CT_UP] = os_cidx(lx + 1, ly, lz)
                    T[r, OS_CT_UP + 1] = os_cidx(lx, ly + 1, lz)
                    T[r, OS_CT_UP + 2] = os_cidx(lx, ly, lz + 1)
                }
                if lz > 0 {
                    T[r, OS_CT_PIV] = 2
                }
                if ly > 0 {
                    T[r, OS_CT_PIV] = 1
                }
                if lx > 0 {
                    T[r, OS_CT_PIV] = 0
                }
            }
        }
    }
    return T
}

#| The 1-D overlap table `S[i, j]`, `i <= la`, `j <= lb`, for exponents `a`
#| and `b` on centers `A` and `B` along one axis. The table is written into
#| `S` (at least `[la+1, lb+1]`).
def os_s1d(a, b, A, B, la, lb, S: Arr) {
    p = a + b
    P = (a * A + b * B) / p
    PA = P - A
    PB = P - B
    h = 0.5 / p
    S[0, 0] = sqrt(OS_PI / p) * exp(-a * b / p * (A - B) * (A - B))
    for i to la {
        v = PA * S[i, 0]
        if i > 0 {
            v = v + h * i * S[i - 1, 0]
        }
        S[i + 1, 0] = v
    }
    for j to lb {
        for i to la + 1 {
            v = PB * S[i, j]
            if i > 0 {
                v = v + h * i * S[i - 1, j]
            }
            if j > 0 {
                v = v + h * j * S[i, j - 1]
            }
            S[i, j + 1] = v
        }
    }
    return 0.0
}

#| The 1-D kinetic table `T[i, j]` from the overlap table `S` of the same
#| pair (filled by `os_s1d`). Written into `T`.
def os_t1d(a, b, A, B, la, lb, S: Arr, T: Arr) {
    p = a + b
    P = (a * A + b * B) / p
    PA = P - A
    PB = P - B
    h = 0.5 / p
    mu = a * b / p
    T[0, 0] = mu * (1.0 - 2.0 * mu * (A - B) * (A - B)) * S[0, 0]
    for i to la {
        v = PA * T[i, 0] + b / p * 2.0 * a * S[i + 1, 0]
        if i > 0 {
            v = v + h * i * T[i - 1, 0] - b / p * i * S[i - 1, 0]
        }
        T[i + 1, 0] = v
    }
    for j to lb {
        for i to la + 1 {
            v = PB * T[i, j] + a / p * 2.0 * b * S[i, j + 1]
            if i > 0 {
                v = v + h * i * T[i - 1, j]
            }
            if j > 0 {
                v = v + h * j * T[i, j - 1] - a / p * j * S[i, j - 1]
            }
            T[i, j + 1] = v
        }
    }
    return 0.0
}

#| Overlap (`kind` 0) or kinetic (`kind` 1) block of shells `sa`, `sb`,
#| `[ncart(la), ncart(lb)]`, unit-Cartesian.
def os_st_block(atoms: Arr, shells: Arr, prims: Arr, cn: Arr, sa, sb, kind) {
    la = shells[sa, SH_L]
    lb = shells[sb, SH_L]
    ia = shells[sa, SH_ATOM]
    ib = shells[sb, SH_ATOM]
    pa0 = shells[sa, SH_P0]
    pb0 = shells[sb, SH_P0]
    Pa = bas_powers(la)
    Pb = bas_powers(lb)
    na = len(Pa)
    nb = len(Pb)
    out = zeros(na, nb)
    Sx = zeros(la + 2, lb + 2)
    Sy = zeros(la + 2, lb + 2)
    Sz = zeros(la + 2, lb + 2)
    Tx = zeros(la + 2, lb + 2)
    Ty = zeros(la + 2, lb + 2)
    Tz = zeros(la + 2, lb + 2)
    for u to shells[sa, SH_NPRIM] {
        a = prims[pa0 + u, PR_EXP]
        for v to shells[sb, SH_NPRIM] {
            b = prims[pb0 + v, PR_EXP]
            cc = cn[pa0 + u] * cn[pb0 + v]
            os_s1d(a, b, atoms[ia, AT_X], atoms[ib, AT_X], la, lb, Sx)
            os_s1d(a, b, atoms[ia, AT_Y], atoms[ib, AT_Y], la, lb, Sy)
            os_s1d(a, b, atoms[ia, AT_Z], atoms[ib, AT_Z], la, lb, Sz)
            if kind == 1 {
                os_t1d(a, b, atoms[ia, AT_X], atoms[ib, AT_X], la, lb, Sx, Tx)
                os_t1d(a, b, atoms[ia, AT_Y], atoms[ib, AT_Y], la, lb, Sy, Ty)
                os_t1d(a, b, atoms[ia, AT_Z], atoms[ib, AT_Z], la, lb, Sz, Tz)
            }
            for r to na {
                x1 = Pa[r, 0]
                y1 = Pa[r, 1]
                z1 = Pa[r, 2]
                for c to nb {
                    x2 = Pb[c, 0]
                    y2 = Pb[c, 1]
                    z2 = Pb[c, 2]
                    if kind == 1 {
                        val = Tx[x1, x2] * Sy[y1, y2] * Sz[z1, z2] + Sx[x1, x2] * Ty[y1, y2] * Sz[z1, z2] + Sx[x1, x2] * Sy[y1, y2] * Tz[z1, z2]
                    } else {
                        val = Sx[x1, x2] * Sy[y1, y2] * Sz[z1, z2]
                    }
                    out[r, c] = out[r, c] + cc * val
                }
            }
        }
    }
    for r to na {
        for c to nb {
            out[r, c] = out[r, c] * bas_nc(Pa[r, 0], Pa[r, 1], Pa[r, 2]) * bas_nc(Pb[c, 0], Pb[c, 1], Pb[c, 2])
        }
    }
    return out
}

#| Overlap block of shells `sa`, `sb`.
def os_overlap_block(atoms: Arr, shells: Arr, prims: Arr, cn: Arr, sa, sb) {
    return os_st_block(atoms, shells, prims, cn, sa, sb, 0)
}

#| Kinetic block of shells `sa`, `sb`.
def os_kinetic_block(atoms: Arr, shells: Arr, prims: Arr, cn: Arr, sa, sb) {
    return os_st_block(atoms, shells, prims, cn, sa, sb, 1)
}

#| Nuclear attraction block of shells `sa`, `sb`: `-sum_C Z_C <a|1/r_C|b>`
#| over every row of `atoms`. The vertical recurrence builds `Theta(e)^(m)`
#| for every `e` up to `la + lb` on A; primitives and nuclei are summed into
#| `acc[e]` for `la <= l_e <= la + lb`, and the horizontal step then moves
#| `lb` units of power to B.
def os_nuclear_block(atoms: Arr, shells: Arr, prims: Arr, cn: Arr, sa, sb) {
    la = shells[sa, SH_L]
    lb = shells[sb, SH_L]
    L = la + lb
    CT = os_ctab(L)
    ne = os_ncum(L + 1)
    A = atoms[shells[sa, SH_ATOM], AT_X:AT_Z + 1]
    B = atoms[shells[sb, SH_ATOM], AT_X:AT_Z + 1]
    AB2 = sum((A - B) * (A - B))
    pa0 = shells[sa, SH_P0]
    pb0 = shells[sb, SH_P0]
    TH = zeros(ne, L + 1)
    acc = zeros(ne)
    for u to shells[sa, SH_NPRIM] {
        a = prims[pa0 + u, PR_EXP]
        for v to shells[sb, SH_NPRIM] {
            b = prims[pb0 + v, PR_EXP]
            p = a + b
            h = 0.5 / p
            P = (a * A + b * B) / p
            PA = P - A
            pref = 2.0 * OS_PI / p * exp(-a * b / p * AB2) * cn[pa0 + u] * cn[pb0 + v]
            for C to len(atoms) {
                Z = atoms[C, AT_CHG]
                if Z != 0.0 {
                    PC = P - atoms[C, AT_X:AT_Z + 1]
                    TH[0, 0:L + 1] = pref * boys_values(L, p * sum(PC * PC))
                    for q to ne - 1 {
                        e = q + 1
                        i = CT[e, OS_CT_PIV]
                        e1 = CT[e, OS_CT_DN + i]
                        e2 = CT[e1, OS_CT_DN + i]
                        ni = CT[e1, i]
                        for m to L + 1 - CT[e, OS_CT_L] {
                            w = PA[i] * TH[e1, m] - PC[i] * TH[e1, m + 1]
                            if ni > 0 {
                                w = w + h * ni * (TH[e2, m] - TH[e2, m + 1])
                            }
                            TH[e, m] = w
                        }
                    }
                    acc[os_ncum(la):ne] = acc[os_ncum(la):ne] - Z * TH[os_ncum(la):ne, 0]
                }
            }
        }
    }
    return os_hrr_pair(acc, CT, la, lb, A - B)
}

#| The horizontal step on a one-electron pair: from `acc[e]` (every `e` with
#| `la <= l_e <= la + lb`, on A) to the block `[ncart(la), ncart(lb)]`,
#| times the component factors. `(a, b+1_i) = (a+1_i, b) + AB_i (a, b)`.
def os_hrr_pair(acc: Arr, CT: Arr, la, lb, AB: Arr) {
    L = la + lb
    ne = os_ncum(L + 1)
    nbc = os_ncum(lb + 1)
    # H[e, f]: e on A, f on B (f indexed from 0 over all powers up to lb)
    H = zeros(ne, nbc)
    for q to ne - os_ncum(la) {
        e = os_ncum(la) + q
        H[e, 0] = acc[e]
    }
    for q to nbc - 1 {
        f = q + 1
        i = CT[f, OS_CT_PIV]
        f1 = CT[f, OS_CT_DN + i]
        lf = CT[f, OS_CT_L]
        for t to os_ncum(L - lf + 1) - os_ncum(la) {
            e = os_ncum(la) + t
            H[e, f] = H[CT[e, OS_CT_UP + i], f1] + AB[i] * H[e, f1]
        }
    }
    na = bas_ncart(la)
    nb = bas_ncart(lb)
    out = zeros(na, nb)
    for r to na {
        e = os_ncum(la) + r
        for c to nb {
            f = os_ncum(lb) + c
            out[r, c] = H[e, f] * bas_nc(CT[e, 0], CT[e, 1], CT[e, 2]) * bas_nc(CT[f, 0], CT[f, 1], CT[f, 2])
        }
    }
    return out
}

#| Fill a whole matrix from per-pair blocks: `kind` 0 overlap, 1 kinetic,
#| 2 nuclear attraction. Only `sa >= sb` is computed; the other triangle is
#| the transpose.
def os_matrix(atoms: Arr, shells: Arr, prims: Arr, kind) {
    cn = bas_contract(shells, prims)
    n = bas_nao_cart(shells)
    M = zeros(n, n)
    for sa to len(shells) {
        for sb to sa + 1 {
            if kind == 2 {
                blk = os_nuclear_block(atoms, shells, prims, cn, sa, sb)
            } else {
                blk = os_st_block(atoms, shells, prims, cn, sa, sb, kind)
            }
            a0 = shells[sa, SH_CART0]
            b0 = shells[sb, SH_CART0]
            for r to len(blk) {
                for c to bas_ncart(shells[sb, SH_L]) {
                    M[a0 + r, b0 + c] = blk[r, c]
                    M[b0 + c, a0 + r] = blk[r, c]
                }
            }
        }
    }
    return M
}

#| Stage member: the overlap matrix, `[nao_cart, nao_cart]`.
def os_overlap(atoms: Arr, shells: Arr, prims: Arr) -> Arr {
    return os_matrix(atoms, shells, prims, 0)
}

#| Stage member: the kinetic energy matrix.
def os_kinetic(atoms: Arr, shells: Arr, prims: Arr) -> Arr {
    return os_matrix(atoms, shells, prims, 1)
}

#| Stage member: the nuclear attraction matrix, over every atom's charge.
def os_nuclear(atoms: Arr, shells: Arr, prims: Arr) -> Arr {
    return os_matrix(atoms, shells, prims, 2)
}

# ---- tests ----------------------------------------------------------------------
# The reference modules used above exist for these tests. An import cannot sit
# inside a test block, so a program that uses this module also compiles their
# literal-returning defs; nothing calls them, and the measured cost of two of
# them (water_631gs_cart and water_tzvp) is 0.55 s of cc -O2.

test "the component table lists every power once, in PySCF's order within a shell" {
    # every shell up to l = 6 (i)
    os_lmax = 6
    os_ct = os_ctab(os_lmax)
    assert_near(len(os_ct), os_ncum(os_lmax + 1), 0.0)
    os_bad = 0
    for os_l to os_lmax + 1 {
        os_pw = bas_powers(os_l)
        for os_c to len(os_pw) {
            os_r = os_ncum(os_l) + os_c
            if os_ct[os_r, OS_CT_LX] != os_pw[os_c, 0] or os_ct[os_r, OS_CT_LY] != os_pw[os_c, 1] or os_ct[os_r, OS_CT_LZ] != os_pw[os_c, 2] {
                os_bad = os_bad + 1
            }
            # one step down and back up along the pivot returns to the same row
            if os_l > 0 {
                os_i = os_ct[os_r, OS_CT_PIV]
                os_dn = os_ct[os_r, OS_CT_DN + os_i]
                if os_ct[os_dn, OS_CT_UP + os_i] != os_r {
                    os_bad = os_bad + 1
                }
            }
        }
    }
    assert_near(os_bad, 0.0, 0.0)
}

# SPDX-License-Identifier: MIT OR Apache-2.0
# SPDX-FileCopyrightText: 2026 The Viv authors