`deepmd.dpmodel.utils`

Submodules

Package Contents

Classes

`EnvMat`	The unit operation of a native model.
`AtomExcludeMask`	Computes the type exclusion mask for atoms.
`PairExcludeMask`	Computes the type exclusion mask for atom pairs.
`NativeLayer`	Native representation of a layer.
`NetworkCollection`	A collection of networks for multiple elements.

Functions

`make_embedding_network`(T_Network, T_NetworkLayer)
`make_fitting_network`(T_EmbeddingNet, T_Network, ...)
`make_multilayer_network`(T_NetworkLayer, ModuleBase)
`build_multiple_neighbor_list`(→ Dict[str, numpy.ndarray])	Input one neighbor list, and produce multiple neighbor lists with
`build_neighbor_list`(→ numpy.ndarray)	Build neightbor list for a single frame. keeps nsel neighbors.
`extend_coord_with_ghosts`(coord, atype, cell, rcut)	Extend the coordinates of the atoms by appending peridoc images.
`get_multiple_nlist_key`(→ str)
`nlist_distinguish_types`(nlist, atype, sel)	Given a nlist that does not distinguish atom types, return a nlist that
`inter2phys`(→ numpy.ndarray)	Convert internal(direct) coordinates to physical coordinates.
`normalize_coord`(→ numpy.ndarray)	Apply PBC according to the atomic coordinates.
`phys2inter`(→ numpy.ndarray)	Convert physical coordinates to internal(direct) coordinates.
`to_face_distance`(→ numpy.ndarray)	Compute the to-face-distance of the simulation cell.
`load_dp_model`(→ dict)	Load a DP model from a file in the native format.
`save_dp_model`(→ None)	Save a DP model to a file in the native format.
`traverse_model_dict`(model_obj, callback[, is_variable])	Traverse a model dict and call callback on each variable.

Attributes

`EmbeddingNet`
`FittingNet`
`NativeNet`

class deepmd.dpmodel.utils.EnvMat(rcut, rcut_smth, protection: float = 0.0)[source]

Bases: deepmd.dpmodel.NativeOP

The unit operation of a native model.

call(coord_ext: numpy.ndarray, atype_ext: numpy.ndarray, nlist: numpy.ndarray, davg: numpy.ndarray | None = None, dstd: numpy.ndarray | None = None, radial_only: bool = False) → numpy.ndarray | numpy.ndarray[source]

Compute the environment matrix.

Parameters:

nlist: The neighbor list. shape: nf x nloc x nnei
coord_ext: The extended coordinates of atoms. shape: nf x (nallx3)
atype_ext: The extended aotm types. shape: nf x nall
davg: The data avg. shape: nt x nnei x (4 or 1)
dstd: The inverse of data std. shape: nt x nnei x (4 or 1)
radial_only: Whether to only compute radial part of the environment matrix. If True, the output will be of shape nf x nloc x nnei x 1. Otherwise, the output will be of shape nf x nloc x nnei x 4. Default: False.

Returns:

env_mat: The environment matrix. shape: nf x nloc x nnei x (4 or 1)
switch: The value of switch function. shape: nf x nloc x nnei

_call(nlist, coord_ext, radial_only)[source]

serialize() → dict[source]

classmethod deserialize(data: dict) → EnvMat[source]

class deepmd.dpmodel.utils.AtomExcludeMask(ntypes: int, exclude_types: List[int] = [])[source]

Computes the type exclusion mask for atoms.

get_exclude_types()[source]

get_type_mask()[source]

build_type_exclude_mask(atype: numpy.ndarray)[source]

Compute type exclusion mask for atoms.

Parameters:

atype: The extended aotm types. shape: nf x natom

Returns:

mask: The type exclusion mask for atoms. shape: nf x natom Element [ff,ii] being 0 if type(ii) is excluded, otherwise being 1.

class deepmd.dpmodel.utils.PairExcludeMask(ntypes: int, exclude_types: List[Tuple[int, int]] = [])[source]

