deepmd.pt.utils.nlist

Functions

extend_input_and_build_neighbor_list(coord, atype, ...)
build_neighbor_list(→ torch.Tensor)	Build neighbor list for a single frame. keeps nsel neighbors.
_trim_mask_distinguish_nlist(→ torch.Tensor)	Trim the size of nlist, mask if any central atom is virtual, distinguish types if necessary.
build_directional_neighbor_list(→ torch.Tensor)	Build directional neighbor list.
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.

Module Contents

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 neighbor 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._trim_mask_distinguish_nlist(is_vir_cntl: torch.Tensor, atype_neig: torch.Tensor, rr: torch.Tensor, nlist: torch.Tensor, rcut: float, sel: list[int], distinguish_types: bool) → torch.Tensor

Trim the size of nlist, mask if any central atom is virtual, distinguish types if necessary.

deepmd.pt.utils.nlist.build_directional_neighbor_list(coord_cntl: torch.Tensor, atype_cntl: torch.Tensor, coord_neig: torch.Tensor, atype_neig: torch.Tensor, rcut: float, sel: int | list[int], distinguish_types: bool = True) → torch.Tensor

Build directional neighbor list.

With each central atom, all the neighbor atoms in the cut-off radius will be recorded in the neighbor list. The maximum neighbors is nsel. If the real number of neighbors is larger than nsel, the neighbors will be sorted with the distance and the first nsel neighbors are kept.

Important: the central and neighboring atoms are assume to be different atoms.

Parameters:

coord_centraltorch.Tensor

coordinates of central atoms. assumed to be local atoms. shape [batch_size, nloc_central x 3]

atype_centraltorch.Tensor

atomic types of central atoms. shape [batch_size, nloc_central] if type < 0 the atom is treated as virtual atoms.

coord_neighbortorch.Tensor

extended coordinates of neighbors atoms. shape [batch_size, nall_neighbor x 3]

atype_centraltorch.Tensor

extended atomic types of neighbors atoms. shape [batch_size, nall_neighbor] if type < 0 the atom is treated as virtual 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_central, 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

mapping extended index to the local index