deepmd.descriptor package

class deepmd.descriptor.Descriptor(*args, **kwargs)[source]

Bases: PluginVariant

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., rcut_smth=0.5, sel=[50])
>>> type(descript)
<class 'deepmd.descriptor.se_a.DescrptSeA'>
Attributes
explicit_ntypes

Explicit ntypes with type embedding.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

pass_tensors_from_frz_model(*tensors)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

abstract build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: Dict[str, Any], reuse: Optional[bool] = None, suffix: str = '') 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.

build_type_exclude_mask(exclude_types: List[Tuple[int, int]], ntypes: int, sel: List[int], ndescrpt: int, atype: Tensor, shape0: Tensor) 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[R08579741114c-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.

abstract compute_input_stats(data_coord: List[ndarray], data_box: List[ndarray], data_atype: List[ndarray], natoms_vec: List[ndarray], mesh: List[ndarray], input_dict: Dict[str, List[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.model.model_stat.make_stat_input()

data_boxlist[np.ndarray]

The box. Can be generated by deepmd.model.model_stat.make_stat_input()

data_atypelist[np.ndarray]

The atom types. Can be generated by deepmd.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.model.model_stat.make_stat_input()

meshlist[np.ndarray]

The mesh for neighbor searching. Can be generated by deepmd.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.

enable_compression(min_nbor_dist: float, graph: Graph, graph_def: 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]

Reveive 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.

enable_mixed_precision(mixed_prec: Optional[dict] = None) None[source]

Reveive 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.

property explicit_ntypes: bool

Explicit ntypes with type embedding.

classmethod get_class_by_input(input: dict)[source]
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.

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

get_nlist() Tuple[Tensor, 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 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_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.

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

init_variables(graph: Graph, graph_def: 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.

pass_tensors_from_frz_model(*tensors: 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().

abstract prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, 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

static register(key: str) Callable[source]

Register a descriptor plugin.

Parameters
keystr

the key of a descriptor

Returns
Descriptor

the registered descriptor

Examples

>>> @Descriptor.register("some_descrpt")
    class SomeDescript(Descriptor):
        pass
abstract classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]

Update the selection and perform neighbor statistics.

Parameters
global_jdatadict

The global data, containing the training section

local_jdatadict

The local data refer to the current class

class deepmd.descriptor.DescrptHybrid(*args, **kwargs)[source]

Bases: Descriptor

Concate a list of descriptors to form a new descriptor.

Parameters
listlist

Build a descriptor from the concatenation of the list of descriptors.

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Get the neighbor information of the descriptor, returns the nlist of the descriptor with the largest cut-off radius.

get_nlist_i(ii)

Get the neighbor information of the ii-th descriptor.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

pass_tensors_from_frz_model(*tensors)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') Tensor[source]

Build the computational graph for the descriptor.

Parameters
coord_

The coordinate of atoms

atype_

The type of atoms

natoms

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 the system

mesh

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_dict

Dictionary for additional inputs

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns
descriptor

The output descriptor

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, mixed_type: bool = False, real_natoms_vec: Optional[list] = None, **kwargs) None[source]

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

Parameters
data_coord

The coordinates. Can be generated by deepmd.model.make_stat_input

data_box

The box. Can be generated by deepmd.model.make_stat_input

data_atype

The atom types. Can be generated by deepmd.model.make_stat_input

natoms_vec

The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input

mesh

The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input

input_dict

Dictionary for additional input

mixed_type

Whether to perform the mixed_type mode. If True, the input data has the mixed_type format (see doc/model/train_se_atten.md), in which frames in a system may have different natoms_vec(s), with the same nloc.

real_natoms_vec

If mixed_type is True, it takes in the real natoms_vec for each frame.

**kwargs

Additional keyword arguments.

enable_compression(min_nbor_dist: float, graph: Graph, graph_def: 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]

Reveive 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_def 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

enable_mixed_precision(mixed_prec: Optional[dict] = None) None[source]

Reveive the mixed precision setting.

Parameters
mixed_prec

The mixed precision setting used in the embedding net

property explicit_ntypes: bool

Explicit ntypes with type embedding.

get_dim_out() int[source]

Returns the output dimension of this descriptor.

get_nlist() Tuple[Tensor, Tensor, List[int], List[int]][source]

Get the neighbor information of the descriptor, returns the nlist of the descriptor with the largest cut-off radius.

Returns
nlist

Neighbor list

rij

The relative distance between the neighbor and the center atom.

sel_a

The number of neighbors with full information

sel_r

The number of neighbors with only radial information

get_nlist_i(ii: int) Tuple[Tensor, Tensor, List[int], List[int]][source]

Get the neighbor information of the ii-th descriptor.

Parameters
iiint

The index of the descriptor

Returns
nlist

Neighbor list

rij

The relative distance between the neighbor and the center atom.

sel_a

The number of neighbors with full information

sel_r

The number of neighbors with only radial information

get_ntypes() int[source]

Returns the number of atom types.

get_rcut() float[source]

Returns the cut-off radius.

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

init_variables(graph: Graph, graph_def: 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

merge_input_stats(stat_dict)[source]

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

Parameters
stat_dict

The dict of statisitcs computed from compute_input_stats, including:

sumr

The sum of radial statisitcs.

suma

The sum of relative coord statisitcs.

sumn

The sum of neighbor numbers.

sumr2

The sum of square of radial statisitcs.

suma2

The sum of square of relative coord statisitcs.

pass_tensors_from_frz_model(*tensors: Tensor) None[source]

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

Parameters
*tensorstf.Tensor

passed tensors

prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters
atom_ener

The atomic energy

natoms

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
force

The force on atoms

virial

The total virial

atom_virial

The atomic virial

classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]

Update the selection and perform neighbor statistics.

Parameters
global_jdatadict

The global data, containing the training section

local_jdatadict

The local data refer to the current class

class deepmd.descriptor.DescrptLocFrame(*args, **kwargs)[source]

Bases: Descriptor

Defines a local frame at each atom, and the compute the descriptor as local coordinates under this frame.

Parameters
rcut

The cut-off radius

sel_alist[str]

The length of the list should be the same as the number of atom types in the system. sel_a[i] gives the selected number of type-i neighbors. The full relative coordinates of the neighbors are used by the descriptor.

sel_rlist[str]

The length of the list should be the same as the number of atom types in the system. sel_r[i] gives the selected number of type-i neighbors. Only relative distance of the neighbors are used by the descriptor. sel_a[i] + sel_r[i] is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius.

axis_rule: list[int]

The length should be 6 times of the number of types. - axis_rule[i*6+0]: class of the atom defining the first axis of type-i atom. 0 for neighbors with full coordinates and 1 for neighbors only with relative distance. - axis_rule[i*6+1]: type of the atom defining the first axis of type-i atom. - axis_rule[i*6+2]: index of the axis atom defining the first axis. Note that the neighbors with the same class and type are sorted according to their relative distance. - axis_rule[i*6+3]: class of the atom defining the second axis of type-i atom. 0 for neighbors with full coordinates and 1 for neighbors only with relative distance. - axis_rule[i*6+4]: type of the atom defining the second axis of type-i atom. - axis_rule[i*6+5]: index of the axis atom defining the second axis. Note that the neighbors with the same class and type are sorted according to their relative distance.

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_rot_mat()

Get rotational matrix.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

pass_tensors_from_frz_model(*tensors)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') Tensor[source]

Build the computational graph for the descriptor.

Parameters
coord_

The coordinate of atoms

atype_

The type of atoms

natoms

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 the system

mesh

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_dict

Dictionary for additional inputs

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns
descriptor

The output descriptor

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, **kwargs) None[source]

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

Parameters
data_coord

The coordinates. Can be generated by deepmd.model.make_stat_input

data_box

The box. Can be generated by deepmd.model.make_stat_input

data_atype

The atom types. Can be generated by deepmd.model.make_stat_input

natoms_vec

The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input

mesh

The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input

input_dict

Dictionary for additional input

**kwargs

Additional keyword arguments.

get_dim_out() int[source]

Returns the output dimension of this descriptor.

get_nlist() Tuple[Tensor, Tensor, List[int], List[int]][source]
Returns
nlist

Neighbor list

rij

The relative distance between the neighbor and the center atom.

sel_a

The number of neighbors with full information

sel_r

The number of neighbors with only radial information

get_ntypes() int[source]

Returns the number of atom types.

get_rcut() float[source]

Returns the cut-off radius.

get_rot_mat() Tensor[source]

Get rotational matrix.

init_variables(graph: Graph, graph_def: 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

prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters
atom_ener

The atomic energy

natoms

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
force

The force on atoms

virial

The total virial

atom_virial

The atomic virial

classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]

Update the selection and perform neighbor statistics.

Parameters
global_jdatadict

The global data, containing the training section

local_jdatadict

The local data refer to the current class

class deepmd.descriptor.DescrptSeA(*args, **kwargs)[source]

Bases: DescrptSe

DeepPot-SE constructed from all information (both angular and radial) of atomic configurations. The embedding takes the distance between atoms as input.

The descriptor \(\mathcal{D}^i \in \mathcal{R}^{M_1 \times M_2}\) is given by [1]

\[\mathcal{D}^i = (\mathcal{G}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \mathcal{G}^i_<\]

where \(\mathcal{R}^i \in \mathbb{R}^{N \times 4}\) is the coordinate matrix, and each row of \(\mathcal{R}^i\) can be constructed as follows

\[(\mathcal{R}^i)_j = [ \begin{array}{c} s(r_{ji}) & \frac{s(r_{ji})x_{ji}}{r_{ji}} & \frac{s(r_{ji})y_{ji}}{r_{ji}} & \frac{s(r_{ji})z_{ji}}{r_{ji}} \end{array} ]\]

where \(\mathbf{R}_{ji}=\mathbf{R}_j-\mathbf{R}_i = (x_{ji}, y_{ji}, z_{ji})\) is the relative coordinate and \(r_{ji}=\lVert \mathbf{R}_{ji} \lVert\) is its norm. The switching function \(s(r)\) is defined as:

\[\begin{split}s(r)= \begin{cases} \frac{1}{r}, & r<r_s \\ \frac{1}{r} \{ {(\frac{r - r_s}{ r_c - r_s})}^3 (-6 {(\frac{r - r_s}{ r_c - r_s})}^2 +15 \frac{r - r_s}{ r_c - r_s} -10) +1 \}, & r_s \leq r<r_c \\ 0, & r \geq r_c \end{cases}\end{split}\]

Each row of the embedding matrix \(\mathcal{G}^i \in \mathbb{R}^{N \times M_1}\) consists of outputs of a embedding network \(\mathcal{N}\) of \(s(r_{ji})\):

\[(\mathcal{G}^i)_j = \mathcal{N}(s(r_{ji}))\]

\(\mathcal{G}^i_< \in \mathbb{R}^{N \times M_2}\) takes first \(M_2\) columns of \(\mathcal{G}^i\). The equation of embedding network \(\mathcal{N}\) can be found at deepmd.utils.network.embedding_net().

Parameters
rcut

The cut-off radius \(r_c\)

rcut_smth

From where the environment matrix should be smoothed \(r_s\)

sellist[str]

sel[i] specifies the maxmum number of type i atoms in the cut-off radius

neuronlist[int]

Number of neurons in each hidden layers of the embedding net \(\mathcal{N}\)

axis_neuron

Number of the axis neuron \(M_2\) (number of columns of the sub-matrix of the embedding matrix)

resnet_dt

Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)

trainable

If the weights of embedding net are trainable.

seed

Random seed for initializing the network parameters.

type_one_side

Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets

exclude_typesList[List[int]]

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

set_davg_zero

Set the shift of embedding net input to zero.

activation_function

The activation function in the embedding net. Supported options are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.

precision

The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.

uniform_seed

Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed

multi_task

If the model has multi fitting nets to train.

References

1

Linfeng Zhang, Jiequn Han, Han Wang, Wissam A. Saidi, Roberto Car, and E. Weinan. 2018. End-to-end symmetry preserving inter-atomic potential energy model for finite and extended systems. In Proceedings of the 32nd International Conference on Neural Information Processing Systems (NIPS’18). Curran Associates Inc., Red Hook, NY, USA, 4441-4451.

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

precision

Precision of filter network.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_rot_mat()

Get rotational matrix.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') Tensor[source]

Build the computational graph for the descriptor.

Parameters
coord_

The coordinate of atoms

atype_

The type of atoms

natoms

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 the system

mesh

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_dict

Dictionary for additional inputs

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns
descriptor

The output descriptor

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, **kwargs) None[source]

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

Parameters
data_coord

The coordinates. Can be generated by deepmd.model.make_stat_input

data_box

The box. Can be generated by deepmd.model.make_stat_input

data_atype

The atom types. Can be generated by deepmd.model.make_stat_input

natoms_vec

The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input

mesh

The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input

input_dict

Dictionary for additional input

**kwargs

Additional keyword arguments.

enable_compression(min_nbor_dist: float, graph: Graph, graph_def: GraphDef, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') None[source]

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

Parameters
min_nbor_dist

The nearest distance between atoms

graphtf.Graph

The graph of the model

graph_deftf.GraphDef

The graph_def of the model

table_extrapolate

The scale of model extrapolation

table_stride_1

The uniform stride of the first table

table_stride_2

The uniform stride of the second table

check_frequency

The overflow check frequency

suffixstr, optional

The suffix of the scope

enable_mixed_precision(mixed_prec: Optional[dict] = None) None[source]

Reveive the mixed precision setting.

Parameters
mixed_prec

The mixed precision setting used in the embedding net

property explicit_ntypes: bool

Explicit ntypes with type embedding.

get_dim_out() int[source]

Returns the output dimension of this descriptor.

get_dim_rot_mat_1() int[source]

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

get_nlist() Tuple[Tensor, Tensor, List[int], List[int]][source]

Returns neighbor information.

Returns
nlist

Neighbor list

rij

The relative distance between the neighbor and the center atom.

sel_a

The number of neighbors with full information

sel_r

The number of neighbors with only radial information

get_ntypes() int[source]

Returns the number of atom types.

get_rcut() float[source]

Returns the cut-off radius.

get_rot_mat() Tensor[source]

Get rotational matrix.

init_variables(graph: Graph, graph_def: 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

merge_input_stats(stat_dict)[source]

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

Parameters
stat_dict

The dict of statisitcs computed from compute_input_stats, including:

sumr

The sum of radial statisitcs.

suma

The sum of relative coord statisitcs.

sumn

The sum of neighbor numbers.

sumr2

The sum of square of radial statisitcs.

suma2

The sum of square of relative coord statisitcs.

prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters
atom_ener

The atomic energy

natoms

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
force

The force on atoms

virial

The total virial

atom_virial

The atomic virial

class deepmd.descriptor.DescrptSeAEbd(*args, **kwargs)[source]

Bases: DescrptSeA

DeepPot-SE descriptor with type embedding approach.

Parameters
rcut

The cut-off radius

rcut_smth

From where the environment matrix should be smoothed

sellist[str]

sel[i] specifies the maxmum number of type i atoms in the cut-off radius

neuronlist[int]

Number of neurons in each hidden layers of the embedding net

axis_neuron

Number of the axis neuron (number of columns of the sub-matrix of the embedding matrix)

resnet_dt

Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)

trainable

If the weights of embedding net are trainable.

seed

Random seed for initializing the network parameters.

type_one_side

Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets

type_nchanl

Number of channels for type representation

type_nlayer

Number of hidden layers for the type embedding net (skip connected).

numb_aparam

Number of atomic parameters. If >0 it will be embedded with atom types.

set_davg_zero

Set the shift of embedding net input to zero.

activation_function

The activation function in the embedding net. Supported options are {0}

precision

The precision of the embedding net parameters. Supported options are {1}

exclude_typesList[List[int]]

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

precision

Precision of filter network.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_rot_mat()

Get rotational matrix.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') Tensor[source]

Build the computational graph for the descriptor.

Parameters
coord_

The coordinate of atoms

atype_

The type of atoms

natoms

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 the system

mesh

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_dict

Dictionary for additional inputs

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns
descriptor

The output descriptor

class deepmd.descriptor.DescrptSeAEbdV2(*args, **kwargs)[source]

Bases: DescrptSeA

A compressible se_a_ebd model.

This model is a warpper for DescriptorSeA, which set stripped_type_embedding=True.

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

precision

Precision of filter network.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_rot_mat()

Get rotational matrix.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

class deepmd.descriptor.DescrptSeAEf(*args, **kwargs)[source]

Bases: DescrptSe

Smooth edition descriptor with Ef.

Parameters
rcut

The cut-off radius

rcut_smth

From where the environment matrix should be smoothed

sellist[str]

sel[i] specifies the maxmum number of type i atoms in the cut-off radius

neuronlist[int]

Number of neurons in each hidden layers of the embedding net

axis_neuron

Number of the axis neuron (number of columns of the sub-matrix of the embedding matrix)

resnet_dt

Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)

trainable

If the weights of embedding net are trainable.

seed

Random seed for initializing the network parameters.

type_one_side

Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets

exclude_typesList[List[int]]

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

set_davg_zero

Set the shift of embedding net input to zero.

activation_function

The activation function in the embedding net. Supported options are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.

precision

The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.

uniform_seed

Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

precision

Precision of filter network.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_rot_mat()

Get rotational matrix.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') Tensor[source]

Build the computational graph for the descriptor.

Parameters
coord_

The coordinate of atoms

atype_

The type of atoms

natoms

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 the system

mesh

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_dict

Dictionary for additional inputs. Should have ‘efield’.

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns
descriptor

The output descriptor

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, **kwargs) None[source]

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

Parameters
data_coord

The coordinates. Can be generated by deepmd.model.make_stat_input

data_box

The box. Can be generated by deepmd.model.make_stat_input

data_atype

The atom types. Can be generated by deepmd.model.make_stat_input

natoms_vec

The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input

mesh

The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input

input_dict

Dictionary for additional input

**kwargs

Additional keyword arguments.

get_dim_out() int[source]

Returns the output dimension of this descriptor.

get_dim_rot_mat_1() int[source]

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

get_nlist() Tuple[Tensor, Tensor, List[int], List[int]][source]

Returns neighbor information.

Returns
nlist

Neighbor list

rij

The relative distance between the neighbor and the center atom.

sel_a

The number of neighbors with full information

sel_r

The number of neighbors with only radial information

get_ntypes() int[source]

Returns the number of atom types.

get_rcut() float[source]

Returns the cut-off radius.

get_rot_mat() Tensor[source]

Get rotational matrix.

prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters
atom_ener

The atomic energy

natoms

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
force

The force on atoms

virial

The total virial

atom_virial

The atomic virial

class deepmd.descriptor.DescrptSeAEfLower(*args, **kwargs)[source]

Bases: DescrptSeA

Helper class for implementing DescrptSeAEf.

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

precision

Precision of filter network.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_rot_mat()

Get rotational matrix.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

build(coord_, atype_, natoms, box_, mesh, input_dict, suffix='', reuse=None)[source]

Build the computational graph for the descriptor.

Parameters
coord_

The coordinate of atoms

atype_

The type of atoms

natoms

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 the system

mesh

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_dict

Dictionary for additional inputs

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns
descriptor

The output descriptor

compute_input_stats(data_coord, data_box, data_atype, natoms_vec, mesh, input_dict, **kwargs)[source]

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

Parameters
data_coord

The coordinates. Can be generated by deepmd.model.make_stat_input

data_box

The box. Can be generated by deepmd.model.make_stat_input

data_atype

The atom types. Can be generated by deepmd.model.make_stat_input

natoms_vec

The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input

mesh

The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input

input_dict

Dictionary for additional input

**kwargs

Additional keyword arguments.

class deepmd.descriptor.DescrptSeAMask(*args, **kwargs)[source]

Bases: DescrptSeA

DeepPot-SE constructed from all information (both angular and radial) of atomic configurations. The embedding takes the distance between atoms as input.

The descriptor \(\mathcal{D}^i \in \mathcal{R}^{M_1 \times M_2}\) is given by [1]

\[\mathcal{D}^i = (\mathcal{G}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \mathcal{G}^i_<\]

where \(\mathcal{R}^i \in \mathbb{R}^{N \times 4}\) is the coordinate matrix, and each row of \(\mathcal{R}^i\) can be constructed as follows

\[(\mathcal{R}^i)_j = [ \begin{array}{c} s(r_{ji}) & \frac{s(r_{ji})x_{ji}}{r_{ji}} & \frac{s(r_{ji})y_{ji}}{r_{ji}} & \frac{s(r_{ji})z_{ji}}{r_{ji}} \end{array} ]\]

where \(\mathbf{R}_{ji}=\mathbf{R}_j-\mathbf{R}_i = (x_{ji}, y_{ji}, z_{ji})\) is the relative coordinate and \(r_{ji}=\lVert \mathbf{R}_{ji} \lVert\) is its norm. The switching function \(s(r)\) is defined as:

\[\begin{split}s(r)= \begin{cases} \frac{1}{r}, & r<r_s \\ \frac{1}{r} \{ {(\frac{r - r_s}{ r_c - r_s})}^3 (-6 {(\frac{r - r_s}{ r_c - r_s})}^2 +15 \frac{r - r_s}{ r_c - r_s} -10) +1 \}, & r_s \leq r<r_c \\ 0, & r \geq r_c \end{cases}\end{split}\]

Each row of the embedding matrix \(\mathcal{G}^i \in \mathbb{R}^{N \times M_1}\) consists of outputs of a embedding network \(\mathcal{N}\) of \(s(r_{ji})\):

\[(\mathcal{G}^i)_j = \mathcal{N}(s(r_{ji}))\]

\(\mathcal{G}^i_< \in \mathbb{R}^{N \times M_2}\) takes first \(M_2\) columns of \(\mathcal{G}^i\). The equation of embedding network \(\mathcal{N}\) can be found at deepmd.utils.network.embedding_net(). Specially for descriptor se_a_mask is a concise implementation of se_a. The difference is that se_a_mask only considered a non-pbc system. And accept a mask matrix to indicate the atom i in frame j is a real atom or not. (1 means real atom, 0 means ghost atom) Thus se_a_mask can accept a variable number of atoms in a frame.

Parameters
sellist[str]

sel[i] specifies the maxmum number of type i atoms in the neighbor list.

neuronlist[int]

Number of neurons in each hidden layers of the embedding net \(\mathcal{N}\)

axis_neuron

Number of the axis neuron \(M_2\) (number of columns of the sub-matrix of the embedding matrix)

resnet_dt

Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)

trainable

If the weights of embedding net are trainable.

seed

Random seed for initializing the network parameters.

type_one_side

Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets

exclude_typesList[List[int]]

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

activation_function

The activation function in the embedding net. Supported options are {0}

precision

The precision of the embedding net parameters. Supported options are {1}

uniform_seed

Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed

References

1

Linfeng Zhang, Jiequn Han, Han Wang, Wissam A. Saidi, Roberto Car, and E. Weinan. 2018. End-to-end symmetry preserving inter-atomic potential energy model for finite and extended systems. In Proceedings of the 32nd International Conference on Neural Information Processing Systems (NIPS’18). Curran Associates Inc., Red Hook, NY, USA, 4441-4451.

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

precision

Precision of filter network.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cutoff radius.

get_rot_mat()

Get rotational matrix.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: Dict[str, Any], reuse: Optional[bool] = None, suffix: str = '') Tensor[source]

Build the computational graph for the descriptor.

Parameters
coord_

The coordinate of atoms

atype_

The type of atoms

natoms

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 the system

mesh

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_dict

Dictionary for additional inputs

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns
descriptor

The output descriptor

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, **kwargs) None[source]

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

