deepmd.tf.descriptor.descriptor

deepmd.tf.descriptor.descriptor#

Classes#

Descriptor

The abstract class for descriptors. All specific descriptors should

Module Contents#

class deepmd.tf.descriptor.descriptor.Descriptor[source]#

Bases: deepmd.tf.utils.PluginVariant, make_plugin_registry('descriptor')

The abstract class for descriptors. All specific descriptors should be based on this class.

The descriptor \(\mathcal{D}\) describes the environment of an atom, which should be a function of coordinates and types of its neighbour atoms.

Notes

Only methods and attributes defined in this class are generally public, that can be called by other classes.

Examples

>>> descript = Descriptor(type="se_e2_a", rcut=6.0, rcut_smth=0.5, sel=[50])
>>> type(descript)
<class 'deepmd.tf.descriptor.se_a.DescrptSeA'>

abstract get_rcut() → float[source]#

Returns the cut-off radius.

Returns:

float: the cut-off radius

Notes

This method must be implemented, as it’s called by other classes.

abstract get_ntypes() → int[source]#

Returns the number of atom types.

Returns:

int: the number of atom types

Notes

This method must be implemented, as it’s called by other classes.

abstract get_dim_out() → int[source]#

Returns the output dimension of this descriptor.

Returns:

int: the output dimension of this descriptor

Notes

This method must be implemented, as it’s called by other classes.

abstract get_dim_rot_mat_1() → int[source]#

Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3.

Returns:

int: the first dimension of the rotation matrix

abstract get_nlist() → tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, list[int], list[int]][source]#

Returns neighbor information.

Returns:

nlisttf.Tensor: Neighbor list
rijtf.Tensor: The relative distance between the neighbor and the center atom.
sel_alist[int]: The number of neighbors with full information
sel_rlist[int]: The number of neighbors with only radial information

abstract compute_input_stats(data_coord: list[numpy.ndarray], data_box: list[numpy.ndarray], data_atype: list[numpy.ndarray], natoms_vec: list[numpy.ndarray], mesh: list[numpy.ndarray], input_dict: dict[str, list[numpy.ndarray]], **kwargs) → None[source]#

Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.

Parameters:

data_coordlist[np.ndarray]: The coordinates. Can be generated by deepmd.tf.model.model_stat.make_stat_input()
data_boxlist[np.ndarray]: The box. Can be generated by deepmd.tf.model.model_stat.make_stat_input()
data_atypelist[np.ndarray]: The atom types. Can be generated by deepmd.tf.model.model_stat.make_stat_input()
natoms_veclist[np.ndarray]: The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.tf.model.model_stat.make_stat_input()
meshlist[np.ndarray]: The mesh for neighbor searching. Can be generated by deepmd.tf.model.model_stat.make_stat_input()
input_dictdict[str, list[np.ndarray]]: Dictionary for additional input
**kwargs: Additional keyword arguments which may contain mixed_type and real_natoms_vec.

Notes

This method must be implemented, as it’s called by other classes.

abstract build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box_: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict[str, Any], reuse: bool | None = None, suffix: str = '') → deepmd.tf.env.tf.Tensor[source]#

Build the computational graph for the descriptor.

Parameters:

coord_tf.Tensor: The coordinate of atoms
atype_tf.Tensor: The type of atoms
natomstf.Tensor: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
box_tf.Tensor: The box of frames
meshtf.Tensor: For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed.
input_dictdict[str, Any]: Dictionary for additional inputs
reusebool, optional: The weights in the networks should be reused when get the variable.
suffixstr, optional: Name suffix to identify this descriptor

Returns:

descriptor: tf.Tensor: The output descriptor

Notes

This method must be implemented, as it’s called by other classes.

abstract enable_compression(min_nbor_dist: float, graph: deepmd.tf.env.tf.Graph, graph_def: deepmd.tf.env.tf.GraphDef, table_extrapolate: float = 5.0, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') → None[source]#

Receive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.

Parameters:

min_nbor_distfloat: The nearest distance between atoms
graphtf.Graph: The graph of the model
graph_deftf.GraphDef: The graph definition of the model
table_extrapolatefloat, default: 5.: The scale of model extrapolation
table_stride_1float, default: 0.01: The uniform stride of the first table
table_stride_2float, default: 0.1: The uniform stride of the second table
check_frequencyint, default: -1: The overflow check frequency
suffixstr, optional: The suffix of the scope

Notes

This method is called by others when the descriptor supported compression.

abstract enable_mixed_precision(mixed_prec: dict | None = None) → None[source]#

Receive the mixed precision setting.

Parameters:

mixed_prec: The mixed precision setting used in the embedding net

Notes

This method is called by others when the descriptor supported compression.

abstract prod_force_virial(atom_ener: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor) → tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor][source]#

