deepmd.pt.model.model

The model that takes the coordinates, cell and atom types as input and predicts some property. The models are automatically generated from atomic models by the deepmd.dpmodel.make_model method.

The make_model method does the reduction, auto-differentiation and communication of the atomic properties according to output variable definition deepmd.dpmodel.OutputVariableDef.

All models should be inherited from deepmd.pt.model.model.model.BaseModel. Models generated by make_model have already done it.

Submodules

• deepmd.pt.model.model.dipole_model
• deepmd.pt.model.model.dos_model
• deepmd.pt.model.model.dp_linear_model
• deepmd.pt.model.model.dp_model
• deepmd.pt.model.model.dp_zbl_model
• deepmd.pt.model.model.ener_model
• deepmd.pt.model.model.frozen
• deepmd.pt.model.model.make_hessian_model
• deepmd.pt.model.model.make_model
• deepmd.pt.model.model.model
• deepmd.pt.model.model.polar_model
• deepmd.pt.model.model.property_model
• deepmd.pt.model.model.spin_model
• deepmd.pt.model.model.transform_output

Classes

DipoleModel	A base class to implement common methods for all the Models.
DOSModel	A base class to implement common methods for all the Models.
LinearEnergyModel	Base class for final exported model that will be directly used for inference.
DPModelCommon	A base class to implement common methods for all the Models.
DPZBLModel	Base class for final exported model that will be directly used for inference.
EnergyModel	A base class to implement common methods for all the Models.
FrozenModel	Load model from a frozen model, which cannot be trained.
BaseModel	Base class for final exported model that will be directly used for inference.
PolarModel	A base class to implement common methods for all the Models.
SpinEnergyModel	A spin model for energy.
SpinModel	A spin model wrapper, with spin input preprocess and output split.

Functions

make_hessian_model(T_Model)	Make a model that can compute Hessian.
make_model(T_AtomicModel)	Make a model as a derived class of an atomic model.
get_model(model_params)

Package Contents

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

Bases: deepmd.pt.model.model.dp_model.DPModelCommon, DPDipoleModel_

A base class to implement common methods for all the Models.

model_type = 'dipole'

translated_output_def()

forward(coord, atype, box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False) → dict[str, torch.Tensor]

forward_lower(extended_coord, extended_atype, nlist, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False)

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

Bases: deepmd.pt.model.model.dp_model.DPModelCommon, DPDOSModel_

A base class to implement common methods for all the Models.

model_type = 'dos'

translated_output_def()

forward(coord, atype, box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False) → dict[str, torch.Tensor]

get_numb_dos() → int

Get the number of DOS for DOSFittingNet.

forward_lower(extended_coord, extended_atype, nlist, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False)

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

Bases: DPLinearModel_

Base class for final exported model that will be directly used for inference.

The class defines some abstractmethods that will be directly called by the inference interface. If the final model class inherits some of those methods from other classes, BaseModel should be inherited as the last class to ensure the correct method resolution order.

This class is backend-indepedent.

See also

deepmd.dpmodel.model.base_model.BaseModel

BaseModel class for DPModel backend.

model_type = 'linear_ener'

translated_output_def()

forward(coord, atype, box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False) → dict[str, torch.Tensor]

forward_lower(extended_coord, extended_atype, nlist, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False)

classmethod update_sel(train_data: deepmd.utils.data_system.DeepmdDataSystem, type_map: list[str] | None, local_jdata: dict) → tuple[dict, float | None]

Update the selection and perform neighbor statistics.

Parameters:

train_dataDeepmdDataSystem

data used to do neighbor statistics

type_maplist[str], optional

The name of each type of atoms

local_jdatadict

The local data refer to the current class

Returns:

dict

The updated local data

float

The minimum distance between two atoms

class deepmd.pt.model.model.DPModelCommon

A base class to implement common methods for all the Models.

classmethod update_sel(train_data: deepmd.utils.data_system.DeepmdDataSystem, type_map: list[str] | None, local_jdata: dict) → tuple[dict, float | None]

