deepmd.jax.jax2tf.nlist

Functions

build_neighbor_list(...)	Build neighbor 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
tf_outer(a, b)
extend_coord_with_ghosts(coord, atype, cell, rcut)	Extend the coordinates of the atoms by appending peridoc images.

Module Contents

deepmd.jax.jax2tf.nlist.build_neighbor_list(coord: tensorflow.experimental.numpy.ndarray, atype: tensorflow.experimental.numpy.ndarray, nloc: int, rcut: float, sel: int | list[int], distinguish_types: bool = True) → tensorflow.experimental.numpy.ndarray

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

Parameters:

coordtnp.ndarray

exptended coordinates of shape [batch_size, nall x 3]

atypetnp.ndarray

extended atomic types of shape [batch_size, nall] 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_listtnp.ndarray

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.jax.jax2tf.nlist.nlist_distinguish_types(nlist: tensorflow.experimental.numpy.ndarray, atype: tensorflow.experimental.numpy.ndarray, sel: list[int])

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

deepmd.jax.jax2tf.nlist.tf_outer(a, b)

deepmd.jax.jax2tf.nlist.extend_coord_with_ghosts(coord: tensorflow.experimental.numpy.ndarray, atype: tensorflow.experimental.numpy.ndarray, cell: tensorflow.experimental.numpy.ndarray, rcut: float)

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:

coordtnp.ndarray

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

atypetnp.ndarray

atom type of shape [-1, nloc].

celltnp.ndarray

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

rcutfloat

the cutoff radius

Returns:

extended_coord: tnp.ndarray

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

extended_atype: tnp.ndarray

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

index_mapping: tnp.ndarray

mapping extended index to the local index