deepmd.model package

class deepmd.model.DOSModel(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
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)

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)

Data staticis.

get_ntypes()

Get the number of types.

get_numb_aparam() → int

Get the number of atomic parameters.

get_numb_dos()

Get the number of gridpoints in energy space.

get_numb_fparam() → int

Get the number of frame parameters.

get_rcut()

Get cutoff radius of the model.

get_type_map()

Get the type map.

init_variables(graph: Graph, graph_def: GraphDef, model_type: str = 'original_model', suffix: str = '') → None

Init the embedding net variables with the given frozen model.

Parameters

graphtf.Graph

The input frozen model graph

graph_deftf.GraphDef

The input frozen model graph_def

model_typestr

the type of the model

suffixstr

suffix to name scope

model_type = 'dos'

class deepmd.model.DipoleModel(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
update_sel(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_out_size
get_sel_type

class deepmd.model.EnerModel(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
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)

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

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)

Data staticis.

get_ntypes()

Get the number of types.

get_numb_aparam() → int

Get the number of atomic parameters.

get_numb_fparam() → int

Get the number of frame parameters.

get_rcut()

Get cutoff radius of the model.

get_type_map()

Get the type map.

init_variables(graph: Graph, graph_def: GraphDef, model_type: str = 'original_model', suffix: str = '') → None

Init the embedding net variables with the given frozen model.

Parameters

graphtf.Graph

The input frozen model graph

graph_deftf.GraphDef

The input frozen model graph_def

model_typestr

the type of the model

suffixstr

suffix to name scope

model_type = 'ener'

natoms_match(force, natoms)

natoms_not_match(force, natoms, atype)

class deepmd.model.GlobalPolarModel(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
update_sel(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_out_size
get_sel_type

class deepmd.model.MultiModel(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
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)

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)

Data staticis.

enable_mixed_precision(mixed_prec: dict)

Enable mixed precision for the model.

Parameters

mixed_precdict

The mixed precision config

get_fitting() → dict

Get the fitting(s).

get_loss(loss: dict, lr: dict) → Dict[str, Loss]

Get the loss function(s).

get_ntypes()

Get the number of types.

get_numb_aparam() → dict

Get the number of atomic parameters.

get_numb_dos() → dict

Get the number of gridpoints in energy space.

get_numb_fparam() → dict

Get the number of frame parameters.

get_rcut()

Get cutoff radius of the model.

get_type_map()

Get the type map.

init_variables(graph: Graph, graph_def: GraphDef, model_type: str = 'original_model', suffix: str = '') → None

Init the embedding net variables with the given frozen model.

Parameters

graphtf.Graph

The input frozen model graph

graph_deftf.GraphDef

The input frozen model graph_def

model_typestr

the type of the model

suffixstr

suffix to name scope

model_type = 'multi_task'

classmethod update_sel(global_jdata: dict, local_jdata: dict)

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.model.PolarModel(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
update_sel(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_out_size
get_sel_type

class deepmd.model.WFCModel(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
update_sel(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_out_size
get_sel_type

Submodules

deepmd.model.dos module

class deepmd.model.dos.DOSModel(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
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)

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)

Data staticis.

get_ntypes()

Get the number of types.

get_numb_aparam() → int

Get the number of atomic parameters.

get_numb_dos()

Get the number of gridpoints in energy space.

get_numb_fparam() → int

Get the number of frame parameters.

get_rcut()

Get cutoff radius of the model.

get_type_map()

Get the type map.

init_variables(graph: Graph, graph_def: GraphDef, model_type: str = 'original_model', suffix: str = '') → None

Init the embedding net variables with the given frozen model.

Parameters

graphtf.Graph

The input frozen model graph

graph_deftf.GraphDef

The input frozen model graph_def

model_typestr

the type of the model

suffixstr

suffix to name scope

model_type = 'dos'

deepmd.model.ener module

class deepmd.model.ener.EnerModel(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
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)

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

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)

Data staticis.

get_ntypes()

Get the number of types.

get_numb_aparam() → int

Get the number of atomic parameters.

get_numb_fparam() → int

Get the number of frame parameters.

get_rcut()

Get cutoff radius of the model.

get_type_map()

Get the type map.

init_variables(graph: Graph, graph_def: GraphDef, model_type: str = 'original_model', suffix: str = '') → None

Init the embedding net variables with the given frozen model.

Parameters

graphtf.Graph

The input frozen model graph

graph_deftf.GraphDef

The input frozen model graph_def

model_typestr

the type of the model

suffixstr

suffix to name scope

model_type = 'ener'

natoms_match(force, natoms)

natoms_not_match(force, natoms, atype)

deepmd.model.frozen module

class deepmd.model.frozen.FrozenModel(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
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

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)

Data staticis.

enable_compression(suffix: str = '') → None

Enable compression.

Parameters

suffixstr

suffix to name scope

get_fitting() → Union[Fitting, dict]

Get the fitting(s).

get_loss(loss: dict, lr) → Optional[Union[Loss, dict]]

Get the loss function(s).

get_ntypes() → int

Get the number of types.

get_rcut()

Get cutoff radius of the model.

get_type_map() → list

Get the type map.

init_variables(graph: Graph, graph_def: GraphDef, model_type: str = 'original_model', suffix: str = '') → None

Init the embedding net variables with the given frozen model.

Parameters

graphtf.Graph

The input frozen model graph

graph_deftf.GraphDef

The input frozen model graph_def

model_typestr

the type of the model

suffixstr

suffix to name scope

classmethod update_sel(global_jdata: dict, local_jdata: dict)

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

class deepmd.model.linear.LinearEnergyModel(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
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

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.model.linear.LinearModel(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
update_sel(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_ntypes

data_stat(data)

Data staticis.

enable_compression(suffix: str = '') → None

Enable compression.

Parameters

suffixstr

suffix to name scope

get_fitting() → Union[Fitting, dict]

Get the fitting(s).

get_loss(loss: dict, lr) → Optional[Union[Loss, dict]]

Get the loss function(s).

get_ntypes() → int

Get the number of types.

get_rcut()

Get cutoff radius of the model.

get_type_map() → list

Get the type map.

init_variables(graph: Graph, graph_def: GraphDef, model_type: str = 'original_model', suffix: str = '') → None

Init the embedding net variables with the given frozen model.

Parameters

graphtf.Graph

The input frozen model graph

graph_deftf.GraphDef

The input frozen model graph_def

model_typestr

the type of the model

suffixstr

suffix to name scope

classmethod update_sel(global_jdata: dict, local_jdata: dict)

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

class deepmd.model.model.Model(*args, **kwargs)

Bases: ABC

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
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)

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)

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

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

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)

Data staticis.

enable_compression(suffix: str = '')

Enable compression.

Parameters

suffixstr

suffix to name scope

enable_mixed_precision(mixed_prec: dict)

Enable mixed precision for the model.

Parameters

mixed_precdict

The mixed precision config

classmethod get_class_by_input(input: dict)

Get the class by input data.

Parameters

inputdict

The input data

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

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.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]

Get the fitting(s).

abstract get_loss(loss: dict, lr) → Optional[Union[Loss, dict]]

Get the loss function(s).

abstract get_ntypes() → int

Get the number of types.

get_numb_aparam() → Union[int, dict]

Get the number of atomic parameters.

get_numb_dos() → Union[int, dict]

Get the number of gridpoints in energy space.

get_numb_fparam() → Union[int, dict]

Get the number of frame parameters.

abstract get_rcut() → float

Get cutoff radius of the model.

get_type_map() → list

Get the type map.

init_variables(graph: Graph, graph_def: GraphDef, model_type: str = 'original_model', suffix: str = '') → None

Init the embedding net variables with the given frozen model.

Parameters

graphtf.Graph

The input frozen model graph

graph_deftf.GraphDef

The input frozen model graph_def

model_typestr

the type of the model

suffixstr

suffix to name scope

abstract classmethod update_sel(global_jdata: dict, local_jdata: dict) → dict

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.model.model.StandardModel(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
update_sel(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

enable_compression(suffix: str = '')

Enable compression.

Parameters

suffixstr

suffix to name scope

enable_mixed_precision(mixed_prec: dict)

Enable mixed precision for the model.

Parameters

mixed_precdict

The mixed precision config

get_fitting() → Union[Fitting, dict]

Get the fitting(s).

get_loss(loss: dict, lr) → Union[Loss, dict]

Get the loss function(s).

get_ntypes() → int

Get the number of types.

get_rcut() → float

Get cutoff radius of the model.

classmethod update_sel(global_jdata: dict, local_jdata: dict)

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

Alias for backward compatibility.

deepmd.model.model_stat.make_stat_input(data, nbatches, merge_sys=True)

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.model.model_stat.merge_sys_stat(all_stat)

deepmd.model.multi module

class deepmd.model.multi.MultiModel(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
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)

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)

Data staticis.

enable_mixed_precision(mixed_prec: dict)

Enable mixed precision for the model.

Parameters

mixed_precdict

The mixed precision config

get_fitting() → dict

Get the fitting(s).

get_loss(loss: dict, lr: dict) → Dict[str, Loss]

Get the loss function(s).

get_ntypes()

Get the number of types.

get_numb_aparam() → dict

Get the number of atomic parameters.

get_numb_dos() → dict

Get the number of gridpoints in energy space.

get_numb_fparam() → dict

Get the number of frame parameters.

get_rcut()

Get cutoff radius of the model.

get_type_map()

Get the type map.

init_variables(graph: Graph, graph_def: GraphDef, model_type: str = 'original_model', suffix: str = '') → None

Init the embedding net variables with the given frozen model.

Parameters

graphtf.Graph

The input frozen model graph

graph_deftf.GraphDef

The input frozen model graph_def

model_typestr

the type of the model

suffixstr

suffix to name scope

model_type = 'multi_task'

classmethod update_sel(global_jdata: dict, local_jdata: dict)

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

class deepmd.model.pairtab.PairTabModel(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
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)

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)

Data staticis.

enable_compression(suffix: str = '') → None

Enable compression.

Parameters

suffixstr

suffix to name scope

get_fitting() → Union[Fitting, dict]

Get the fitting(s).

get_loss(loss: dict, lr) → Optional[Union[Loss, dict]]

Get the loss function(s).

get_ntypes() → int

Get the number of types.

get_rcut() → float

Get cutoff radius of the model.

init_variables(graph: Graph, graph_def: GraphDef, model_type: str = 'original_model', suffix: str = '') → None

Init the embedding net variables with the given frozen model.

Parameters

graphtf.Graph

The input frozen model graph

graph_deftf.GraphDef

The input frozen model graph_def

model_typestr

the type of the model

suffixstr

suffix to name scope

model_type = 'ener'

classmethod update_sel(global_jdata: dict, local_jdata: dict) → dict

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.model.pairwise_dprc module

class deepmd.model.pairwise_dprc.PairwiseDPRc(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
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)

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)

Data staticis.

enable_compression(suffix: str = '') → None

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]

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.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]