Update the selection and perform neighbor statistics.

Parameters:

train_dataDeepmdDataSystem

data used to do neighbor statistics

type_maplist[str], optional

The name of each type of atoms

local_jdatadict

The local data refer to the current class

Returns:

dict

The updated local data

float

The minimum distance between two atoms

get_fitting_net()

Get the fitting network.

get_descriptor()

Get the descriptor.

set_eval_descriptor_hook(enable: bool) → None

Set the hook for evaluating descriptor and clear the cache for descriptor list.

eval_descriptor() → torch.Tensor

Evaluate the descriptor.

class deepmd.pt.model.model.DPZBLModel(*args, **kwargs)

Bases: DPZBLModel_

Base class for final exported model that will be directly used for inference.

The class defines some abstractmethods that will be directly called by the inference interface. If the final model class inherits some of those methods from other classes, BaseModel should be inherited as the last class to ensure the correct method resolution order.

This class is backend-indepedent.

See also

deepmd.dpmodel.model.base_model.BaseModel

BaseModel class for DPModel backend.

model_type = 'zbl'

translated_output_def()

forward(coord, atype, box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False) → dict[str, torch.Tensor]

forward_lower(extended_coord, extended_atype, nlist, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False)

classmethod update_sel(train_data: deepmd.utils.data_system.DeepmdDataSystem, type_map: list[str] | None, local_jdata: dict) → tuple[dict, float | None]

Update the selection and perform neighbor statistics.

Parameters:

train_dataDeepmdDataSystem

data used to do neighbor statistics

type_maplist[str], optional

The name of each type of atoms

local_jdatadict

The local data refer to the current class

Returns:

dict

The updated local data

float

The minimum distance between two atoms

class deepmd.pt.model.model.EnergyModel(*args, **kwargs)

Bases: deepmd.pt.model.model.dp_model.DPModelCommon, DPEnergyModel_

A base class to implement common methods for all the Models.

model_type = 'ener'

translated_output_def()

forward(coord, atype, box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False) → dict[str, torch.Tensor]

forward_lower(extended_coord, extended_atype, nlist, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None)

class deepmd.pt.model.model.FrozenModel(model_file: str, **kwargs)

Bases: deepmd.pt.model.model.model.BaseModel

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

Parameters:

model_filestr

The path to the frozen model

model_file

fitting_output_def() → deepmd.dpmodel.output_def.FittingOutputDef

Get the output def of developer implemented atomic models.

get_rcut() → float

Get the cut-off radius.

get_type_map() → list[str]

Get the type map.

get_sel() → list[int]

Returns the number of selected atoms for each type.

get_dim_fparam() → int

Get the number (dimension) of frame parameters of this atomic model.

get_dim_aparam() → int

Get the number (dimension) of atomic parameters of this atomic model.

get_sel_type() → list[int]

Get the selected atom types of this model.

Only atoms with selected atom types have atomic contribution to the result of the model. If returning an empty list, all atom types are selected.

is_aparam_nall() → bool

Check whether the shape of atomic parameters is (nframes, nall, ndim).

If False, the shape is (nframes, nloc, ndim).

mixed_types() → bool

If true, the model 1. assumes total number of atoms aligned across frames; 2. uses a neighbor list that does not distinguish different atomic types.

If false, the model 1. assumes total number of atoms of each atom type aligned across frames; 2. uses a neighbor list that distinguishes different atomic types.

has_message_passing() → bool

Returns whether the descriptor has message passing.

need_sorted_nlist_for_lower() → bool

Returns whether the model needs sorted nlist when using forward_lower.

forward(coord, atype, box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False) → dict[str, torch.Tensor]

get_model_def_script() → str

Get the model definition script.

get_min_nbor_dist() → float | None

Get the minimum neighbor distance.

serialize() → dict

Serialize the model.

Returns:

dict

The serialized data

classmethod deserialize(data: dict) → NoReturn

Deserialize the model.

