get_fully_connected_mapping

torch_sim.transforms.get_fully_connected_mapping(*, i_ids, shifts_idx, self_interaction)[source]

Generate a fully connected mapping of atom indices with optional cell shifts.

This function computes a mapping of atom indices for a fully connected graph, considering periodic boundary conditions through cell shifts. It can also exclude self-interactions based on the provided flag.

Parameters:
  • i_ids (Tensor) – A tensor of shape (n_atoms,) containing the indices of the atoms.

  • shifts_idx (Tensor) – A tensor of shape (n_shifts, 3) representing the shifts to apply for periodic boundary conditions.

  • self_interaction (bool) – A flag indicating whether to include self-interactions in the mapping.

Returns:

A tuple containing:
  • mapping (torch.Tensor): A tensor of shape (n_pairs, 2)

    representing the pairs of indices for which distances will be computed.

  • shifts_idx (torch.Tensor): A tensor of shape (n_pairs, 3)

    representing the corresponding shifts for the computed pairs.

Return type:

tuple[Tensor, Tensor]