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 checkingorder (
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 computedindex_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 relationcell_shifts (
Optional
[Tensor
]) – A [n_edge, 3] tensor giving the periodic cell shiftrcut (
Optional
[float
]) – cutoff radius used to compute the connectivityself_interaction (
Optional
[bool
]) – whether the mapping includes self referring mappings, e.g. mappings where i == j.
- Returns:
Neighborlist dictionary
- Return type:
Dict
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 instancercut (
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 representationsnum_workers (
int
) – Number of threads to spawn and use for computing the neighbor listmax_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 instancercut (
float
) – Upper distance cutoff for determining neighbor edgesself_interaction (
Optional
[bool
]) – If True, self edges will added for each atomnum_workers (
Optional
[int
]) – Number of threads to use for neighbor enumerationmax_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 instanceeps (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 neighborsself_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