Parameters
data_coord

The coordinates. Can be generated by deepmd.model.make_stat_input

data_box

The box. Can be generated by deepmd.model.make_stat_input

data_atype

The atom types. Can be generated by deepmd.model.make_stat_input

natoms_vec

The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input

mesh

The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input

input_dict

Dictionary for additional input

**kwargs

Additional keyword arguments.

get_rcut() float[source]

Returns the cutoff radius.

prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters
atom_ener

The atomic energy

natoms

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
force

The force on atoms

virial

None for se_a_mask op

atom_virial

None for se_a_mask op

classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]

Update the selection and perform neighbor statistics.

Parameters
global_jdatadict

The global data, containing the training section

local_jdatadict

The local data refer to the current class

class deepmd.descriptor.DescrptSeAtten(*args, **kwargs)[source]

Bases: DescrptSeA

Smooth version descriptor with attention.

Parameters
rcut

The cut-off radius \(r_c\)

rcut_smth

From where the environment matrix should be smoothed \(r_s\)

sellist[str]

sel[i] specifies the maxmum number of type i atoms in the cut-off radius

neuronlist[int]

Number of neurons in each hidden layers of the embedding net \(\mathcal{N}\)

axis_neuron

Number of the axis neuron \(M_2\) (number of columns of the sub-matrix of the embedding matrix)

