deepmd.jax.utils.neighbor_stat

Classes

NeighborStat

Neighbor statistics using JAX.

Module Contents

class deepmd.jax.utils.neighbor_stat.NeighborStat(ntypes: int, rcut: float, mixed_type: bool = False)

Bases: deepmd.utils.neighbor_stat.NeighborStat

Neighbor statistics using JAX.

Parameters:

ntypesint

The num of atom types

rcutfloat

The cut-off radius

mixed_typebool, optional, default=False

Treat all types as a single type.

op

auto_batch_size

iterator(data: deepmd.utils.data_system.DeepmdDataSystem) → collections.abc.Iterator[tuple[numpy.ndarray, float, str]]

Iterator method for producing neighbor statistics data.

Yields:

np.ndarray

The maximal number of neighbors

float

The squared minimal distance between two atoms

str

The directory of the data system

_execute(coord: numpy.ndarray, atype: numpy.ndarray, cell: numpy.ndarray | None)

Execute the operation.

Parameters:

coord

The coordinates of atoms.

atype

The atom types.

cell

The cell.