Compute force and virial.

Parameters:

atom_enertf.Tensor: The atomic energy
natomstf.Tensor: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms

Returns:

forcetf.Tensor: The force on atoms
virialtf.Tensor: The total virial
atom_virialtf.Tensor: The atomic virial

abstract init_variables(graph: deepmd.tf.env.tf.Graph, graph_def: deepmd.tf.env.tf.GraphDef, suffix: str = '') → None[source]#

Init the embedding net variables with the given dict.

Parameters:

graphtf.Graph: The input frozen model graph
graph_deftf.GraphDef: The input frozen model graph_def
suffixstr, optional: The suffix of the scope

Notes

This method is called by others when the descriptor supported initialization from the given variables.

abstract get_tensor_names(suffix: str = '') → tuple[str][source]#

Get names of tensors.

Parameters:

suffixstr: The suffix of the scope

Returns:

tuple[str]: Names of tensors

abstract pass_tensors_from_frz_model(*tensors: deepmd.tf.env.tf.Tensor) → None[source]#

Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.

Parameters:

*tensorstf.Tensor: passed tensors

Notes

The number of parameters in the method must be equal to the numbers of returns in get_tensor_names().

build_type_exclude_mask(exclude_types: set[tuple[int, int]], ntypes: int, sel: list[int], ndescrpt: int, atype: deepmd.tf.env.tf.Tensor, shape0: deepmd.tf.env.tf.Tensor) → deepmd.tf.env.tf.Tensor[source]#

Build the type exclude mask for the descriptor.

Parameters:

exclude_typeslist[tuple[int, int]]: The list of excluded types, e.g. [(0, 1), (1, 0)] means the interaction between type 0 and type 1 is excluded.
ntypesint: The number of types.
sellist[int]: The list of the number of selected neighbors for each type.
ndescrptint: The number of descriptors for each atom.
atypetf.Tensor: The type of atoms, with the size of shape0.
shape0tf.Tensor: The shape of the first dimension of the inputs, which is equal to nsamples * natoms.

Returns:

tf.Tensor: The type exclude mask, with the shape of (shape0, ndescrpt), and the precision of GLOBAL_TF_FLOAT_PRECISION. The mask has the value of 1 if the interaction between two types is not excluded, and 0 otherwise.

Notes

To exclude the interaction between two types, the derivative of energy with respect to distances (or angles) between two atoms should be zero[Rdfe82bee38f2-1]_, i.e.

\[\forall i \in \text{type 1}, j \in \text{type 2}, \frac{\partial{E}}{\partial{r_{ij}}} = 0\]

When embedding networks between every two types are built, we can just remove that network. But when type_one_side is enabled, a network may be built for multiple pairs of types. In this case, we need to build a mask to exclude the interaction between two types.

The mask assumes the descriptors are sorted by neighbro type with the fixed number of given sel and each neighbor has the same number of descriptors (for example 4).

References

[1]

Jinzhe Zeng, Timothy J. Giese, ̧Sölen Ekesan, Darrin M. York, Development of Range-Corrected Deep Learning Potentials for Fast, Accurate Quantum Mechanical/molecular Mechanical Simulations of Chemical Reactions in Solution, J. Chem. Theory Comput., 2021, 17 (11), 6993-7009.

property explicit_ntypes: bool[source]#: Explicit ntypes with type embedding.

classmethod update_sel(train_data: deepmd.utils.data_system.DeepmdDataSystem, type_map: list[str] | None, local_jdata: dict) → tuple[dict, float | None][source]#

Abstractmethod:

Update the selection and perform neighbor statistics.

Parameters:

train_dataDeepmdDataSystem: data used to do neighbor statistics
type_maplist[str], optional: The name of each type of atoms
local_jdatadict: The local data refer to the current class

Returns:

dict: The updated local data
float: The minimum distance between two atoms

classmethod deserialize(data: dict, suffix: str = '') → Descriptor[source]#

Deserialize the model.

There is no suffix in a native DP model, but it is important for the TF backend.

Parameters:

datadict: The serialized data
suffixstr, optional: Name suffix to identify this descriptor

Returns:

Descriptor: The deserialized descriptor

abstract serialize(suffix: str = '') → dict[source]#

Serialize the model.

There is no suffix in a native DP model, but it is important for the TF backend.

Returns:

dict: The serialized data
suffixstr, optional: Name suffix to identify this descriptor

property input_requirement: list[deepmd.utils.data.DataRequirementItem][source]#: Return data requirements needed for the model input.