resnet_dt

Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)

trainable

If the weights of embedding net are trainable.

seed

Random seed for initializing the network parameters.

type_one_side

Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets

exclude_typesList[List[int]]

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

set_davg_zero

Set the shift of embedding net input to zero.

activation_function

The activation function in the embedding net. Supported options are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.

precision

The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.

uniform_seed

Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed

attn

The length of hidden vector during scale-dot attention computation.

attn_layer

The number of layers in attention mechanism.

attn_dotr

Whether to dot the relative coordinates on the attention weights as a gated scheme.

attn_mask

Whether to mask the diagonal in the attention weights.

multi_task

If the model has multi fitting nets to train.

stripped_type_embedding

Whether to strip the type embedding into a separated embedding network. Default value will be True in se_atten_v2 descriptor.

smooth_type_embdding

When using stripped type embedding, whether to dot smooth factor on the network output of type embedding to keep the network smooth, instead of setting set_davg_zero to be True. Default value will be True in se_atten_v2 descriptor.

Raises
ValueError

if ntypes is 0.

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

precision

Precision of filter network.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the attention descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_rot_mat()

Get rotational matrix.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') Tensor[source]

Build the computational graph for the descriptor.

Parameters
coord_

The coordinate of atoms

atype_

The type of atoms

natoms

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 the system

mesh

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_dict

Dictionary for additional inputs

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns
descriptor

The output descriptor

build_type_exclude_mask(exclude_types: List[Tuple[int, int]], ntypes: int, sel: List[int], ndescrpt: int, atype: Tensor, shape0: Tensor, nei_type_vec: Tensor) Tensor[source]

Build the type exclude mask for the attention 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.

nei_type_vectf.Tensor

The type of neighbors, with the size of (shape0, nnei).

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

This method has the similiar way to build the type exclude mask as deepmd.descriptor.descriptor.Descriptor.build_type_exclude_mask(). The mathmatical expression has been explained in that method. The difference is that the attention descriptor has provided the type of the neighbors (idx_j) that is not in order, so we use it from an extra input.

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, mixed_type: bool = False, real_natoms_vec: Optional[list] = None, **kwargs) None[source]

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

Parameters
data_coord

The coordinates. Can be generated by deepmd.model.make_stat_input

data_box

The box. Can be generated by deepmd.model.make_stat_input

data_atype

The atom types. Can be generated by deepmd.model.make_stat_input

natoms_vec

The vector for the number of atoms of the system and different types of atoms. If mixed_type is True, this para is blank. See real_natoms_vec.

mesh

The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input

input_dict

Dictionary for additional input

mixed_type

Whether to perform the mixed_type mode. If True, the input data has the mixed_type format (see doc/model/train_se_atten.md), in which frames in a system may have different natoms_vec(s), with the same nloc.

real_natoms_vec

If mixed_type is True, it takes in the real natoms_vec for each frame.

**kwargs

Additional keyword arguments.

enable_compression(min_nbor_dist: float, graph: Graph, graph_def: GraphDef, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') None[source]

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

Parameters
min_nbor_dist

The nearest distance between atoms

graphtf.Graph

The graph of the model

graph_deftf.GraphDef

The graph_def of the model

table_extrapolate

The scale of model extrapolation

table_stride_1

The uniform stride of the first table

table_stride_2

The uniform stride of the second table

check_frequency

The overflow check frequency

suffixstr, optional

The suffix of the scope

property explicit_ntypes: bool

Explicit ntypes with type embedding.

init_variables(graph: Graph, graph_def: 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

classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]

Update the selection and perform neighbor statistics.

Parameters
global_jdatadict

The global data, containing the training section

local_jdatadict

The local data refer to the current class

class deepmd.descriptor.DescrptSeAttenV2(*args, **kwargs)[source]

Bases: DescrptSeAtten

Smooth version 2.0 descriptor with attention.

Parameters
rcut

The cut-off radius \(r_c\)

rcut_smth

From where the environment matrix should be smoothed \(r_s\)

sellist[str]

sel[i] specifies the maxmum number of type i atoms in the cut-off radius

neuronlist[int]

Number of neurons in each hidden layers of the embedding net \(\mathcal{N}\)

axis_neuron

Number of the axis neuron \(M_2\) (number of columns of the sub-matrix of the embedding matrix)

resnet_dt

Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)

trainable

If the weights of embedding net are trainable.

seed

Random seed for initializing the network parameters.

type_one_side

Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets

exclude_typesList[List[int]]

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

set_davg_zero

Set the shift of embedding net input to zero.

activation_function

The activation function in the embedding net. Supported options are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.

precision

The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.

uniform_seed

Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed

attn

The length of hidden vector during scale-dot attention computation.

attn_layer

The number of layers in attention mechanism.

attn_dotr

Whether to dot the relative coordinates on the attention weights as a gated scheme.

attn_mask

Whether to mask the diagonal in the attention weights.

multi_task

If the model has multi fitting nets to train.

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

precision

Precision of filter network.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the attention descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_rot_mat()

Get rotational matrix.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

class deepmd.descriptor.DescrptSeR(*args, **kwargs)[source]

Bases: DescrptSe

DeepPot-SE constructed from radial information of atomic configurations.

The embedding takes the distance between atoms as input.

Parameters
rcut

The cut-off radius

rcut_smth

From where the environment matrix should be smoothed

sellist[str]

sel[i] specifies the maxmum number of type i atoms in the cut-off radius

neuronlist[int]

Number of neurons in each hidden layers of the embedding net

resnet_dt

Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)

trainable

If the weights of embedding net are trainable.

seed

Random seed for initializing the network parameters.

type_one_side

Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets

exclude_typesList[List[int]]

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

activation_function

The activation function in the embedding net. Supported options are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.

precision

The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.

uniform_seed

Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

precision

Precision of filter network.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') Tensor[source]

Build the computational graph for the descriptor.

Parameters
coord_

The coordinate of atoms

atype_

The type of atoms

natoms

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 the system

mesh

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_dict

Dictionary for additional inputs

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns
descriptor

The output descriptor

compute_input_stats(data_coord, data_box, data_atype, natoms_vec, mesh, input_dict, **kwargs)[source]

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

Parameters
data_coord

The coordinates. Can be generated by deepmd.model.make_stat_input

data_box

The box. Can be generated by deepmd.model.make_stat_input

data_atype

The atom types. Can be generated by deepmd.model.make_stat_input

natoms_vec

The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input

mesh

The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input

input_dict

Dictionary for additional input

**kwargs

Additional keyword arguments.

enable_compression(min_nbor_dist: float, graph: Graph, graph_def: GraphDef, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') None[source]

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

Parameters
min_nbor_dist

The nearest distance between atoms

graphtf.Graph

The graph of the model

graph_deftf.GraphDef

The graph_def of the model

table_extrapolate

The scale of model extrapolation

table_stride_1

The uniform stride of the first table

table_stride_2

The uniform stride of the second table

check_frequency

The overflow check frequency

suffixstr, optional

The suffix of the scope

get_dim_out()[source]

Returns the output dimension of this descriptor.

get_nlist()[source]

Returns neighbor information.

Returns
nlist

Neighbor list

rij

The relative distance between the neighbor and the center atom.

sel_a

The number of neighbors with full information

sel_r

The number of neighbors with only radial information

get_ntypes()[source]

Returns the number of atom types.

get_rcut()[source]

Returns the cut-off radius.

merge_input_stats(stat_dict)[source]

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

Parameters
stat_dict

The dict of statisitcs computed from compute_input_stats, including:

sumr

The sum of radial statisitcs.

sumn

The sum of neighbor numbers.

sumr2

The sum of square of radial statisitcs.

prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters
atom_ener

The atomic energy

natoms

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
force

The force on atoms

virial

The total virial

atom_virial

The atomic virial

class deepmd.descriptor.DescrptSeT(*args, **kwargs)[source]

Bases: DescrptSe

DeepPot-SE constructed from all information (both angular and radial) of atomic configurations.

The embedding takes angles between two neighboring atoms as input.

Parameters
rcut

The cut-off radius

rcut_smth

From where the environment matrix should be smoothed

sellist[str]

sel[i] specifies the maxmum number of type i atoms in the cut-off radius

neuronlist[int]

Number of neurons in each hidden layers of the embedding net

resnet_dt

Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)

trainable

If the weights of embedding net are trainable.

seed

Random seed for initializing the network parameters.

set_davg_zero

Set the shift of embedding net input to zero.

activation_function

The activation function in the embedding net. Supported options are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.

precision

The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.

uniform_seed

Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

precision

Precision of filter network.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') Tensor[source]

Build the computational graph for the descriptor.

Parameters
coord_

The coordinate of atoms

atype_

The type of atoms

natoms

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 the system

mesh

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_dict

Dictionary for additional inputs

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns
descriptor

The output descriptor

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, **kwargs) None[source]

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

Parameters
data_coord

The coordinates. Can be generated by deepmd.model.make_stat_input

data_box

The box. Can be generated by deepmd.model.make_stat_input

data_atype

The atom types. Can be generated by deepmd.model.make_stat_input

natoms_vec

The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input

mesh

The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input

input_dict

Dictionary for additional input

**kwargs

Additional keyword arguments.

enable_compression(min_nbor_dist: float, graph: Graph, graph_def: GraphDef, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') None[source]

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

Parameters
min_nbor_dist

The nearest distance between atoms

graphtf.Graph

The graph of the model

graph_deftf.GraphDef

The graph_def of the model

table_extrapolate

The scale of model extrapolation

table_stride_1

The uniform stride of the first table

table_stride_2

The uniform stride of the second table

check_frequency

The overflow check frequency

suffixstr, optional

The suffix of the scope

get_dim_out() int[source]

Returns the output dimension of this descriptor.

get_nlist() Tuple[Tensor, Tensor, List[int], List[int]][source]

Returns neighbor information.

Returns
nlist

Neighbor list

rij

The relative distance between the neighbor and the center atom.

sel_a

The number of neighbors with full information

sel_r

The number of neighbors with only radial information

get_ntypes() int[source]

Returns the number of atom types.

get_rcut() float[source]

Returns the cut-off radius.

merge_input_stats(stat_dict)[source]

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

Parameters
stat_dict

The dict of statisitcs computed from compute_input_stats, including:

sumr

The sum of radial statisitcs.

suma

The sum of relative coord statisitcs.

sumn

The sum of neighbor numbers.

sumr2

The sum of square of radial statisitcs.

suma2

The sum of square of relative coord statisitcs.

prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters
atom_ener

The atomic energy

natoms

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
force

The force on atoms

virial

The total virial

atom_virial

The atomic virial

Submodules

deepmd.descriptor.descriptor module

class deepmd.descriptor.descriptor.Descriptor(*args, **kwargs)[source]

Bases: PluginVariant

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., rcut_smth=0.5, sel=[50])
>>> type(descript)
<class 'deepmd.descriptor.se_a.DescrptSeA'>
Attributes
explicit_ntypes

Explicit ntypes with type embedding.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

pass_tensors_from_frz_model(*tensors)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

abstract build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: Dict[str, Any], reuse: Optional[bool] = None, suffix: str = '') 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.

