xnn.common.models.dispersion_fast.atm.atm_energy#
- xnn.common.models.dispersion_fast.atm.atm_energy(edge_index, edge_vec, feat_a, feat_b, consts, n_atoms, reference=None)[source]#
Per-atom ATM energies
(N,)in hartree, with the Triton kernels.- Parameters:
edge_index (Tensor) – Directed edges
[src, dst]within the three-body cutoff,(2, E), both directions of every pair.edge_vec (Tensor) – Their vectors in bohr,
(E, 3)(differentiable).feat_a (Tensor) – Pair-C6 factors
(N, K),C6_xy = <feat_a[x], feat_b[y]>(differentiable; may be the same tensor).feat_b (Tensor) – Pair-C6 factors
(N, K),C6_xy = <feat_a[x], feat_b[y]>(differentiable; may be the same tensor).consts (ATMConstants) – Species, the damping-radius table and the scalars.
n_atoms (int) – Number of atoms.
reference (callable, optional) –
reference(edge_vec, feat_a, feat_b) -> (N,): the same energy in differentiable torch ops, run by the backward pass when a second derivative is needed; without it the kernel’s (first-order) gradient is used throughout.
- Return type: