deepmd.jax.jax2tf.format_nlist

Functions

format_nlist(extended_coord, nlist, nsel, rcut)

Format neighbor list.

Module Contents

deepmd.jax.jax2tf.format_nlist.format_nlist(extended_coord: tensorflow.experimental.numpy.ndarray, nlist: tensorflow.experimental.numpy.ndarray, nsel: int, rcut: float)

Format neighbor list.

If nnei == nsel, do nothing; If nnei < nsel, pad -1; If nnei > nsel, sort by distance and truncate.

Parameters:

extended_coord

The extended coordinates of the atoms. shape: nf x nall x 3

nlist

The neighbor list. shape: nf x nloc x nnei

nsel

The number of selected neighbors.

rcut

The cutoff radius.

Returns:

nlist

The formatted neighbor list. shape: nf x nloc x nsel