xnn.gnn.models.dimenet.directed_triplets#

xnn.gnn.models.dimenet.directed_triplets(edge_index, edge_vec, num_nodes)[source]#

Enumerate the message pairs (k -> j, j -> i) of directional message passing.

For every directed edge j -> i (the message m_ji) all edges k -> j arriving at its source are listed, except the reverse of the edge itself (k = i in a molecule; in a periodic system only the same image of i, found by the exact antisymmetry of the two edge vectors, so other images of i count as neighbors k). Vectorized: the edges are grouped by their destination node and every edge j -> i is paired with the whole group of j.

Parameters:

edge_index (Tensor) – Edge index of shape (2, E); row 0 is the source j and row 1 the destination i of each edge.
edge_vec (Tensor) – Edge vectors x_i - x_j (cell shifts included), shape (E, 3).
num_nodes (int) – Number of atoms.

Returns:

(kj, ji), long tensors of shape (T,) indexing the edge dimension: kj[t] is the edge k -> j whose message feeds the update of the edge ji[t] (j -> i). Empty when there is no triplet. TorchScript-compatible.

Return type:

tuple of Tensor