build_type_exclude_mask(exclude_types: List[Tuple[int, int]], ntypes: int, sel: List[int], ndescrpt: int, atype: Tensor, shape0: Tensor) 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[Rafc1ae60e195-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.

abstract compute_input_stats(data_coord: List[ndarray], data_box: List[ndarray], data_atype: List[ndarray], natoms_vec: List[ndarray], mesh: List[ndarray], input_dict: Dict[str, List[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.model.model_stat.make_stat_input()

data_boxlist[np.ndarray]

The box. Can be generated by deepmd.model.model_stat.make_stat_input()

data_atypelist[np.ndarray]

The atom types. Can be generated by deepmd.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.model.model_stat.make_stat_input()

meshlist[np.ndarray]

The mesh for neighbor searching. Can be generated by deepmd.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.

enable_compression(min_nbor_dist: float, graph: Graph, graph_def: 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]

Reveive 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.

enable_mixed_precision(mixed_prec: Optional[dict] = None) None[source]

Reveive 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.

property explicit_ntypes: bool

Explicit ntypes with type embedding.

classmethod get_class_by_input(input: dict)[source]
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.

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

get_nlist() Tuple[Tensor, 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 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_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.

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

init_variables(graph: Graph, graph_def: 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.

pass_tensors_from_frz_model(*tensors: 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().

abstract prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, 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

static register(key: str) Callable[source]

Register a descriptor plugin.

Parameters
keystr

the key of a descriptor

Returns
Descriptor

the registered descriptor

Examples

>>> @Descriptor.register("some_descrpt")
    class SomeDescript(Descriptor):
        pass
abstract classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]

Update the selection and perform neighbor statistics.

Parameters
global_jdatadict

The global data, containing the training section

local_jdatadict

The local data refer to the current class

deepmd.descriptor.hybrid module

class deepmd.descriptor.hybrid.DescrptHybrid(*args, **kwargs)[source]

Bases: Descriptor

Concate a list of descriptors to form a new descriptor.

Parameters
listlist

Build a descriptor from the concatenation of the list of descriptors.

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Get the neighbor information of the descriptor, returns the nlist of the descriptor with the largest cut-off radius.

get_nlist_i(ii)

Get the neighbor information of the ii-th descriptor.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

pass_tensors_from_frz_model(*tensors)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') Tensor[source]

Build the computational graph for the descriptor.

Parameters
coord_

The coordinate of atoms

atype_

The type of atoms

natoms

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 the system

mesh

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_dict

Dictionary for additional inputs

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns
descriptor

The output descriptor

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, mixed_type: bool = False, real_natoms_vec: Optional[list] = None, **kwargs) None[source]

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

Parameters
data_coord

The coordinates. Can be generated by deepmd.model.make_stat_input

data_box

The box. Can be generated by deepmd.model.make_stat_input

data_atype

The atom types. Can be generated by deepmd.model.make_stat_input

natoms_vec

The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input

mesh

The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input

input_dict

Dictionary for additional input

mixed_type

Whether to perform the mixed_type mode. If True, the input data has the mixed_type format (see doc/model/train_se_atten.md), in which frames in a system may have different natoms_vec(s), with the same nloc.

real_natoms_vec

If mixed_type is True, it takes in the real natoms_vec for each frame.

**kwargs

Additional keyword arguments.

enable_compression(min_nbor_dist: float, graph: Graph, graph_def: 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]

Reveive 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_def 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

enable_mixed_precision(mixed_prec: Optional[dict] = None) None[source]

Reveive the mixed precision setting.

Parameters
mixed_prec

The mixed precision setting used in the embedding net

property explicit_ntypes: bool

Explicit ntypes with type embedding.

get_dim_out() int[source]

Returns the output dimension of this descriptor.

get_nlist() Tuple[Tensor, Tensor, List[int], List[int]][source]

Get the neighbor information of the descriptor, returns the nlist of the descriptor with the largest cut-off radius.

Returns
nlist

Neighbor list

rij

The relative distance between the neighbor and the center atom.

sel_a

The number of neighbors with full information

sel_r

The number of neighbors with only radial information

get_nlist_i(ii: int) Tuple[Tensor, Tensor, List[int], List[int]][source]

Get the neighbor information of the ii-th descriptor.

Parameters
iiint

The index of the descriptor

Returns
nlist

Neighbor list

rij

The relative distance between the neighbor and the center atom.

sel_a

The number of neighbors with full information

sel_r

The number of neighbors with only radial information

get_ntypes() int[source]

Returns the number of atom types.

get_rcut() float[source]

Returns the cut-off radius.

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

init_variables(graph: Graph, graph_def: 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

merge_input_stats(stat_dict)[source]

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

Parameters
stat_dict

The dict of statisitcs computed from compute_input_stats, including:

sumr

The sum of radial statisitcs.

suma

The sum of relative coord statisitcs.

sumn

The sum of neighbor numbers.

sumr2

The sum of square of radial statisitcs.

suma2

The sum of square of relative coord statisitcs.

pass_tensors_from_frz_model(*tensors: Tensor) None[source]

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

Parameters
*tensorstf.Tensor

passed tensors

prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters
atom_ener

The atomic energy

natoms

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
force

The force on atoms

virial

The total virial

atom_virial

The atomic virial

classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]

Update the selection and perform neighbor statistics.

Parameters
global_jdatadict

The global data, containing the training section

local_jdatadict

The local data refer to the current class

deepmd.descriptor.loc_frame module

class deepmd.descriptor.loc_frame.DescrptLocFrame(*args, **kwargs)[source]

Bases: Descriptor

Defines a local frame at each atom, and the compute the descriptor as local coordinates under this frame.

Parameters
rcut

The cut-off radius

sel_alist[str]

The length of the list should be the same as the number of atom types in the system. sel_a[i] gives the selected number of type-i neighbors. The full relative coordinates of the neighbors are used by the descriptor.

sel_rlist[str]

The length of the list should be the same as the number of atom types in the system. sel_r[i] gives the selected number of type-i neighbors. Only relative distance of the neighbors are used by the descriptor. sel_a[i] + sel_r[i] is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius.

axis_rule: list[int]

The length should be 6 times of the number of types. - axis_rule[i*6+0]: class of the atom defining the first axis of type-i atom. 0 for neighbors with full coordinates and 1 for neighbors only with relative distance. - axis_rule[i*6+1]: type of the atom defining the first axis of type-i atom. - axis_rule[i*6+2]: index of the axis atom defining the first axis. Note that the neighbors with the same class and type are sorted according to their relative distance. - axis_rule[i*6+3]: class of the atom defining the second axis of type-i atom. 0 for neighbors with full coordinates and 1 for neighbors only with relative distance. - axis_rule[i*6+4]: type of the atom defining the second axis of type-i atom. - axis_rule[i*6+5]: index of the axis atom defining the second axis. Note that the neighbors with the same class and type are sorted according to their relative distance.

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_rot_mat()

Get rotational matrix.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

pass_tensors_from_frz_model(*tensors)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') Tensor[source]

Build the computational graph for the descriptor.

Parameters
coord_

The coordinate of atoms

atype_

The type of atoms

natoms

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 the system

mesh

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_dict

Dictionary for additional inputs

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns
descriptor

The output descriptor

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, **kwargs) None[source]

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

Parameters
data_coord

The coordinates. Can be generated by deepmd.model.make_stat_input

data_box

The box. Can be generated by deepmd.model.make_stat_input

data_atype

The atom types. Can be generated by deepmd.model.make_stat_input

natoms_vec

The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input

mesh

The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input

input_dict

Dictionary for additional input

**kwargs

Additional keyword arguments.

get_dim_out() int[source]

Returns the output dimension of this descriptor.

get_nlist() Tuple[Tensor, Tensor, List[int], List[int]][source]
Returns
nlist

Neighbor list

rij

The relative distance between the neighbor and the center atom.

sel_a

The number of neighbors with full information

sel_r

The number of neighbors with only radial information

get_ntypes() int[source]

Returns the number of atom types.

get_rcut() float[source]

Returns the cut-off radius.

get_rot_mat() Tensor[source]

Get rotational matrix.

init_variables(graph: Graph, graph_def: 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

prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters
atom_ener

The atomic energy

natoms

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
force

The force on atoms

virial

The total virial

atom_virial

The atomic virial

classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]

Update the selection and perform neighbor statistics.

Parameters
global_jdatadict

The global data, containing the training section

local_jdatadict

The local data refer to the current class

deepmd.descriptor.se module

class deepmd.descriptor.se.DescrptSe(*args, **kwargs)[source]

Bases: Descriptor

A base class for smooth version of descriptors.

Notes

All of these descriptors have an environmental matrix and an embedding network (deepmd.utils.network.embedding_net()), so they can share some similiar methods without defining them twice.

Attributes
embedding_net_variablesdict

initial embedding network variables

descrpt_reshapetf.Tensor

the reshaped descriptor

descrpt_derivtf.Tensor

the descriptor derivative

rijtf.Tensor

distances between two atoms

nlisttf.Tensor

the neighbor list

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

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

init_variables(graph: Graph, graph_def: 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

pass_tensors_from_frz_model(descrpt_reshape: Tensor, descrpt_deriv: Tensor, rij: Tensor, nlist: Tensor)[source]

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

Parameters
descrpt_reshape

The passed descrpt_reshape tensor

descrpt_deriv

The passed descrpt_deriv tensor

rij

The passed rij tensor

nlist

The passed nlist tensor

property precision: DType

Precision of filter network.

classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]

Update the selection and perform neighbor statistics.

Parameters
global_jdatadict

The global data, containing the training section

local_jdatadict

The local data refer to the current class

deepmd.descriptor.se_a module

class deepmd.descriptor.se_a.DescrptSeA(*args, **kwargs)[source]

Bases: DescrptSe

DeepPot-SE constructed from all information (both angular and radial) of atomic configurations. The embedding takes the distance between atoms as input.

The descriptor \(\mathcal{D}^i \in \mathcal{R}^{M_1 \times M_2}\) is given by [1]

\[\mathcal{D}^i = (\mathcal{G}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \mathcal{G}^i_<\]

where \(\mathcal{R}^i \in \mathbb{R}^{N \times 4}\) is the coordinate matrix, and each row of \(\mathcal{R}^i\) can be constructed as follows

\[(\mathcal{R}^i)_j = [ \begin{array}{c} s(r_{ji}) & \frac{s(r_{ji})x_{ji}}{r_{ji}} & \frac{s(r_{ji})y_{ji}}{r_{ji}} & \frac{s(r_{ji})z_{ji}}{r_{ji}} \end{array} ]\]

where \(\mathbf{R}_{ji}=\mathbf{R}_j-\mathbf{R}_i = (x_{ji}, y_{ji}, z_{ji})\) is the relative coordinate and \(r_{ji}=\lVert \mathbf{R}_{ji} \lVert\) is its norm. The switching function \(s(r)\) is defined as:

\[\begin{split}s(r)= \begin{cases} \frac{1}{r}, & r<r_s \\ \frac{1}{r} \{ {(\frac{r - r_s}{ r_c - r_s})}^3 (-6 {(\frac{r - r_s}{ r_c - r_s})}^2 +15 \frac{r - r_s}{ r_c - r_s} -10) +1 \}, & r_s \leq r<r_c \\ 0, & r \geq r_c \end{cases}\end{split}\]

Each row of the embedding matrix \(\mathcal{G}^i \in \mathbb{R}^{N \times M_1}\) consists of outputs of a embedding network \(\mathcal{N}\) of \(s(r_{ji})\):

\[(\mathcal{G}^i)_j = \mathcal{N}(s(r_{ji}))\]

\(\mathcal{G}^i_< \in \mathbb{R}^{N \times M_2}\) takes first \(M_2\) columns of \(\mathcal{G}^i\). The equation of embedding network \(\mathcal{N}\) can be found at deepmd.utils.network.embedding_net().

Parameters
rcut

The cut-off radius \(r_c\)

rcut_smth

From where the environment matrix should be smoothed \(r_s\)

sellist[str]

sel[i] specifies the maxmum number of type i atoms in the cut-off radius

neuronlist[int]

Number of neurons in each hidden layers of the embedding net \(\mathcal{N}\)

axis_neuron

Number of the axis neuron \(M_2\) (number of columns of the sub-matrix of the embedding matrix)

resnet_dt

Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)

trainable

If the weights of embedding net are trainable.

seed

Random seed for initializing the network parameters.

type_one_side

Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets

exclude_typesList[List[int]]

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

set_davg_zero

Set the shift of embedding net input to zero.

activation_function

The activation function in the embedding net. Supported options are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.

precision

The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.

uniform_seed

Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed

multi_task

If the model has multi fitting nets to train.

References

1

Linfeng Zhang, Jiequn Han, Han Wang, Wissam A. Saidi, Roberto Car, and E. Weinan. 2018. End-to-end symmetry preserving inter-atomic potential energy model for finite and extended systems. In Proceedings of the 32nd International Conference on Neural Information Processing Systems (NIPS’18). Curran Associates Inc., Red Hook, NY, USA, 4441-4451.

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

precision

Precision of filter network.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_rot_mat()

Get rotational matrix.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') Tensor[source]

Build the computational graph for the descriptor.

Parameters
coord_

The coordinate of atoms

atype_

The type of atoms

natoms

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 the system

mesh

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_dict

Dictionary for additional inputs

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns
descriptor

The output descriptor

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, **kwargs) None[source]

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

Parameters
data_coord

The coordinates. Can be generated by deepmd.model.make_stat_input

data_box

The box. Can be generated by deepmd.model.make_stat_input

data_atype

The atom types. Can be generated by deepmd.model.make_stat_input

natoms_vec

The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input

mesh

The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input

input_dict

Dictionary for additional input

**kwargs

Additional keyword arguments.

enable_compression(min_nbor_dist: float, graph: Graph, graph_def: GraphDef, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') None[source]

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

Parameters
min_nbor_dist

The nearest distance between atoms

graphtf.Graph

The graph of the model

graph_deftf.GraphDef

The graph_def of the model

table_extrapolate

The scale of model extrapolation

table_stride_1

The uniform stride of the first table

table_stride_2

The uniform stride of the second table

check_frequency

The overflow check frequency

suffixstr, optional

The suffix of the scope

enable_mixed_precision(mixed_prec: Optional[dict] = None) None[source]

Reveive the mixed precision setting.

Parameters
mixed_prec

The mixed precision setting used in the embedding net

property explicit_ntypes: bool

Explicit ntypes with type embedding.

get_dim_out() int[source]

Returns the output dimension of this descriptor.

get_dim_rot_mat_1() int[source]

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

get_nlist() Tuple[Tensor, Tensor, List[int], List[int]][source]

Returns neighbor information.

Returns
nlist

Neighbor list

rij

The relative distance between the neighbor and the center atom.

sel_a

The number of neighbors with full information

sel_r

The number of neighbors with only radial information

get_ntypes() int[source]

Returns the number of atom types.

get_rcut() float[source]

Returns the cut-off radius.

get_rot_mat() Tensor[source]

Get rotational matrix.

init_variables(graph: Graph, graph_def: 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

merge_input_stats(stat_dict)[source]

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

Parameters
stat_dict

The dict of statisitcs computed from compute_input_stats, including:

sumr

The sum of radial statisitcs.

suma

The sum of relative coord statisitcs.

sumn

The sum of neighbor numbers.

sumr2

The sum of square of radial statisitcs.

suma2

The sum of square of relative coord statisitcs.

prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters
atom_ener

The atomic energy

natoms

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
force

The force on atoms

virial

The total virial

atom_virial

The atomic virial

deepmd.descriptor.se_a_ebd module

class deepmd.descriptor.se_a_ebd.DescrptSeAEbd(*args, **kwargs)[source]

Bases: DescrptSeA

DeepPot-SE descriptor with type embedding approach.

Parameters
rcut

The cut-off radius

rcut_smth

From where the environment matrix should be smoothed

sellist[str]

sel[i] specifies the maxmum number of type i atoms in the cut-off radius

neuronlist[int]

Number of neurons in each hidden layers of the embedding net

axis_neuron

Number of the axis neuron (number of columns of the sub-matrix of the embedding matrix)

resnet_dt

Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)

trainable

If the weights of embedding net are trainable.

seed

Random seed for initializing the network parameters.

type_one_side

Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets

type_nchanl

Number of channels for type representation

type_nlayer

Number of hidden layers for the type embedding net (skip connected).

numb_aparam

Number of atomic parameters. If >0 it will be embedded with atom types.

set_davg_zero

Set the shift of embedding net input to zero.

activation_function

The activation function in the embedding net. Supported options are {0}

precision

The precision of the embedding net parameters. Supported options are {1}

exclude_typesList[List[int]]

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

precision

Precision of filter network.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_rot_mat()

Get rotational matrix.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') Tensor[source]

Build the computational graph for the descriptor.

Parameters
coord_

The coordinate of atoms

atype_

The type of atoms

natoms

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 the system

mesh

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_dict

Dictionary for additional inputs

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns
descriptor

The output descriptor

deepmd.descriptor.se_a_ebd_v2 module

class deepmd.descriptor.se_a_ebd_v2.DescrptSeAEbdV2(*args, **kwargs)[source]

Bases: DescrptSeA

A compressible se_a_ebd model.

This model is a warpper for DescriptorSeA, which set stripped_type_embedding=True.

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

precision

Precision of filter network.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_rot_mat()

Get rotational matrix.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

deepmd.descriptor.se_a_ef module

class deepmd.descriptor.se_a_ef.DescrptSeAEf(*args, **kwargs)[source]

Bases: DescrptSe

Smooth edition descriptor with Ef.

Parameters
rcut

The cut-off radius

rcut_smth

From where the environment matrix should be smoothed

sellist[str]

sel[i] specifies the maxmum number of type i atoms in the cut-off radius

neuronlist[int]

Number of neurons in each hidden layers of the embedding net

axis_neuron

Number of the axis neuron (number of columns of the sub-matrix of the embedding matrix)

resnet_dt

Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)

trainable

If the weights of embedding net are trainable.

seed

Random seed for initializing the network parameters.

type_one_side

Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets

exclude_typesList[List[int]]

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

set_davg_zero

Set the shift of embedding net input to zero.

activation_function

The activation function in the embedding net. Supported options are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.

precision

The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.

uniform_seed

Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

precision

Precision of filter network.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_rot_mat()

Get rotational matrix.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') Tensor[source]

Build the computational graph for the descriptor.

Parameters
coord_

The coordinate of atoms

atype_

The type of atoms

natoms

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 the system

mesh

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_dict

Dictionary for additional inputs. Should have ‘efield’.

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns
descriptor

The output descriptor

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, **kwargs) None[source]

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

Parameters
data_coord

The coordinates. Can be generated by deepmd.model.make_stat_input

data_box

The box. Can be generated by deepmd.model.make_stat_input

data_atype

The atom types. Can be generated by deepmd.model.make_stat_input

natoms_vec

The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input

mesh

The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input

input_dict

Dictionary for additional input

**kwargs

Additional keyword arguments.

get_dim_out() int[source]

Returns the output dimension of this descriptor.

get_dim_rot_mat_1() int[source]

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

get_nlist() Tuple[Tensor, Tensor, List[int], List[int]][source]

Returns neighbor information.

Returns
nlist

Neighbor list

rij

The relative distance between the neighbor and the center atom.

sel_a

The number of neighbors with full information

sel_r

The number of neighbors with only radial information

get_ntypes() int[source]

Returns the number of atom types.

get_rcut() float[source]

Returns the cut-off radius.

get_rot_mat() Tensor[source]

Get rotational matrix.

prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters
atom_ener

The atomic energy

natoms

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
force

The force on atoms

virial

The total virial

atom_virial

The atomic virial

class deepmd.descriptor.se_a_ef.DescrptSeAEfLower(*args, **kwargs)[source]

Bases: DescrptSeA

Helper class for implementing DescrptSeAEf.

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

precision

Precision of filter network.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_rot_mat()

Get rotational matrix.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

build(coord_, atype_, natoms, box_, mesh, input_dict, suffix='', reuse=None)[source]

Build the computational graph for the descriptor.

Parameters
coord_

The coordinate of atoms

atype_

The type of atoms

natoms

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 the system

mesh

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_dict

Dictionary for additional inputs

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns
descriptor

The output descriptor

compute_input_stats(data_coord, data_box, data_atype, natoms_vec, mesh, input_dict, **kwargs)[source]

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

Parameters
data_coord

The coordinates. Can be generated by deepmd.model.make_stat_input

data_box

The box. Can be generated by deepmd.model.make_stat_input

data_atype

The atom types. Can be generated by deepmd.model.make_stat_input

natoms_vec

The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input

mesh

The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input

input_dict

Dictionary for additional input

**kwargs

Additional keyword arguments.

deepmd.descriptor.se_a_mask module

class deepmd.descriptor.se_a_mask.DescrptSeAMask(*args, **kwargs)[source]

Bases: DescrptSeA

DeepPot-SE constructed from all information (both angular and radial) of atomic configurations. The embedding takes the distance between atoms as input.

The descriptor \(\mathcal{D}^i \in \mathcal{R}^{M_1 \times M_2}\) is given by [1]

\[\mathcal{D}^i = (\mathcal{G}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \mathcal{G}^i_<\]

where \(\mathcal{R}^i \in \mathbb{R}^{N \times 4}\) is the coordinate matrix, and each row of \(\mathcal{R}^i\) can be constructed as follows

\[(\mathcal{R}^i)_j = [ \begin{array}{c} s(r_{ji}) & \frac{s(r_{ji})x_{ji}}{r_{ji}} & \frac{s(r_{ji})y_{ji}}{r_{ji}} & \frac{s(r_{ji})z_{ji}}{r_{ji}} \end{array} ]\]

where \(\mathbf{R}_{ji}=\mathbf{R}_j-\mathbf{R}_i = (x_{ji}, y_{ji}, z_{ji})\) is the relative coordinate and \(r_{ji}=\lVert \mathbf{R}_{ji} \lVert\) is its norm. The switching function \(s(r)\) is defined as:

\[\begin{split}s(r)= \begin{cases} \frac{1}{r}, & r<r_s \\ \frac{1}{r} \{ {(\frac{r - r_s}{ r_c - r_s})}^3 (-6 {(\frac{r - r_s}{ r_c - r_s})}^2 +15 \frac{r - r_s}{ r_c - r_s} -10) +1 \}, & r_s \leq r<r_c \\ 0, & r \geq r_c \end{cases}\end{split}\]

Each row of the embedding matrix \(\mathcal{G}^i \in \mathbb{R}^{N \times M_1}\) consists of outputs of a embedding network \(\mathcal{N}\) of \(s(r_{ji})\):

\[(\mathcal{G}^i)_j = \mathcal{N}(s(r_{ji}))\]

\(\mathcal{G}^i_< \in \mathbb{R}^{N \times M_2}\) takes first \(M_2\) columns of \(\mathcal{G}^i\). The equation of embedding network \(\mathcal{N}\) can be found at deepmd.utils.network.embedding_net(). Specially for descriptor se_a_mask is a concise implementation of se_a. The difference is that se_a_mask only considered a non-pbc system. And accept a mask matrix to indicate the atom i in frame j is a real atom or not. (1 means real atom, 0 means ghost atom) Thus se_a_mask can accept a variable number of atoms in a frame.

Parameters
sellist[str]

sel[i] specifies the maxmum number of type i atoms in the neighbor list.

neuronlist[int]

Number of neurons in each hidden layers of the embedding net \(\mathcal{N}\)

axis_neuron

Number of the axis neuron \(M_2\) (number of columns of the sub-matrix of the embedding matrix)

resnet_dt

Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)

trainable

If the weights of embedding net are trainable.

seed

Random seed for initializing the network parameters.

type_one_side

Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets

exclude_typesList[List[int]]

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

activation_function

The activation function in the embedding net. Supported options are {0}

precision

The precision of the embedding net parameters. Supported options are {1}

uniform_seed

Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed

References

1

Linfeng Zhang, Jiequn Han, Han Wang, Wissam A. Saidi, Roberto Car, and E. Weinan. 2018. End-to-end symmetry preserving inter-atomic potential energy model for finite and extended systems. In Proceedings of the 32nd International Conference on Neural Information Processing Systems (NIPS’18). Curran Associates Inc., Red Hook, NY, USA, 4441-4451.

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

precision

Precision of filter network.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cutoff radius.

get_rot_mat()

Get rotational matrix.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: Dict[str, Any], reuse: Optional[bool] = None, suffix: str = '') Tensor[source]

Build the computational graph for the descriptor.

Parameters
coord_

The coordinate of atoms

atype_

The type of atoms

natoms

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 the system

mesh

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_dict

Dictionary for additional inputs

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns
descriptor

The output descriptor

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, **kwargs) None[source]

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

