xnn.common.models.dispersion#

Shared machinery of the dispersion add-ons (DFT-D3, DFT-D4).

The two Grimme dispersion models share most of their structure: a real-space cutoff with an optional quintic switching window, Gaussian coordination-number weights over tabulated reference systems (D3 paper eq 16, D4 paper eq 8), an Axilrod-Teller-Muto three-body term with zero damping over the atom triples of a neighbor list (D3 paper eqs 11-14, D4 paper eqs 22-27), and the same way of attaching themselves to a short-range model. Those pieces live here, once; d3 and d4 hold what is specific to each model (the coordination-number counting function, the reference data, the damping functions, D4’s charge scaling and EEQ charges).

DispersionCorrection is the model-agnostic wrapper both add-ons derive from: standalone it is the dispersion energy, given a short-range model it adds the dispersion energy to that model’s prediction while handing the model only the edges within its own cutoff. The deploy channels (torchscript, lammps) peel the wrapper off through the same interface.

Functions

`auto_triplet_chunk`(like)	Triplets per recompute block from the free device memory.
`edge_cell_shifts`(pos, edge_index, edge_vec, ...)	Integer lattice shifts `S` of directed edges, `(E, 3)` int64.
`evaluate_on_graph`(term, data)	Run a dispersion term's `evaluate` on a (batched) `AtomicGraph`.
`gaussian_reference_weights`(cn, refcn, valid, wf)	Normalized Gaussian weights of the reference systems in the coordination number.
`options_from_extra`(extra, keys)	Pick the recognized keyword arguments out of a config `extra` dict.
`pair_edge_keys`(edge_index, shifts, n_atoms)	One int64 key per directed edge from `(dst, src, shift)`, unique per image.
`switching_function`(r, cutoff, width)	Quintic switching window: 1 below `cutoff - width`, 0 above `cutoff`.
`three_body_energy`(z, edge_index, edge_vec, ...)	Per-atom Axilrod-Teller-Muto three-body energy with zero damping.
`three_body_energy_chunked`(z, edge_index, ...)	`three_body_energy()` with memory bounded by one block (eager only).
`triplet_cache_budget`(like, budget)	Bytes the triple cache may hold: `budget` in GB, or `None` for a quarter of the free device memory (CPU: 1 GB); 0 disables the cache.
`triplet_energy`(a, b, c, j, k, edge_vec, r, ...)	ATM energy of one block of triplets, a third per center visit, `(T,)`.

Classes

DispersionCorrection(term[, model])

A dispersion term as an xnn potential, standalone or wrapped around a model.