Notebooks · lib/hf/md_direct

Direct J and K, recomputing the integrals instead of storing them

The Fock matrix needs the Coulomb and exchange matrices of a density D:

J[i, j] = sum_kl (ij|kl) D[k, l]
K[i, j] = sum_kl (ik|jl) D[k, l]

With the integrals stored in packed-8 form this is a loop over the vector. The vector grows as n^4 / 8: 18.4 MB at 65 AOs and 172 MB at 114. Direct J/K stores nothing. Every SCF iteration recomputes each shell quartet with mderi_block and adds it into J and K straight away, then throws it away.

One unique integral, eight contributions. mderi_block gives each integral once, for i >= j, k >= l, ij >= kl. Its 8 index orders all carry the same value, and each order adds to J and K in its own place:

order (p q | r s)   adds  g D[r, s] to J[p, q]   and  g D[q, s] to K[p, r]

When indices coincide (i = j, k = l, or ij = kl), some of the 8 orders are the same order counted twice, so g is halved once for each coincidence before the 8 contributions are added. The result is exact for any D, symmetric or not.

Parallel, with the same answer on any thread count. The quartets come from the same flat list as mderi_build, and they are dealt round-robin into MDJK_NCHUNK = 32 chunks. A parallel for runs over the chunks. Each chunk adds into its own copy of J and K, and after the loop the copies are summed in chunk order. A chunk handles the same quartets in the same order whichever thread runs it, so the default build, the OpenMP build and any OMP_NUM_THREADS give the same bits.

The signature is the frozen J/K one, (X1, X2, X3, D, opts) -> [2 ndens, n, n] with atoms, shells and prims in X1 .. X3, so this member can sit in the same registry as the in-core ones. opts[OPT_SCREEN_TOL] is the Schwarz threshold. Everything is over unit-normalized Cartesian AOs.

Density screening (mdjk_direct_screened, the production member). A quartet adds g D terms to J and K, with |g| <= Q_AB Q_CD, so a quartet whose bound times the largest density element it meets is small adds nothing that matters. With Dm[A, B] the largest |D| over the functions of shells A and B (and every density), the quartet (AB|CD) is skipped when

Q_AB Q_CD max(4 Dm[A,B], 4 Dm[C,D], Dm[A,C], Dm[A,D], Dm[B,C], Dm[B,D])

is at or below opts[OPT_DIRECT_TOL], 10⁻¹³: the rule and the threshold of PySCF's direct SCF (CVHFnrs8_prescreen with direct_scf_tol = 1e-13; the 4 is J's weight). On a full density this drops few quartets. On the change of density between two SCF iterations, which the incremental Fock build of lib/hf/scf_loop hands it, it drops most of them once the SCF is close to converged. The Schwarz factors Q are a property of the basis: a caller that builds J and K many times computes them once (mderi_schwarz) and passes them (mdjk_screened_q). The quartets are dealt into the chunks by bra pair, so no quartet list is built per call; mdjk_direct, which deals the Schwarz list round-robin and screens by nothing else, checks it.

Used with import hf.md_direct.

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 61–383

from hf.opts import OPT_DIRECT_TOL, OPT_SCREEN_TOL
from hf.md_eri import (
    mderi_block, MDERI_D_GA, MDERI_D_GB, MDERI_D_NA, MDERI_D_NB, MDERI_D_SA, MDERI_D_SB,
    MDERI_DIR_LEN, mderi_grid, mderi_groups, mderi_qlist, mderi_schwarz, mderi_store,
    mderi_store_seg, MDERI_TOP,
)
from hf.basis import bas_nao_cart, SH_CART0, SH_NCART
from hf.eri_store import es_get, es_rng

MDJK_NCHUNK = 32