Parameters
data_coord

The coordinates. Can be generated by deepmd.model.make_stat_input

data_box

The box. Can be generated by deepmd.model.make_stat_input

data_atype

The atom types. Can be generated by deepmd.model.make_stat_input

natoms_vec

The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input

mesh

The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input

input_dict

Dictionary for additional input

**kwargs

Additional keyword arguments.

get_rcut() float[source]

Returns the cutoff radius.

prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters
atom_ener

The atomic energy

natoms

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
force

The force on atoms

virial

None for se_a_mask op

atom_virial

None for se_a_mask op

classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]

Update the selection and perform neighbor statistics.

Parameters
global_jdatadict

The global data, containing the training section

local_jdatadict

The local data refer to the current class

deepmd.descriptor.se_atten module

class deepmd.descriptor.se_atten.DescrptSeAtten(*args, **kwargs)[source]

Bases: DescrptSeA

Smooth version descriptor with attention.

Parameters
rcut

The cut-off radius \(r_c\)

rcut_smth

From where the environment matrix should be smoothed \(r_s\)

sellist[str]

sel[i] specifies the maxmum number of type i atoms in the cut-off radius

neuronlist[int]

Number of neurons in each hidden layers of the embedding net \(\mathcal{N}\)

axis_neuron

Number of the axis neuron \(M_2\) (number of columns of the sub-matrix of the embedding matrix)