Get the fitting(s).

get_loss(loss: dict, lr) → Union[Loss, dict]

Get the loss function(s).

get_ntypes() → int

Get the number of types.

get_rcut()

Get cutoff radius of the model.

init_variables(graph: Graph, graph_def: GraphDef, model_type: str = 'original_model', suffix: str = '') → None

Init the embedding net variables with the given frozen model.

Parameters

graphtf.Graph

The input frozen model graph

graph_deftf.GraphDef

The input frozen model graph_def

model_typestr

the type of the model

suffixstr

suffix to name scope

model_type = 'ener'

classmethod update_sel(global_jdata: dict, local_jdata: dict)

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.model.pairwise_dprc.gather_placeholder(params: Tensor, indices: Tensor, placeholder: float = 0.0, **kwargs) → Tensor

Call tf.gather but allow indices to contain placeholders (-1).

deepmd.model.tensor module

class deepmd.model.tensor.DipoleModel(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
update_sel(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_out_size
get_sel_type

class deepmd.model.tensor.GlobalPolarModel(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
update_sel(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_out_size
get_sel_type

class deepmd.model.tensor.PolarModel(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
update_sel(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_out_size
get_sel_type

class deepmd.model.tensor.TensorModel(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
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)

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)

Data staticis.

get_ntypes()

Get the number of types.

get_out_size()

get_rcut()

Get cutoff radius of the model.

get_sel_type()

get_type_map()

Get the type map.

init_variables(graph: Graph, graph_def: GraphDef, model_type: str = 'original_model', suffix: str = '') → None

Init the embedding net variables with the given frozen model.

Parameters

graphtf.Graph

The input frozen model graph

graph_deftf.GraphDef

The input frozen model graph_def

model_typestr

the type of the model

suffixstr

suffix to name scope

class deepmd.model.tensor.WFCModel(*args, **kwargs)

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.
enable_compression([suffix])	Enable compression.
enable_mixed_precision(mixed_prec)	Enable mixed precision for the model.
get_class_by_input(input)	Get the class by input data.
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_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.
update_sel(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

get_out_size
get_sel_type