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
minimum edges of the neighbor graph per GPU model. |
Modules