xnn.common.models.dispersion.three_body_energy_chunked#

xnn.common.models.dispersion.three_body_energy_chunked(z, edge_index, edge_vec, r, r0_table, s9, alp3, cutoff, width, n_atoms, c6_mat=None, alpha_a=None, alpha_b=None, chunk=None, r0_atom=None, a1=None, a2=None, c6_edge=None, edge_shift=None, once=True, analytic=True, triplet_cache=None)[source]#

three_body_energy() with memory bounded by one block (eager only).

The edges are sorted by center once, the centers are cut into blocks holding about chunk triplets each (read off the cumulative pair counts; None sizes the block from the free device memory, between 2^20 and 2^24, since every block pays a fixed set-up in both passes), and every block – triplet enumeration included – runs as a recompute() block, so the retained state is the neighbor list and the per-atom sums at every derivative order (force training included); the block is recomputed in the backward passes. Pair C6 values come from a dense matrix c6_mat (N, N) or, without one, from the polarizability factors alpha_a, alpha_b (N, K) as sum_k alpha_a[i, k] alpha_b[j, k] (D4: the Casimir-Polder weights folded into alpha_a), which removes the (N, N) table altogether. Faster still, c6_edge (E,) gives the pair C6 per directed edge (computed once, differentiable in whatever it depends on) together with the edges’ integer cell shifts edge_shift (E, 3): the two sides through the center are then scalar gathers and the third side is found by a binary search over packed pair keys, which removes the (T, 23) gathers and, above all, their contended index-accumulate backward (energy + forces 53 -> 15 ns per triplet on an A100). The pair critical radii come from the element table r0_table or, for BJ radii a1 sqrt(3 Q_A Q_B) + a2, from the per-atom factors r0_atom = 3^(1/4) sqrt(Q) with the scalars a1, a2 (cheap gradients when those are trainable). With once (default) each triangle of distinct atoms is evaluated from its smallest-index corner only. With analytic (default) and the per-edge C6 / per-atom radius inputs, the block’s first derivative is evaluated in closed form (_triplet_block_grad()) in one pass instead of re-running the block under autograd. On that path triplet_cache keeps each block’s triples from the forward to the backward pass (_TripletCache; GB, None for a quarter of the free device memory, 0 to disable), so they are enumerated once per step. Results equal the scripted loop to rounding.

Parameters:
Return type:

Tensor

On this page
Show Source