deepmd.dpmodel.utils

deepmd.dpmodel.utils#

Submodules#

Attributes#

`EmbeddingNet`
`FittingNet`
`NativeNet`

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#

`aggregate`(data, owners[, average, num_owner])	Aggregate rows in data by specifying the owners.
`get_graph_index`(nlist, nlist_mask, a_nlist_mask, nall)	Get the index mapping for edge graph and angle graph, ready in aggregate or index_select.
`make_embedding_network`(→ type)
`make_fitting_network`(→ type)
`make_multilayer_network`(→ type)
`build_multiple_neighbor_list`(→ dict[str, ...)	Input one neighbor list, and produce multiple neighbor lists with
`build_neighbor_list`(→ deepmd.dpmodel.array_api.Array)	Build neighbor list for a single frame. keeps nsel neighbors.
`extend_coord_with_ghosts`(...)	Extend the coordinates of the atoms by appending peridoc images.
`get_multiple_nlist_key`(→ str)
`nlist_distinguish_types`(→ deepmd.dpmodel.array_api.Array)	Given a nlist that does not distinguish atom types, return a nlist that
`inter2phys`(→ deepmd.dpmodel.array_api.Array)	Convert internal(direct) coordinates to physical coordinates.
`normalize_coord`(→ deepmd.dpmodel.array_api.Array)	Apply PBC according to the atomic coordinates.
`phys2inter`(→ deepmd.dpmodel.array_api.Array)	Convert physical coordinates to internal(direct) coordinates.
`to_face_distance`(→ deepmd.dpmodel.array_api.Array)	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`(→ Any)	Traverse a model dict and call callback on each variable.

Package Contents#

class deepmd.dpmodel.utils.EnvMat(rcut: float, rcut_smth: float, protection: float = 0.0, use_exp_switch: bool = False)[source]#

Bases: deepmd.dpmodel.NativeOP

The unit operation of a native model.

rcut#

rcut_smth#

protection = 0.0#

use_exp_switch = False#

call(coord_ext: deepmd.dpmodel.array_api.Array, atype_ext: deepmd.dpmodel.array_api.Array, nlist: deepmd.dpmodel.array_api.Array, davg: deepmd.dpmodel.array_api.Array | None = None, dstd: deepmd.dpmodel.array_api.Array | None = None, radial_only: bool = False) → tuple[deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array][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)
diff: The relative coordinate of neighbors. shape: nf x nloc x nnei x 3
switch: The value of switch function. shape: nf x nloc x nnei

_call(nlist: Any, coord_ext: Any, radial_only: bool) → tuple[Any, Any, Any][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.

ntypes#

exclude_types = []#

type_mask#

get_exclude_types() → list[int][source]#

get_type_mask() → deepmd.dpmodel.array_api.Array[source]#

build_type_exclude_mask(atype: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array[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.

ntypes#

exclude_types#

type_mask#

get_exclude_types() → list[tuple[int, int]][source]#

build_type_exclude_mask(nlist: deepmd.dpmodel.array_api.Array, atype_ext: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array[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: tuple[int, int]) → bool[source]#

deepmd.dpmodel.utils.EmbeddingNet[source]#

deepmd.dpmodel.utils.FittingNet[source]#

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

Bases: deepmd.dpmodel.NativeOP

Native representation of a layer.

Parameters:

wArray, optional: The weights of the layer.
bArray, optional: The biases of the layer.
idtArray, 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.
precisionstr, optional: The precision of the layer.
seedint, optional: Random seed.
trainablebool, default=True: Whether the layer is trainable.

trainable = True#

precision = 'float64'#

w#

b#

idt#

activation_function = None#

resnet = False#

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() → None[source]#

check_type_consistency() → None[source]#

__setitem__(key: str, value: Any) → None[source]#

__getitem__(key: str) → Any[source]#

dim_in() → int[source]#

dim_out() → int[source]#

call(x)[source]#

Forward pass.

Parameters:

xArray: 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 dimensions 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]]#

ndim#

ntypes#

network_type#

_networks#

check_completeness() → None[source]#

Check whether the collection is complete.

Raises:

RuntimeError: If the collection is incomplete.

_convert_key(key: int | tuple) → int[source]#

__getitem__(key: int | tuple) → Any[source]#

__setitem__(key: int | tuple, value: Any) → None[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.aggregate(data, owners, average: bool = True, num_owner: int | None = None)[source]#

Aggregate rows in data by specifying the owners.

Parameters:

datadata tensor to aggregate [n_row, feature_dim]
ownersspecify the owner of each row [n_row, 1]
averageif True, average the rows, if False, sum the rows.: Default = True
num_ownerthe number of owners, this is needed if the: max idx of owner is not presented in owners tensor Default = None

Returns:

output: [num_owner, feature_dim]

deepmd.dpmodel.utils.get_graph_index(nlist, nlist_mask, a_nlist_mask, nall: int, use_loc_mapping: bool = True)[source]#

Get the index mapping for edge graph and angle graph, ready in aggregate or index_select.

Parameters:

nlistnf x nloc x nnei: Neighbor list. (padded neis are set to 0)
nlist_masknf x nloc x nnei: Masks of the neighbor list. real nei 1 otherwise 0
a_nlist_masknf x nloc x a_nnei: Masks of the neighbor list for angle. real nei 1 otherwise 0
nall: The number of extended atoms.
use_loc_mapping: Whether to use local atom index mapping in training or non-parallel inference. When True, local indexing and mapping are applied to neighbor lists and embeddings during descriptor computation.

Returns:

edge_index2 x n_edge

n2e_indexn_edge: Broadcast indices from node(i) to edge(ij), or reduction indices from edge(ij) to node(i).
n_ext2e_indexn_edge: Broadcast indices from extended node(j) to edge(ij).

angle_index3 x n_angle

n2a_indexn_angle: Broadcast indices from extended node(j) to angle(ijk).
eij2a_indexn_angle: Broadcast indices from extended edge(ij) to angle(ijk), or reduction indices from angle(ijk) to edge(ij).
eik2a_indexn_angle: Broadcast indices from extended edge(ik) to angle(ijk).

deepmd.dpmodel.utils.make_embedding_network(T_Network: type, T_NetworkLayer: type) → type[source]#

deepmd.dpmodel.utils.make_fitting_network(T_EmbeddingNet: type, T_Network: type, T_NetworkLayer: type) → type[source]#

deepmd.dpmodel.utils.make_multilayer_network(T_NetworkLayer: type, ModuleBase: type) → type[source]#

deepmd.dpmodel.utils.build_multiple_neighbor_list(coord: deepmd.dpmodel.array_api.Array, nlist: deepmd.dpmodel.array_api.Array, rcuts: list[float], nsels: list[int]) → dict[str, deepmd.dpmodel.array_api.Array][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:

coordArray: exptended coordinates of shape [batch_size, nall x 3]
nlistArray: 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, Array]: 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: deepmd.dpmodel.array_api.Array, atype: deepmd.dpmodel.array_api.Array, nloc: int, rcut: float, sel: int | list[int], distinguish_types: bool = True) → deepmd.dpmodel.array_api.Array[source]#

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

Parameters:

coordArray: exptended coordinates of shape [batch_size, nall x 3]
atypeArray: 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_listArray: 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: deepmd.dpmodel.array_api.Array, atype: deepmd.dpmodel.array_api.Array, cell: deepmd.dpmodel.array_api.Array | None, rcut: float) → tuple[deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array][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:

coordArray: original coordinates of shape [-1, nloc*3].
atypeArray: atom type of shape [-1, nloc].
cellArray: simulation cell tensor of shape [-1, 9].
rcutfloat: the cutoff radius

Returns:

extended_coord: Array: extended coordinates of shape [-1, nall*3].
extended_atype: Array: extended atom type of shape [-1, nall].
index_mapping: Array: mapping 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: deepmd.dpmodel.array_api.Array, atype: deepmd.dpmodel.array_api.Array, sel: list[int]) → deepmd.dpmodel.array_api.Array[source]#: Given a nlist that does not distinguish atom types, return a nlist that distinguish atom types.

deepmd.dpmodel.utils.inter2phys(coord: deepmd.dpmodel.array_api.Array, cell: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array[source]#

Convert internal(direct) coordinates to physical coordinates.

Parameters:

coordArray: internal coordinates of shape [*, na, 3].
cellArray: simulation cell tensor of shape [*, 3, 3].

Returns:

phys_coord: Array: the physical coordinates

deepmd.dpmodel.utils.normalize_coord(coord: deepmd.dpmodel.array_api.Array, cell: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array[source]#

Apply PBC according to the atomic coordinates.

Parameters:

coordArray: original coordinates of shape [*, na, 3].
cellArray: simulation cell shape [*, 3, 3].

Returns:

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

deepmd.dpmodel.utils.phys2inter(coord: deepmd.dpmodel.array_api.Array, cell: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array[source]#

Convert physical coordinates to internal(direct) coordinates.

Parameters:

coordArray: physical coordinates of shape [*, na, 3].
cellArray: simulation cell tensor of shape [*, 3, 3].

Returns:

inter_coord: Array: the internal coordinates

deepmd.dpmodel.utils.to_face_distance(cell: deepmd.dpmodel.array_api.Array) → deepmd.dpmodel.array_api.Array[source]#

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

Parameters:

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

Returns:

dist: Array: 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: Any, callback: Callable, is_variable: bool = False) → Any[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.