deepmd.tf.model package

class deepmd.tf.model.DOSModel(*args, **kwargs)[source]

Bases: StandardModel

DOS model.

Parameters

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

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

build(coord_, atype_, natoms, box, mesh, input_dict, frz_model=None, ckpt_meta: Optional[str] = None, suffix='', reuse=None)[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

data_stat(data)[source]: Data staticis.

get_ntypes()[source]: Get the number of types.

get_numb_aparam() → int[source]: Get the number of atomic parameters.

get_numb_dos()[source]: Get the number of gridpoints in energy space.

get_numb_fparam() → int[source]: Get the number of frame parameters.

get_rcut()[source]: Get cutoff radius of the model.

get_type_map()[source]: Get the type map.

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

model_type = 'dos'

class deepmd.tf.model.DipoleModel(*args, **kwargs)[source]

Bases: TensorModel

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_out_size
get_sel_type

class deepmd.tf.model.EnerModel(*args, **kwargs)[source]

Bases: StandardModel

Energy model.

Parameters

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
use_srtab: The table for the short-range pairwise interaction added on top of DP. The table is a text data file with (N_t + 1) * N_t / 2 + 1 columes. The first colume is the distance between atoms. The second to the last columes are energies for pairs of certain types. For example we have two atom types, 0 and 1. The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly.
smin_alpha: The short-range tabulated interaction will be swithed according to the distance of the nearest neighbor. This distance is calculated by softmin. This parameter is the decaying parameter in the softmin. It is only required when use_srtab is provided.
sw_rmin: The lower boundary of the interpolation between short-range tabulated interaction and DP. It is only required when use_srtab is provided.
sw_rmin: The upper boundary of the interpolation between short-range tabulated interaction and DP. It is only required when use_srtab is provided.
srtab_add_biasbool: Whether add energy bias from the statistics of the data to short-range tabulated atomic energy. It only takes effect when use_srtab is provided.
spin: spin
data_stat_nsample: The number of training samples in a system to compute and change the energy bias.

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

natoms_match
natoms_not_match

build(coord_, atype_, natoms, box, mesh, input_dict, frz_model=None, ckpt_meta: Optional[str] = None, suffix='', reuse=None)[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

change_energy_bias(data: DeepmdDataSystem, frozen_model: str, origin_type_map: list, full_type_map: str, bias_shift: str = 'delta') → None[source]

Change the energy bias according to the input data and the pretrained model.

Parameters

dataDeepmdDataSystem

The training data.

frozen_modelstr

The path file of frozen model.

origin_type_maplist

The original type_map in dataset, they are targets to change the energy bias.

full_type_mapstr

The full type_map in pretrained model

bias_shiftstr

The mode for changing energy bias : [‘delta’, ‘statistic’] ‘delta’ : perform predictions on energies of target dataset,

and do least sqaure on the errors to obtain the target shift as bias.

‘statistic’ : directly use the statistic energy bias in the target dataset.

data_stat(data)[source]: Data staticis.

get_ntypes()[source]: Get the number of types.

get_numb_aparam() → int[source]: Get the number of atomic parameters.

get_numb_fparam() → int[source]: Get the number of frame parameters.

get_rcut()[source]: Get cutoff radius of the model.

get_type_map()[source]: Get the type map.

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

model_type = 'ener'

natoms_match(force, natoms)[source]

natoms_not_match(force, natoms, atype)[source]

class deepmd.tf.model.FrozenModel(*args, **kwargs)[source]

Bases: Model

Load model from a frozen model, which cannot be trained.

Parameters

model_filestr: The path to the frozen model

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box: Tensor, mesh: Tensor, input_dict: dict, frz_model: Optional[str] = None, ckpt_meta: Optional[str] = None, suffix: str = '', reuse: Optional[Union[bool, Enum]] = 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

data_stat(data)[source]: Data staticis.

enable_compression(suffix: str = '') → None[source]

Enable compression.

Parameters

suffixstr: suffix to name scope

get_fitting() → Union[Fitting, dict][source]: Get the fitting(s).

get_loss(loss: dict, lr) → Optional[Union[Loss, dict]][source]: Get the loss function(s).

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

get_rcut()[source]: Get cutoff radius of the model.

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

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

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.tf.model.GlobalPolarModel(*args, **kwargs)[source]

Bases: TensorModel

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_out_size
get_sel_type

class deepmd.tf.model.LinearEnergyModel(*args, **kwargs)[source]

Bases: LinearModel

Linear energy model make linear combinations of several existing energy models.

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_ntypes

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box: Tensor, mesh: Tensor, input_dict: dict, frz_model: Optional[str] = None, ckpt_meta: Optional[str] = None, suffix: str = '', reuse: Optional[Union[bool, Enum]] = 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

model_type = 'ener'

class deepmd.tf.model.MultiModel(*args, **kwargs)[source]

Bases: Model

Multi-task model.

Parameters

descriptor: Descriptor
fitting_net_dict: Dictionary of fitting nets
fitting_type_dict: deprecated argument
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
use_srtab: The table for the short-range pairwise interaction added on top of DP. The table is a text data file with (N_t + 1) * N_t / 2 + 1 columes. The first colume is the distance between atoms. The second to the last columes are energies for pairs of certain types. For example we have two atom types, 0 and 1. The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly.
smin_alpha: The short-range tabulated interaction will be swithed according to the distance of the nearest neighbor. This distance is calculated by softmin. This parameter is the decaying parameter in the softmin. It is only required when use_srtab is provided.
sw_rmin: The lower boundary of the interpolation between short-range tabulated interaction and DP. It is only required when use_srtab is provided.
sw_rmin: The upper boundary of the interpolation between short-range tabulated interaction and DP. It is only required when use_srtab is provided.

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

build(coord_, atype_, natoms, box, mesh, input_dict, frz_model=None, ckpt_meta: Optional[str] = None, suffix='', reuse=None)[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

data_stat(data)[source]: Data staticis.

enable_mixed_precision(mixed_prec: dict)[source]

Enable mixed precision for the model.

Parameters

mixed_precdict: The mixed precision config

get_fitting() → dict[source]: Get the fitting(s).

get_loss(loss: dict, lr: dict) → Dict[str, Loss][source]: Get the loss function(s).

get_ntypes()[source]: Get the number of types.

get_numb_aparam() → dict[source]: Get the number of atomic parameters.

get_numb_dos() → dict[source]: Get the number of gridpoints in energy space.

get_numb_fparam() → dict[source]: Get the number of frame parameters.

get_rcut()[source]: Get cutoff radius of the model.

get_type_map()[source]: Get the type map.

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

model_type = 'multi_task'

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.tf.model.PairTabModel(*args, **kwargs)[source]

Bases: Model

Pairwise tabulation energy model.

This model can be used to tabulate the pairwise energy between atoms for either short-range or long-range interactions, such as D3, LJ, ZBL, etc. It should not be used alone, but rather as one submodel of a linear (sum) model, such as DP+D3.

Do not put the model on the first model of a linear model, since the linear model fetches the type map from the first model.

At this moment, the model does not smooth the energy at the cutoff radius, so one needs to make sure the energy has been smoothed to zero.

Parameters

tab_filestr: The path to the tabulation file.
rcutfloat: The cutoff radius
selint or list[int]: The maxmum number of atoms in the cut-off radius

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box: Tensor, mesh: Tensor, input_dict: dict, frz_model: Optional[str] = None, ckpt_meta: Optional[str] = None, suffix: str = '', reuse: Optional[Union[bool, Enum]] = None)[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

data_stat(data: dict)[source]: Data staticis.

enable_compression(suffix: str = '') → None[source]

Enable compression.

Parameters

suffixstr: suffix to name scope

get_fitting() → Union[Fitting, dict][source]: Get the fitting(s).

get_loss(loss: dict, lr) → Optional[Union[Loss, dict]][source]: Get the loss function(s).

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

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

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

model_type = 'ener'

classmethod update_sel(global_jdata: dict, local_jdata: dict) → 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

Returns

dict: The updated local data

Notes

Do not modify the input data without copying it.

class deepmd.tf.model.PairwiseDPRc(*args, **kwargs)[source]

Bases: Model

Pairwise Deep Potential - Range Correction.

Methods

`build`(coord_, atype_, natoms, box_, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

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

data_stat(data)[source]: Data staticis.

enable_compression(suffix: str = '') → None[source]

Enable compression.

Parameters

suffixstr: suffix to name scope

get_feed_dict(coord_: Tensor, atype_: Tensor, natoms: Tensor, box: Tensor, mesh: Tensor, **kwargs) → Dict[str, Tensor][source]

Generate the feed_dict for current 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
boxtf.Tensor: The box. Can be generated by deepmd.tf.model.make_stat_input
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.
aparamtf.Tensor: The parameters of the descriptor
**kwargsdict: The keyword arguments

Returns

feed_dictdict[str, tf.Tensor]: The output feed_dict of current descriptor

get_fitting() → Union[str, dict][source]: Get the fitting(s).

get_loss(loss: dict, lr) → Union[Loss, dict][source]: Get the loss function(s).

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

get_rcut()[source]: Get cutoff radius of the model.

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

model_type = 'ener'

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.tf.model.PolarModel(*args, **kwargs)[source]

Bases: TensorModel

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_out_size
get_sel_type

class deepmd.tf.model.WFCModel(*args, **kwargs)[source]

Bases: TensorModel

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_out_size
get_sel_type

Submodules

deepmd.tf.model.dos module

class deepmd.tf.model.dos.DOSModel(*args, **kwargs)[source]

Bases: StandardModel

DOS model.

Parameters

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

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

build(coord_, atype_, natoms, box, mesh, input_dict, frz_model=None, ckpt_meta: Optional[str] = None, suffix='', reuse=None)[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

data_stat(data)[source]: Data staticis.

get_ntypes()[source]: Get the number of types.

get_numb_aparam() → int[source]: Get the number of atomic parameters.

get_numb_dos()[source]: Get the number of gridpoints in energy space.

get_numb_fparam() → int[source]: Get the number of frame parameters.

get_rcut()[source]: Get cutoff radius of the model.

get_type_map()[source]: Get the type map.

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

model_type = 'dos'

deepmd.tf.model.ener module

class deepmd.tf.model.ener.EnerModel(*args, **kwargs)[source]

Bases: StandardModel

Energy model.

Parameters

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
use_srtab: The table for the short-range pairwise interaction added on top of DP. The table is a text data file with (N_t + 1) * N_t / 2 + 1 columes. The first colume is the distance between atoms. The second to the last columes are energies for pairs of certain types. For example we have two atom types, 0 and 1. The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly.
smin_alpha: The short-range tabulated interaction will be swithed according to the distance of the nearest neighbor. This distance is calculated by softmin. This parameter is the decaying parameter in the softmin. It is only required when use_srtab is provided.
sw_rmin: The lower boundary of the interpolation between short-range tabulated interaction and DP. It is only required when use_srtab is provided.
sw_rmin: The upper boundary of the interpolation between short-range tabulated interaction and DP. It is only required when use_srtab is provided.
srtab_add_biasbool: Whether add energy bias from the statistics of the data to short-range tabulated atomic energy. It only takes effect when use_srtab is provided.
spin: spin
data_stat_nsample: The number of training samples in a system to compute and change the energy bias.

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

natoms_match
natoms_not_match

build(coord_, atype_, natoms, box, mesh, input_dict, frz_model=None, ckpt_meta: Optional[str] = None, suffix='', reuse=None)[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

change_energy_bias(data: DeepmdDataSystem, frozen_model: str, origin_type_map: list, full_type_map: str, bias_shift: str = 'delta') → None[source]

Change the energy bias according to the input data and the pretrained model.

Parameters

dataDeepmdDataSystem

The training data.

frozen_modelstr

The path file of frozen model.

origin_type_maplist

The original type_map in dataset, they are targets to change the energy bias.

full_type_mapstr

The full type_map in pretrained model

bias_shiftstr

The mode for changing energy bias : [‘delta’, ‘statistic’] ‘delta’ : perform predictions on energies of target dataset,

and do least sqaure on the errors to obtain the target shift as bias.

‘statistic’ : directly use the statistic energy bias in the target dataset.

data_stat(data)[source]: Data staticis.

get_ntypes()[source]: Get the number of types.

get_numb_aparam() → int[source]: Get the number of atomic parameters.

get_numb_fparam() → int[source]: Get the number of frame parameters.

get_rcut()[source]: Get cutoff radius of the model.

get_type_map()[source]: Get the type map.

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

model_type = 'ener'

natoms_match(force, natoms)[source]

natoms_not_match(force, natoms, atype)[source]

deepmd.tf.model.frozen module

class deepmd.tf.model.frozen.FrozenModel(*args, **kwargs)[source]

Bases: Model

Load model from a frozen model, which cannot be trained.

Parameters

model_filestr: The path to the frozen model

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box: Tensor, mesh: Tensor, input_dict: dict, frz_model: Optional[str] = None, ckpt_meta: Optional[str] = None, suffix: str = '', reuse: Optional[Union[bool, Enum]] = 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

data_stat(data)[source]: Data staticis.

enable_compression(suffix: str = '') → None[source]

Enable compression.

Parameters

suffixstr: suffix to name scope

get_fitting() → Union[Fitting, dict][source]: Get the fitting(s).

get_loss(loss: dict, lr) → Optional[Union[Loss, dict]][source]: Get the loss function(s).

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

get_rcut()[source]: Get cutoff radius of the model.

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

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

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.tf.model.linear module

class deepmd.tf.model.linear.LinearEnergyModel(*args, **kwargs)[source]

Bases: LinearModel

Linear energy model make linear combinations of several existing energy models.

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_ntypes

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box: Tensor, mesh: Tensor, input_dict: dict, frz_model: Optional[str] = None, ckpt_meta: Optional[str] = None, suffix: str = '', reuse: Optional[Union[bool, Enum]] = 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

model_type = 'ener'

class deepmd.tf.model.linear.LinearModel(*args, **kwargs)[source]

Bases: Model

Linear model make linear combinations of several existing models.

Parameters

modelslist[dict]: A list of models to be combined.
weightslist[float] or str: If the type is list[float], a list of weights for each model. If “mean”, the weights are set to be 1 / len(models). If “sum”, the weights are set to be 1.

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_ntypes

data_stat(data)[source]: Data staticis.

enable_compression(suffix: str = '') → None[source]

Enable compression.

Parameters

suffixstr: suffix to name scope

get_fitting() → Union[Fitting, dict][source]: Get the fitting(s).

get_loss(loss: dict, lr) → Optional[Union[Loss, dict]][source]: Get the loss function(s).

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

get_rcut()[source]: Get cutoff radius of the model.

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

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

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.tf.model.model module

class deepmd.tf.model.model.Model(*args, **kwargs)[source]

Bases: ABC, PR

Abstract base model.

Parameters

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_bias_nsample: The number of training samples in a system to compute and change the energy bias.
data_stat_protect: Protect parameter for atomic energy regression
use_srtab: The table for the short-range pairwise interaction added on top of DP. The table is a text data file with (N_t + 1) * N_t / 2 + 1 columes. The first colume is the distance between atoms. The second to the last columes are energies for pairs of certain types. For example we have two atom types, 0 and 1. The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly.
smin_alpha: The short-range tabulated interaction will be swithed according to the distance of the nearest neighbor. This distance is calculated by softmin. This parameter is the decaying parameter in the softmin. It is only required when use_srtab is provided.
sw_rmin: The lower boundary of the interpolation between short-range tabulated interaction and DP. It is only required when use_srtab is provided.
sw_rmin: The upper boundary of the interpolation between short-range tabulated interaction and DP. It is only required when use_srtab is provided.
srtab_add_biasbool: Whether add energy bias from the statistics of the data to short-range tabulated atomic energy. It only takes effect when use_srtab is provided.
spin: spin
compress: Compression information for internal use

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

abstract build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box: Tensor, mesh: Tensor, input_dict: dict, frz_model: Optional[str] = None, ckpt_meta: Optional[str] = None, suffix: str = '', reuse: Optional[Union[bool, Enum]] = None)[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

build_descrpt(coord_: Tensor, atype_: Tensor, natoms: Tensor, box: Tensor, mesh: Tensor, input_dict: dict, frz_model: Optional[str] = None, ckpt_meta: Optional[str] = None, suffix: str = '', reuse: Optional[Union[bool, Enum]] = None)[source]

Build the descriptor part of 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

tf.Tensor: The descriptor tensor

build_type_embedding(ntypes: int, frz_model: Optional[str] = None, ckpt_meta: Optional[str] = None, suffix: str = '', reuse: Optional[Union[bool, Enum]] = None) → Tensor[source]

Build the type embedding part of the model.

Parameters

ntypesint: The number of types
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

tf.Tensor: The type embedding tensor

change_energy_bias(data: DeepmdDataSystem, frozen_model: str, origin_type_map: list, full_type_map: str, bias_shift: str = 'delta') → None[source]

Change the energy bias according to the input data and the pretrained model.

Parameters

dataDeepmdDataSystem

The training data.

frozen_modelstr

The path file of frozen model.

origin_type_maplist

The original type_map in dataset, they are targets to change the energy bias.

full_type_mapstr

The full type_map in pretrained model

bias_shiftstr

The mode for changing energy bias : [‘delta’, ‘statistic’] ‘delta’ : perform predictions on energies of target dataset,

and do least sqaure on the errors to obtain the target shift as bias.

‘statistic’ : directly use the statistic energy bias in the target dataset.

abstract data_stat(data: dict)[source]: Data staticis.

classmethod deserialize(data: dict, suffix: str = '') → Model[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 model

Returns

Model: The deserialized Model

enable_compression(suffix: str = '')[source]

Enable compression.

Parameters

suffixstr: suffix to name scope

enable_mixed_precision(mixed_prec: dict)[source]

Enable mixed precision for the model.

Parameters

mixed_precdict: The mixed precision config

get_feed_dict(coord_: Tensor, atype_: Tensor, natoms: Tensor, box: Tensor, mesh: Tensor, **kwargs) → Dict[str, Tensor][source]

Generate the feed_dict for current 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
boxtf.Tensor: The box. Can be generated by deepmd.tf.model.make_stat_input
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.
**kwargsdict: The additional arguments

Returns

feed_dictdict[str, tf.Tensor]: The output feed_dict of current descriptor

abstract get_fitting() → Union[Fitting, dict][source]: Get the fitting(s).

abstract get_loss(loss: dict, lr) → Optional[Union[Loss, dict]][source]: Get the loss function(s).

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

get_numb_aparam() → Union[int, dict][source]: Get the number of atomic parameters.

get_numb_dos() → Union[int, dict][source]: Get the number of gridpoints in energy space.

get_numb_fparam() → Union[int, dict][source]: Get the number of frame parameters.

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

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

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

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

abstract classmethod update_sel(global_jdata: dict, local_jdata: dict) → 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

Returns

dict: The updated local data

Notes

Do not modify the input data without copying it.

class deepmd.tf.model.model.StandardModel(*args, **kwargs)[source]

Bases: Model

Standard model, which must contain a descriptor and a fitting.

Parameters

descriptorUnion[dict, Descriptor]: The descriptor
fitting_netUnion[dict, Fitting]: The fitting network
type_embeddingdict, optional: The type embedding
type_maplist of dict, optional: The type map

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

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

enable_compression(suffix: str = '')[source]

Enable compression.

Parameters

suffixstr: suffix to name scope

enable_mixed_precision(mixed_prec: dict)[source]

Enable mixed precision for the model.

Parameters

mixed_precdict: The mixed precision config

get_fitting() → Union[Fitting, dict][source]: Get the fitting(s).

get_loss(loss: dict, lr) → Union[Loss, dict][source]: Get the loss function(s).

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

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

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 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.tf.model.model_stat module

Alias for backward compatibility.

deepmd.tf.model.model_stat.make_stat_input(data, nbatches, merge_sys=True)[source]

Pack data for statistics.

Parameters

data: The data
nbatchesint: The number of batches
merge_sysbool (True): Merge system data

Returns

all_stat:

A dictionary of list of list storing data for stat. if merge_sys == False data can be accessed by

all_stat[key][sys_idx][batch_idx][frame_idx]

else merge_sys == True can be accessed by: all_stat[key][batch_idx][frame_idx]

deepmd.tf.model.model_stat.merge_sys_stat(all_stat)[source]

deepmd.tf.model.multi module

class deepmd.tf.model.multi.MultiModel(*args, **kwargs)[source]

Bases: Model

Multi-task model.

Parameters

descriptor: Descriptor
fitting_net_dict: Dictionary of fitting nets
fitting_type_dict: deprecated argument
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
use_srtab: The table for the short-range pairwise interaction added on top of DP. The table is a text data file with (N_t + 1) * N_t / 2 + 1 columes. The first colume is the distance between atoms. The second to the last columes are energies for pairs of certain types. For example we have two atom types, 0 and 1. The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly.
smin_alpha: The short-range tabulated interaction will be swithed according to the distance of the nearest neighbor. This distance is calculated by softmin. This parameter is the decaying parameter in the softmin. It is only required when use_srtab is provided.
sw_rmin: The lower boundary of the interpolation between short-range tabulated interaction and DP. It is only required when use_srtab is provided.
sw_rmin: The upper boundary of the interpolation between short-range tabulated interaction and DP. It is only required when use_srtab is provided.

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

build(coord_, atype_, natoms, box, mesh, input_dict, frz_model=None, ckpt_meta: Optional[str] = None, suffix='', reuse=None)[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

data_stat(data)[source]: Data staticis.

enable_mixed_precision(mixed_prec: dict)[source]

Enable mixed precision for the model.

Parameters

mixed_precdict: The mixed precision config

get_fitting() → dict[source]: Get the fitting(s).

get_loss(loss: dict, lr: dict) → Dict[str, Loss][source]: Get the loss function(s).

get_ntypes()[source]: Get the number of types.

get_numb_aparam() → dict[source]: Get the number of atomic parameters.

get_numb_dos() → dict[source]: Get the number of gridpoints in energy space.

get_numb_fparam() → dict[source]: Get the number of frame parameters.

get_rcut()[source]: Get cutoff radius of the model.

get_type_map()[source]: Get the type map.

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

model_type = 'multi_task'

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.tf.model.pairtab module

class deepmd.tf.model.pairtab.PairTabModel(*args, **kwargs)[source]

Bases: Model

Pairwise tabulation energy model.

This model can be used to tabulate the pairwise energy between atoms for either short-range or long-range interactions, such as D3, LJ, ZBL, etc. It should not be used alone, but rather as one submodel of a linear (sum) model, such as DP+D3.

Do not put the model on the first model of a linear model, since the linear model fetches the type map from the first model.

At this moment, the model does not smooth the energy at the cutoff radius, so one needs to make sure the energy has been smoothed to zero.

Parameters

tab_filestr: The path to the tabulation file.
rcutfloat: The cutoff radius
selint or list[int]: The maxmum number of atoms in the cut-off radius

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box: Tensor, mesh: Tensor, input_dict: dict, frz_model: Optional[str] = None, ckpt_meta: Optional[str] = None, suffix: str = '', reuse: Optional[Union[bool, Enum]] = None)[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

data_stat(data: dict)[source]: Data staticis.

enable_compression(suffix: str = '') → None[source]

Enable compression.

Parameters

suffixstr: suffix to name scope

get_fitting() → Union[Fitting, dict][source]: Get the fitting(s).

get_loss(loss: dict, lr) → Optional[Union[Loss, dict]][source]: Get the loss function(s).

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

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

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

model_type = 'ener'

classmethod update_sel(global_jdata: dict, local_jdata: dict) → 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

Returns

dict: The updated local data

Notes

Do not modify the input data without copying it.

deepmd.tf.model.pairwise_dprc module

class deepmd.tf.model.pairwise_dprc.PairwiseDPRc(*args, **kwargs)[source]

Bases: Model

Pairwise Deep Potential - Range Correction.

Methods

`build`(coord_, atype_, natoms, box_, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

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

data_stat(data)[source]: Data staticis.

enable_compression(suffix: str = '') → None[source]

Enable compression.

Parameters

suffixstr: suffix to name scope

get_feed_dict(coord_: Tensor, atype_: Tensor, natoms: Tensor, box: Tensor, mesh: Tensor, **kwargs) → Dict[str, Tensor][source]

Generate the feed_dict for current 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
boxtf.Tensor: The box. Can be generated by deepmd.tf.model.make_stat_input
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.
aparamtf.Tensor: The parameters of the descriptor
**kwargsdict: The keyword arguments

Returns

feed_dictdict[str, tf.Tensor]: The output feed_dict of current descriptor

get_fitting() → Union[str, dict][source]: Get the fitting(s).

get_loss(loss: dict, lr) → Union[Loss, dict][source]: Get the loss function(s).

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

get_rcut()[source]: Get cutoff radius of the model.

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

model_type = 'ener'

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.tf.model.pairwise_dprc.gather_placeholder(params: Tensor, indices: Tensor, placeholder: float = 0.0, **kwargs) → Tensor[source]: Call tf.gather but allow indices to contain placeholders (-1).

deepmd.tf.model.tensor module

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

Bases: TensorModel

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_out_size
get_sel_type

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

Bases: TensorModel

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_out_size
get_sel_type

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

Bases: TensorModel

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_out_size
get_sel_type

class deepmd.tf.model.tensor.TensorModel(*args, **kwargs)[source]

Bases: 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

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_out_size
get_sel_type

build(coord_, atype_, natoms, box, mesh, input_dict, frz_model=None, ckpt_meta: Optional[str] = None, suffix='', reuse=None)[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

data_stat(data)[source]: Data staticis.

get_ntypes()[source]: Get the number of types.

get_out_size()[source]

get_rcut()[source]: Get cutoff radius of the model.

get_sel_type()[source]

get_type_map()[source]: Get the type map.

init_variables(graph: Graph, graph_def: 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, **kwargs)[source]

Bases: TensorModel

Methods

`build`(coord_, atype_, natoms, box, mesh, ...)	Build the model.
`build_descrpt`(coord_, atype_, natoms, box, ...)	Build the descriptor part of the model.
`build_type_embedding`(ntypes[, frz_model, ...])	Build the type embedding part of the model.
`change_energy_bias`(data, frozen_model, ...)	Change the energy bias according to the input data and the pretrained model.
`data_stat`(data)	Data staticis.
`deserialize`(data[, suffix])	Deserialize the model.
`enable_compression`([suffix])	Enable compression.
`enable_mixed_precision`(mixed_prec)	Enable mixed precision for the model.
`get_class_by_type`(class_type)	Get the class by the plugin type.
`get_feed_dict`(coord_, atype_, natoms, box, ...)	Generate the feed_dict for current descriptor.
`get_fitting`()	Get the fitting(s).
`get_loss`(loss, lr)	Get the loss function(s).
`get_ntypes`()	Get the number of types.
`get_numb_aparam`()	Get the number of atomic parameters.
`get_numb_dos`()	Get the number of gridpoints in energy space.
`get_numb_fparam`()	Get the number of frame parameters.
`get_plugins`()	Get all the registered plugins.
`get_rcut`()	Get cutoff radius of the model.
`get_type_map`()	Get the type map.
`init_variables`(graph, graph_def[, ...])	Init the embedding net variables with the given frozen model.
`register`(key)	Register a descriptor plugin.
`serialize`([suffix])	Serialize the model.
`update_sel`(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_out_size
get_sel_type