Computes the type exclusion mask for atom pairs.

get_exclude_types()[source]

build_type_exclude_mask(nlist: numpy.ndarray, atype_ext: numpy.ndarray)[source]

Compute type exclusion mask for atom pairs.

Parameters:

nlist: The neighbor list. shape: nf x nloc x nnei
atype_ext: The extended aotm types. shape: nf x nall

Returns:

mask: The type exclusion mask for pair atoms of shape: nf x nloc x nnei. Element [ff,ii,jj] being 0 if type(ii), type(nlist[ff,ii,jj]) is excluded, otherwise being 1.

__contains__(item)[source]

deepmd.dpmodel.utils.EmbeddingNet[source]

deepmd.dpmodel.utils.FittingNet[source]

class deepmd.dpmodel.utils.NativeLayer(num_in, num_out, bias: bool = True, use_timestep: bool = False, activation_function: str | None = None, resnet: bool = False, precision: str = DEFAULT_PRECISION)[source]

Bases: deepmd.dpmodel.NativeOP

Native representation of a layer.

Parameters:

wnp.ndarray, optional: The weights of the layer.
bnp.ndarray, optional: The biases of the layer.
idtnp.ndarray, optional: The identity matrix of the layer.
activation_functionstr, optional: The activation function of the layer.
resnetbool, optional: Whether the layer is a residual layer.

serialize() → dict[source]

Serialize the layer to a dict.

Returns:

dict: The serialized layer.

classmethod deserialize(data: dict) → NativeLayer[source]

Deserialize the layer from a dict.

Parameters:

datadict: The dict to deserialize from.

check_shape_consistency()[source]

check_type_consistency()[source]

__setitem__(key, value)[source]

__getitem__(key)[source]

dim_in() → int[source]

dim_out() → int[source]

call(x: numpy.ndarray) → numpy.ndarray[source]

Forward pass.

Parameters:

xnp.ndarray: The input.

Returns:

np.ndarray: The output.

deepmd.dpmodel.utils.NativeNet[source]

class deepmd.dpmodel.utils.NetworkCollection(ndim: int, ntypes: int, network_type: str = 'network', networks: List[NativeNet | dict] = [])[source]

A collection of networks for multiple elements.

The number of dimesions for types might be 0, 1, or 2. - 0: embedding or fitting with type embedding, in () - 1: embedding with type_one_side, or fitting, in (type_i) - 2: embedding without type_one_side, in (type_i, type_j)

Parameters:

ndimint: The number of dimensions.
network_typestr, optional: The type of the network.
networksdict, optional: The networks to initialize with.

NETWORK_TYPE_MAP: ClassVar[Dict[str, type]]

check_completeness()[source]

Check whether the collection is complete.

Raises:

RuntimeError: If the collection is incomplete.

_convert_key(key)[source]

__getitem__(key)[source]

__setitem__(key, value)[source]

serialize() → dict[source]

Serialize the networks to a dict.

Returns:

dict: The serialized networks.

classmethod deserialize(data: dict) → NetworkCollection[source]

Deserialize the networks from a dict.

Parameters:

datadict: The dict to deserialize from.

deepmd.dpmodel.utils.make_embedding_network(T_Network, T_NetworkLayer)[source]

deepmd.dpmodel.utils.make_fitting_network(T_EmbeddingNet, T_Network, T_NetworkLayer)[source]

deepmd.dpmodel.utils.make_multilayer_network(T_NetworkLayer, ModuleBase)[source]

deepmd.dpmodel.utils.build_multiple_neighbor_list(coord: numpy.ndarray, nlist: numpy.ndarray, rcuts: List[float], nsels: List[int]) → Dict[str, numpy.ndarray][source]

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:

coordnp.ndarray: exptended coordinates of shape [batch_size, nall x 3]
nlistnp.ndarray: 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, np.ndarray]: A dict of nlists, key given by get_multiple_nlist_key(rc, nsel) value being the corresponding nlist.

