deepmd.tf.descriptor

deepmd.tf.descriptor#

Submodules#

Classes#

`Descriptor`	The abstract class for descriptors. All specific descriptors should
`DescrptHybrid`	Concate a list of descriptors to form a new descriptor.
`DescrptLocFrame`	Defines a local frame at each atom, and the compute the descriptor as local
`DescrptSeA`	DeepPot-SE constructed from all information (both angular and radial) of
`DescrptSeAEbd`	DeepPot-SE descriptor with type embedding approach.
`DescrptSeAEbdV2`	A compressible se_a_ebd model.
`DescrptSeAEf`	Smooth edition descriptor with Ef.
`DescrptSeAEfLower`	Helper class for implementing DescrptSeAEf.
`DescrptSeAMask`	DeepPot-SE constructed from all information (both angular and radial) of
`DescrptSeAtten`	Smooth version descriptor with attention.
`DescrptSeAttenV2`	Smooth version 2.0 descriptor with attention.
`DescrptSeR`	DeepPot-SE constructed from radial information of atomic configurations.
`DescrptSeT`	DeepPot-SE constructed from all information (both angular and radial) of atomic

Package Contents#

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

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

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

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

Notes

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

Examples

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

abstract get_rcut() → float[source]#

Returns the cut-off radius.

Returns:

float: the cut-off radius

Notes

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

abstract get_ntypes() → int[source]#

Returns the number of atom types.

Returns:

int: the number of atom types

Notes

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

abstract get_dim_out() → int[source]#

Returns the output dimension of this descriptor.

Returns:

int: the output dimension of this descriptor

Notes

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

abstract get_dim_rot_mat_1() → int[source]#

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

Returns:

int: the first dimension of the rotation matrix

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

Returns neighbor information.

Returns:

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

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

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

Parameters:

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

Notes

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

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

Build the computational graph for the descriptor.

Parameters:

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

Returns:

descriptor: tf.Tensor: The output descriptor

Notes

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

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

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

Parameters:

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

Notes

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

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

Receive the mixed precision setting.

Parameters:

mixed_prec: The mixed precision setting used in the embedding net

Notes

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

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

Compute force and virial.

Parameters:

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

Returns:

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

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

Init the embedding net variables with the given dict.

Parameters:

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

Notes

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

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

Get names of tensors.

Parameters:

suffixstr: The suffix of the scope

Returns:

tuple[str]: Names of tensors

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

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

Parameters:

*tensorstf.Tensor: passed tensors

Notes

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

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

Build the type exclude mask for the descriptor.

Parameters:

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

Returns:

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

Notes

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

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

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

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

References

[1]

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

property explicit_ntypes: bool#: Explicit ntypes with type embedding.

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

Abstractmethod:

Update the selection and perform neighbor statistics.

Parameters:

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

Returns:

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

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

Deserialize the model.

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

Parameters:

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

Returns:

Descriptor: The deserialized descriptor

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

Serialize the model.

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

Returns:

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

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

class deepmd.tf.descriptor.DescrptHybrid(list: DescrptHybrid.__init__.list[deepmd.tf.descriptor.descriptor.Descriptor | dict[str, Any]], ntypes: int | None = None, spin: deepmd.tf.utils.spin.Spin | None = None, **kwargs)[source]#

Bases: deepmd.tf.descriptor.descriptor.Descriptor

Concate a list of descriptors to form a new descriptor.

Parameters:

listlistlist[Union[Descriptor, dict[str, Any]]]: Build a descriptor from the concatenation of the list of descriptors. The descriptor can be either an object or a dictionary.

descrpt_list#

numb_descrpt#

get_rcut() → float[source]#: Returns the cut-off radius.

get_ntypes() → int[source]#: Returns the number of atom types.

get_dim_out() → int[source]#: Returns the output dimension of this descriptor.

