xnn.common.models.ops.build_triplets#

xnn.common.models.ops.build_triplets(edge_index, num_nodes)[source]#

Enumerate neighbour pairs (j, k) sharing a centre i.

For every centre atom, all unordered pairs of its incoming edges are returned. Fully vectorised (no Python loop over atoms): edges are grouped by their destination (centre) node and, within each group of c edges, the c*(c-1)/2 pairs are enumerated with a shared lower-triangular index template – the same scheme torchani uses. Used by the angular descriptors of the dnn family (ANI / HDNNP) and by the valence-angle enumeration of the ffnn family (ReaxFF).

Parameters:

edge_index (Tensor) – Edge index of shape (2, E); row 0 is the source (neighbour) and row 1 the destination (centre) node of each edge.
num_nodes (int) – Number of atoms (nodes) in the graph.

Returns:

(edge_jk_first, edge_jk_second, center), each of shape (T,) where T is the number of triplets. The first two index into the edge dimension (the two edges forming a triplet, first < second in the per-centre ordering) and center is the shared centre node. All three are empty long tensors when no triplet exists.

Return type:

tuple[Tensor, Tensor, Tensor]