deepmd.tf.model.tensor

Module Contents

Classes

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

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)

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

get_rcut()

Get cutoff radius of the model.

get_ntypes()

Get the number of types.

get_type_map()

Get the type map.

get_sel_type()

get_out_size()

data_stat(data)

Data staticis.

_compute_input_stat(all_stat, protection=0.01)

_compute_output_stat(all_stat)

build(coord_, atype_, natoms, box, mesh, input_dict, frz_model=None, ckpt_meta: str | None = None, suffix='', reuse=None)

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

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, **kwargs)

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, **kwargs)

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, **kwargs)

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, **kwargs)

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