vesin_nl_ts

torch_sim.neighbors.vesin.vesin_nl_ts(positions, cell, pbc, cutoff, system_idx, self_interaction=False)[source]

Compute neighbor lists using TorchScript-compatible Vesin.

This function provides a TorchScript-compatible interface to the Vesin neighbor list algorithm using VesinNeighborListTorch.

Parameters:

  • positions (Tensor) – Atomic positions tensor [n_atoms, 3]

  • cell (Tensor) – Unit cell vectors [n_systems, 3, 3] or [3, 3]

  • pbc (Tensor) – Boolean tensor [n_systems, 3] or [3]

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

  • system_idx (Tensor) – Tensor [n_atoms] indicating which system each atom belongs to

  • self_interaction (bool) – If True, include self-pairs. Default: False

Returns:

  • mapping: Tensor [2, num_neighbors] - pairs of atom indices

  • system_mapping: Tensor [num_neighbors] - system assignment for each pair

  • shifts_idx: Tensor [num_neighbors, 3] - periodic shift indices

Return type:

tuple containing

Example

>>> # Single system
>>> positions = torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]])
>>> system_idx = torch.zeros(2, dtype=torch.long)
>>> mapping, sys_map, shifts = vesin_nl_ts(
...     positions, cell, pbc, cutoff, system_idx
... )

Notes

  • Uses VesinNeighborListTorch for TorchScript compatibility

  • Requires CPU tensors in float64 precision internally

  • Returns tensors on the same device as input with original precision

  • For non-periodic systems, shifts will be zero vectors

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

References

https://github.com/Luthaf/vesin