`deepmd.tf.model`

Submodules

Package Contents

Classes

`FrozenModel`	Load model from a frozen model, which cannot be trained.
`LinearEnergyModel`	Linear energy model make linear combinations of several existing energy models.
`PairTabModel`	Pairwise tabulation energy model.
`PairwiseDPRc`	Pairwise Deep Potential - Range Correction.
`DOSModel`	DOS model.
`EnerModel`	Energy model.
`MultiModel`	Multi-task model.
`DipoleModel`	Tensor model.
`GlobalPolarModel`	Tensor model.
`PolarModel`	Tensor model.
`WFCModel`	Tensor model.

class deepmd.tf.model.FrozenModel(model_file: str, **kwargs)[source]

Bases: deepmd.tf.model.model.Model

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

Parameters:

model_filestr: The path to the frozen model

build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict, frz_model: str | None = None, ckpt_meta: str | None = None, suffix: str = '', reuse: bool | enum.Enum | None = None) → dict[source]

Build the model.

Parameters:

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

Returns:

dict: The output dict

get_fitting() → deepmd.tf.fit.fitting.Fitting | dict[source]: Get the fitting(s).

get_loss(loss: dict, lr) → deepmd.tf.loss.loss.Loss | dict | None[source]: Get the loss function(s).

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

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

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

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

Init the embedding net variables with the given frozen model.

Parameters:

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

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

Enable compression.

Parameters:

suffixstr: suffix to name scope

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

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

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

Serialize the model.

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

Returns:

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

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

class deepmd.tf.model.LinearEnergyModel(models: List[dict], weights: List[float], **kwargs)[source]

Bases: LinearModel

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

model_type = 'ener'

build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict, frz_model: str | None = None, ckpt_meta: str | None = None, suffix: str = '', reuse: bool | enum.Enum | None = None) → dict[source]

Build the model.

Parameters:

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

Returns:

dict: The output dict

class deepmd.tf.model.PairTabModel(tab_file: str, rcut: float, sel: int | List[int], **kwargs)[source]

Bases: deepmd.tf.model.model.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

model_type = 'ener'

build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict, frz_model: str | None = None, ckpt_meta: str | None = None, suffix: str = '', reuse: bool | enum.Enum | None = 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

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

Init the embedding net variables with the given frozen model.

Parameters:

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

get_fitting() → deepmd.tf.fit.fitting.Fitting | dict[source]: Get the fitting(s).

get_loss(loss: dict, lr) → deepmd.tf.loss.loss.Loss | dict | None[source]: Get the loss function(s).

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

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

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

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

Enable compression.

Parameters:

suffixstr: suffix to name scope

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(qm_model: dict, qmmm_model: dict, type_embedding: dict | deepmd.tf.utils.type_embed.TypeEmbedNet, type_map: List[str], data_stat_nbatch: int = 10, data_stat_nsample: int = 10, data_stat_protect: float = 0.01, use_srtab: str | None = None, smin_alpha: float | None = None, sw_rmin: float | None = None, sw_rmax: float | None = None, spin: deepmd.tf.utils.spin.Spin | None = None, compress: dict | None = None, **kwargs)[source]

Bases: deepmd.tf.model.model.Model

Pairwise Deep Potential - Range Correction.

model_type = 'ener'

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

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

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

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

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

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

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

Init the embedding net variables with the given frozen model.

Parameters:

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

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

Enable compression.

Parameters:

suffixstr: suffix to name scope

get_feed_dict(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, **kwargs) → Dict[str, deepmd.tf.env.tf.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

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.DOSModel(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)[source]

Bases: deepmd.tf.model.model.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

model_type = 'dos'

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

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

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

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

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

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

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

_compute_input_stat(all_stat, protection=0.01, mixed_type=False)[source]

_compute_output_stat(all_stat, mixed_type=False)[source]

build(coord_, atype_, natoms, box, mesh, input_dict, frz_model=None, ckpt_meta: str | None = 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

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

Init the embedding net variables with the given frozen model.

Parameters:

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

class deepmd.tf.model.EnerModel(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, use_srtab: str | None = None, smin_alpha: float | None = None, sw_rmin: float | None = None, sw_rmax: float | None = None, srtab_add_bias: bool = True, spin: deepmd.tf.utils.spin.Spin | None = None, data_bias_nsample: int = 10, **kwargs)[source]

Bases: deepmd.tf.model.model.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.

model_type = 'ener'

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

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

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

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

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

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

_compute_input_stat(all_stat, protection=0.01, mixed_type=False)[source]

_compute_output_stat(all_stat, mixed_type=False)[source]

build(coord_, atype_, natoms, box, mesh, input_dict, frz_model=None, ckpt_meta: str | None = 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

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

Init the embedding net variables with the given frozen model.

Parameters:

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

natoms_match(force, natoms)[source]

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

change_energy_bias(data: deepmd.tf.utils.data_system.DeepmdDataSystem, frozen_model: str, origin_type_map: list, full_type_map: str, bias_adjust_mode: str = 'change-by-statistic') → 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_adjust_modestr

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

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

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

class deepmd.tf.model.MultiModel(descriptor: dict, fitting_net_dict: dict, fitting_type_dict: dict | None = None, type_embedding=None, type_map: List[str] | None = None, data_stat_nbatch: int = 10, data_stat_protect: float = 0.01, use_srtab: str | None = None, smin_alpha: float | None = None, sw_rmin: float | None = None, sw_rmax: float | None = None, **kwargs)[source]

Bases: deepmd.tf.model.model.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.

model_type = 'multi_task'

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

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

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

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

_compute_input_stat(all_stat, protection=0.01, mixed_type=False, fitting_key='')[source]

_compute_output_stat(all_stat, mixed_type=False, fitting_key='')[source]

build(coord_, atype_, natoms, box, mesh, input_dict, frz_model=None, ckpt_meta: str | None = 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

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

Init the embedding net variables with the given frozen model.

Parameters:

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

enable_mixed_precision(mixed_prec: dict)[source]

Enable mixed precision for the model.

Parameters:

mixed_precdict: The mixed precision config

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

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_fitting() → dict[source]: Get the fitting(s).

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

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

Bases: TensorModel

Tensor model.

Parameters:

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

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

Bases: TensorModel

Tensor model.

Parameters:

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

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

Bases: TensorModel

Tensor model.

Parameters:

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

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

Bases: TensorModel

Tensor model.

Parameters:

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

deepmd.tf.model

Submodules

Package Contents

Classes

`deepmd.tf.model`