deepmd.dpmodel.utils.build_neighbor_list(coord: numpy.ndarray, atype: numpy.ndarray, nloc: int, rcut: float, sel: int | List[int], distinguish_types: bool = True) → numpy.ndarray[source]

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

Parameters:

coordnp.ndarray: exptended coordinates of shape [batch_size, nall x 3]
atypenp.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_listnp.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.dpmodel.utils.extend_coord_with_ghosts(coord: numpy.ndarray, atype: numpy.ndarray, cell: numpy.ndarray | None, rcut: float)[source]

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:

coordnp.ndarray: original coordinates of shape [-1, nloc*3].
atypenp.ndarray: atom type of shape [-1, nloc].
cellnp.ndarray: simulation cell tensor of shape [-1, 9].
rcutfloat: the cutoff radius

Returns:

extended_coord: np.ndarray: extended coordinates of shape [-1, nall*3].
extended_atype: np.ndarray: extended atom type of shape [-1, nall].
index_mapping: np.ndarray: maping extended index to the local index

deepmd.dpmodel.utils.get_multiple_nlist_key(rcut: float, nsel: int) → str[source]

deepmd.dpmodel.utils.nlist_distinguish_types(nlist: numpy.ndarray, atype: numpy.ndarray, sel: List[int])[source]: Given a nlist that does not distinguish atom types, return a nlist that distinguish atom types.

deepmd.dpmodel.utils.inter2phys(coord: numpy.ndarray, cell: numpy.ndarray) → numpy.ndarray[source]

Convert internal(direct) coordinates to physical coordinates.

Parameters:

coordnp.ndarray: internal coordinates of shape [*, na, 3].
cellnp.ndarray: simulation cell tensor of shape [*, 3, 3].

Returns:

phys_coord: np.ndarray: the physical coordinates

deepmd.dpmodel.utils.normalize_coord(coord: numpy.ndarray, cell: numpy.ndarray) → numpy.ndarray[source]

Apply PBC according to the atomic coordinates.

Parameters:

coordnp.ndarray: orignal coordinates of shape [*, na, 3].
cellnp.ndarray: simulation cell shape [*, 3, 3].

Returns:

wrapped_coord: np.ndarray: wrapped coordinates of shape [*, na, 3].

deepmd.dpmodel.utils.phys2inter(coord: numpy.ndarray, cell: numpy.ndarray) → numpy.ndarray[source]

Convert physical coordinates to internal(direct) coordinates.

Parameters:

coordnp.ndarray: physical coordinates of shape [*, na, 3].
cellnp.ndarray: simulation cell tensor of shape [*, 3, 3].

Returns:

inter_coord: np.ndarray: the internal coordinates

deepmd.dpmodel.utils.to_face_distance(cell: numpy.ndarray) → numpy.ndarray[source]

Compute the to-face-distance of the simulation cell.

Parameters:

cellnp.ndarray: simulation cell tensor of shape [*, 3, 3].

Returns:

dist: np.ndarray: the to face distances of shape [*, 3]

deepmd.dpmodel.utils.load_dp_model(filename: str) → dict[source]

Load a DP model from a file in the native format.

Parameters:

filenamestr: The filename to load from.

Returns:

dict: The loaded model dict, including meta information.

deepmd.dpmodel.utils.save_dp_model(filename: str, model_dict: dict) → None[source]

Save a DP model to a file in the native format.

Parameters:

filenamestr: The filename to save to.
model_dictdict: The model dict to save.

deepmd.dpmodel.utils.traverse_model_dict(model_obj, callback: Callable, is_variable: bool = False)[source]

Traverse a model dict and call callback on each variable.

Parameters:

model_objobject: The model object to traverse.
callbackcallable(): The callback function to call on each variable.
is_variablebool, optional: Whether the current node is a variable.

Returns:

object: The model object after traversing.

deepmd.dpmodel.utils

Submodules

Package Contents

Classes

Functions

Attributes

`deepmd.dpmodel.utils`