deepmd.pt.utils.nlist

Module Contents

Functions

extend_input_and_build_neighbor_list(coord, atype, ...)
build_neighbor_list(→ torch.Tensor)	Build neightbor list for a single frame. keeps nsel neighbors.
nlist_distinguish_types(nlist, atype, sel)	Given a nlist that does not distinguish atom types, return a nlist that
get_multiple_nlist_key(→ str)
build_multiple_neighbor_list(→ Dict[str, torch.Tensor])	Input one neighbor list, and produce multiple neighbor lists with
extend_coord_with_ghosts(coord, atype, cell, rcut[, ...])	Extend the coordinates of the atoms by appending peridoc images.

deepmd.pt.utils.nlist.extend_input_and_build_neighbor_list(coord, atype, rcut: float, sel: List[int], mixed_types: bool = False, box: torch.Tensor | None = None)

deepmd.pt.utils.nlist.build_neighbor_list(coord: torch.Tensor, atype: torch.Tensor, nloc: int, rcut: float, sel: int | List[int], distinguish_types: bool = True) → torch.Tensor

Build neightbor list for a single frame. keeps nsel neighbors.

Parameters:

coordtorch.Tensor

exptended coordinates of shape [batch_size, nall x 3]

atypetorch.Tensor

extended atomic types of shape [batch_size, nall] if type < 0 the atom is treat as virtual atoms.

nlocint

number of local atoms.

rcutfloat

cut-off radius

selint or List[int]

maximal number of neighbors (of each type). if distinguish_types==True, nsel should be list and the length of nsel should be equal to number of types.

distinguish_typesbool

distinguish different types.

Returns:

neighbor_listtorch.Tensor

Neighbor list of shape [batch_size, nloc, nsel], the neighbors are stored in an ascending order. If the number of neighbors is less than nsel, the positions are masked with -1. The neighbor list of an atom looks like |------ nsel ------| xx xx xx xx -1 -1 -1 if distinguish_types==True and we have two types |---- nsel[0] -----| |---- nsel[1] -----| xx xx xx xx -1 -1 -1 xx xx xx -1 -1 -1 -1 For virtual atoms all neighboring positions are filled with -1.

deepmd.pt.utils.nlist.nlist_distinguish_types(nlist: torch.Tensor, atype: torch.Tensor, sel: List[int])

Given a nlist that does not distinguish atom types, return a nlist that distinguish atom types.

deepmd.pt.utils.nlist.get_multiple_nlist_key(rcut: float, nsel: int) → str

deepmd.pt.utils.nlist.build_multiple_neighbor_list(coord: torch.Tensor, nlist: torch.Tensor, rcuts: List[float], nsels: List[int]) → Dict[str, torch.Tensor]

Input one neighbor list, and produce multiple neighbor lists with different cutoff radius and numbers of selection out of it. The required rcuts and nsels should be smaller or equal to the input nlist.

Parameters:

coordtorch.Tensor

exptended coordinates of shape [batch_size, nall x 3]

nlisttorch.Tensor

Neighbor list of shape [batch_size, nloc, nsel], the neighbors should be stored in an ascending order.

rcutsList[float]

list of cut-off radius in ascending order.

nselsList[int]

maximal number of neighbors in ascending order.

Returns:

nlist_dictDict[str, torch.Tensor]

A dict of nlists, key given by get_multiple_nlist_key(rc, nsel) value being the corresponding nlist.

deepmd.pt.utils.nlist.extend_coord_with_ghosts(coord: torch.Tensor, atype: torch.Tensor, cell: torch.Tensor | None, rcut: float, cell_cpu: torch.Tensor | None = None)

Extend the coordinates of the atoms by appending peridoc images. The number of images is large enough to ensure all the neighbors within rcut are appended.

Parameters:

coordtorch.Tensor

original coordinates of shape [-1, nloc*3].

atypetorch.Tensor

atom type of shape [-1, nloc].

celltorch.Tensor

simulation cell tensor of shape [-1, 9].

rcutfloat

the cutoff radius

cell_cputorch.Tensor

cell on cpu for performance

Returns:

extended_coord: torch.Tensor

extended coordinates of shape [-1, nall*3].

extended_atype: torch.Tensor

extended atom type of shape [-1, nall].

index_mapping: torch.Tensor

maping extended index to the local index