build_naive_neighborhood

torch_sim.transforms.build_naive_neighborhood(positions, cell, pbc, cutoff, n_atoms, self_interaction)[source]

Build a vectorized O(N^2) neighborhood list for batched atomic systems.

All systems are padded to a common size and processed simultaneously using batched tensor operations. Pairs within the cutoff are returned with global atom indices.

NOTE: due to the use of pad_sequence, this function is best used when all the systems being batched have a similar number of atoms as this reduces the memory overhead of the padding.

Parameters:

  • positions (torch.Tensor) – A tensor of shape (n_total_atoms, 3) representing the positions of atoms.

  • cell (torch.Tensor) – A tensor of shape (n_systems, 3, 3) representing the unit cell matrices.

  • pbc (torch.Tensor) – A tensor of shape (n_systems, 3) indicating whether periodic boundary conditions are applied.

  • cutoff (float) – The cutoff distance beyond which atoms are not considered neighbors.

  • n_atoms (torch.Tensor) – A tensor containing the number of atoms in each structure.

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

Returns:

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

    representing the pairs of global indices for neighboring atoms.

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

    indicating the structure index for each pair.

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

    representing the integer lattice shifts for each pair.

Return type:

tuple[Tensor, Tensor, Tensor]

References