Coulomb backends: exact, density-fitted, streamed, screened

The two-electron Coulomb (and exact-exchange) term is the engine's cost center. dftax offers a ladder of backends with the same energy but different memory/compute trade-offs. Each backend is a value passed to the builder as coulomb=; invalid combinations raise at the factory.

Density fitting (RI, the default)

from dftax import KS, Molecule, df, scf
from dftax.energy.xc import PBE

water = Molecule.from_xyz("O 0 0 0; H 0.7586 0 0.5043; H 0.7586 0 -0.5043", "sto-3g")
scf(KS(water, PBE()))                          # same as coulomb=df()
scf(KS(water, PBE(), coulomb=df()))            # def2-universal-jkfit, chunk="auto"

Resolution-of-identity reduces the 4-center integral to 2- and 3-center pieces against an auxiliary basis (functions up to i, so transition metals work): O(N³) compute, O(N²·N_aux) memory for the materialized 3-center tensor, and the chunk="auto" policy switches to the streamed contraction when that tensor would exceed a device-aware memory budget (a fraction of the GPU pool; a fixed fallback on CPU). The RI error against the exact path is sub-mHa with the JK-fitting set.

On the materialized paths (single-device and mesh-sharded alike) the auxiliary basis is built in spherical harmonics (fewer functions, a tighter fit, and better cross-backend force reproducibility); only the streamed backend keeps the cartesian span, since it contracts cartesian auxiliary elements on the fly. df(spherical=False) pins a common fit space when comparing against the streamed backend, and df(screen=...) without a chunk runs a shell-pair Schwarz compact gather on the materialized build (screened pairs stay exactly zero), so extended systems build O(N) rather than O(N²) shell pairs.

Exact 4-center ERIs

from dftax import exact

scf(KS(water, PBE(), coulomb=exact()))         # builds the (nao⁴) ERI tensor

Exact (μν|λσ) via McMurchie-Davidson, using the 8-fold permutational symmetry. O(N⁴) memory: small systems and reference comparisons (it is also the fallback when building from a raw System, which has no element symbols to resolve an auxiliary basis name). Two variants trade compute for memory:

exact(screen=1e-10)   # Cauchy-Schwarz-screened materialized ERI tensor
exact(stream=True)    # J/K contracted on the fly, O(N²) memory, no ERI tensor

Streaming + screening (large systems)

For memory-light runs, stream the RI-J contraction over auxiliary chunks (never forming the nao²×naux tensor) and Schwarz-screen negligible bra pairs; the XC integral streams over grid chunks the same way via the grid spec:

from dftax import KS, becke, df, scf

ks = KS(water, PBE(),
        coulomb=df("def2-universal-jkfit", chunk=64, screen=1e-10),
        grid=becke(chunk=20_000))
scf(ks).e_tot

df(..., chunk=...) sets the auxiliary streaming chunk; screen= drops bra pairs whose self-Schwarz factor is below the relative cutoff (O(N) survivors for extended systems; it requires chunk). becke(chunk=...) (or points(coords, weights, chunk=...) for an explicit grid) streams the XC integral over grid chunks of that many points, O(chunk·nao) memory. Hybrids stream RI-K the same way.

Hybrids and exact exchange

PBE0 / B3LYP add a fraction of exact exchange K. On the exact path it is contracted on the fly (exchange_k_4c); under density fitting it uses streamed RI-K. The L ≥ 2 exact path is compile- and memory-bounded by scanning the bra primitive pair in the ERI kernel (so cc-pVDZ compiles in seconds rather than OOMing).

Range-separated hybrids (CAMB3LYP, WB97X, WB97XV) split the exchange into short- and long-range parts via erf(ω·r₁₂)/r₁₂ integrals and run on every DF backend: the materialized path holds the attenuated tensors, the streamed path recomputes attenuated 3-center elements on the fly (only the small attenuated metric is stored), and the mesh-sharded path runs the long-range exchange as slab-wise all-to-all rounds. exact(stream=True) is the one backend without an attenuated variant. WB97XV additionally carries VV10 nonlocal correlation, which needs the materialized XC grid (no becke(chunk=...) or mesh=).

Choosing a backend

System size	Backend
anything, by default	df() (auto materialize/stream)
small / validation	exact()
small, memory-tight	exact(stream=True) or exact(screen=1e-10)
large, hand-tuned	df(chunk=64, screen=1e-10) + becke(chunk=20_000)

Multi-GPU: mesh=

KS(..., mesh=mesh()) shards the calculation across a 1-D device mesh (all local devices by default; pass mesh(devices=[...]) to choose). The XC quadrature is sharded over grid points, and the density-fitted backend builds and holds its 3-center tensor in per-device aux slabs (no device ever materializes more than naux/ndev of it), with hybrid exact exchange computed slab-wise. The dense nao² matrices (S, hcore, P, Fock) stay replicated, and every collective differentiates, so scf, minimize, and property workflows run unchanged. A one-device mesh is a no-op.

The streamed df(chunk=...) backend shards too, and it is the one that reaches protein scale: the materialized slabs still cost nao²·naux/ndev per device (15 TiB for insulin at triple zeta), while the streamed path holds no tensor at all and each device takes its own contiguous slice of the auxiliary range. RI-J, RI-K and range-separated exchange all shard, dense or Schwarz-screened, so the combination covers the whole functional range with one exception: WB97XV's VV10 nonlocal correlation, whose double-grid pair quadrature is nonlocal across shards and cannot be evaluated shard by shard at all. That one raises at build time.

Across nodes

The same mesh spans several nodes once the processes have joined a group. JAX is multi-controller, so the launcher starts the same program once per node and distributed() does the rest:

from dftax import KS, df, distributed, mesh, scf

distributed()                    # reads the SLURM environment; no-op on one node
ks = KS(mol, PBE0(), coulomb=df("def2-universal-jkfit"), mesh=mesh())
res = scf(ks)                    # every process runs this and gets the same answer

Call it before any other JAX work (importing dftax starts no backend, so import order does not matter). Each process builds only the auxiliary slabs its own GPUs hold, which is what makes the capacity scale with the whole allocation rather than with one node. Every process must run the same sequence of collectives, so use is_coordinator() to guard printing and file writes, never to skip a build; barrier() synchronizes the group where they stop being symmetric. scripts/gpu/distributed.sbatch is a working two-node template, and scripts/gpu/validate_distributed.py checks a distributed solve against the single-device one in the same run.