build_quadruplets

torch_sim.neighbors.nbody.build_quadruplets(main_edge_index, internal_edge_index, n_atoms, main_cell_offsets, internal_cell_offsets)[source]

Build quadruplet interaction indices d→b→a←c from two edge sets.

For each internal (short-cutoff) bond b→a, pairs every main-graph neighbour d of b with every main-graph neighbour c of a, excluding c == d in the same periodic image. The resulting four-atom chains have a short central bond flanked by longer outer bonds:

d ——(main)——> b ===(internal)===> a <——(main)—— c

Pure-PyTorch equivalent of GemNet-OC get_quadruplets.

Parameters:

  • main_edge_index (Tensor) – [2, n_main] — long-range (outer) graph edges.

  • internal_edge_index (Tensor) – [2, n_internal] — short-range (central) graph edges.

  • n_atoms (int) – Total number of atoms.

  • main_cell_offsets (Tensor) – [n_main, 3] periodic cell offsets for main graph.

  • internal_cell_offsets (Tensor) – [n_internal, 3] periodic cell offsets for internal graph.

Returns:

  • "d_to_b_edge" — main-edge indices for d→b, shape [n_trip_in].

  • "b_to_a_edge" — internal-edge indices for the central bond b→a, shape [n_trip_in].

  • "b_to_a_edge_agg" — local aggregation index within each b→a edge, shape [n_trip_in].

  • "c_to_a_edge" — main-edge indices for c→a, shape [n_trip_out].

  • "c_to_a_edge_agg" — local aggregation index within each c→a edge, shape [n_trip_out].

  • "quad_c_to_a_edge" — main-edge index of the c→a bond for each quadruplet, shape [n_quads].

  • "quad_d_to_b_trip_idx" — index into d_to_b_edge / b_to_a_edge for each quadruplet, shape [n_quads].

  • "quad_c_to_a_trip_idx" — index into c_to_a_edge for each quadruplet, shape [n_quads].

  • "quad_c_to_a_agg" — local aggregation index within each c→a main edge across quadruplets, shape [n_quads].

Return type:

Dict with keys describing the quadruplet d→b→a←c