build_mixed_triplets

torch_sim.neighbors.nbody.build_mixed_triplets(edge_index_in, edge_index_out, n_atoms, to_outedge=False, cell_offsets_in=None, cell_offsets_out=None)[source]

Build triplet indices across two different edge sets sharing the same atoms.

For each edge in edge_index_out, finds all edges in edge_index_in that share the same atom (target or source depending on to_outedge), filtering self-loops via cell offsets when provided.

This is the pure-PyTorch equivalent of GemNet-OC get_mixed_triplets.

Parameters:

  • edge_index_in (Tensor) – [2, n_edges_in] — input graph edges.

  • edge_index_out (Tensor) – [2, n_edges_out] — output graph edges.

  • n_atoms (int) – Total number of atoms.

  • to_outedge (bool) – If True, match on the source atom of out-edges (a→c style); otherwise match on the target atom (c→a style).

  • cell_offsets_in (Tensor | None) – [n_edges_in, 3] periodic offsets for input graph.

  • cell_offsets_out (Tensor | None) – [n_edges_out, 3] periodic offsets for output graph.

Returns:

Dict with keys "trip_in", "trip_out", "trip_out_agg".

Return type:

dict[str, Tensor]