Parameters:

datadict

The serialized data

Returns:

BaseModel

The deserialized model

get_nnei() → int

Returns the total number of selected neighboring atoms in the cut-off radius.

get_nsel() → int

Returns the total number of selected neighboring atoms in the cut-off radius.

classmethod update_sel(train_data: deepmd.utils.data_system.DeepmdDataSystem, type_map: list[str] | None, local_jdata: dict) → tuple[dict, float | None]

Update the selection and perform neighbor statistics.

Parameters:

train_dataDeepmdDataSystem

data used to do neighbor statistics

type_maplist[str], optional

The name of each type of atoms

local_jdatadict

The local data refer to the current class

Returns:

dict

The updated local data

float

The minimum distance between two atoms

model_output_type() → str

Get the output type for the model.

deepmd.pt.model.model.make_hessian_model(T_Model)

Make a model that can compute Hessian.

LIMITATION: this model is not jitable due to the restrictions of torch jit script.

LIMITATION: only the hessian of forward_common is available.

Parameters:

T_Model

The model. Should provide the forward_common and atomic_output_def methods

Returns:

The model computes hessian.

deepmd.pt.model.model.make_model(T_AtomicModel: type[deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel])

Make a model as a derived class of an atomic model.

The model provide two interfaces.

1. the forward_common_lower, that takes extended coordinates, atyps and neighbor list, and outputs the atomic and property and derivatives (if required) on the extended region.

2. the forward_common, that takes coordinates, atypes and cell and predicts the atomic and reduced property, and derivatives (if required) on the local region.

Parameters:

T_AtomicModel

The atomic model.

Returns:

CM

The model.

class deepmd.pt.model.model.BaseModel(*args, **kwargs)

Bases: torch.nn.Module, make_base_model()

Base class for final exported model that will be directly used for inference.

The class defines some abstractmethods that will be directly called by the inference interface. If the final model class inherits some of those methods from other classes, BaseModel should be inherited as the last class to ensure the correct method resolution order.

This class is backend-indepedent.

See also

deepmd.dpmodel.model.base_model.BaseModel

BaseModel class for DPModel backend.

model_def_script = ''

The model definition script.

abstract compute_or_load_stat(sampled_func, stat_file_path: deepmd.utils.path.DPPath | None = None) → NoReturn

Compute or load the statistics parameters of the model, such as mean and standard deviation of descriptors or the energy bias of the fitting net. When sampled is provided, all the statistics parameters will be calculated (or re-calculated for update), and saved in the stat_file_path`(s). When `sampled is not provided, it will check the existence of `stat_file_path`(s) and load the calculated statistics parameters.

Parameters:

sampled_func

The sampled data frames from different data systems.

stat_file_path

The path to the statistics files.

get_model_def_script() → str

Get the model definition script.

get_min_nbor_dist() → float | None

Get the minimum distance between two atoms.

get_ntypes()

Returns the number of element types.

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

Bases: deepmd.pt.model.model.dp_model.DPModelCommon, DPPolarModel_

A base class to implement common methods for all the Models.

model_type = 'polar'

translated_output_def()

forward(coord, atype, box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False) → dict[str, torch.Tensor]

forward_lower(extended_coord, extended_atype, nlist, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False)

class deepmd.pt.model.model.SpinEnergyModel(backbone_model, spin: deepmd.utils.spin.Spin)

Bases: SpinModel

A spin model for energy.

model_type = 'ener'

translated_output_def()

forward(coord, atype, spin, box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False) → dict[str, torch.Tensor]

forward_lower(extended_coord, extended_atype, extended_spin, nlist, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None)

class deepmd.pt.model.model.SpinModel(backbone_model, spin: deepmd.utils.spin.Spin)

Bases: torch.nn.Module

A spin model wrapper, with spin input preprocess and output split.

backbone_model

spin

ntypes_real

virtual_scale_mask

spin_mask

process_spin_input(coord, atype, spin)

Generate virtual coordinates and types, concat into the input.

