deepmd.pt.model.model package

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.

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

Bases: BaseBaseModel

Methods

__call__(*args, **kwds)	Inference method.
compute_or_load_stat(sampled[, stat_file_path])	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.
deserialize(data)	Deserialize the model.
get_class_by_type(class_type)	Get the class by the plugin type.
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_model_def_script()	Get the model definition script.
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_plugins()	Get all the registered plugins.
get_rcut()	Get the cut-off radius.
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).
model_output_type()	Get the output type for the model.
register(key)	Register a descriptor plugin.
serialize()	Serialize the model.
update_sel(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

compute_or_load_stat(sampled, stat_file_path: Optional[DPPath] = 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

The sampled data frames from different data systems.

stat_file_path

The path to the statistics files.

model_def_script: str

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

Bases: CM, BaseModel

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.
deserialize(data)	Deserialize the 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.
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.
format_nlist(extended_coord, extended_atype, ...)	Format the neighbor list.
forward(*input)	Define the computation performed at every call.
forward_atomic(extended_coord, ...[, ...])	Return atomic prediction.
forward_common(coord, atype[, box, fparam, ...])	Return model prediction.
forward_common_lower(extended_coord, ...[, ...])	Return model prediction.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
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_plugins()	Get all the registered plugins.
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.
input_type_cast(coord[, box, fparam, aparam])	Cast the input data to global float type.
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.
model_output_def()	Get the output def for the model.
model_output_type()	Get the output type for the model.
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.
output_type_cast(model_ret, input_prec)	Convert the model output to the input prec.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
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.
update_sel(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__
forward_common_atomic
get_model_def_script
reinit_atom_exclude
reinit_pair_exclude
serialize

training: bool

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

Bases: CM, BaseModel

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[, stat_file_path])	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.
deserialize(data)	Deserialize the 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.
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.
format_nlist(extended_coord, extended_atype, ...)	Format the neighbor list.
forward(coord, atype[, box, fparam, aparam, ...])	Define the computation performed at every call.
forward_atomic(extended_coord, ...[, ...])	Return atomic prediction.
forward_common(coord, atype[, box, fparam, ...])	Return model prediction.
forward_common_lower(extended_coord, ...[, ...])	Return model prediction.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
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_plugins()	Get all the registered plugins.
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.
input_type_cast(coord[, box, fparam, aparam])	Cast the input data to global float type.
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.
model_output_def()	Get the output def for the model.
model_output_type()	Get the output type for the model.
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.
output_type_cast(model_ret, input_prec)	Convert the model output to the input prec.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
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.
update_sel(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__
forward_common_atomic
forward_lower
get_model_def_script
reinit_atom_exclude
reinit_pair_exclude
serialize

forward(coord, atype, box: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None, do_atomic_virial: bool = False) → Dict[str, Tensor]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

forward_lower(extended_coord, extended_atype, nlist, mapping: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None, do_atomic_virial: bool = False)

model_type = 'ener'

training: bool

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

Bases: DPModel

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.
deserialize(data)	Deserialize the 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.
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.
format_nlist(extended_coord, extended_atype, ...)	Format the neighbor list.
forward(coord, atype[, box, fparam, aparam, ...])	Define the computation performed at every call.
forward_atomic(extended_coord, ...[, ...])	Return atomic prediction.
forward_common(coord, atype[, box, fparam, ...])	Return model prediction.
forward_common_lower(extended_coord, ...[, ...])	Return model prediction.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
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_plugins()	Get all the registered plugins.
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.
input_type_cast(coord[, box, fparam, aparam])	Cast the input data to global float type.
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.
model_output_def()	Get the output def for the model.
model_output_type()	Get the output type for the model.
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.
output_type_cast(model_ret, input_prec)	Convert the model output to the input prec.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
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.
update_sel(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__
forward_common_atomic
forward_lower
get_model_def_script
reinit_atom_exclude
reinit_pair_exclude
serialize

forward(coord, atype, box: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None, do_atomic_virial: bool = False) → Dict[str, Tensor]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

forward_lower(extended_coord, extended_atype, nlist, mapping: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None, do_atomic_virial: bool = False)

model_type = 'ener'

training: bool

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

deepmd.pt.model.model.get_zbl_model(model_params)

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)

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.

Submodules

deepmd.pt.model.model.dipole_model module

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

Bases: DPModel

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.
deserialize(data)	Deserialize the 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.
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.
format_nlist(extended_coord, extended_atype, ...)	Format the neighbor list.
forward(coord, atype[, box, fparam, aparam, ...])	Define the computation performed at every call.
forward_atomic(extended_coord, ...[, ...])	Return atomic prediction.
forward_common(coord, atype[, box, fparam, ...])	Return model prediction.
forward_common_lower(extended_coord, ...[, ...])	Return model prediction.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
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_plugins()	Get all the registered plugins.
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.
input_type_cast(coord[, box, fparam, aparam])	Cast the input data to global float type.
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.
model_output_def()	Get the output def for the model.
model_output_type()	Get the output type for the model.
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.
output_type_cast(model_ret, input_prec)	Convert the model output to the input prec.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
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.
update_sel(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__
forward_common_atomic
forward_lower
get_model_def_script
reinit_atom_exclude
reinit_pair_exclude
serialize

forward(coord, atype, box: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None, do_atomic_virial: bool = False) → Dict[str, Tensor]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

forward_lower(extended_coord, extended_atype, nlist, mapping: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None, do_atomic_virial: bool = False)

model_type = 'dipole'

training: bool

deepmd.pt.model.model.dp_model module

class deepmd.pt.model.model.dp_model.DPModel(*args, **kwargs)

Bases: CM, BaseModel

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.
deserialize(data)	Deserialize the 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.
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.
format_nlist(extended_coord, extended_atype, ...)	Format the neighbor list.
forward(*input)	Define the computation performed at every call.
forward_atomic(extended_coord, ...[, ...])	Return atomic prediction.
forward_common(coord, atype[, box, fparam, ...])	Return model prediction.
forward_common_lower(extended_coord, ...[, ...])	Return model prediction.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
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_plugins()	Get all the registered plugins.
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.
input_type_cast(coord[, box, fparam, aparam])	Cast the input data to global float type.
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.
model_output_def()	Get the output def for the model.
model_output_type()	Get the output type for the model.
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.
output_type_cast(model_ret, input_prec)	Convert the model output to the input prec.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
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.
update_sel(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__
forward_common_atomic
get_model_def_script
reinit_atom_exclude
reinit_pair_exclude
serialize

training: bool

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

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

Bases: CM, BaseModel

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[, stat_file_path])	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.
deserialize(data)	Deserialize the 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.
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.
format_nlist(extended_coord, extended_atype, ...)	Format the neighbor list.
forward(coord, atype[, box, fparam, aparam, ...])	Define the computation performed at every call.
forward_atomic(extended_coord, ...[, ...])	Return atomic prediction.
forward_common(coord, atype[, box, fparam, ...])	Return model prediction.
forward_common_lower(extended_coord, ...[, ...])	Return model prediction.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
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_plugins()	Get all the registered plugins.
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.
input_type_cast(coord[, box, fparam, aparam])	Cast the input data to global float type.
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.
model_output_def()	Get the output def for the model.
model_output_type()	Get the output type for the model.
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.
output_type_cast(model_ret, input_prec)	Convert the model output to the input prec.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
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.
update_sel(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__
forward_common_atomic
forward_lower
get_model_def_script
reinit_atom_exclude
reinit_pair_exclude
serialize

forward(coord, atype, box: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None, do_atomic_virial: bool = False) → Dict[str, Tensor]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

forward_lower(extended_coord, extended_atype, nlist, mapping: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None, do_atomic_virial: bool = False)

model_type = 'ener'

training: bool

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

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

Bases: DPModel

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.
deserialize(data)	Deserialize the 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.
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.
format_nlist(extended_coord, extended_atype, ...)	Format the neighbor list.
forward(coord, atype[, box, fparam, aparam, ...])	Define the computation performed at every call.
forward_atomic(extended_coord, ...[, ...])	Return atomic prediction.
forward_common(coord, atype[, box, fparam, ...])	Return model prediction.
forward_common_lower(extended_coord, ...[, ...])	Return model prediction.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
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_plugins()	Get all the registered plugins.
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.
input_type_cast(coord[, box, fparam, aparam])	Cast the input data to global float type.
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.
model_output_def()	Get the output def for the model.
model_output_type()	Get the output type for the model.
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.
output_type_cast(model_ret, input_prec)	Convert the model output to the input prec.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
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.
update_sel(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__
forward_common_atomic
forward_lower
get_model_def_script
reinit_atom_exclude
reinit_pair_exclude
serialize

forward(coord, atype, box: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None, do_atomic_virial: bool = False) → Dict[str, Tensor]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

forward_lower(extended_coord, extended_atype, nlist, mapping: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None, do_atomic_virial: bool = False)

model_type = 'ener'

training: bool

deepmd.pt.model.model.make_hessian_model module

deepmd.pt.model.model.make_hessian_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 module

deepmd.pt.model.model.make_model.make_model(T_AtomicModel)

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.

deepmd.pt.model.model.model module

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

Bases: BaseBaseModel

Methods

__call__(*args, **kwds)	Inference method.
compute_or_load_stat(sampled[, stat_file_path])	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.
deserialize(data)	Deserialize the model.
get_class_by_type(class_type)	Get the class by the plugin type.
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_model_def_script()	Get the model definition script.
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_plugins()	Get all the registered plugins.
get_rcut()	Get the cut-off radius.
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).
model_output_type()	Get the output type for the model.
register(key)	Register a descriptor plugin.
serialize()	Serialize the model.
update_sel(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.

compute_or_load_stat(sampled, stat_file_path: Optional[DPPath] = 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

The sampled data frames from different data systems.

stat_file_path

The path to the statistics files.

model_def_script: str

deepmd.pt.model.model.polar_model module

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

Bases: DPModel

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.
deserialize(data)	Deserialize the 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.
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.
format_nlist(extended_coord, extended_atype, ...)	Format the neighbor list.
forward(coord, atype[, box, fparam, aparam, ...])	Define the computation performed at every call.
forward_atomic(extended_coord, ...[, ...])	Return atomic prediction.
forward_common(coord, atype[, box, fparam, ...])	Return model prediction.
forward_common_lower(extended_coord, ...[, ...])	Return model prediction.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_class_by_type(class_type)	Get the class by the plugin type.
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_plugins()	Get all the registered plugins.
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.
input_type_cast(coord[, box, fparam, aparam])	Cast the input data to global float type.
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.
model_output_def()	Get the output def for the model.
model_output_type()	Get the output type for the model.
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.
output_type_cast(model_ret, input_prec)	Convert the model output to the input prec.
parameters([recurse])	Return an iterator over module parameters.
register(key)	Register a descriptor plugin.
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.
update_sel(global_jdata, local_jdata)	Update the selection and perform neighbor statistics.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__
forward_common_atomic
forward_lower
get_model_def_script
reinit_atom_exclude
reinit_pair_exclude
serialize

forward(coord, atype, box: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None, do_atomic_virial: bool = False) → Dict[str, Tensor]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

forward_lower(extended_coord, extended_atype, nlist, mapping: Optional[Tensor] = None, fparam: Optional[Tensor] = None, aparam: Optional[Tensor] = None, do_atomic_virial: bool = False)

model_type = 'polar'

training: bool

deepmd.pt.model.model.transform_output module

deepmd.pt.model.model.transform_output.atomic_virial_corr(extended_coord: Tensor, atom_energy: Tensor)

deepmd.pt.model.model.transform_output.communicate_extended_output(model_ret: Dict[str, Tensor], model_output_def: ModelOutputDef, mapping: Tensor, do_atomic_virial: bool = False) → Dict[str, Tensor]

Transform the output of the model network defined on local and ghost (extended) atoms to local atoms.

deepmd.pt.model.model.transform_output.fit_output_to_model_output(fit_ret: Dict[str, Tensor], fit_output_def: FittingOutputDef, coord_ext: Tensor, do_atomic_virial: bool = False) → Dict[str, Tensor]

Transform the output of the fitting network to the model output.

deepmd.pt.model.model.transform_output.get_atom_axis(vdef: Tensor)

Get the axis of atoms.

deepmd.pt.model.model.transform_output.get_leading_dims(vv: Tensor, vdef: OutputVariableDef)

Get the dimensions of nf x nloc.

deepmd.pt.model.model.transform_output.take_deriv(vv: Tensor, svv: Tensor, vdef: OutputVariableDef, coord_ext: Tensor, do_virial: bool = False, do_atomic_virial: bool = False)

deepmd.pt.model.model.transform_output.task_deriv_one(atom_energy: Tensor, energy: Tensor, extended_coord: Tensor, do_virial: bool = True, do_atomic_virial: bool = False)