torchsim_nl

torch_sim.neighbors.torchsim_nl(positions, cell, pbc, cutoff, sort_id=False)[source]

Compute neighbor lists with automatic fallback for AMD ROCm compatibility.

This function automatically selects the best available neighbor list implementation. When vesin is available, it uses vesin_nl_ts for optimal performance. When vesin is not available (e.g., on AMD ROCm systems), it falls back to standard_nl.

Parameters:

  • positions (Tensor) – Atomic positions tensor of shape (num_atoms, 3)

  • cell (Tensor) – Unit cell vectors according to the row vector convention, i.e. [[a1, a2, a3], [b1, b2, b3], [c1, c2, c3]].

  • pbc (Tensor) – Boolean tensor of shape (3,) indicating periodic boundary conditions in each axis.

  • cutoff (Tensor) – Maximum distance (scalar tensor) for considering atoms as neighbors

  • sort_id (bool) – If True, sort neighbors by first atom index for better memory access patterns

Returns:

  • mapping: Tensor of shape (2, num_neighbors) containing pairs of atom indices that are neighbors. Each column (i,j) represents a neighbor pair.

  • shifts: Tensor of shape (num_neighbors, 3) containing the periodic shift vectors needed to get the correct periodic image for each neighbor pair.

Return type:

tuple containing

Notes

  • Automatically uses vesin_nl_ts when vesin is available

  • Falls back to standard_nl when vesin is unavailable (AMD ROCm)

  • Fallback works on NVIDIA CUDA, AMD ROCm, and CPU

  • For non-periodic systems (pbc=False), shifts will be zero vectors

  • The neighbor list includes both (i,j) and (j,i) pairs