process_spin_input_lower(extended_coord, extended_atype, extended_spin, nlist, mapping: torch.Tensor | None = None)

Add extended_spin into extended_coord to generate virtual atoms, and extend nlist and mapping. Note that the final extended_coord_updated with shape [nframes, nall + nall, 3] has the following order: - [:, :nloc]: original nloc real atoms. - [:, nloc: nloc + nloc]: virtual atoms corresponding to nloc real atoms. - [:, nloc + nloc: nloc + nall]: ghost real atoms. - [:, nloc + nall: nall + nall]: virtual atoms corresponding to ghost real atoms.

process_spin_output(atype, out_tensor, add_mag: bool = True, virtual_scale: bool = True)

Split the output both real and virtual atoms, and scale the latter. add_mag: whether to add magnetic tensor onto the real tensor.

Default: True. e.g. True for forces and False for atomic virials on real atoms.

virtual_scale: whether to scale the magnetic tensor with virtual scale factor.

Default: True. e.g. True for forces and False for atomic virials on virtual atoms.

process_spin_output_lower(extended_atype, extended_out_tensor, nloc: int, add_mag: bool = True, virtual_scale: bool = True)

Split the extended output of both real and virtual atoms with switch, and scale the latter. add_mag: whether to add magnetic tensor onto the real tensor.

Default: True. e.g. True for forces and False for atomic virials on real atoms.

virtual_scale: whether to scale the magnetic tensor with virtual scale factor.

Default: True. e.g. True for forces and False for atomic virials on virtual atoms.

static extend_nlist(extended_atype, nlist)

static expand_aparam(aparam, nloc: int)

Expand the atom parameters for virtual atoms if necessary.

get_type_map() → list[str]

Get the type map.

get_ntypes()

Returns the number of element types.

get_rcut()

Get the cut-off radius.

get_dim_fparam()

Get the number (dimension) of frame parameters of this atomic model.

get_dim_aparam()

Get the number (dimension) of atomic parameters of this atomic model.

get_sel_type() → list[int]

Get the selected atom types of this model. Only atoms with selected atom types have atomic contribution to the result of the model. If returning an empty list, all atom types are selected.

is_aparam_nall() → bool

Check whether the shape of atomic parameters is (nframes, nall, ndim). If False, the shape is (nframes, nloc, ndim).

model_output_type() → list[str]

Get the output type for the model.

get_model_def_script() → str

Get the model definition script.

get_min_nbor_dist() → float | None

Get the minimum neighbor distance.

get_nnei() → int

Returns the total number of selected neighboring atoms in the cut-off radius.

get_nsel() → int

Returns the total number of selected neighboring atoms in the cut-off radius.

has_spin() → bool

Returns whether it has spin input and output.

has_message_passing() → bool

Returns whether the model has message passing.

need_sorted_nlist_for_lower() → bool

Returns whether the model needs sorted nlist when using forward_lower.

model_output_def()

Get the output def for the model.

__getattr__(name)

Get attribute from the wrapped model.

compute_or_load_stat(sampled_func, stat_file_path: deepmd.utils.path.DPPath | None = None) → None

Compute or load the statistics parameters of the model, such as mean and standard deviation of descriptors or the energy bias of the fitting net. When sampled is provided, all the statistics parameters will be calculated (or re-calculated for update), and saved in the stat_file_path`(s). When `sampled is not provided, it will check the existence of `stat_file_path`(s) and load the calculated statistics parameters.

Parameters:

sampled_func

The lazy sampled function to get data frames from different data systems.

stat_file_path

The dictionary of paths to the statistics files.

forward_common(coord, atype, spin, box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False) → dict[str, torch.Tensor]

forward_common_lower(extended_coord, extended_atype, extended_spin, nlist, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, comm_dict: dict[str, torch.Tensor] | None = None, extra_nlist_sort: bool = False)

serialize() → dict

classmethod deserialize(data) → SpinModel

deepmd.pt.model.model.get_model(model_params)