get_nlist() → tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.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[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.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

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: list | None = 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.tf.model.make_stat_input
data_box: The box. Can be generated by deepmd.tf.model.make_stat_input
data_atype: The atom types. Can be generated by deepmd.tf.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.tf.model.make_stat_input
mesh: The mesh for neighbor searching. Can be generated by deepmd.tf.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.

merge_input_stats(stat_dict) → None[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.

build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box_: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') → deepmd.tf.env.tf.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

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

Compute force and virial.

Parameters:

atom_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

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

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

Parameters:

min_nbor_distfloat: The nearest distance between atoms
graphtf.Graph: The graph of the model
graph_deftf.GraphDef: The graph_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: dict | None = None) → None[source]#

Receive the mixed precision setting.

Parameters:

mixed_prec: The mixed precision setting used in the embedding net

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

Init the embedding net variables with the given dict.

Parameters:

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

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

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

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

Parameters:

*tensorstf.Tensor: passed tensors

property explicit_ntypes: bool#: Explicit ntypes with type embedding.

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

Update the selection and perform neighbor statistics.

Parameters:

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

Returns:

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

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

Serialize the model.

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

Returns:

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

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

Deserialize the model.

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

Parameters:

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

Returns:

Descriptor: The deserialized descriptor

class deepmd.tf.descriptor.DescrptLocFrame(rcut: float, sel_a: list[int], sel_r: list[int], axis_rule: list[int], **kwargs)[source]#

Bases: deepmd.tf.descriptor.descriptor.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[int]: 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[int]: 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.

sel_a#

sel_r#

axis_rule#

rcut_r#

ntypes#

rcut_a#

nnei_a#

nnei_r#

nnei#

ndescrpt_a#

ndescrpt_r#

ndescrpt#

davg = None#

dstd = None#

place_holders#

sub_sess#

get_rcut() → float[source]#: Returns the cut-off radius.

get_ntypes() → int[source]#: Returns the number of atom types.

get_dim_out() → int[source]#: Returns the output dimension of this descriptor.

get_nlist() → tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.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

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.tf.model.make_stat_input
data_box: The box. Can be generated by deepmd.tf.model.make_stat_input
data_atype: The atom types. Can be generated by deepmd.tf.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.tf.model.make_stat_input
mesh: The mesh for neighbor searching. Can be generated by deepmd.tf.model.make_stat_input
input_dict: Dictionary for additional input
**kwargs: Additional keyword arguments.

build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box_: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') → deepmd.tf.env.tf.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

get_rot_mat() → deepmd.tf.env.tf.Tensor[source]#: Get rotational matrix.

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

Compute force and virial.

Parameters:

atom_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

_compute_dstats_sys_nonsmth(data_coord, data_box, data_atype, natoms_vec, mesh)[source]#

_compute_std(sumv2, sumv, sumn)[source]#

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

Init the embedding net variables with the given dict.

Parameters:

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

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

Update the selection and perform neighbor statistics.

Parameters:

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

Returns:

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

class deepmd.tf.descriptor.DescrptSeA(rcut: float, rcut_smth: float, sel: list[int], neuron: list[int] = [24, 48, 96], axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, seed: int | None = None, type_one_side: bool = True, exclude_types: list[list[int]] = [], set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = 'default', uniform_seed: bool = False, spin: deepmd.tf.utils.spin.Spin | None = None, tebd_input_mode: str = 'concat', type_map: list[str] | None = None, env_protection: float = 0.0, **kwargs)[source]#

Bases: deepmd.tf.descriptor.se.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.tf.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[int]: 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 “none”, “gelu_tf”, “linear”, “relu6”, “sigmoid”, “tanh”, “gelu”, “relu”, “softplus”.
precision: The precision of the embedding net parameters. Supported options are “float16”, “float64”, “default”, “float32”.
uniform_seed: Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
env_protection: float: Protection parameter to prevent division by zero errors during environment matrix calculations.
type_map: list[str], Optional: A list of strings. Give the name to each type of atoms.

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.

sel_a#

rcut_r#

rcut_r_smth#

filter_neuron#

n_axis_neuron#

filter_resnet_dt#

seed#

uniform_seed#

seed_shift#

trainable#

compress_activation_fn#

filter_activation_fn#

activation_function_name#

filter_precision#

filter_np_precision#

orig_exclude_types#

exclude_types#

env_protection#

type_map#

set_davg_zero#

type_one_side#

spin#

stripped_type_embedding#

extra_embedding_net_variables = None#

layer_size#

sel_r#

ntypes#

rcut_a#

nnei_a#

nnei_r#

nnei#

ndescrpt_a#

ndescrpt_r#

ndescrpt#

useBN = False#

dstd = None#

davg = None#

compress = False#

embedding_net_variables = None#

mixed_prec = None#

place_holders#

nei_type#

sub_sess#

original_sel = None#

use_tebd: bool = False#

get_rcut() → float[source]#: Returns the cut-off radius.

get_ntypes() → int[source]#: Returns the number of atom types.

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[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.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

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.tf.model.make_stat_input
data_box: The box. Can be generated by deepmd.tf.model.make_stat_input
data_atype: The atom types. Can be generated by deepmd.tf.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.tf.model.make_stat_input
mesh: The mesh for neighbor searching. Can be generated by deepmd.tf.model.make_stat_input
input_dict: Dictionary for additional input
**kwargs: Additional keyword arguments.

merge_input_stats(stat_dict) → None[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.

enable_compression(min_nbor_dist: float, graph: deepmd.tf.env.tf.Graph, graph_def: deepmd.tf.env.tf.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]#

Receive 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: dict | None = None) → None[source]#

Receive the mixed precision setting.

Parameters:

mixed_prec: The mixed precision setting used in the embedding net

build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box_: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') → deepmd.tf.env.tf.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

get_rot_mat() → deepmd.tf.env.tf.Tensor[source]#: Get rotational matrix.

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

Compute force and virial.

Parameters:

atom_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

_pass_filter(inputs, atype, natoms, input_dict, reuse=None, suffix='', trainable=True)[source]#

_compute_dstats_sys_smth(data_coord, data_box, data_atype, natoms_vec, mesh)[source]#

_compute_std(sumv2, sumv, sumn)[source]#

_concat_type_embedding(xyz_scatter, nframes, natoms, type_embedding)[source]#

Concatenate type_embedding of neighbors and xyz_scatter. If not self.type_one_side, concatenate type_embedding of center atoms as well.

Parameters:

xyz_scatter:: shape is [nframes*natoms[0]*self.nnei, 1]
nframes:: shape is []
natoms:: shape is [1+1+self.ntypes]
type_embedding:: shape is [self.ntypes, Y] where Y=jdata[‘type_embedding’][‘neuron’][-1]

Returns:

embedding:: environment of each atom represented by embedding.

_filter_lower(type_i, type_input, start_index, incrs_index, inputs, nframes, natoms, type_embedding=None, is_exclude=False, activation_fn=None, bavg=0.0, stddev=1.0, trainable=True, suffix='')[source]#: Input env matrix, returns R.G.

_filter(inputs, type_input, natoms, type_embedding=None, activation_fn=tf.nn.tanh, stddev=1.0, bavg=0.0, name='linear', reuse=None, trainable=True)[source]#

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

Init the embedding net variables with the given dict.

Parameters:

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

property explicit_ntypes: bool#: Explicit ntypes with type embedding.

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

Deserialize the model.

Parameters:

datadict: The serialized data

Returns:

Model: The deserialized model

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

Serialize the model.

Parameters:

suffixstr, optional: The suffix of the scope

Returns:

dict: The serialized data

class deepmd.tf.descriptor.DescrptSeAEbd(rcut: float, rcut_smth: float, sel: list[int], neuron: list[int] = [24, 48, 96], axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, seed: int | None = None, type_one_side: bool = True, type_nchanl: int = 2, type_nlayer: int = 1, numb_aparam: int = 0, set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = 'default', exclude_types: list[list[int]] = [], **kwargs)[source]#

Bases: deepmd.tf.descriptor.se_a.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[int]: 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.

type_nchanl#

type_nlayer#

type_one_side#

numb_aparam#

build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box_: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') → deepmd.tf.env.tf.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

_type_embed(atype, ndim=1, reuse=None, suffix='', trainable=True)[source]#

_embedding_net(inputs, natoms, filter_neuron, activation_fn=tf.nn.tanh, stddev=1.0, bavg=0.0, name='linear', reuse=None, seed=None, trainable=True)[source]#: inputs: nf x na x (nei x 4) outputs: nf x na x nei x output_size.

_type_embedding_net_two_sides(mat_g, atype, natoms, name='', reuse=None, seed=None, trainable=True)[source]#

_type_embedding_net_one_side(mat_g, atype, natoms, name='', reuse=None, seed=None, trainable=True)[source]#

_type_embedding_net_one_side_aparam(mat_g, atype, natoms, aparam, name='', reuse=None, seed=None, trainable=True)[source]#

_pass_filter(inputs, atype, natoms, input_dict, reuse=None, suffix='', trainable=True)[source]#

_ebd_filter(inputs, atype, natoms, input_dict, activation_fn=tf.nn.tanh, stddev=1.0, bavg=0.0, name='linear', reuse=None, seed=None, trainable=True)[source]#

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

class deepmd.tf.descriptor.DescrptSeAEbdV2(rcut: float, rcut_smth: float, sel: list[int], neuron: list[int] = [24, 48, 96], axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, seed: int | None = None, type_one_side: bool = True, exclude_types: list[list[int]] = [], set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = 'default', uniform_seed: bool = False, spin: deepmd.tf.utils.spin.Spin | None = None, **kwargs)[source]#

Bases: deepmd.tf.descriptor.se_a.DescrptSeA

A compressible se_a_ebd model.

This model is a wrapper for DescriptorSeA, which set tebd_input_mode=’strip’.

class deepmd.tf.descriptor.DescrptSeAEf(rcut: float, rcut_smth: float, sel: list[int], neuron: list[int] = [24, 48, 96], axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, seed: int | None = None, type_one_side: bool = True, exclude_types: list[list[int]] = [], set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = 'default', uniform_seed=False, **kwargs)[source]#

Bases: deepmd.tf.descriptor.se.DescrptSe

Smooth edition descriptor with Ef.

Parameters:

rcut: The cut-off radius
rcut_smth: From where the environment matrix should be smoothed
sellist[int]: 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 “none”, “gelu_tf”, “linear”, “relu6”, “sigmoid”, “tanh”, “gelu”, “relu”, “softplus”.
precision: The precision of the embedding net parameters. Supported options are “float16”, “float64”, “default”, “float32”.
uniform_seed: Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed

descrpt_para#

descrpt_vert#

get_rcut() → float[source]#: Returns the cut-off radius.

get_ntypes() → int[source]#: Returns the number of atom types.

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_rot_mat() → deepmd.tf.env.tf.Tensor[source]#: Get rotational matrix.

get_nlist() → tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.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

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.tf.model.make_stat_input
data_box: The box. Can be generated by deepmd.tf.model.make_stat_input
data_atype: The atom types. Can be generated by deepmd.tf.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.tf.model.make_stat_input
mesh: The mesh for neighbor searching. Can be generated by deepmd.tf.model.make_stat_input
input_dict: Dictionary for additional input
**kwargs: Additional keyword arguments.

build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box_: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') → deepmd.tf.env.tf.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

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

Compute force and virial.

Parameters:

atom_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.tf.descriptor.DescrptSeAEfLower(op, rcut: float, rcut_smth: float, sel: list[int], neuron: list[int] = [24, 48, 96], axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, seed: int | None = None, type_one_side: bool = True, exclude_types: list[list[int]] = [], set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = 'default', uniform_seed: bool = False)[source]#

Bases: deepmd.tf.descriptor.se_a.DescrptSeA

Helper class for implementing DescrptSeAEf.

sel_a#

rcut_r#

rcut_r_smth#

filter_neuron#

n_axis_neuron#

filter_resnet_dt#

seed#

trainable#

op#

sel_r#

ntypes#

rcut_a#

nnei_a#

nnei_r#

nnei#

ndescrpt_a#

ndescrpt_r#

ndescrpt#

useBN = False#

dstd = None#

davg = None#

place_holders#

sub_sess#

compute_input_stats(data_coord, data_box, data_atype, natoms_vec, mesh, input_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.tf.model.make_stat_input
data_box: The box. Can be generated by deepmd.tf.model.make_stat_input
data_atype: The atom types. Can be generated by deepmd.tf.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.tf.model.make_stat_input
mesh: The mesh for neighbor searching. Can be generated by deepmd.tf.model.make_stat_input
input_dict: Dictionary for additional input
**kwargs: Additional keyword arguments.

_normalize_3d(a)[source]#

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_dstats_sys_smth(data_coord, data_box, data_atype, natoms_vec, mesh, data_efield)[source]#

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

class deepmd.tf.descriptor.DescrptSeAMask(sel: list[int], neuron: list[int] = [24, 48, 96], axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, type_one_side: bool = False, exclude_types: list[list[int]] = [], seed: int | None = None, activation_function: str = 'tanh', precision: str = 'default', uniform_seed: bool = False, tebd_input_mode: str = 'concat', **kwargs)[source]#