TorchScript API reference

class vesin_torch.NeighborList(cutoff: float, full_list: bool, sorted: bool = False, algorithm: str = 'auto', skin: float = 0.0, n_threads: int = 0)

A neighbor list calculator that can be used with TorchScript.

Parameters:

  • cutoff – spherical cutoff for this neighbor list

  • full_list – should we return each pair twice (as i-j and j-i) or only once

  • sorted – Should the neighbor list be sorted? If True, the returned pairs will be sorted by the first point index (i). The order of the second point index (j) and shifts in the list of pairs is unspecified.

  • algorithm – algorithm to use when computing the neighbor list. One of "auto", "brute_force", or "cell_list".

  • skin – skin size for Verlet caching. A positive value enables caching the neighbor list until any atom moves farther than skin/2 from its reference coordinates. A negative value sets the skin to an unspecified fraction of the cutoff that is optimal for the current implementation. Zero disables caching (the default).

  • n_threads – number of CPU threads to use. Must be 0 or a positive integer. If set to 0, Vesin uses OMP_NUM_THREADS when set to a positive integer, and otherwise the number of available CPU cores.

property cutoff: float

The cutoff distance for this neighbor list.

property full_list: bool

Whether this neighbor list returns both i-j and j-i pairs.

property sorted: bool

Whether this neighbor list is sorted by the first index i.

property algorithm: str

Algorithm used to compute this neighbor list.

property skin: float

Skin size used for Verlet caching.

property n_threads: int

Number of CPU threads requested for the calculation.

compute(points: Tensor, box: Tensor, periodic: bool | Tensor, quantities: str = 'ij', copy: bool = True) → List[Tensor]

Compute the neighbor list for the system defined by positions, box, and periodic; returning the requested quantities.

quantities can contain any combination of the following values:

  • "i" to get the index of the first point in the pair

  • "j" to get the index of the second point in the pair

  • "P" to get the indexes of the two points in the pair simultaneously

  • "S" to get the periodic shift of the pair

  • "d" to get the distance between points in the pair

  • "D" to get the distance vector between points in the pair

Parameters:

  • points – positions of all points in the system

  • box – bounding box of the system

  • periodic – should we use periodic boundary conditions? This can be a single boolean to enable/disable periodic boundary conditions in all directions, or a tensor of three booleans (one for each direction).

  • quantities – quantities to return, defaults to “ij”

  • copy – should we copy the returned quantities, defaults to True. Setting this to False might be a bit faster, but the returned tensors are view inside this class, and will be invalidated whenever this class is garbage collected or used to run a new calculation.

Returns:

list of torch.Tensor as indicated by quantities.