build_triplets

torch_sim.neighbors.nbody.build_triplets(edge_index, n_atoms)[source]

Build triplet interaction indices from an edge list.

For every pair of edges (b→a) and (c→a) that share the same target atom a with edge_ba ≠ edge_ca, produces a triplet b→a←c.

Uses only ops that are JIT/AOTInductor safe: argsort, bincount, repeat_interleave, and boolean indexing.

Parameters:

  • edge_index (Tensor) – [2, n_edges] tensor where edge_index[0] are sources and edge_index[1] are targets.

  • n_atoms (int) – Total number of atoms (used for bincount sizing).

Returns:

  • "trip_in" — edge indices of the incoming edge b→a, shape [n_triplets].

  • "trip_out" — edge indices of the outgoing edge c→a, shape [n_triplets].

  • "trip_out_agg" — per-segment local index for aggregation, shape [n_triplets].

  • "center_atom" — atom index a for each triplet, shape [n_triplets].

Return type:

Dict with keys