resnet_dt

Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)

trainable

If the weights of embedding net are trainable.

seed

Random seed for initializing the network parameters.

type_one_side

Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets

exclude_typesList[List[int]]

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

set_davg_zero

Set the shift of embedding net input to zero.

activation_function

The activation function in the embedding net. Supported options are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.

precision

The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.

uniform_seed

Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed

attn

The length of hidden vector during scale-dot attention computation.

attn_layer

The number of layers in attention mechanism.

attn_dotr

Whether to dot the relative coordinates on the attention weights as a gated scheme.

attn_mask

Whether to mask the diagonal in the attention weights.

multi_task

If the model has multi fitting nets to train.

stripped_type_embedding

Whether to strip the type embedding into a separated embedding network. Default value will be True in se_atten_v2 descriptor.

smooth_type_embdding

When using stripped type embedding, whether to dot smooth factor on the network output of type embedding to keep the network smooth, instead of setting set_davg_zero to be True. Default value will be True in se_atten_v2 descriptor.

Raises
ValueError

if ntypes is 0.

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

precision

Precision of filter network.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the attention descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_rot_mat()

Get rotational matrix.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') Tensor[source]

Build the computational graph for the descriptor.

Parameters
coord_

The coordinate of atoms

atype_

The type of atoms

natoms

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 the system

mesh

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_dict

Dictionary for additional inputs

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns
descriptor

The output descriptor

build_type_exclude_mask(exclude_types: List[Tuple[int, int]], ntypes: int, sel: List[int], ndescrpt: int, atype: Tensor, shape0: Tensor, nei_type_vec: Tensor) Tensor[source]

Build the type exclude mask for the attention 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.

nei_type_vectf.Tensor

The type of neighbors, with the size of (shape0, nnei).

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

This method has the similiar way to build the type exclude mask as deepmd.descriptor.descriptor.Descriptor.build_type_exclude_mask(). The mathmatical expression has been explained in that method. The difference is that the attention descriptor has provided the type of the neighbors (idx_j) that is not in order, so we use it from an extra input.

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, mixed_type: bool = False, real_natoms_vec: Optional[list] = None, **kwargs) None[source]

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

Parameters
data_coord

The coordinates. Can be generated by deepmd.model.make_stat_input

data_box

The box. Can be generated by deepmd.model.make_stat_input

data_atype

The atom types. Can be generated by deepmd.model.make_stat_input

natoms_vec

The vector for the number of atoms of the system and different types of atoms. If mixed_type is True, this para is blank. See real_natoms_vec.

mesh

The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input

input_dict

Dictionary for additional input

mixed_type

Whether to perform the mixed_type mode. If True, the input data has the mixed_type format (see doc/model/train_se_atten.md), in which frames in a system may have different natoms_vec(s), with the same nloc.

real_natoms_vec

