xnn.common.models.dispersion.three_body_energy#

xnn.common.models.dispersion.three_body_energy(z, edge_index, edge_vec, r, c6_mat, r0_table, s9, alp3, cutoff, width, n_atoms)[source]#

Per-atom Axilrod-Teller-Muto three-body energy with zero damping.

E = s9 sum_ABC C9 (3 cos cos cos + 1) / (R_AB R_BC R_CA)^3 / (1 + 6 (R_0 / R)^alp3) with C9 = sqrt(C6_AB C6_BC C6_CA), R_0 the geometric product of the three pair critical radii and R that of the three distances. Triples are enumerated as pairs of edges sharing a center; each geometric triangle is visited once per corner (three times for distinct atoms, and correspondingly for the self-image triangles of a periodic cell), so a third of the triangle energy is assigned to the center at every visit – which reproduces the atom-resolved bookkeeping of the reference codes exactly. The triplet arithmetic runs in chunks so memory stays bounded for long cutoffs.

Parameters:

z (Tensor) – Atomic numbers, shape (N,).
edge_index (Tensor) – Edges [src, dst] within the three-body cutoff, shape (2, E).
edge_vec (Tensor) – Edge vectors in bohr, shape (E, 3).
r (Tensor) – Edge lengths in bohr, shape (E,).
c6_mat (Tensor) – Dense pair C6 matrix, shape (N, N).
r0_table (Tensor) – Pair critical radii by element, shape (Z_max + 1, Z_max + 1), bohr.
s9 (Tensor) – Three-body scaling (scalar tensor, may be learnable).
alp3 (float) – Damping exponent (alp / 3 in D4, (alp + 2) / 3 in D3).
cutoff (float) – Three-body cutoff and switching width, bohr.
width (float) – Three-body cutoff and switching width, bohr.
n_atoms (int) – Number of atoms.

Returns:

Per-atom three-body energies in hartree, shape (N,).

Return type:

Tensor