deepmd.pt.model.atomic_model package

The atomic model provides the prediction of some property on each atom. All the atomic models are not supposed to be directly accessed by users, but it provides a convenient interface for the implementation of models.

Taking the energy models for example, the developeres only needs to implement the atomic energy prediction via an atomic model, and the model can be automatically made by the deepmd.dpmodel.make_model method. The DPModel is made by ` DPModel = make_model(DPAtomicModel) `

class deepmd.pt.model.atomic_model.BaseAtomicModel(atom_exclude_types: List[int] = [], pair_exclude_types: List[Tuple[int, int]] = [])[source]

Bases: BAM

Methods

`atomic_output_def`()	Get the output def of the atomic model.
`do_grad_`(var_name, base)	Tell if the output variable var_name is differentiable.
`do_grad_c`([var_name])	Tell if the output variable var_name is c_differentiable.
`do_grad_r`([var_name])	Tell if the output variable var_name is r_differentiable.
`fitting_output_def`()	Get the output def of developer implemented atomic models.
`get_dim_aparam`()	Get the number (dimension) of atomic parameters of this atomic model.
`get_dim_fparam`()	Get the number (dimension) of frame parameters of this atomic model.
`get_nnei`()	Returns the total number of selected neighboring atoms in the cut-off radius.
`get_nsel`()	Returns the total number of selected neighboring atoms in the cut-off radius.
`get_ntypes`()	Get the number of atom types.
`get_rcut`()	Get the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_sel_type`()	Get the selected atom types of this model.
`get_type_map`()	Get the type map.
`is_aparam_nall`()	Check whether the shape of atomic parameters is (nframes, nall, ndim).
`mixed_types`()	If true, the model 1.

deserialize
forward_atomic
forward_common_atomic
get_model_def_script
reinit_atom_exclude
reinit_pair_exclude
serialize

atomic_output_def() → FittingOutputDef[source]

Get the output def of the atomic model.

By default it is the same as FittingOutputDef, but it allows model level wrapper of the output defined by the developer.

forward_common_atomic(extended_coord: Tensor, extended_atype: Tensor, nlist: Tensor, mapping: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None) → Dict[str, Tensor][source]

get_model_def_script() → str[source]

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.

get_ntypes() → int: Get the number of atom types.

reinit_atom_exclude(exclude_types: List[int] = [])[source]

reinit_pair_exclude(exclude_types: List[Tuple[int, int]] = [])[source]

serialize() → dict[source]

class deepmd.pt.model.atomic_model.DPAtomicModel(descriptor, fitting, type_map: Optional[List[str]], **kwargs)[source]

Bases: Module, BaseAtomicModel

Model give atomic prediction of some physical property.

Parameters

descriptor: Descriptor
fitting_net: Fitting 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.

Methods

`add_module`(name, module)	Add a child module to the current module.
`apply`(fn)	Apply `fn` recursively to every submodule (as returned by `.children()`) as well as self.
`atomic_output_def`()	Get the output def of the atomic model.
`bfloat16`()	Casts all floating point parameters and buffers to `bfloat16` datatype.
`buffers`([recurse])	Return an iterator over module buffers.
`children`()	Return an iterator over immediate children modules.
`compile`(args, *kwargs)	Compile this Module's forward using `torch.compile()`.
`compute_or_load_stat`(sampled_func[, ...])	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.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`do_grad_`(var_name, base)	Tell if the output variable var_name is differentiable.
`do_grad_c`([var_name])	Tell if the output variable var_name is c_differentiable.
`do_grad_r`([var_name])	Tell if the output variable var_name is r_differentiable.
`double`()	Casts all floating point parameters and buffers to `double` datatype.
`eval`()	Set the module in evaluation mode.
`extra_repr`()	Set the extra representation of the module.
`fitting_output_def`()	Get the output def of the fitting net.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(*input)	Define the computation performed at every call.
`forward_atomic`(extended_coord, ...[, ...])	Return atomic prediction.
`get_buffer`(target)	Return the buffer given by `target` if it exists, otherwise throw an error.
`get_dim_aparam`()	Get the number (dimension) of atomic parameters of this atomic model.
`get_dim_fparam`()	Get the number (dimension) of frame parameters of this atomic model.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_nnei`()	Returns the total number of selected neighboring atoms in the cut-off radius.
`get_nsel`()	Returns the total number of selected neighboring atoms in the cut-off radius.
`get_ntypes`()	Get the number of atom types.
`get_parameter`(target)	Return the parameter given by `target` if it exists, otherwise throw an error.
`get_rcut`()	Get the cut-off radius.
`get_sel`()	Get the neighbor selection.
`get_sel_type`()	Get the selected atom types of this model.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`get_type_map`()	Get the type map.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`is_aparam_nall`()	Check whether the shape of atomic parameters is (nframes, nall, ndim).
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`mixed_types`()	If true, the model 1.
`modules`()	Return an iterator over all modules in the network.
`named_buffers`([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
`named_children`()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
`named_modules`([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
`named_parameters`([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
`parameters`([recurse])	Return an iterator over module parameters.
`register_backward_hook`(hook)	Register a backward hook on the module.
`register_buffer`(name, tensor[, persistent])	Add a buffer to the module.
`register_forward_hook`(hook, *[, prepend, ...])	Register a forward hook on the module.
`register_forward_pre_hook`(hook, *[, ...])	Register a forward pre-hook on the module.
`register_full_backward_hook`(hook[, prepend])	Register a backward hook on the module.
`register_full_backward_pre_hook`(hook[, prepend])	Register a backward pre-hook on the module.
`register_load_state_dict_post_hook`(hook)	Register a post hook to be run after module's `load_state_dict` is called.
`register_module`(name, module)	Alias for `add_module()`.
`register_parameter`(name, param)	Add a parameter to the module.
`register_state_dict_pre_hook`(hook)	Register a pre-hook for the `load_state_dict()` method.
`requires_grad_`([requires_grad])	Change if autograd should record operations on parameters in this module.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`state_dict`(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
`to`(args, *kwargs)	Move and/or cast the parameters and buffers.
`to_empty`(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
`train`([mode])	Set the module in training mode.
`type`(dst_type)	Casts all parameters and buffers to `dst_type`.
`xpu`([device])	Move all model parameters and buffers to the XPU.
`zero_grad`([set_to_none])	Reset gradients of all model parameters.

__call__
deserialize
forward_common_atomic
get_model_def_script
reinit_atom_exclude
reinit_pair_exclude
serialize

compute_or_load_stat(sampled_func, stat_file_path: Optional[DPPath] = None)[source]

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.

classmethod deserialize(data) → DPAtomicModel[source]

fitting_output_def() → FittingOutputDef[source]: Get the output def of the fitting net.

forward_atomic(extended_coord, extended_atype, nlist, mapping: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None) → Dict[str, Tensor][source]

Return atomic prediction.

Parameters

extended_coord: coodinates in extended region
extended_atype: atomic type in extended region
nlist: neighbor list. nf x nloc x nsel
mapping: mapps the extended indices to local indices
fparam: frame parameter. nf x ndf
aparam: atomic parameter. nf x nloc x nda

Returns

result_dict: the result dict, defined by the FittingOutputDef.

get_dim_aparam() → int[source]: Get the number (dimension) of atomic parameters of this atomic model.

get_dim_fparam() → int[source]: Get the number (dimension) of frame parameters of this atomic model.

get_rcut() → float[source]: Get the cut-off radius.

get_sel() → List[int][source]: Get the neighbor selection.

get_sel_type() → List[int][source]

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.

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

is_aparam_nall() → bool[source]

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

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

mixed_types() → bool[source]

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.

serialize() → dict[source]

training: bool

class deepmd.pt.model.atomic_model.DPZBLLinearAtomicModel(dp_model: DPAtomicModel, zbl_model: PairTabAtomicModel, sw_rmin: float, sw_rmax: float, smin_alpha: Optional[float] = 0.1, **kwargs)[source]

Bases: LinearAtomicModel

Model linearly combine a list of AtomicModels.

Parameters

models: This linear model should take a DPAtomicModel and a PairTable model.

Methods

`add_module`(name, module)	Add a child module to the current module.
`apply`(fn)	Apply `fn` recursively to every submodule (as returned by `.children()`) as well as self.
`atomic_output_def`()	Get the output def of the atomic model.
`bfloat16`()	Casts all floating point parameters and buffers to `bfloat16` datatype.
`buffers`([recurse])	Return an iterator over module buffers.
`children`()	Return an iterator over immediate children modules.
`compile`(args, *kwargs)	Compile this Module's forward using `torch.compile()`.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`do_grad_`(var_name, base)	Tell if the output variable var_name is differentiable.
`do_grad_c`([var_name])	Tell if the output variable var_name is c_differentiable.
`do_grad_r`([var_name])	Tell if the output variable var_name is r_differentiable.
`double`()	Casts all floating point parameters and buffers to `double` datatype.
`eval`()	Set the module in evaluation mode.
`extra_repr`()	Set the extra representation of the module.
`fitting_output_def`()	Get the output def of developer implemented atomic models.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(*input)	Define the computation performed at every call.
`forward_atomic`(extended_coord, ...[, ...])	Return atomic prediction.
`get_buffer`(target)	Return the buffer given by `target` if it exists, otherwise throw an error.
`get_dim_aparam`()	Get the number (dimension) of atomic parameters of this atomic model.
`get_dim_fparam`()	Get the number (dimension) of frame parameters of this atomic model.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_model_nsels`()	Get the processed sels for each individual models.
`get_model_rcuts`()	Get the cut-off radius for each individual models.
`get_model_sels`()	Get the sels for each individual models.
`get_nnei`()	Returns the total number of selected neighboring atoms in the cut-off radius.
`get_nsel`()	Returns the total number of selected neighboring atoms in the cut-off radius.
`get_ntypes`()	Get the number of atom types.
`get_parameter`(target)	Return the parameter given by `target` if it exists, otherwise throw an error.
`get_rcut`()	Get the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_sel_type`()	Get the selected atom types of this model.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`get_type_map`()	Get the type map.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`is_aparam_nall`()	Check whether the shape of atomic parameters is (nframes, nall, ndim).
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`mixed_types`()	If true, the model 1.
`modules`()	Return an iterator over all modules in the network.
`named_buffers`([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
`named_children`()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
`named_modules`([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
`named_parameters`([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
`parameters`([recurse])	Return an iterator over module parameters.
`register_backward_hook`(hook)	Register a backward hook on the module.
`register_buffer`(name, tensor[, persistent])	Add a buffer to the module.
`register_forward_hook`(hook, *[, prepend, ...])	Register a forward hook on the module.
`register_forward_pre_hook`(hook, *[, ...])	Register a forward pre-hook on the module.
`register_full_backward_hook`(hook[, prepend])	Register a backward hook on the module.
`register_full_backward_pre_hook`(hook[, prepend])	Register a backward pre-hook on the module.
`register_load_state_dict_post_hook`(hook)	Register a post hook to be run after module's `load_state_dict` is called.
`register_module`(name, module)	Alias for `add_module()`.
`register_parameter`(name, param)	Add a parameter to the module.
`register_state_dict_pre_hook`(hook)	Register a pre-hook for the `load_state_dict()` method.
`requires_grad_`([requires_grad])	Change if autograd should record operations on parameters in this module.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`state_dict`(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
`to`(args, *kwargs)	Move and/or cast the parameters and buffers.
`to_empty`(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
`train`([mode])	Set the module in training mode.
`type`(dst_type)	Casts all parameters and buffers to `dst_type`.
`xpu`([device])	Move all model parameters and buffers to the XPU.
`zero_grad`([set_to_none])	Reset gradients of all model parameters.

__call__
deserialize
forward_common_atomic
get_model_def_script
reinit_atom_exclude
reinit_pair_exclude
serialize

classmethod deserialize(data) → DPZBLLinearAtomicModel[source]

serialize() → dict[source]

training: bool

class deepmd.pt.model.atomic_model.LinearAtomicModel(models: List[BaseAtomicModel], **kwargs)[source]

Bases: Module, BaseAtomicModel

Linear model make linear combinations of several existing models.

Parameters

modelslist[DPAtomicModel or PairTabAtomicModel]: A list of models to be combined. PairTabAtomicModel must be used together with a DPAtomicModel.

Methods

`add_module`(name, module)	Add a child module to the current module.
`apply`(fn)	Apply `fn` recursively to every submodule (as returned by `.children()`) as well as self.
`atomic_output_def`()	Get the output def of the atomic model.
`bfloat16`()	Casts all floating point parameters and buffers to `bfloat16` datatype.
`buffers`([recurse])	Return an iterator over module buffers.
`children`()	Return an iterator over immediate children modules.
`compile`(args, *kwargs)	Compile this Module's forward using `torch.compile()`.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`do_grad_`(var_name, base)	Tell if the output variable var_name is differentiable.
`do_grad_c`([var_name])	Tell if the output variable var_name is c_differentiable.
`do_grad_r`([var_name])	Tell if the output variable var_name is r_differentiable.
`double`()	Casts all floating point parameters and buffers to `double` datatype.
`eval`()	Set the module in evaluation mode.
`extra_repr`()	Set the extra representation of the module.
`fitting_output_def`()	Get the output def of developer implemented atomic models.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(*input)	Define the computation performed at every call.
`forward_atomic`(extended_coord, ...[, ...])	Return atomic prediction.
`get_buffer`(target)	Return the buffer given by `target` if it exists, otherwise throw an error.
`get_dim_aparam`()	Get the number (dimension) of atomic parameters of this atomic model.
`get_dim_fparam`()	Get the number (dimension) of frame parameters of this atomic model.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_model_nsels`()	Get the processed sels for each individual models.
`get_model_rcuts`()	Get the cut-off radius for each individual models.
`get_model_sels`()	Get the sels for each individual models.
`get_nnei`()	Returns the total number of selected neighboring atoms in the cut-off radius.
`get_nsel`()	Returns the total number of selected neighboring atoms in the cut-off radius.
`get_ntypes`()	Get the number of atom types.
`get_parameter`(target)	Return the parameter given by `target` if it exists, otherwise throw an error.
`get_rcut`()	Get the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_sel_type`()	Get the selected atom types of this model.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`get_type_map`()	Get the type map.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`is_aparam_nall`()	Check whether the shape of atomic parameters is (nframes, nall, ndim).
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`mixed_types`()	If true, the model 1.
`modules`()	Return an iterator over all modules in the network.
`named_buffers`([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
`named_children`()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
`named_modules`([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
`named_parameters`([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
`parameters`([recurse])	Return an iterator over module parameters.
`register_backward_hook`(hook)	Register a backward hook on the module.
`register_buffer`(name, tensor[, persistent])	Add a buffer to the module.
`register_forward_hook`(hook, *[, prepend, ...])	Register a forward hook on the module.
`register_forward_pre_hook`(hook, *[, ...])	Register a forward pre-hook on the module.
`register_full_backward_hook`(hook[, prepend])	Register a backward hook on the module.
`register_full_backward_pre_hook`(hook[, prepend])	Register a backward pre-hook on the module.
`register_load_state_dict_post_hook`(hook)	Register a post hook to be run after module's `load_state_dict` is called.
`register_module`(name, module)	Alias for `add_module()`.
`register_parameter`(name, param)	Add a parameter to the module.
`register_state_dict_pre_hook`(hook)	Register a pre-hook for the `load_state_dict()` method.
`requires_grad_`([requires_grad])	Change if autograd should record operations on parameters in this module.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`state_dict`(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
`to`(args, *kwargs)	Move and/or cast the parameters and buffers.
`to_empty`(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
`train`([mode])	Set the module in training mode.
`type`(dst_type)	Casts all parameters and buffers to `dst_type`.
`xpu`([device])	Move all model parameters and buffers to the XPU.
`zero_grad`([set_to_none])	Reset gradients of all model parameters.

__call__
deserialize
forward_common_atomic
get_model_def_script
reinit_atom_exclude
reinit_pair_exclude
serialize

static deserialize(data) → List[BaseAtomicModel][source]

fitting_output_def() → FittingOutputDef[source]: Get the output def of developer implemented atomic models.

forward_atomic(extended_coord: Tensor, extended_atype: Tensor, nlist: Tensor, mapping: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None) → Dict[str, Tensor][source]

Return atomic prediction.

Parameters

extended_coord: coodinates in extended region, (nframes, nall * 3)
extended_atype: atomic type in extended region, (nframes, nall)
nlist: neighbor list, (nframes, nloc, nsel).
mapping: mapps the extended indices to local indices.
fparam: frame parameter. (nframes, ndf)
aparam: atomic parameter. (nframes, nloc, nda)

Returns

result_dict: the result dict, defined by the fitting net output def.

get_dim_aparam() → int[source]: Get the number (dimension) of atomic parameters of this atomic model.

get_dim_fparam() → int[source]: Get the number (dimension) of frame parameters of this atomic model.

get_model_nsels() → List[int][source]: Get the processed sels for each individual models. Not distinguishing types.

get_model_rcuts() → List[float][source]: Get the cut-off radius for each individual models.

get_model_sels() → List[List[int]][source]: Get the sels for each individual models.

get_rcut() → float[source]: Get the cut-off radius.

get_sel() → List[int][source]: Returns the number of selected atoms for each type.

get_sel_type() → List[int][source]

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.

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

is_aparam_nall() → bool[source]

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

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

mixed_types() → bool[source]

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.

static serialize(models) → dict[source]

training: bool

class deepmd.pt.model.atomic_model.PairTabAtomicModel(tab_file: str, rcut: float, sel: Union[int, List[int]], **kwargs)[source]

Bases: Module, BaseAtomicModel

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

`add_module`(name, module)	Add a child module to the current module.
`apply`(fn)	Apply `fn` recursively to every submodule (as returned by `.children()`) as well as self.
`atomic_output_def`()	Get the output def of the atomic model.
`bfloat16`()	Casts all floating point parameters and buffers to `bfloat16` datatype.
`buffers`([recurse])	Return an iterator over module buffers.
`children`()	Return an iterator over immediate children modules.
`compile`(args, *kwargs)	Compile this Module's forward using `torch.compile()`.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`do_grad_`(var_name, base)	Tell if the output variable var_name is differentiable.
`do_grad_c`([var_name])	Tell if the output variable var_name is c_differentiable.
`do_grad_r`([var_name])	Tell if the output variable var_name is r_differentiable.
`double`()	Casts all floating point parameters and buffers to `double` datatype.
`eval`()	Set the module in evaluation mode.
`extra_repr`()	Set the extra representation of the module.
`fitting_output_def`()	Get the output def of developer implemented atomic models.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(*input)	Define the computation performed at every call.
`get_buffer`(target)	Return the buffer given by `target` if it exists, otherwise throw an error.
`get_dim_aparam`()	Get the number (dimension) of atomic parameters of this atomic model.
`get_dim_fparam`()	Get the number (dimension) of frame parameters of this atomic model.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_nnei`()	Returns the total number of selected neighboring atoms in the cut-off radius.
`get_nsel`()	Returns the total number of selected neighboring atoms in the cut-off radius.
`get_ntypes`()	Get the number of atom types.
`get_parameter`(target)	Return the parameter given by `target` if it exists, otherwise throw an error.
`get_rcut`()	Get the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_sel_type`()	Get the selected atom types of this model.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`get_type_map`()	Get the type map.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`is_aparam_nall`()	Check whether the shape of atomic parameters is (nframes, nall, ndim).
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`mixed_types`()	If true, the model 1.
`modules`()	Return an iterator over all modules in the network.
`named_buffers`([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
`named_children`()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
`named_modules`([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
`named_parameters`([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
`parameters`([recurse])	Return an iterator over module parameters.
`register_backward_hook`(hook)	Register a backward hook on the module.
`register_buffer`(name, tensor[, persistent])	Add a buffer to the module.
`register_forward_hook`(hook, *[, prepend, ...])	Register a forward hook on the module.
`register_forward_pre_hook`(hook, *[, ...])	Register a forward pre-hook on the module.
`register_full_backward_hook`(hook[, prepend])	Register a backward hook on the module.
`register_full_backward_pre_hook`(hook[, prepend])	Register a backward pre-hook on the module.
`register_load_state_dict_post_hook`(hook)	Register a post hook to be run after module's `load_state_dict` is called.
`register_module`(name, module)	Alias for `add_module()`.
`register_parameter`(name, param)	Add a parameter to the module.
`register_state_dict_pre_hook`(hook)	Register a pre-hook for the `load_state_dict()` method.
`requires_grad_`([requires_grad])	Change if autograd should record operations on parameters in this module.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`state_dict`(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
`to`(args, *kwargs)	Move and/or cast the parameters and buffers.
`to_empty`(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
`train`([mode])	Set the module in training mode.
`type`(dst_type)	Casts all parameters and buffers to `dst_type`.
`xpu`([device])	Move all model parameters and buffers to the XPU.
`zero_grad`([set_to_none])	Reset gradients of all model parameters.

__call__
deserialize
forward_atomic
forward_common_atomic
get_model_def_script
reinit_atom_exclude
reinit_pair_exclude
serialize

classmethod deserialize(data) → PairTabAtomicModel[source]

fitting_output_def() → FittingOutputDef[source]: Get the output def of developer implemented atomic models.

forward_atomic(extended_coord: Tensor, extended_atype: Tensor, nlist: Tensor, mapping: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None, do_atomic_virial: bool = False) → Dict[str, Tensor][source]

get_dim_aparam() → int[source]: Get the number (dimension) of atomic parameters of this atomic model.

get_dim_fparam() → int[source]: Get the number (dimension) of frame parameters of this atomic model.

get_nsel() → int[source]: Returns the total number of selected neighboring atoms in the cut-off radius.

get_rcut() → float[source]: Get the cut-off radius.

get_sel() → List[int][source]: Returns the number of selected atoms for each type.

get_sel_type() → List[int][source]

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.

get_type_map() → Optional[List[str]][source]: Get the type map.

is_aparam_nall() → bool[source]

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

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

mixed_types() → bool[source]

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.

serialize() → dict[source]

training: bool

Submodules

deepmd.pt.model.atomic_model.base_atomic_model module

class deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel(atom_exclude_types: List[int] = [], pair_exclude_types: List[Tuple[int, int]] = [])[source]

Bases: BAM

Methods

`atomic_output_def`()	Get the output def of the atomic model.
`do_grad_`(var_name, base)	Tell if the output variable var_name is differentiable.
`do_grad_c`([var_name])	Tell if the output variable var_name is c_differentiable.
`do_grad_r`([var_name])	Tell if the output variable var_name is r_differentiable.
`fitting_output_def`()	Get the output def of developer implemented atomic models.
`get_dim_aparam`()	Get the number (dimension) of atomic parameters of this atomic model.
`get_dim_fparam`()	Get the number (dimension) of frame parameters of this atomic model.
`get_nnei`()	Returns the total number of selected neighboring atoms in the cut-off radius.
`get_nsel`()	Returns the total number of selected neighboring atoms in the cut-off radius.
`get_ntypes`()	Get the number of atom types.
`get_rcut`()	Get the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_sel_type`()	Get the selected atom types of this model.
`get_type_map`()	Get the type map.
`is_aparam_nall`()	Check whether the shape of atomic parameters is (nframes, nall, ndim).
`mixed_types`()	If true, the model 1.

deserialize
forward_atomic
forward_common_atomic
get_model_def_script
reinit_atom_exclude
reinit_pair_exclude
serialize

atomic_output_def() → FittingOutputDef[source]

Get the output def of the atomic model.

By default it is the same as FittingOutputDef, but it allows model level wrapper of the output defined by the developer.

forward_common_atomic(extended_coord: Tensor, extended_atype: Tensor, nlist: Tensor, mapping: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None) → Dict[str, Tensor][source]

get_model_def_script() → str[source]

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.

get_ntypes() → int: Get the number of atom types.

reinit_atom_exclude(exclude_types: List[int] = [])[source]

reinit_pair_exclude(exclude_types: List[Tuple[int, int]] = [])[source]

serialize() → dict[source]

deepmd.pt.model.atomic_model.dp_atomic_model module

class deepmd.pt.model.atomic_model.dp_atomic_model.DPAtomicModel(descriptor, fitting, type_map: Optional[List[str]], **kwargs)[source]

Bases: Module, BaseAtomicModel

Model give atomic prediction of some physical property.

Parameters

descriptor: Descriptor
fitting_net: Fitting 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.

Methods

`add_module`(name, module)	Add a child module to the current module.
`apply`(fn)	Apply `fn` recursively to every submodule (as returned by `.children()`) as well as self.
`atomic_output_def`()	Get the output def of the atomic model.
`bfloat16`()	Casts all floating point parameters and buffers to `bfloat16` datatype.
`buffers`([recurse])	Return an iterator over module buffers.
`children`()	Return an iterator over immediate children modules.
`compile`(args, *kwargs)	Compile this Module's forward using `torch.compile()`.
`compute_or_load_stat`(sampled_func[, ...])	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.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`do_grad_`(var_name, base)	Tell if the output variable var_name is differentiable.
`do_grad_c`([var_name])	Tell if the output variable var_name is c_differentiable.
`do_grad_r`([var_name])	Tell if the output variable var_name is r_differentiable.
`double`()	Casts all floating point parameters and buffers to `double` datatype.
`eval`()	Set the module in evaluation mode.
`extra_repr`()	Set the extra representation of the module.
`fitting_output_def`()	Get the output def of the fitting net.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(*input)	Define the computation performed at every call.
`forward_atomic`(extended_coord, ...[, ...])	Return atomic prediction.
`get_buffer`(target)	Return the buffer given by `target` if it exists, otherwise throw an error.
`get_dim_aparam`()	Get the number (dimension) of atomic parameters of this atomic model.
`get_dim_fparam`()	Get the number (dimension) of frame parameters of this atomic model.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_nnei`()	Returns the total number of selected neighboring atoms in the cut-off radius.
`get_nsel`()	Returns the total number of selected neighboring atoms in the cut-off radius.
`get_ntypes`()	Get the number of atom types.
`get_parameter`(target)	Return the parameter given by `target` if it exists, otherwise throw an error.
`get_rcut`()	Get the cut-off radius.
`get_sel`()	Get the neighbor selection.
`get_sel_type`()	Get the selected atom types of this model.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`get_type_map`()	Get the type map.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`is_aparam_nall`()	Check whether the shape of atomic parameters is (nframes, nall, ndim).
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`mixed_types`()	If true, the model 1.
`modules`()	Return an iterator over all modules in the network.
`named_buffers`([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
`named_children`()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
`named_modules`([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
`named_parameters`([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
`parameters`([recurse])	Return an iterator over module parameters.
`register_backward_hook`(hook)	Register a backward hook on the module.
`register_buffer`(name, tensor[, persistent])	Add a buffer to the module.
`register_forward_hook`(hook, *[, prepend, ...])	Register a forward hook on the module.
`register_forward_pre_hook`(hook, *[, ...])	Register a forward pre-hook on the module.
`register_full_backward_hook`(hook[, prepend])	Register a backward hook on the module.
`register_full_backward_pre_hook`(hook[, prepend])	Register a backward pre-hook on the module.
`register_load_state_dict_post_hook`(hook)	Register a post hook to be run after module's `load_state_dict` is called.
`register_module`(name, module)	Alias for `add_module()`.
`register_parameter`(name, param)	Add a parameter to the module.
`register_state_dict_pre_hook`(hook)	Register a pre-hook for the `load_state_dict()` method.
`requires_grad_`([requires_grad])	Change if autograd should record operations on parameters in this module.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`state_dict`(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
`to`(args, *kwargs)	Move and/or cast the parameters and buffers.
`to_empty`(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
`train`([mode])	Set the module in training mode.
`type`(dst_type)	Casts all parameters and buffers to `dst_type`.
`xpu`([device])	Move all model parameters and buffers to the XPU.
`zero_grad`([set_to_none])	Reset gradients of all model parameters.

__call__
deserialize
forward_common_atomic
get_model_def_script
reinit_atom_exclude
reinit_pair_exclude
serialize

compute_or_load_stat(sampled_func, stat_file_path: Optional[DPPath] = None)[source]

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.

classmethod deserialize(data) → DPAtomicModel[source]

fitting_output_def() → FittingOutputDef[source]: Get the output def of the fitting net.

forward_atomic(extended_coord, extended_atype, nlist, mapping: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None) → Dict[str, Tensor][source]

Return atomic prediction.

Parameters

extended_coord: coodinates in extended region
extended_atype: atomic type in extended region
nlist: neighbor list. nf x nloc x nsel
mapping: mapps the extended indices to local indices
fparam: frame parameter. nf x ndf
aparam: atomic parameter. nf x nloc x nda

Returns

result_dict: the result dict, defined by the FittingOutputDef.

get_dim_aparam() → int[source]: Get the number (dimension) of atomic parameters of this atomic model.

get_dim_fparam() → int[source]: Get the number (dimension) of frame parameters of this atomic model.

get_rcut() → float[source]: Get the cut-off radius.

get_sel() → List[int][source]: Get the neighbor selection.

get_sel_type() → List[int][source]

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.

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

is_aparam_nall() → bool[source]

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

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

mixed_types() → bool[source]

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.

serialize() → dict[source]

training: bool

deepmd.pt.model.atomic_model.linear_atomic_model module

class deepmd.pt.model.atomic_model.linear_atomic_model.DPZBLLinearAtomicModel(dp_model: DPAtomicModel, zbl_model: PairTabAtomicModel, sw_rmin: float, sw_rmax: float, smin_alpha: Optional[float] = 0.1, **kwargs)[source]

Bases: LinearAtomicModel

Model linearly combine a list of AtomicModels.

Parameters

models: This linear model should take a DPAtomicModel and a PairTable model.

Methods

`add_module`(name, module)	Add a child module to the current module.
`apply`(fn)	Apply `fn` recursively to every submodule (as returned by `.children()`) as well as self.
`atomic_output_def`()	Get the output def of the atomic model.
`bfloat16`()	Casts all floating point parameters and buffers to `bfloat16` datatype.
`buffers`([recurse])	Return an iterator over module buffers.
`children`()	Return an iterator over immediate children modules.
`compile`(args, *kwargs)	Compile this Module's forward using `torch.compile()`.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`do_grad_`(var_name, base)	Tell if the output variable var_name is differentiable.
`do_grad_c`([var_name])	Tell if the output variable var_name is c_differentiable.
`do_grad_r`([var_name])	Tell if the output variable var_name is r_differentiable.
`double`()	Casts all floating point parameters and buffers to `double` datatype.
`eval`()	Set the module in evaluation mode.
`extra_repr`()	Set the extra representation of the module.
`fitting_output_def`()	Get the output def of developer implemented atomic models.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(*input)	Define the computation performed at every call.
`forward_atomic`(extended_coord, ...[, ...])	Return atomic prediction.
`get_buffer`(target)	Return the buffer given by `target` if it exists, otherwise throw an error.
`get_dim_aparam`()	Get the number (dimension) of atomic parameters of this atomic model.
`get_dim_fparam`()	Get the number (dimension) of frame parameters of this atomic model.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_model_nsels`()	Get the processed sels for each individual models.
`get_model_rcuts`()	Get the cut-off radius for each individual models.
`get_model_sels`()	Get the sels for each individual models.
`get_nnei`()	Returns the total number of selected neighboring atoms in the cut-off radius.
`get_nsel`()	Returns the total number of selected neighboring atoms in the cut-off radius.
`get_ntypes`()	Get the number of atom types.
`get_parameter`(target)	Return the parameter given by `target` if it exists, otherwise throw an error.
`get_rcut`()	Get the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_sel_type`()	Get the selected atom types of this model.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`get_type_map`()	Get the type map.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`is_aparam_nall`()	Check whether the shape of atomic parameters is (nframes, nall, ndim).
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`mixed_types`()	If true, the model 1.
`modules`()	Return an iterator over all modules in the network.
`named_buffers`([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
`named_children`()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
`named_modules`([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
`named_parameters`([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
`parameters`([recurse])	Return an iterator over module parameters.
`register_backward_hook`(hook)	Register a backward hook on the module.
`register_buffer`(name, tensor[, persistent])	Add a buffer to the module.
`register_forward_hook`(hook, *[, prepend, ...])	Register a forward hook on the module.
`register_forward_pre_hook`(hook, *[, ...])	Register a forward pre-hook on the module.
`register_full_backward_hook`(hook[, prepend])	Register a backward hook on the module.
`register_full_backward_pre_hook`(hook[, prepend])	Register a backward pre-hook on the module.
`register_load_state_dict_post_hook`(hook)	Register a post hook to be run after module's `load_state_dict` is called.
`register_module`(name, module)	Alias for `add_module()`.
`register_parameter`(name, param)	Add a parameter to the module.
`register_state_dict_pre_hook`(hook)	Register a pre-hook for the `load_state_dict()` method.
`requires_grad_`([requires_grad])	Change if autograd should record operations on parameters in this module.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`state_dict`(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
`to`(args, *kwargs)	Move and/or cast the parameters and buffers.
`to_empty`(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
`train`([mode])	Set the module in training mode.
`type`(dst_type)	Casts all parameters and buffers to `dst_type`.
`xpu`([device])	Move all model parameters and buffers to the XPU.
`zero_grad`([set_to_none])	Reset gradients of all model parameters.

__call__
deserialize
forward_common_atomic
get_model_def_script
reinit_atom_exclude
reinit_pair_exclude
serialize

classmethod deserialize(data) → DPZBLLinearAtomicModel[source]

serialize() → dict[source]

training: bool

class deepmd.pt.model.atomic_model.linear_atomic_model.LinearAtomicModel(models: List[BaseAtomicModel], **kwargs)[source]

Bases: Module, BaseAtomicModel

Linear model make linear combinations of several existing models.

Parameters

modelslist[DPAtomicModel or PairTabAtomicModel]: A list of models to be combined. PairTabAtomicModel must be used together with a DPAtomicModel.

Methods

`add_module`(name, module)	Add a child module to the current module.
`apply`(fn)	Apply `fn` recursively to every submodule (as returned by `.children()`) as well as self.
`atomic_output_def`()	Get the output def of the atomic model.
`bfloat16`()	Casts all floating point parameters and buffers to `bfloat16` datatype.
`buffers`([recurse])	Return an iterator over module buffers.
`children`()	Return an iterator over immediate children modules.
`compile`(args, *kwargs)	Compile this Module's forward using `torch.compile()`.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`do_grad_`(var_name, base)	Tell if the output variable var_name is differentiable.
`do_grad_c`([var_name])	Tell if the output variable var_name is c_differentiable.
`do_grad_r`([var_name])	Tell if the output variable var_name is r_differentiable.
`double`()	Casts all floating point parameters and buffers to `double` datatype.
`eval`()	Set the module in evaluation mode.
`extra_repr`()	Set the extra representation of the module.
`fitting_output_def`()	Get the output def of developer implemented atomic models.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(*input)	Define the computation performed at every call.
`forward_atomic`(extended_coord, ...[, ...])	Return atomic prediction.
`get_buffer`(target)	Return the buffer given by `target` if it exists, otherwise throw an error.
`get_dim_aparam`()	Get the number (dimension) of atomic parameters of this atomic model.
`get_dim_fparam`()	Get the number (dimension) of frame parameters of this atomic model.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_model_nsels`()	Get the processed sels for each individual models.
`get_model_rcuts`()	Get the cut-off radius for each individual models.
`get_model_sels`()	Get the sels for each individual models.
`get_nnei`()	Returns the total number of selected neighboring atoms in the cut-off radius.
`get_nsel`()	Returns the total number of selected neighboring atoms in the cut-off radius.
`get_ntypes`()	Get the number of atom types.
`get_parameter`(target)	Return the parameter given by `target` if it exists, otherwise throw an error.
`get_rcut`()	Get the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_sel_type`()	Get the selected atom types of this model.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`get_type_map`()	Get the type map.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`is_aparam_nall`()	Check whether the shape of atomic parameters is (nframes, nall, ndim).
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`mixed_types`()	If true, the model 1.
`modules`()	Return an iterator over all modules in the network.
`named_buffers`([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
`named_children`()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
`named_modules`([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
`named_parameters`([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
`parameters`([recurse])	Return an iterator over module parameters.
`register_backward_hook`(hook)	Register a backward hook on the module.
`register_buffer`(name, tensor[, persistent])	Add a buffer to the module.
`register_forward_hook`(hook, *[, prepend, ...])	Register a forward hook on the module.
`register_forward_pre_hook`(hook, *[, ...])	Register a forward pre-hook on the module.
`register_full_backward_hook`(hook[, prepend])	Register a backward hook on the module.
`register_full_backward_pre_hook`(hook[, prepend])	Register a backward pre-hook on the module.
`register_load_state_dict_post_hook`(hook)	Register a post hook to be run after module's `load_state_dict` is called.
`register_module`(name, module)	Alias for `add_module()`.
`register_parameter`(name, param)	Add a parameter to the module.
`register_state_dict_pre_hook`(hook)	Register a pre-hook for the `load_state_dict()` method.
`requires_grad_`([requires_grad])	Change if autograd should record operations on parameters in this module.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`state_dict`(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
`to`(args, *kwargs)	Move and/or cast the parameters and buffers.
`to_empty`(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
`train`([mode])	Set the module in training mode.
`type`(dst_type)	Casts all parameters and buffers to `dst_type`.
`xpu`([device])	Move all model parameters and buffers to the XPU.
`zero_grad`([set_to_none])	Reset gradients of all model parameters.

__call__
deserialize
forward_common_atomic
get_model_def_script
reinit_atom_exclude
reinit_pair_exclude
serialize

static deserialize(data) → List[BaseAtomicModel][source]

fitting_output_def() → FittingOutputDef[source]: Get the output def of developer implemented atomic models.

forward_atomic(extended_coord: Tensor, extended_atype: Tensor, nlist: Tensor, mapping: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None) → Dict[str, Tensor][source]

Return atomic prediction.

Parameters

extended_coord: coodinates in extended region, (nframes, nall * 3)
extended_atype: atomic type in extended region, (nframes, nall)
nlist: neighbor list, (nframes, nloc, nsel).
mapping: mapps the extended indices to local indices.
fparam: frame parameter. (nframes, ndf)
aparam: atomic parameter. (nframes, nloc, nda)

Returns

result_dict: the result dict, defined by the fitting net output def.

get_dim_aparam() → int[source]: Get the number (dimension) of atomic parameters of this atomic model.

get_dim_fparam() → int[source]: Get the number (dimension) of frame parameters of this atomic model.

get_model_nsels() → List[int][source]: Get the processed sels for each individual models. Not distinguishing types.

get_model_rcuts() → List[float][source]: Get the cut-off radius for each individual models.

get_model_sels() → List[List[int]][source]: Get the sels for each individual models.

get_rcut() → float[source]: Get the cut-off radius.

get_sel() → List[int][source]: Returns the number of selected atoms for each type.

get_sel_type() → List[int][source]

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.

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

is_aparam_nall() → bool[source]

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

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

mixed_types() → bool[source]

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.

static serialize(models) → dict[source]

training: bool

deepmd.pt.model.atomic_model.pairtab_atomic_model module

class deepmd.pt.model.atomic_model.pairtab_atomic_model.PairTabAtomicModel(tab_file: str, rcut: float, sel: Union[int, List[int]], **kwargs)[source]

Bases: Module, BaseAtomicModel

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

`add_module`(name, module)	Add a child module to the current module.
`apply`(fn)	Apply `fn` recursively to every submodule (as returned by `.children()`) as well as self.
`atomic_output_def`()	Get the output def of the atomic model.
`bfloat16`()	Casts all floating point parameters and buffers to `bfloat16` datatype.
`buffers`([recurse])	Return an iterator over module buffers.
`children`()	Return an iterator over immediate children modules.
`compile`(args, *kwargs)	Compile this Module's forward using `torch.compile()`.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`do_grad_`(var_name, base)	Tell if the output variable var_name is differentiable.
`do_grad_c`([var_name])	Tell if the output variable var_name is c_differentiable.
`do_grad_r`([var_name])	Tell if the output variable var_name is r_differentiable.
`double`()	Casts all floating point parameters and buffers to `double` datatype.
`eval`()	Set the module in evaluation mode.
`extra_repr`()	Set the extra representation of the module.
`fitting_output_def`()	Get the output def of developer implemented atomic models.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(*input)	Define the computation performed at every call.
`get_buffer`(target)	Return the buffer given by `target` if it exists, otherwise throw an error.
`get_dim_aparam`()	Get the number (dimension) of atomic parameters of this atomic model.
`get_dim_fparam`()	Get the number (dimension) of frame parameters of this atomic model.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_nnei`()	Returns the total number of selected neighboring atoms in the cut-off radius.
`get_nsel`()	Returns the total number of selected neighboring atoms in the cut-off radius.
`get_ntypes`()	Get the number of atom types.
`get_parameter`(target)	Return the parameter given by `target` if it exists, otherwise throw an error.
`get_rcut`()	Get the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_sel_type`()	Get the selected atom types of this model.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`get_type_map`()	Get the type map.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`is_aparam_nall`()	Check whether the shape of atomic parameters is (nframes, nall, ndim).
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`mixed_types`()	If true, the model 1.
`modules`()	Return an iterator over all modules in the network.
`named_buffers`([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
`named_children`()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
`named_modules`([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
`named_parameters`([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
`parameters`([recurse])	Return an iterator over module parameters.
`register_backward_hook`(hook)	Register a backward hook on the module.
`register_buffer`(name, tensor[, persistent])	Add a buffer to the module.
`register_forward_hook`(hook, *[, prepend, ...])	Register a forward hook on the module.
`register_forward_pre_hook`(hook, *[, ...])	Register a forward pre-hook on the module.
`register_full_backward_hook`(hook[, prepend])	Register a backward hook on the module.
`register_full_backward_pre_hook`(hook[, prepend])	Register a backward pre-hook on the module.
`register_load_state_dict_post_hook`(hook)	Register a post hook to be run after module's `load_state_dict` is called.
`register_module`(name, module)	Alias for `add_module()`.
`register_parameter`(name, param)	Add a parameter to the module.
`register_state_dict_pre_hook`(hook)	Register a pre-hook for the `load_state_dict()` method.
`requires_grad_`([requires_grad])	Change if autograd should record operations on parameters in this module.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`state_dict`(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
`to`(args, *kwargs)	Move and/or cast the parameters and buffers.
`to_empty`(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
`train`([mode])	Set the module in training mode.
`type`(dst_type)	Casts all parameters and buffers to `dst_type`.
`xpu`([device])	Move all model parameters and buffers to the XPU.
`zero_grad`([set_to_none])	Reset gradients of all model parameters.

__call__
deserialize
forward_atomic
forward_common_atomic
get_model_def_script
reinit_atom_exclude
reinit_pair_exclude
serialize

classmethod deserialize(data) → PairTabAtomicModel[source]

fitting_output_def() → FittingOutputDef[source]: Get the output def of developer implemented atomic models.

forward_atomic(extended_coord: Tensor, extended_atype: Tensor, nlist: Tensor, mapping: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None, do_atomic_virial: bool = False) → Dict[str, Tensor][source]

get_dim_aparam() → int[source]: Get the number (dimension) of atomic parameters of this atomic model.

get_dim_fparam() → int[source]: Get the number (dimension) of frame parameters of this atomic model.

get_nsel() → int[source]: Returns the total number of selected neighboring atoms in the cut-off radius.

get_rcut() → float[source]: Get the cut-off radius.

get_sel() → List[int][source]: Returns the number of selected atoms for each type.

get_sel_type() → List[int][source]

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.

get_type_map() → Optional[List[str]][source]: Get the type map.

is_aparam_nall() → bool[source]

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

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

mixed_types() → bool[source]

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.

serialize() → dict[source]

training: bool