xnn.common.models.dispersion_fast.atm

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

The reference (three_body_energy_chunked()) enumerates the triplets of every block of centers into index tensors, gathers per-triplet inputs, evaluates the summand and accumulates the results, in a forward and a closed-form backward pass. Here one kernel does all of it in registers. Each program takes one center and a block of its edges (sorted by center, then neighbor) and walks the pairs of that block with the later edges of the center: a BLOCK x BLOCK tile of pairs at a time, with the geometry, the three pair C6, the damping, the switching functions and the energy (and, in the backward kernel, every derivative) evaluated per pair and accumulated without ever storing a triplet.

Bookkeeping, as in the reference:

• A triangle of three distinct atoms is evaluated once, from the corner with the smallest index (the edges of a center are sorted by neighbor, so these are the pairs among the edges to larger-index neighbors), and a third of its energy goes to each corner.

• A triangle with a repeated atom (periodic self-images: a center with its own image, or two images of one neighbor) is evaluated at every visit with a third to the center. These exist only when a cell is narrower than twice the cutoff; a second launch covers the centers that have them.

The pair C6 enters through per-atom factors: C6_xy = <f_x, g_y> with f, g (N, K). For D4 f = g = alpha_x(i w) sqrt(3 w / pi) over the 23 Casimir-Polder frequencies; for D3 f holds the reference-weighted C6 rows of each species and g the reference weights placed at the atom’s own species. No (N, N) table and no per-triplet gather of 23 values. The pair damping radii come from a small table over the species present.

Derivatives: the backward kernel returns the gradients with respect to the edge vectors and both factor matrices in closed form; they are not themselves differentiable. When a second derivative is needed (create_graph=True, e.g. force training) the backward pass runs the reference implementation instead, which is differentiable to second order.

Module Attributes

BLOCK

a program holds BLOCK x BLOCK pairs in registers

Functions

atm_energy(edge_index, edge_vec, feat_a, ...)	Per-atom ATM energies (N,) in hartree, with the Triton kernels.
plan_triplets(edge_index, n_atoms)	Sort the edges by (center, neighbor) and lay out the kernel launches.

Classes

ATMConstants(species, table, s9, alp3, ...)	The scalars and tables of one ATM term (species-indexed).
TripletPlan(order, src, dst, launches, n_atoms)	The center-sorted edge list of one evaluation and its kernel launches.