xnn.common.models.dispersion_fast#

Fast paths of the dispersion terms (DFT-D3, DFT-D4).

Profiled on periodic water boxes (5k-24k atoms, energy + forces + stress), the reference D4 spends its time in two places: the large-regime EEQ solve above about 10k atoms (85-90% of a step, nearly all of it recomputing structure factors in the conjugate-gradient loop) and the three-body ATM term below that and in molecular dynamics with the EEQ reuse (about 80% of a step); D3 spends it in its three-body term and in the dense (N, N) C6 matrix that term reads. The fast paths replace exactly those:

  • FastEEQSystem: the EEQ operator with factorized reciprocal phases (no trigonometry per product), a sparse real-space matrix, a fused molecular kernel and lockstep solves;

  • atm_energy(): the ATM energy and its gradient as fused Triton kernels, with the pair C6 from per-atom factors (no (N, N) matrix, so D3’s three-body term has no atom limit).

They follow the use_fast convention of xnn.common.models.fast; the parameters, buffers and outputs of the modules are unchanged, except that the fast D3 path does not form the dense "c6_matrix" output (it is returned empty, as above 20000 atoms; c6_matrix() computes it on request).

Module Attributes

AUTO_POLICY

minimum edges of the neighbor graph per GPU model.

Modules

atm

The Axilrod-Teller-Muto three-body term of D3 and D4 as fused Triton kernels.

coulomb

The molecular EEQ interaction erf(gamma_ij r_ij) / r_ij applied without storing it.

eeq

The large-regime EEQ operator, evaluated without trigonometry in the solver loop.