mlipops.NeighborList#

class mlipops.NeighborList(cutoff: float | None = None, include_self: bool = False, include_symmetric: bool = False, padding: float | None = None, device: str = 'cpu')#

Identifies neighboring particles that can interact with each other.

Neighbors are usually defined by a distance cutoff: any pairs that are closer than the cutoff, possibly taking periodic boundary conditions into account, can interact and therefore are returned. It is also possible to omit the cutoff, in which case all pairs are returned regardless of distance.

Two additional options further restrict the result. include_self determines whether a self interaction pair (i, i) should be included for each particle. include_symmetric whether the result should include both of the symmetric pairs (i, j) and (j, i), or whether only one should be included.

You can optionally specify a padding value, in which case all pairs that are within the distance cutoff+padding are returned. This allows saving computation be reusing cached results. If you call the neighbor list again and no particle has moved by more than half the padding distance, the previous result can be returned. The disadvantage is that more pairs need to be included in the neighbor list. This option is generally useful only when the cost of finding neighbors is large compared to the cost of computing interactions based on them.

Warning

Because NeighborList caches its inputs and outputs, you must never modify them in place after calling the NeighborList. For example, if you call it once with a Tensor of positions, modify the Tensor in place to contain different positions, and then pass it to the NeighborList again, it will incorrectly think the positions have not changed.

__init__(cutoff: float | None = None, include_self: bool = False, include_symmetric: bool = False, padding: float | None = None, device: str = 'cpu')#

Create a NeighborList for identifying neighbors.

Parameters:

  • cutoff (float | None) – the cutoff distance to use when identifying neighbors. If None, all pairs are returned regardless of distance.

  • include_self (bool) – if True, include self interaction pairs of the form (i, i).

  • include_symmetric (bool) – if True, include both of the symmetric pairs (i, j) and (j, i) for particles i and j. If False, only one of the two is included.

  • padding (float | None) – the padding distance to add to the cutoff, allowing cached results to be reused. If None, caching is disabled.

  • device (str) – the PyTorch device to perform calculation on.

forward(positions: Tensor, box_vectors: Tensor | None = None, batch: Tensor | None = None) → Tensor#

Compute the neighbor list.

Parameters:

  • positions (torch.Tensor) – a Tensor of shape (particles, 3) containing the cartesian coordinates of each particle

  • box_vectors (torch.Tensor | None) – if batch is None, a Tensor of shape (3, 3) containing box vectors defining the periodic box. If batch is not None, a Tensor of shape (systems, 3, 3) containing the box vectors for each system. If None, periodic boundary conditions are not used.

  • batch (torch.Tensor | None) – a Tensor of shape (particles,) containing the index of the system each particle belongs to. This must be sorted in ascending order, and every system must contain at least one particle. If None, the neighbor list is computed for a single system instead of a batch of systems.

Returns:

a Tensor of shape (pairs, 2). Each row contains the indices of two particles that can interact.

Return type:

torch.Tensor