Neighbor Lists

mlcg.neighborlist contains several functions that can be used to create neighbor lists using a finite cutoff with mixed periodic boundary conditions using pytorch. There are also tools to interface with ASE tools. The neighbor lists are dictionaries containing meta-data (user defined tag, body order, cutoff, self_interaction) and the actual indices, index_mapping, and cell_shifts.

Neighbor List Utilities

These utilities are meant to format and validate neighbor list dictionaries.

mlcg.neighbor_list.neighbor_list.make_neighbor_list(tag, order, index_mapping, mapping_batch=None, cell_shifts=None, rcut=None, self_interaction=None)[source]

Build a neighbor_list dictionary.

Parameters:

  • tag (str) – quick identifier for compatibility checking

  • order (int) – an int providing the order of the neighborlist, e.g. order == 2 means that central atoms i have 1 neighbor j so distances can be computed, order == 3 means that central atoms i have 2 neighbors j and k so angles can be computed

  • index_mapping (Tensor) – The [order, n_edge] index tensor giving center -> neighbor relations. 1st column refers to the central atom index and the 2nd column to the neighbor atom index in the list of atoms (so it has to be shifted by a cell_shift to get the actual position of the neighboring atoms)

  • mapping_batch (Optional[Tensor]) – [n_edge] map for the neighbor -> structure index relation

  • cell_shifts (Optional[Tensor]) – A [n_edge, 3] tensor giving the periodic cell shift

  • rcut (Optional[float]) – cutoff radius used to compute the connectivity

  • self_interaction (Optional[bool]) – whether the mapping includes self referring mappings, e.g. mappings where i == j.

Returns:

Neighborlist dictionary

Return type:

Dict

mlcg.neighbor_list.neighbor_list.validate_neighborlist(inp)[source]

Tool to validate that the neighborlist dictionary has the required fields

Parameters:

inp (Dict) – Input neighborlist to be validated

Returns:

True if the supplied neighborlist is valid, false otherwise

Return type:

bool

Torch Implementation

mlcg.neighbor_list.torch_impl.torch_neighbor_list(data, rcut, self_interaction=True, num_workers=1, max_num_neighbors=1000)[source]

Function for computing neighbor lists from pytorch geometric data instances that may or may not use periodic boundary conditions

Parameters:

  • data (Data) – Pytorch geometric data instance

  • rcut (float) – upper distance cutoff, in which neighbors with distances larger than this cutoff will be excluded.

  • self_interaction (bool) – If True, each atom will be considered it’s own neighbor. This can be convenient for certain bases and representations

  • num_workers (int) – Number of threads to spawn and use for computing the neighbor list

  • max_number_neighbors – kwarg for radius_graph function from torch_cluster package, specifying the maximum number of neighbors for each atom

Return type:

Tuple[Tensor, Tensor, Tensor, Tensor]

Returns:

  • torch.Tensor – The atom indices of the first atoms in each neighbor pair

  • torch.Tensor – The atom indices of the second atoms in each neighbor pair

  • torch.Tensor – The cell shifts associated with minimum image distances in the presence of periodic boundary conditions

  • torch.Tensor – Mask for excluding self interactions

mlcg.neighbor_list.torch_impl.torch_neighbor_list_no_pbc(data, rcut, self_interaction=True, num_workers=1, max_num_neighbors=1000)[source]

Function for producing torch neighborlists without periodic boundary conditions

Parameters:

  • data (AtomicData) – AtomicData instance

  • rcut (float) – Upper distance cutoff for determining neighbor edges

  • self_interaction (Optional[bool]) – If True, self edges will added for each atom

  • num_workers (Optional[int]) – Number of threads to use for neighbor enumeration

  • max_num_neighbors (Optional[int]) – The maximum number of neighbors to return for each atom. For larger systems, it is important to make sure that this number is sufficiently large.

Return type:

Tuple[Tensor, Tensor, Tensor]

Returns:

  • torch.Tensor – The first atoms in each edge

  • torch.Tensor – The second atoms in each edge

  • torch.Tensor – Boolean tensor identifying self_interacting edges

mlcg.neighbor_list.torch_impl.compute_images(positions, cell, pbc, cutoff, batch, n_atoms)[source]

TODO: add doc

Return type:

Tuple[Tensor, Tensor, Tensor, Tensor]

mlcg.neighbor_list.torch_impl.wrap_positions(data, eps=1e-07)[source]

Wrap positions to unit cell.

Returns positions changed by a multiple of the unit cell vectors to fit inside the space spanned by these vectors.

Parameters:

  • data (Data) – torch_geometric.Data instance

  • eps (float) – Small number to prevent slightly negative coordinates from being wrapped.

Return type:

None

ASE Implementation

mlcg.neighbor_list.ase_impl.ase_neighbor_list(data, rcut, self_interaction=False)[source]

Function for converting a list of neighbor edges from an input AtomicData instance, to an ASE neighborlist as output by ase.neighborlist.neighbor_list. Return documenation taken from https://wiki.fysik.dtu.dk/ase/_modules/ase/neighborlist.html#neighbor_list.

Parameters:

  • data (AtomicData) – Input AtomicData instance. Must contain only one structure.

  • rcut (float) – Upper distance cover for determining neighbors

  • self_interaction (bool) – If True, self edges will be added.

Return type:

Tuple[Tensor, Tensor, Tensor]

Returns:

  • torch.Tensor – first atom indices, of shape (n_atoms)

  • torch.Tensor – second atom index, of shape (n_atoms)

  • torch.Tensor – Dot product of the periodic shift vectors with the system unit cell vectors