#| Add block `B` of store pairs `(k1, k2)` into the J and K copies of one
#| chunk, `W[w0 + 2d]` (J) and `W[w0 + 2d + 1]` (K) for each density `d`.
#| Rows and columns of `B` are numbered as lib/hf/md_eri's packs number
#| them, a shell pair of the two groups at a time.
#@ 69e0ba72
def mdjk_put(W: Arr, w0, B: Arr, D: Arr, nd, ST: Arr, shells: Arr, k1, k2) {
    d1 = MDERI_TOP + MDERI_DIR_LEN * k1
    d2 = MDERI_TOP + MDERI_DIR_LEN * k2
    sa = ST[d1 + MDERI_D_SA]
    sb = ST[d1 + MDERI_D_SB]
    sc = ST[d2 + MDERI_D_SA]
    sd = ST[d2 + MDERI_D_SB]
    gna = ST[d1 + MDERI_D_NA]
    gnb = ST[d1 + MDERI_D_NB]
    gnc = ST[d2 + MDERI_D_NA]
    gnd = ST[d2 + MDERI_D_NB]
    na = shells[sa, SH_NCART]
    nb = shells[sb, SH_NCART]
    nc = shells[sc, SH_NCART]
    ndd = shells[sd, SH_NCART]
    for ga to gna {
        for ia to na {
            i = shells[sa + ga, SH_CART0] + ia
            for gb to gnb {
                for ib to nb {
                    j = shells[sb + gb, SH_CART0] + ib
                    if sa == sb and j > i {
                        continue
                    }
                    ij = i * (i + 1) / 2 + j
                    row = ((ga * gnb + gb) * na + ia) * nb + ib
                    for gc to gnc {
                        for ic to nc {
                            k = shells[sc + gc, SH_CART0] + ic
                            for gd to gnd {
                                for id to ndd {
                                    l = shells[sd + gd, SH_CART0] + id
                                    if sc == sd and l > k {
                                        continue
                                    }
                                    kl = k * (k + 1) / 2 + l
                                    if k1 == k2 and kl > ij {
                                        continue
                                    }
                                    #: halve once per equal index pair, so the eight orders count each integral once
                                    f = B[row, ((gc * gnd + gd) * nc + ic) * ndd + id]
                                    if i == j {
                                        f = 0.5 * f
                                    }
                                    if k == l {
                                        f = 0.5 * f
                                    }
                                    if ij == kl {
                                        f = 0.5 * f
                                    }
                                    for d to nd {
                                        wj = w0 + 2 * d
                                        wk = wj + 1
                                        dkl = f * (D[d, k, l] + D[d, l, k])
                                        dij = f * (D[d, i, j] + D[d, j, i])
                                        W[wj, i, j] = W[wj, i, j] + dkl
                                        W[wj, j, i] = W[wj, j, i] + dkl
                                        W[wj, k, l] = W[wj, k, l] + dij
                                        W[wj, l, k] = W[wj, l, k] + dij
                                        W[wk, i, k] = W[wk, i, k] + f * D[d, j, l]
                                        W[wk, j, k] = W[wk, j, k] + f * D[d, i, l]
                                        W[wk, i, l] = W[wk, i, l] + f * D[d, j, k]
                                        W[wk, j, l] = W[wk, j, l] + f * D[d, i, k]
                                        W[wk, k, i] = W[wk, k, i] + f * D[d, l, j]
                                        W[wk, l, i] = W[wk, l, i] + f * D[d, k, j]
                                        W[wk, k, j] = W[wk, k, j] + f * D[d, l, i]
                                        W[wk, l, j] = W[wk, l, j] + f * D[d, k, i]
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    return 0
}

#| One quartet of the list into chunk `c`'s copies; a def of its own so the
#| block is freed as soon as it has been added.
#@ 005d2ee4
def mdjk_one(W: Arr, D: Arr, nd, ST: Arr, G: Arr, shells: Arr, QL: Arr, q, c) {
    k1 = QL[q, 0]
    k2 = QL[q, 1]
    mdjk_put(W, 2 * nd * c, mderi_block(ST, k1, k2, G), D, nd, ST, shells, k1, k2)
    return 0
}

#| Chunk `c` of `nch`: quartets `c, c + nch, c + 2 nch, ...` of the list,
#| added into `W`, the chunk's own copies of J and K.
#@ 39acc072
def mdjk_chunk(W: Arr, D: Arr, nd, ST: Arr, G: Arr, shells: Arr, QL: Arr, nq, c, nch) {
    q = c
    while q < nq {
        mdjk_one(W, D, nd, ST, G, shells, QL, q, 0)
        q = q + nch
    }
    return 0
}

#| J and K with the quartets dealt into `nch` chunks. The result does not
#| depend on the thread count; it depends on `nch` only through the order of
#| the final sums.
#@ 53fe5f5d
def mdjk_direct_nch(atoms: Arr, shells: Arr, prims: Arr, D: Arr, opts: Arr, nch) {
    n = bas_nao_cart(shells)
    nd = len(D)
    ST = mderi_store_seg(atoms, shells, prims)
    G = mderi_grid(ST)
    npl = ST[0]
    tol = opts[OPT_SCREEN_TOL]
    Q = mderi_schwarz(ST, G)
    QL = mderi_qlist(Q, npl, tol)
    nq = len(QL)
    # pass c fills a copy of its own and files it as slab c of W, so the
    # passes write disjoint slabs
    W = zeros(nch, 2 * nd, n, n)
    parallel for c to nch {
        Wc = zeros(2 * nd, n, n)
        mdjk_chunk(Wc, D, nd, ST, G, shells, QL, nq, c, nch)
        setslab(W, c, Wc)
    }
    #: summed in chunk order, so the result does not depend on threads
    out = zeros(2 * nd, n, n)
    for c to nch {
        Wc = slab(W, c)
        for r to 2 * nd {
            setslab(out, r, slab(out, r) + slab(Wc, r))
        }
    }
    return out
}

#| The J/K registry member: `X1 .. X3` = atoms, shells, prims; `D` is
#| `[ndens, n, n]`; the result is `[2 ndens, n, n]`, J then K per density.
#@ d52d5288
def mdjk_direct(X1: Arr, X2: Arr, X3: Arr, D: Arr, opts: Arr) -> Arr {
    return mdjk_direct_nch(X1, X2, X3, D, opts, MDJK_NCHUNK)
}

#| The largest `|D|` over the functions of each pair of shell groups
#| (`GT`, lib/hf/md_eri's `mderi_groups`) and every density, `[ng, ng]`,
#| symmetric. A group's functions are consecutive.
#@ 6ff6be83
def mdjk_dmax(D: Arr, shells: Arr, GT: Arr) {
    ng = len(GT)
    A = abs(slab(D, 0))
    for d from 1 to len(D) {
        A = max(A, abs(slab(D, d)))
    }
    lo = zeros(ng)
    hi = zeros(ng)
    for g to ng {
        lo[g] = shells[GT[g, 0], SH_CART0]
        hi[g] = lo[g] + GT[g, 1] * shells[GT[g, 0], SH_NCART]
    }
    Dm = zeros(ng, ng)
    for a to ng {
        for b to a + 1 {
            v = max(max(A[lo[a]:hi[a], lo[b]:hi[b]]), max(A[lo[b]:hi[b], lo[a]:hi[a]]))
            Dm[a, b] = v
            Dm[b, a] = v
        }
    }
    return Dm
}

#| One quartet computed and added into `W`; a def of its own so the block is
#| freed as soon as it has been added.
#@ 27d38d0b
def mdjk_pair(W: Arr, D: Arr, nd, ST: Arr, G: Arr, shells: Arr, k1, k2) {
    mdjk_put(W, 0, mderi_block(ST, k1, k2, G), D, nd, ST, shells, k1, k2)
    return 0
}

#| Chunk `c` of `nch` of the screened build: bra pairs `c, c + nch, ...`,
#| each against every ket pair up to it, a quartet kept when it passes the
#| Schwarz test (`Q_AB Q_CD >= tol_s`) and the density test (header), with
#| `Dm` over the store's shell groups.
#@ e3033706
def mdjk_chunk_screened(W: Arr, D: Arr, nd, ST: Arr, G: Arr, shells: Arr, Q: Arr, Dm: Arr, tol_s, tol_d, c, nch) {
    #: bounds for the early bra test
    npl = ST[0]
    qmax = max(Q)
    dmax = max(Dm)
    #: skip a bra pair no ket can lift over the density test
    k1 = c
    while k1 < npl {
        ga = ST[MDERI_TOP + MDERI_DIR_LEN * k1 + MDERI_D_GA]
        gb = ST[MDERI_TOP + MDERI_DIR_LEN * k1 + MDERI_D_GB]
        q1 = Q[k1]
        # no ket pair can lift this bra pair over the density test
        if q1 * qmax * 4.0 * dmax > tol_d {
            dab = 4.0 * Dm[ga, gb]
            for k2 to k1 + 1 {
                q = q1 * Q[k2]
                if q >= tol_s {
                    gc = ST[MDERI_TOP + MDERI_DIR_LEN * k2 + MDERI_D_GA]
                    gd = ST[MDERI_TOP + MDERI_DIR_LEN * k2 + MDERI_D_GB]
                    dm = max(max(dab, 4.0 * Dm[gc, gd]), max(max(Dm[ga, gc], Dm[ga, gd]), max(Dm[gb, gc], Dm[gb, gd])))
                    if q * dm > tol_d {
                        mdjk_pair(W, D, nd, ST, G, shells, k1, k2)
                    }
                }
            }
        }
        k1 = k1 + nch
    }
    return 0
}

#| J and K of `D` with Schwarz and density screening, from the Schwarz
#| factors `Q` of the store's pairs (`mderi_schwarz`); `[2 ndens, n, n]`.
#@ 58c7efa4
def mdjk_screened_q(atoms: Arr, shells: Arr, prims: Arr, D: Arr, opts: Arr, Q: Arr) -> Arr {
    n = bas_nao_cart(shells)
    nd = len(D)
    nch = MDJK_NCHUNK
    ST = mderi_store(atoms, shells, prims)
    G = mderi_grid(ST)
    Dm = mdjk_dmax(D, shells, mderi_groups(shells, prims))
    tol_s = opts[OPT_SCREEN_TOL]
    tol_d = opts[OPT_DIRECT_TOL]
    #: one J/K copy per chunk: parallel without shared writes
    W = zeros(nch, 2 * nd, n, n)
    parallel for c to nch {
        Wc = zeros(2 * nd, n, n)
        mdjk_chunk_screened(Wc, D, nd, ST, G, shells, Q, Dm, tol_s, tol_d, c, nch)
        setslab(W, c, Wc)
    }
    #: summed in chunk order, so the result does not depend on threads
    out = zeros(2 * nd, n, n)
    for c to nch {
        out = out + slab(W, c)
    }
    return out
}

#| The density-screened J/K registry member: `X1 .. X3` = atoms, shells,
#| prims. It computes the Schwarz factors itself; the SCF's direct member
#| (lib/hf/jk_direct) takes them from the basis blob instead.
#@ 50faef85
def mdjk_direct_screened(X1: Arr, X2: Arr, X3: Arr, D: Arr, opts: Arr) -> Arr {
    ST = mderi_store(X1, X2, X3)
    return mdjk_screened_q(X1, X2, X3, D, opts, mderi_schwarz(ST, mderi_grid(ST)))
}

#| Test helper: `nd` seeded symmetric densities `[nd, n, n]` with entries in
#| (-1, 1), from the Park-Miller generator of `lib/hf/eri_store`.
def mdjk_rand_dens(nd, n, seed) {
    st = zeros(1)
    st[0] = seed
    D = zeros(nd, n, n)
    for d to nd {
        for i to n {
            for j to i + 1 {
                v = 2.0 * es_rng(st) - 1.0
                D[d, i, j] = v
                D[d, j, i] = v
            }
        }
    }
    return D
}

#| Test helper: J and K of `D` by dense contraction of the packed vector,
#| as two `[n^2, n^2]` matrix-vector products: `J = Gj @ vec(D)` with
#| `Gj[(i,j),(k,l)] = (ij|kl)`, and `K = Gk @ vec(D)` with
#| `Gk[(i,j),(k,l)] = (ik|jl)`. The largest gap to `JK` over every entry is
#| returned.
def mdjk_dense_gap(Gp: Arr, JK: Arr, D: Arr, n) {
    #: (ij|kl) for J, (ik|jl) for K, as dense matrices
    Gj = zeros(n * n, n * n)
    Gk = zeros(n * n, n * n)
    for i to n {
        for j to n {
            for k to n {
                for l to n {
                    Gj[i * n + j, k * n + l] = es_get(Gp, i, j, k, l)
                    Gk[i * n + j, k * n + l] = es_get(Gp, i, k, j, l)
                }
            }
        }
    }
    w = 0.0
    for d to len(D) {
        v = zeros(n * n)
        for k to n {
            for l to n {
                v[k * n + l] = D[d, k, l]
            }
        }
        J = Gj @ v
        K = Gk @ v
        for i to n {
            for j to n {
                w = max(w, abs(JK[2 * d, i, j] - J[i * n + j]))
                w = max(w, abs(JK[2 * d + 1, i, j] - K[i * n + j]))
            }
        }
    }
    return w
}

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

from hf.opts import OPT_DIRECT_TOL, OPT_SCREEN_TOL
from hf.md_eri import (
    mderi_block, MDERI_D_GA, MDERI_D_GB, MDERI_D_NA, MDERI_D_NB, MDERI_D_SA, MDERI_D_SB,
    MDERI_DIR_LEN, mderi_grid, mderi_groups, mderi_qlist, mderi_schwarz, mderi_store,
    mderi_store_seg, MDERI_TOP,
)
from hf.basis import bas_nao_cart, SH_CART0, SH_NCART
from hf.eri_store import es_get, es_rng

MDJK_NCHUNK = 32

#| Add block `B` of store pairs `(k1, k2)` into the J and K copies of one
#| chunk, `W[w0 + 2d]` (J) and `W[w0 + 2d + 1]` (K) for each density `d`.
#| Rows and columns of `B` are numbered as lib/hf/md_eri's packs number
#| them, a shell pair of the two groups at a time.
#@ 69e0ba72
def mdjk_put(W: Arr, w0, B: Arr, D: Arr, nd, ST: Arr, shells: Arr, k1, k2) {
    d1 = MDERI_TOP + MDERI_DIR_LEN * k1
    d2 = MDERI_TOP + MDERI_DIR_LEN * k2
    sa = ST[d1 + MDERI_D_SA]
    sb = ST[d1 + MDERI_D_SB]
    sc = ST[d2 + MDERI_D_SA]
    sd = ST[d2 + MDERI_D_SB]
    gna = ST[d1 + MDERI_D_NA]
    gnb = ST[d1 + MDERI_D_NB]
    gnc = ST[d2 + MDERI_D_NA]
    gnd = ST[d2 + MDERI_D_NB]
    na = shells[sa, SH_NCART]
    nb = shells[sb, SH_NCART]
    nc = shells[sc, SH_NCART]
    ndd = shells[sd, SH_NCART]
    for ga to gna {
        for ia to na {
            i = shells[sa + ga, SH_CART0] + ia
            for gb to gnb {
                for ib to nb {
                    j = shells[sb + gb, SH_CART0] + ib
                    if sa == sb and j > i {
                        continue
                    }
                    ij = i * (i + 1) / 2 + j
                    row = ((ga * gnb + gb) * na + ia) * nb + ib
                    for gc to gnc {
                        for ic to nc {
                            k = shells[sc + gc, SH_CART0] + ic
                            for gd to gnd {
                                for id to ndd {
                                    l = shells[sd + gd, SH_CART0] + id
                                    if sc == sd and l > k {
                                        continue
                                    }
                                    kl = k * (k + 1) / 2 + l
                                    if k1 == k2 and kl > ij {
                                        continue
                                    }
                                    #: halve once per equal index pair, so the eight orders count each integral once
                                    f = B[row, ((gc * gnd + gd) * nc + ic) * ndd + id]
                                    if i == j {
                                        f = 0.5 * f
                                    }
                                    if k == l {
                                        f = 0.5 * f
                                    }
                                    if ij == kl {
                                        f = 0.5 * f
                                    }
                                    for d to nd {
                                        wj = w0 + 2 * d
                                        wk = wj + 1
                                        dkl = f * (D[d, k, l] + D[d, l, k])
                                        dij = f * (D[d, i, j] + D[d, j, i])
                                        W[wj, i, j] = W[wj, i, j] + dkl
                                        W[wj, j, i] = W[wj, j, i] + dkl
                                        W[wj, k, l] = W[wj, k, l] + dij
                                        W[wj, l, k] = W[wj, l, k] + dij
                                        W[wk, i, k] = W[wk, i, k] + f * D[d, j, l]
                                        W[wk, j, k] = W[wk, j, k] + f * D[d, i, l]
                                        W[wk, i, l] = W[wk, i, l] + f * D[d, j, k]
                                        W[wk, j, l] = W[wk, j, l] + f * D[d, i, k]
                                        W[wk, k, i] = W[wk, k, i] + f * D[d, l, j]
                                        W[wk, l, i] = W[wk, l, i] + f * D[d, k, j]
                                        W[wk, k, j] = W[wk, k, j] + f * D[d, l, i]
                                        W[wk, l, j] = W[wk, l, j] + f * D[d, k, i]
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    return 0
}

#| One quartet of the list into chunk `c`'s copies; a def of its own so the
#| block is freed as soon as it has been added.
#@ 005d2ee4
def mdjk_one(W: Arr, D: Arr, nd, ST: Arr, G: Arr, shells: Arr, QL: Arr, q, c) {
    k1 = QL[q, 0]
    k2 = QL[q, 1]
    mdjk_put(W, 2 * nd * c, mderi_block(ST, k1, k2, G), D, nd, ST, shells, k1, k2)
    return 0
}

#| Chunk `c` of `nch`: quartets `c, c + nch, c + 2 nch, ...` of the list,
#| added into `W`, the chunk's own copies of J and K.
#@ 39acc072
def mdjk_chunk(W: Arr, D: Arr, nd, ST: Arr, G: Arr, shells: Arr, QL: Arr, nq, c, nch) {
    q = c
    while q < nq {
        mdjk_one(W, D, nd, ST, G, shells, QL, q, 0)
        q = q + nch
    }
    return 0
}

#| J and K with the quartets dealt into `nch` chunks. The result does not
#| depend on the thread count; it depends on `nch` only through the order of
#| the final sums.
#@ 53fe5f5d
def mdjk_direct_nch(atoms: Arr, shells: Arr, prims: Arr, D: Arr, opts: Arr, nch) {
    n = bas_nao_cart(shells)
    nd = len(D)
    ST = mderi_store_seg(atoms, shells, prims)
    G = mderi_grid(ST)
    npl = ST[0]
    tol = opts[OPT_SCREEN_TOL]
    Q = mderi_schwarz(ST, G)
    QL = mderi_qlist(Q, npl, tol)
    nq = len(QL)
    # pass c fills a copy of its own and files it as slab c of W, so the
    # passes write disjoint slabs
    W = zeros(nch, 2 * nd, n, n)
    parallel for c to nch {
        Wc = zeros(2 * nd, n, n)
        mdjk_chunk(Wc, D, nd, ST, G, shells, QL, nq, c, nch)
        setslab(W, c, Wc)
    }
    #: summed in chunk order, so the result does not depend on threads
    out = zeros(2 * nd, n, n)
    for c to nch {
        Wc = slab(W, c)
        for r to 2 * nd {
            setslab(out, r, slab(out, r) + slab(Wc, r))
        }
    }
    return out
}

#| The J/K registry member: `X1 .. X3` = atoms, shells, prims; `D` is
#| `[ndens, n, n]`; the result is `[2 ndens, n, n]`, J then K per density.
#@ d52d5288
def mdjk_direct(X1: Arr, X2: Arr, X3: Arr, D: Arr, opts: Arr) -> Arr {
    return mdjk_direct_nch(X1, X2, X3, D, opts, MDJK_NCHUNK)
}

#| The largest `|D|` over the functions of each pair of shell groups
#| (`GT`, lib/hf/md_eri's `mderi_groups`) and every density, `[ng, ng]`,
#| symmetric. A group's functions are consecutive.
#@ 6ff6be83
def mdjk_dmax(D: Arr, shells: Arr, GT: Arr) {
    ng = len(GT)
    A = abs(slab(D, 0))
    for d from 1 to len(D) {
        A = max(A, abs(slab(D, d)))
    }
    lo = zeros(ng)
    hi = zeros(ng)
    for g to ng {
        lo[g] = shells[GT[g, 0], SH_CART0]
        hi[g] = lo[g] + GT[g, 1] * shells[GT[g, 0], SH_NCART]
    }
    Dm = zeros(ng, ng)
    for a to ng {
        for b to a + 1 {
            v = max(max(A[lo[a]:hi[a], lo[b]:hi[b]]), max(A[lo[b]:hi[b], lo[a]:hi[a]]))
            Dm[a, b] = v
            Dm[b, a] = v
        }
    }
    return Dm
}

#| One quartet computed and added into `W`; a def of its own so the block is
#| freed as soon as it has been added.
#@ 27d38d0b
def mdjk_pair(W: Arr, D: Arr, nd, ST: Arr, G: Arr, shells: Arr, k1, k2) {
    mdjk_put(W, 0, mderi_block(ST, k1, k2, G), D, nd, ST, shells, k1, k2)
    return 0
}

#| Chunk `c` of `nch` of the screened build: bra pairs `c, c + nch, ...`,
#| each against every ket pair up to it, a quartet kept when it passes the
#| Schwarz test (`Q_AB Q_CD >= tol_s`) and the density test (header), with
#| `Dm` over the store's shell groups.
#@ e3033706
def mdjk_chunk_screened(W: Arr, D: Arr, nd, ST: Arr, G: Arr, shells: Arr, Q: Arr, Dm: Arr, tol_s, tol_d, c, nch) {
    #: bounds for the early bra test
    npl = ST[0]
    qmax = max(Q)
    dmax = max(Dm)
    #: skip a bra pair no ket can lift over the density test
    k1 = c
    while k1 < npl {
        ga = ST[MDERI_TOP + MDERI_DIR_LEN * k1 + MDERI_D_GA]
        gb = ST[MDERI_TOP + MDERI_DIR_LEN * k1 + MDERI_D_GB]
        q1 = Q[k1]
        # no ket pair can lift this bra pair over the density test
        if q1 * qmax * 4.0 * dmax > tol_d {
            dab = 4.0 * Dm[ga, gb]
            for k2 to k1 + 1 {
                q = q1 * Q[k2]
                if q >= tol_s {
                    gc = ST[MDERI_TOP + MDERI_DIR_LEN * k2 + MDERI_D_GA]
                    gd = ST[MDERI_TOP + MDERI_DIR_LEN * k2 + MDERI_D_GB]
                    dm = max(max(dab, 4.0 * Dm[gc, gd]), max(max(Dm[ga, gc], Dm[ga, gd]), max(Dm[gb, gc], Dm[gb, gd])))
                    if q * dm > tol_d {
                        mdjk_pair(W, D, nd, ST, G, shells, k1, k2)
                    }
                }
            }
        }
        k1 = k1 + nch
    }
    return 0
}

#| J and K of `D` with Schwarz and density screening, from the Schwarz
#| factors `Q` of the store's pairs (`mderi_schwarz`); `[2 ndens, n, n]`.
#@ 58c7efa4
def mdjk_screened_q(atoms: Arr, shells: Arr, prims: Arr, D: Arr, opts: Arr, Q: Arr) -> Arr {
    n = bas_nao_cart(shells)
    nd = len(D)
    nch = MDJK_NCHUNK
    ST = mderi_store(atoms, shells, prims)
    G = mderi_grid(ST)
    Dm = mdjk_dmax(D, shells, mderi_groups(shells, prims))
    tol_s = opts[OPT_SCREEN_TOL]
    tol_d = opts[OPT_DIRECT_TOL]
    #: one J/K copy per chunk: parallel without shared writes
    W = zeros(nch, 2 * nd, n, n)
    parallel for c to nch {
        Wc = zeros(2 * nd, n, n)
        mdjk_chunk_screened(Wc, D, nd, ST, G, shells, Q, Dm, tol_s, tol_d, c, nch)
        setslab(W, c, Wc)
    }
    #: summed in chunk order, so the result does not depend on threads
    out = zeros(2 * nd, n, n)
    for c to nch {
        out = out + slab(W, c)
    }
    return out
}

#| The density-screened J/K registry member: `X1 .. X3` = atoms, shells,
#| prims. It computes the Schwarz factors itself; the SCF's direct member
#| (lib/hf/jk_direct) takes them from the basis blob instead.
#@ 50faef85
def mdjk_direct_screened(X1: Arr, X2: Arr, X3: Arr, D: Arr, opts: Arr) -> Arr {
    ST = mderi_store(X1, X2, X3)
    return mdjk_screened_q(X1, X2, X3, D, opts, mderi_schwarz(ST, mderi_grid(ST)))
}

#| Test helper: `nd` seeded symmetric densities `[nd, n, n]` with entries in
#| (-1, 1), from the Park-Miller generator of `lib/hf/eri_store`.
def mdjk_rand_dens(nd, n, seed) {
    st = zeros(1)
    st[0] = seed
    D = zeros(nd, n, n)
    for d to nd {
        for i to n {
            for j to i + 1 {
                v = 2.0 * es_rng(st) - 1.0
                D[d, i, j] = v
                D[d, j, i] = v
            }
        }
    }
    return D
}

#| Test helper: J and K of `D` by dense contraction of the packed vector,
#| as two `[n^2, n^2]` matrix-vector products: `J = Gj @ vec(D)` with
#| `Gj[(i,j),(k,l)] = (ij|kl)`, and `K = Gk @ vec(D)` with
#| `Gk[(i,j),(k,l)] = (ik|jl)`. The largest gap to `JK` over every entry is
#| returned.
def mdjk_dense_gap(Gp: Arr, JK: Arr, D: Arr, n) {
    #: (ij|kl) for J, (ik|jl) for K, as dense matrices
    Gj = zeros(n * n, n * n)
    Gk = zeros(n * n, n * n)
    for i to n {
        for j to n {
            for k to n {
                for l to n {
                    Gj[i * n + j, k * n + l] = es_get(Gp, i, j, k, l)
                    Gk[i * n + j, k * n + l] = es_get(Gp, i, k, j, l)
                }
            }
        }
    }
    w = 0.0
    for d to len(D) {
        v = zeros(n * n)
        for k to n {
            for l to n {
                v[k * n + l] = D[d, k, l]
            }
        }
        J = Gj @ v
        K = Gk @ v
        for i to n {
            for j to n {
                w = max(w, abs(JK[2 * d, i, j] - J[i * n + j]))
                w = max(w, abs(JK[2 * d + 1, i, j] - K[i * n + j]))
            }
        }
    }
    return w
}

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