deepmd.tf.model.tensor

deepmd.tf.model.tensor#

Classes#

`TensorModel`	Tensor model.
`WFCModel`	Tensor model.
`DipoleModel`	Tensor model.
`PolarModel`	Tensor model.
`GlobalPolarModel`	Tensor model.

Module Contents#

class deepmd.tf.model.tensor.TensorModel(tensor_name: str, descriptor: dict, fitting_net: dict, type_embedding: dict | deepmd.tf.utils.type_embed.TypeEmbedNet | None = None, type_map: list[str] | None = None, data_stat_nbatch: int = 10, data_stat_protect: float = 0.01, **kwargs: Any)[source]#

Bases: deepmd.tf.model.model.StandardModel

Tensor model.

Parameters:

tensor_name: Name of the tensor.
descriptor: Descriptor
fitting_net: Fitting net
type_embedding: Type embedding net
type_map: Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.
data_stat_nbatch: Number of frames used for data statistic
data_stat_protect: Protect parameter for atomic energy regression

model_type[source]#

get_rcut() → float[source]#: Get cutoff radius of the model.

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

get_type_map() → list[str][source]#: Get the type map.

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

get_out_size() → int[source]#

data_stat(data: deepmd.utils.data_system.DeepmdDataSystem) → None[source]#: Data staticis.

_compute_input_stat(all_stat: dict, protection: float = 0.01) → None[source]#

_compute_output_stat(all_stat: dict) → None[source]#

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, frz_model: str | None = None, ckpt_meta: str | None = None, suffix: str = '', reuse: bool | None = None) → dict[source]#

Build the model.

Parameters:

coord_tf.Tensor: The coordinates of atoms
atype_tf.Tensor: The atom types of atoms
natomstf.Tensor: The number of atoms
boxtf.Tensor: The box vectors
meshtf.Tensor: The mesh vectors
input_dictdict: The input dict
frz_modelstr, optional: The path to the frozen model
ckpt_metastr, optional: The path prefix of the checkpoint and meta files
suffixstr, optional: The suffix of the scope
reusebool or tf.AUTO_REUSE, optional: Whether to reuse the variables

Returns:

dict: The output dict

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

Init the embedding net variables with the given frozen model.

Parameters:

graphtf.Graph: The input frozen model graph
graph_deftf.GraphDef: The input frozen model graph_def
model_typestr: the type of the model
suffixstr: suffix to name scope

class deepmd.tf.model.tensor.WFCModel(*args: Any, **kwargs: Any)[source]#

Bases: TensorModel

Tensor model.

Parameters:

tensor_name: Name of the tensor.
descriptor: Descriptor
fitting_net: Fitting net
type_embedding: Type embedding net
type_map: Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.
data_stat_nbatch: Number of frames used for data statistic
data_stat_protect: Protect parameter for atomic energy regression

class deepmd.tf.model.tensor.DipoleModel(*args: Any, **kwargs: Any)[source]#

Bases: TensorModel

Tensor model.

Parameters:

tensor_name: Name of the tensor.
descriptor: Descriptor
fitting_net: Fitting net
type_embedding: Type embedding net
type_map: Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.
data_stat_nbatch: Number of frames used for data statistic
data_stat_protect: Protect parameter for atomic energy regression

class deepmd.tf.model.tensor.PolarModel(*args: Any, **kwargs: Any)[source]#

Bases: TensorModel

Tensor model.

Parameters:

tensor_name: Name of the tensor.
descriptor: Descriptor
fitting_net: Fitting net
type_embedding: Type embedding net
type_map: Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.
data_stat_nbatch: Number of frames used for data statistic
data_stat_protect: Protect parameter for atomic energy regression

class deepmd.tf.model.tensor.GlobalPolarModel(*args: Any, **kwargs: Any)[source]#

Bases: TensorModel

Tensor model.

Parameters:

tensor_name: Name of the tensor.
descriptor: Descriptor
fitting_net: Fitting net
type_embedding: Type embedding net
type_map: Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.
data_stat_nbatch: Number of frames used for data statistic
data_stat_protect: Protect parameter for atomic energy regression