If mixed_type is True, it takes in the real natoms_vec for each frame.

**kwargs

Additional keyword arguments.

enable_compression(min_nbor_dist: float, graph: Graph, graph_def: GraphDef, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') None[source]

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

Parameters
min_nbor_dist

The nearest distance between atoms

graphtf.Graph

The graph of the model

graph_deftf.GraphDef

The graph_def of the model

table_extrapolate

The scale of model extrapolation

table_stride_1

The uniform stride of the first table

table_stride_2

The uniform stride of the second table

check_frequency

The overflow check frequency

suffixstr, optional

The suffix of the scope

property explicit_ntypes: bool

Explicit ntypes with type embedding.

init_variables(graph: Graph, graph_def: 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

classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]

Update the selection and perform neighbor statistics.

Parameters
global_jdatadict

The global data, containing the training section

local_jdatadict

The local data refer to the current class

deepmd.descriptor.se_atten_v2 module

class deepmd.descriptor.se_atten_v2.DescrptSeAttenV2(*args, **kwargs)[source]

Bases: DescrptSeAtten

Smooth version 2.0 descriptor with attention.

Parameters
rcut

The cut-off radius \(r_c\)

rcut_smth

From where the environment matrix should be smoothed \(r_s\)

sellist[str]

sel[i] specifies the maxmum number of type i atoms in the cut-off radius

neuronlist[int]

Number of neurons in each hidden layers of the embedding net \(\mathcal{N}\)

axis_neuron

Number of the axis neuron \(M_2\) (number of columns of the sub-matrix of the embedding matrix)

resnet_dt

Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)

trainable

If the weights of embedding net are trainable.

seed

Random seed for initializing the network parameters.

type_one_side

Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets

exclude_typesList[List[int]]

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

set_davg_zero

Set the shift of embedding net input to zero.

activation_function

The activation function in the embedding net. Supported options are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.

precision

The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.

uniform_seed

Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed

attn

The length of hidden vector during scale-dot attention computation.

attn_layer

The number of layers in attention mechanism.

attn_dotr

Whether to dot the relative coordinates on the attention weights as a gated scheme.

attn_mask

Whether to mask the diagonal in the attention weights.

multi_task

If the model has multi fitting nets to train.

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

precision

Precision of filter network.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the attention descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_rot_mat()

Get rotational matrix.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

deepmd.descriptor.se_r module

class deepmd.descriptor.se_r.DescrptSeR(*args, **kwargs)[source]

Bases: DescrptSe

DeepPot-SE constructed from radial information of atomic configurations.

The embedding takes the distance between atoms as input.

Parameters
rcut

The cut-off radius

rcut_smth

From where the environment matrix should be smoothed

sellist[str]

sel[i] specifies the maxmum number of type i atoms in the cut-off radius

neuronlist[int]

Number of neurons in each hidden layers of the embedding net

resnet_dt

Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)

trainable

If the weights of embedding net are trainable.

seed

Random seed for initializing the network parameters.

type_one_side

Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets

exclude_typesList[List[int]]

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

activation_function

The activation function in the embedding net. Supported options are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.

precision

The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.

uniform_seed

Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

precision

Precision of filter network.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') Tensor[source]

Build the computational graph for the descriptor.

Parameters
coord_

The coordinate of atoms

atype_

The type of atoms

natoms

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 the system

mesh

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_dict

Dictionary for additional inputs

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns
descriptor

The output descriptor

compute_input_stats(data_coord, data_box, data_atype, natoms_vec, mesh, input_dict, **kwargs)[source]

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

Parameters
data_coord

The coordinates. Can be generated by deepmd.model.make_stat_input

data_box

The box. Can be generated by deepmd.model.make_stat_input

data_atype

The atom types. Can be generated by deepmd.model.make_stat_input

natoms_vec

The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input

mesh

The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input

input_dict

Dictionary for additional input

**kwargs

Additional keyword arguments.

enable_compression(min_nbor_dist: float, graph: Graph, graph_def: GraphDef, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') None[source]

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

Parameters
min_nbor_dist

The nearest distance between atoms

graphtf.Graph

The graph of the model

graph_deftf.GraphDef

The graph_def of the model

table_extrapolate

The scale of model extrapolation

table_stride_1

The uniform stride of the first table

table_stride_2

The uniform stride of the second table

check_frequency

The overflow check frequency

suffixstr, optional

The suffix of the scope

get_dim_out()[source]

Returns the output dimension of this descriptor.

get_nlist()[source]

Returns neighbor information.

Returns
nlist

Neighbor list

rij

The relative distance between the neighbor and the center atom.

sel_a

The number of neighbors with full information

sel_r

The number of neighbors with only radial information

get_ntypes()[source]

Returns the number of atom types.

get_rcut()[source]

Returns the cut-off radius.

merge_input_stats(stat_dict)[source]

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

Parameters
stat_dict

The dict of statisitcs computed from compute_input_stats, including:

sumr

The sum of radial statisitcs.

sumn

The sum of neighbor numbers.

sumr2

The sum of square of radial statisitcs.

prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters
atom_ener

The atomic energy

natoms

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
force

The force on atoms

virial

The total virial

atom_virial

The atomic virial

deepmd.descriptor.se_t module

class deepmd.descriptor.se_t.DescrptSeT(*args, **kwargs)[source]

Bases: DescrptSe

DeepPot-SE constructed from all information (both angular and radial) of atomic configurations.

The embedding takes angles between two neighboring atoms as input.

Parameters
rcut

The cut-off radius

rcut_smth

From where the environment matrix should be smoothed

sellist[str]

sel[i] specifies the maxmum number of type i atoms in the cut-off radius

neuronlist[int]

Number of neurons in each hidden layers of the embedding net

resnet_dt

Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)

trainable

If the weights of embedding net are trainable.

seed

Random seed for initializing the network parameters.

set_davg_zero

Set the shift of embedding net input to zero.

activation_function

The activation function in the embedding net. Supported options are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.

precision

The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.

uniform_seed

Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed

Attributes
explicit_ntypes

Explicit ntypes with type embedding.

precision

Precision of filter network.

Methods

build(coord_, atype_, natoms, box_, mesh, ...)

Build the computational graph for the descriptor.

build_type_exclude_mask(exclude_types, ...)

Build the type exclude mask for the descriptor.

compute_input_stats(data_coord, data_box, ...)

Compute the statisitcs (avg and std) of the training data.

enable_compression(min_nbor_dist, graph, ...)

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

enable_mixed_precision([mixed_prec])

Reveive the mixed precision setting.

get_dim_out()

Returns the output dimension of this descriptor.

get_dim_rot_mat_1()

Returns the first dimension of the rotation matrix.

get_nlist()

Returns neighbor information.

get_ntypes()

Returns the number of atom types.

get_rcut()

Returns the cut-off radius.

get_tensor_names([suffix])

Get names of tensors.

init_variables(graph, graph_def[, suffix])

Init the embedding net variables with the given dict.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

pass_tensors_from_frz_model(descrpt_reshape, ...)

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

prod_force_virial(atom_ener, natoms)

Compute force and virial.

register(key)

Register a descriptor plugin.

update_sel(global_jdata, local_jdata)

Update the selection and perform neighbor statistics.

get_class_by_input

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: Optional[bool] = None, suffix: str = '') Tensor[source]

Build the computational graph for the descriptor.

Parameters
coord_

The coordinate of atoms

atype_

The type of atoms

natoms

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 the system

mesh

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_dict

Dictionary for additional inputs

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns
descriptor

The output descriptor

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, **kwargs) None[source]

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

Parameters
data_coord

The coordinates. Can be generated by deepmd.model.make_stat_input

data_box

The box. Can be generated by deepmd.model.make_stat_input

data_atype

The atom types. Can be generated by deepmd.model.make_stat_input

natoms_vec

The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input

mesh

The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input

input_dict

Dictionary for additional input

**kwargs

Additional keyword arguments.

enable_compression(min_nbor_dist: float, graph: Graph, graph_def: GraphDef, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') None[source]

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

Parameters
min_nbor_dist

The nearest distance between atoms

graphtf.Graph

The graph of the model

graph_deftf.GraphDef

The graph_def of the model

table_extrapolate

The scale of model extrapolation

table_stride_1

The uniform stride of the first table

table_stride_2

The uniform stride of the second table

check_frequency

The overflow check frequency

suffixstr, optional

The suffix of the scope

get_dim_out() int[source]

Returns the output dimension of this descriptor.

get_nlist() Tuple[Tensor, Tensor, List[int], List[int]][source]

Returns neighbor information.

Returns
nlist

Neighbor list

rij

The relative distance between the neighbor and the center atom.

sel_a

The number of neighbors with full information

sel_r

The number of neighbors with only radial information

get_ntypes() int[source]

Returns the number of atom types.

get_rcut() float[source]

Returns the cut-off radius.

merge_input_stats(stat_dict)[source]

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

Parameters
stat_dict

The dict of statisitcs computed from compute_input_stats, including:

sumr

The sum of radial statisitcs.

suma

The sum of relative coord statisitcs.

sumn

The sum of neighbor numbers.

sumr2

The sum of square of radial statisitcs.

suma2

The sum of square of relative coord statisitcs.

prod_force_virial(atom_ener: Tensor, natoms: Tensor) Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters
atom_ener

The atomic energy

natoms

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
force

The force on atoms

virial

The total virial

atom_virial

The atomic virial