deepmd.pt.model.descriptor package

deepmd.pt.model.descriptor.BaseDescriptor: alias of BD

class deepmd.pt.model.descriptor.DescriptorBlock(*args, **kwargs)[source]

Bases: Module, ABC, PR

The building block of descriptor. Given the input descriptor, provide with the atomic coordinates, atomic types and neighbor list, calculate the new descriptor.

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.
`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_input_stats`(merged[, path])	Compute the input statistics (e.g.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(nlist, extended_coord, extended_atype)	Calculate DescriptorBlock.
`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_emb`()	Returns the embedding dimension.
`get_dim_in`()	Returns the output dimension.
`get_dim_out`()	Returns the output dimension.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_nsel`()	Returns the number of selected atoms in the cut-off radius.
`get_ntypes`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_stats`()	Get the statistics of the descriptor.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`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`(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_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__

compute_input_stats(merged: Union[Callable[[], List[dict]], List[dict]], path: Optional[DPPath] = None)[source]

Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

List[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.
Callable[[], List[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

pathOptional[DPPath]

The path to the stat file.

abstract forward(nlist: Tensor, extended_coord: Tensor, extended_atype: Tensor, extended_atype_embd: Optional[Tensor] = None, mapping: Optional[Tensor] = None)[source]: Calculate DescriptorBlock.

abstract get_dim_emb() → int[source]: Returns the embedding dimension.

abstract get_dim_in() → int[source]: Returns the output dimension.

abstract get_dim_out() → int[source]: Returns the output dimension.

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

abstract get_ntypes() → int[source]: Returns the number of element types.

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

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

get_stats() → Dict[str, StatItem][source]: Get the statistics of the descriptor.

local_cluster = False

share_params(base_class, shared_level, resume=False)[source]: Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.

training: bool

class deepmd.pt.model.descriptor.DescrptBlockHybrid(*args, **kwargs)[source]

Bases: DescriptorBlock

Attributes

dim_emb: Returns the output dimension of embedding.
dim_emb_list: Returns the output dimension list of embeddings.
dim_out: Returns the output dimension of this descriptor.

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.
`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_input_stats`(merged[, path])	Compute the input statistics (e.g.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(nlist, extended_coord, extended_atype)	Calculate decoded embedding for each atom.
`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_emb`()	Returns the embedding dimension.
`get_dim_in`()	Returns the input dimension.
`get_dim_out`()	Returns the output dimension.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_nsel`()	Returns the number of selected atoms in the cut-off radius.
`get_ntypes`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_stats`()	Get the statistics of the descriptor.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`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 discriptor 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`(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_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__

compute_input_stats(merged: Union[Callable[[], List[dict]], List[dict]], path: Optional[DPPath] = None)[source]

Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

List[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.
Callable[[], List[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

pathOptional[DPPath]

The path to the stat file.

property dim_emb: Returns the output dimension of embedding.

property dim_emb_list: List[int]: Returns the output dimension list of embeddings.

property dim_out: Returns the output dimension of this descriptor.

forward(nlist: Tensor, extended_coord: Tensor, extended_atype: Tensor, extended_atype_embd: Optional[Tensor] = None, mapping: Optional[Tensor] = None)[source]

Calculate decoded embedding for each atom.

Args: - extended_coord: Tell atom coordinates with shape [nframes, natoms[1]*3]. - nlist: Tell atom types with shape [nframes, natoms[1]]. - atype: Tell atom count and element count. Its shape is [2+self.ntypes]. - nlist_type: Tell simulation box with shape [nframes, 9]. - atype_tebd: Tell simulation box with shape [nframes, 9]. - nlist_tebd: Tell simulation box with shape [nframes, 9].

Returns

result: descriptor with shape [nframes, nloc, self.filter_neuron[-1] * self.axis_neuron].
ret: environment matrix with shape [nframes, nloc, self.neei, out_size]

get_dim_emb()[source]: Returns the embedding dimension.

get_dim_in() → int[source]: Returns the input dimension.

get_dim_out() → int[source]: Returns the output dimension.

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

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

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

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

mixed_types() → bool[source]

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

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

share_params(base_class, shared_level, resume=False)[source]: Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.

training: bool

class deepmd.pt.model.descriptor.DescrptBlockRepformers(*args, **kwargs)[source]

Bases: DescriptorBlock

Attributes

dim_emb: Returns the embedding dimension g2.
dim_in: Returns the atomic input dimension of this descriptor.
dim_out: Returns the output dimension of this descriptor.

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.
`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_input_stats`(merged[, path])	Compute the input statistics (e.g.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(nlist, extended_coord, extended_atype)	Calculate DescriptorBlock.
`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_emb`()	Returns the embedding dimension g2.
`get_dim_in`()	Returns the input dimension.
`get_dim_out`()	Returns the output dimension.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_nsel`()	Returns the number of selected atoms in the cut-off radius.
`get_ntypes`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_stats`()	Get the statistics of the descriptor.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`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 discriptor 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`(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_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__

compute_input_stats(merged: Union[Callable[[], List[dict]], List[dict]], path: Optional[DPPath] = None)[source]

Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

List[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.
Callable[[], List[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

pathOptional[DPPath]

The path to the stat file.

property dim_emb: Returns the embedding dimension g2.

property dim_in: Returns the atomic input dimension of this descriptor.

property dim_out: Returns the output dimension of this descriptor.

forward(nlist: Tensor, extended_coord: Tensor, extended_atype: Tensor, extended_atype_embd: Optional[Tensor] = None, mapping: Optional[Tensor] = None)[source]: Calculate DescriptorBlock.

get_dim_emb() → int[source]: Returns the embedding dimension g2.

get_dim_in() → int[source]: Returns the input dimension.

get_dim_out() → int[source]: Returns the output dimension.

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

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

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

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

get_stats() → Dict[str, StatItem][source]: Get the statistics of the descriptor.

mixed_types() → bool[source]

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

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

training: bool

class deepmd.pt.model.descriptor.DescrptBlockSeA(*args, **kwargs)[source]

Bases: DescriptorBlock

Attributes

dim_in: Returns the atomic input dimension of this descriptor.
dim_out: Returns the output dimension of this descriptor.

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.
`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_input_stats`(merged[, path])	Compute the input statistics (e.g.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(nlist, extended_coord, extended_atype)	Calculate decoded embedding for each atom.
`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_emb`()	Returns the output dimension.
`get_dim_in`()	Returns the input dimension.
`get_dim_out`()	Returns the output dimension.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_nsel`()	Returns the number of selected atoms in the cut-off radius.
`get_ntypes`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_stats`()	Get the statistics of the descriptor.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`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 discriptor 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`(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_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__
reinit_exclude

compute_input_stats(merged: Union[Callable[[], List[dict]], List[dict]], path: Optional[DPPath] = None)[source]

Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

List[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.
Callable[[], List[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

pathOptional[DPPath]

The path to the stat file.

property dim_in: Returns the atomic input dimension of this descriptor.

property dim_out: Returns the output dimension of this descriptor.

forward(nlist: Tensor, extended_coord: Tensor, extended_atype: Tensor, extended_atype_embd: Optional[Tensor] = None, mapping: Optional[Tensor] = None)[source]

Calculate decoded embedding for each atom.

Args: - coord: Tell atom coordinates with shape [nframes, natoms[1]*3]. - atype: Tell atom types with shape [nframes, natoms[1]]. - natoms: Tell atom count and element count. Its shape is [2+self.ntypes]. - box: Tell simulation box with shape [nframes, 9].

Returns

torch.Tensor: descriptor matrix with shape [nframes, natoms[0]*self.filter_neuron[-1]*self.axis_neuron].

get_dim_emb() → int[source]: Returns the output dimension.

get_dim_in() → int[source]: Returns the input dimension.

get_dim_out() → int[source]: Returns the output dimension.

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

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

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

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

get_stats() → Dict[str, StatItem][source]: Get the statistics of the descriptor.

mixed_types() → bool[source]

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

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

ndescrpt: Final[int]

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

class deepmd.pt.model.descriptor.DescrptBlockSeAtten(*args, **kwargs)[source]

Bases: DescriptorBlock

Attributes

dim_emb: Returns the output dimension of embedding.
dim_in: Returns the atomic input dimension of this descriptor.
dim_out: Returns the output dimension of this descriptor.

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.
`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_input_stats`(merged[, path])	Compute the input statistics (e.g.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(nlist, extended_coord, extended_atype)	Calculate decoded embedding for each atom.
`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_emb`()	Returns the output dimension of embedding.
`get_dim_in`()	Returns the output dimension.
`get_dim_out`()	Returns the output dimension.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_nsel`()	Returns the number of selected atoms in the cut-off radius.
`get_ntypes`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_stats`()	Get the statistics of the descriptor.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`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 discriptor 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`(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_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__

compute_input_stats(merged: Union[Callable[[], List[dict]], List[dict]], path: Optional[DPPath] = None)[source]

Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

List[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.
Callable[[], List[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

pathOptional[DPPath]

The path to the stat file.

property dim_emb: Returns the output dimension of embedding.

property dim_in: Returns the atomic input dimension of this descriptor.

property dim_out: Returns the output dimension of this descriptor.

forward(nlist: Tensor, extended_coord: Tensor, extended_atype: Tensor, extended_atype_embd: Optional[Tensor] = None, mapping: Optional[Tensor] = None) → List[Tensor][source]

Calculate decoded embedding for each atom.

Args: - coord: Tell atom coordinates with shape [nframes, natoms[1]*3]. - atype: Tell atom types with shape [nframes, natoms[1]]. - natoms: Tell atom count and element count. Its shape is [2+self.ntypes]. - box: Tell simulation box with shape [nframes, 9].

Returns

result: descriptor with shape [nframes, nloc, self.filter_neuron[-1] * self.axis_neuron].
ret: environment matrix with shape [nframes, nloc, self.neei, out_size]

get_dim_emb() → int[source]: Returns the output dimension of embedding.

get_dim_in() → int[source]: Returns the output dimension.

get_dim_out() → int[source]: Returns the output dimension.

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

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

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

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

get_stats() → Dict[str, StatItem][source]: Get the statistics of the descriptor.

mixed_types() → bool[source]

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

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

training: bool

class deepmd.pt.model.descriptor.DescrptDPA1(*args, **kwargs)[source]

Bases: BD, Module

Attributes

dim_emb
dim_out

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.
`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_input_stats`(merged[, path])	Compute the input statistics (e.g.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`deserialize`()	Deserialize from a dict.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(extended_coord, extended_atype, nlist)	Compute the descriptor.
`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_emb`()	Returns the embedding dimension of g2.
`get_dim_out`()	Returns the output dimension.
`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 number of selected atoms in the cut-off radius.
`get_ntypes`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`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 discriptor 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`(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.
`serialize`()	Serialize the obj to dict.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__

compute_input_stats(merged: Union[Callable[[], List[dict]], List[dict]], path: Optional[DPPath] = None)[source]

Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

List[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.
Callable[[], List[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

pathOptional[DPPath]

The path to the stat file.

classmethod deserialize() → DescrptDPA1[source]: Deserialize from a dict.

property dim_emb

property dim_out

forward(extended_coord: Tensor, extended_atype: Tensor, nlist: Tensor, mapping: Optional[Tensor] = None)[source]

Compute the descriptor.

Parameters

coord_ext: The extended coordinates of atoms. shape: nf x (nallx3)
atype_ext: The extended aotm types. shape: nf x nall
nlist: The neighbor list. shape: nf x nloc x nnei
mapping: The index mapping, not required by this descriptor.

Returns

descriptor: The descriptor. shape: nf x nloc x (ng x axis_neuron)
gr: The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3
g2: The rotationally invariant pair-partical representation. shape: nf x nloc x nnei x ng
h2: The rotationally equivariant pair-partical representation. shape: nf x nloc x nnei x 3
sw: The smooth switch function. shape: nf x nloc x nnei

get_dim_emb() → int[source]: Returns the embedding dimension of g2.

get_dim_out() → int[source]: Returns the output dimension.

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

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

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

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

mixed_types() → bool[source]

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

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

serialize() → dict[source]: Serialize the obj to dict.

share_params(base_class, shared_level, resume=False)[source]: Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.

training: bool

classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]

Update the selection and perform neighbor statistics.

Parameters

global_jdatadict: The global data, containing the training section
local_jdatadict: The local data refer to the current class

class deepmd.pt.model.descriptor.DescrptDPA2(*args, **kwargs)[source]

Bases: Module, BD

Attributes

dim_emb: Returns the embedding dimension g2.
dim_out

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.
`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_input_stats`(merged[, path])	Compute the input statistics (e.g.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`deserialize`()	Deserialize from a dict.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(extended_coord, extended_atype, nlist)	Compute the descriptor.
`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_emb`()	Returns the embedding dimension of this descriptor.
`get_dim_out`()	Returns the output dimension of this descriptor.
`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 number of selected atoms in the cut-off radius.
`get_ntypes`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`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 discriptor 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`(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.
`serialize`()	Serialize the obj to dict.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__

compute_input_stats(merged: Union[Callable[[], List[dict]], List[dict]], path: Optional[DPPath] = None)[source]

Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

List[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.
Callable[[], List[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

pathOptional[DPPath]

The path to the stat file.

classmethod deserialize() → DescrptDPA2[source]: Deserialize from a dict.

property dim_emb: Returns the embedding dimension g2.

property dim_out

forward(extended_coord: Tensor, extended_atype: Tensor, nlist: Tensor, mapping: Optional[Tensor] = None)[source]

Compute the descriptor.

Parameters

coord_ext: The extended coordinates of atoms. shape: nf x (nallx3)
atype_ext: The extended aotm types. shape: nf x nall
nlist: The neighbor list. shape: nf x nloc x nnei
mapping: The index mapping, mapps extended region index to local region.

Returns

descriptor: The descriptor. shape: nf x nloc x (ng x axis_neuron)
gr: The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3
g2: The rotationally invariant pair-partical representation. shape: nf x nloc x nnei x ng
h2: The rotationally equivariant pair-partical representation. shape: nf x nloc x nnei x 3
sw: The smooth switch function. shape: nf x nloc x nnei

get_dim_emb() → int[source]: Returns the embedding dimension of this descriptor.

get_dim_out() → int[source]: Returns the output dimension of this descriptor.

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

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

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

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

mixed_types() → bool[source]

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

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

serialize() → dict[source]: Serialize the obj to dict.

share_params(base_class, shared_level, resume=False)[source]: Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.

training: bool

classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]

Update the selection and perform neighbor statistics.

Parameters

global_jdatadict: The global data, containing the training section
local_jdatadict: The local data refer to the current class

class deepmd.pt.model.descriptor.DescrptGaussianLcc(*args, **kwargs)[source]

Bases: Module, BD

Attributes

dim_emb: Returns the output dimension of pair representation.
dim_in: Returns the atomic input dimension of this descriptor.
dim_out: Returns the output dimension of atomic representation.

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.
`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_input_stats`(merged[, path])	Update mean and stddev for descriptor elements.
`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.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(extended_coord, nlist, atype, nlist_type)	Calculate the atomic and pair representations of this descriptor.
`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_emb`()	Returns the embedding dimension of g2.
`get_dim_out`()	Returns the output descriptor dimension.
`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`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected neighboring atoms for each type.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`mixed_types`()	Returns if the descriptor requires a neighbor list that distinguish different atomic types or not.
`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`(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.
`serialize`()	Serialize the obj to dict.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__

compute_input_stats(merged: List[dict], path: Optional[DPPath] = None)[source]: Update mean and stddev for descriptor elements.

property dim_emb: Returns the output dimension of pair representation.

property dim_in: Returns the atomic input dimension of this descriptor.

property dim_out: Returns the output dimension of atomic representation.

forward(extended_coord, nlist, atype, nlist_type, nlist_loc=None, atype_tebd=None, nlist_tebd=None, seq_input=None)[source]

Calculate the atomic and pair representations of this descriptor.

Args: - extended_coord: Copied atom coordinates with shape [nframes, nall, 3]. - nlist: Neighbor list with shape [nframes, nloc, nnei]. - atype: Atom type with shape [nframes, nloc]. - nlist_type: Atom type of neighbors with shape [nframes, nloc, nnei]. - nlist_loc: Local index of neighbor list with shape [nframes, nloc, nnei]. - atype_tebd: Atomic type embedding with shape [nframes, nloc, tebd_dim]. - nlist_tebd: Type embeddings of neighbor with shape [nframes, nloc, nnei, tebd_dim]. - seq_input: The sequential input from other descriptor with

shape [nframes, nloc, tebd_dim] or [nframes * nloc, 1 + nnei, tebd_dim]

Returns

result: descriptor with shape [nframes, nloc, self.filter_neuron[-1] * self.axis_neuron].
ret: environment matrix with shape [nframes, nloc, self.neei, out_size]

training: bool

class deepmd.pt.model.descriptor.DescrptHybrid(*args, **kwargs)[source]

Bases: BD, Module

Concate a list of descriptors to form a new descriptor.

Parameters

listlistList[Union[BaseDescriptor, Dict[str, Any]]]: Build a descriptor from the concatenation of the list of descriptors. The descriptor can be either an object or a dictionary.

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.
`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_input_stats`(merged[, path])	Update mean and stddev for descriptor elements.
`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.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(coord_ext, atype_ext, nlist[, mapping])	Compute the descriptor.
`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_emb`()	Returns the output dimension.
`get_dim_out`()	Returns the output dimension.
`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`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`mixed_types`()	Returns if the descriptor requires a neighbor list that distinguish different atomic types or not.
`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`(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.
`serialize`()	Serialize the obj to dict.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__

compute_input_stats(merged: List[dict], path: Optional[DPPath] = None)[source]: Update mean and stddev for descriptor elements.

classmethod deserialize(data: dict) → DescrptHybrid[source]

Deserialize the model.

Parameters

datadict: The serialized data

Returns

BD: The deserialized descriptor

forward(coord_ext: Tensor, atype_ext: Tensor, nlist: Tensor, mapping: Optional[Tensor] = None)[source]

Compute the descriptor.

Parameters

coord_ext: The extended coordinates of atoms. shape: nf x (nallx3)
atype_ext: The extended aotm types. shape: nf x nall
nlist: The neighbor list. shape: nf x nloc x nnei
mapping: The index mapping, not required by this descriptor.

Returns

descriptor: The descriptor. shape: nf x nloc x (ng x axis_neuron)
gr: The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3. This descriptor returns None
g2: The rotationally invariant pair-partical representation. this descriptor returns None
h2: The rotationally equivariant pair-partical representation. this descriptor returns None
sw: The smooth switch function. this descriptor returns None

get_dim_emb() → int[source]: Returns the output dimension.

get_dim_out() → int[source]: Returns the output dimension.

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

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

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

mixed_types()[source]: Returns if the descriptor requires a neighbor list that distinguish different atomic types or not.

serialize() → dict[source]: Serialize the obj to dict.

share_params(base_class, shared_level, resume=False)[source]: Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.

training: bool

classmethod update_sel(global_jdata: dict, local_jdata: dict) → dict[source]

Update the selection and perform neighbor statistics.

Parameters

global_jdatadict: The global data, containing the training section
local_jdatadict: The local data refer to the current class

class deepmd.pt.model.descriptor.DescrptSeA(*args, **kwargs)[source]

Bases: BD, Module

Attributes

dim_out: Returns the output dimension of this descriptor.

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.
`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_input_stats`(merged[, path])	Compute the input statistics (e.g.
`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.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(coord_ext, atype_ext, nlist[, mapping])	Compute the descriptor.
`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_emb`()	Returns the output dimension.
`get_dim_out`()	Returns the output dimension.
`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 number of selected atoms in the cut-off radius.
`get_ntypes`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`mixed_types`()	Returns if the descriptor requires a neighbor list that distinguish different atomic types or not.
`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`(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.
`reinit_exclude`([exclude_types])	Update the type exclusions.
`requires_grad_`([requires_grad])	Change if autograd should record operations on parameters in this module.
`serialize`()	Serialize the obj to dict.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__
set_stat_mean_and_stddev

compute_input_stats(merged: Union[Callable[[], List[dict]], List[dict]], path: Optional[DPPath] = None)[source]

Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

List[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.
Callable[[], List[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

pathOptional[DPPath]

The path to the stat file.

classmethod deserialize(data: dict) → DescrptSeA[source]

Deserialize the model.

Parameters

datadict: The serialized data

Returns

BD: The deserialized descriptor

property dim_out: Returns the output dimension of this descriptor.

forward(coord_ext: Tensor, atype_ext: Tensor, nlist: Tensor, mapping: Optional[Tensor] = None)[source]

Compute the descriptor.

Parameters

coord_ext: The extended coordinates of atoms. shape: nf x (nallx3)
atype_ext: The extended aotm types. shape: nf x nall
nlist: The neighbor list. shape: nf x nloc x nnei
mapping: The index mapping, not required by this descriptor.

Returns

descriptor: The descriptor. shape: nf x nloc x (ng x axis_neuron)
gr: The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3
g2: The rotationally invariant pair-partical representation. this descriptor returns None
h2: The rotationally equivariant pair-partical representation. this descriptor returns None
sw: The smooth switch function.

get_dim_emb() → int[source]: Returns the output dimension.

get_dim_out() → int[source]: Returns the output dimension.

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

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

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

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

mixed_types()[source]: Returns if the descriptor requires a neighbor list that distinguish different atomic types or not.

reinit_exclude(exclude_types: List[Tuple[int, int]] = [])[source]: Update the type exclusions.

serialize() → dict[source]: Serialize the obj to dict.

set_stat_mean_and_stddev(mean: Tensor, stddev: Tensor) → None[source]

share_params(base_class, shared_level, resume=False)[source]: Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.

training: bool

classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]

Update the selection and perform neighbor statistics.

Parameters

global_jdatadict: The global data, containing the training section
local_jdatadict: The local data refer to the current class

class deepmd.pt.model.descriptor.DescrptSeR(*args, **kwargs)[source]

Bases: BD, Module

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.
`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_input_stats`(merged[, path])	Compute the input statistics (e.g.
`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.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(coord_ext, atype_ext, nlist[, mapping])	Compute the descriptor.
`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_emb`()	Returns the output dimension.
`get_dim_in`()	Returns the input dimension.
`get_dim_out`()	Returns the output dimension.
`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 number of selected atoms in the cut-off radius.
`get_ntypes`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_stats`()	Get the statistics of the descriptor.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`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 discriptor 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`(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.
`serialize`()	Serialize the obj to dict.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__
set_stat_mean_and_stddev

compute_input_stats(merged: Union[Callable[[], List[dict]], List[dict]], path: Optional[DPPath] = None)[source]

Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

List[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.
Callable[[], List[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

pathOptional[DPPath]

The path to the stat file.

classmethod deserialize(data: dict) → DescrptSeR[source]

Deserialize the model.

Parameters

datadict: The serialized data

Returns

BD: The deserialized descriptor

forward(coord_ext: Tensor, atype_ext: Tensor, nlist: Tensor, mapping: Optional[Tensor] = None)[source]

Compute the descriptor.

Parameters

coord_ext: The extended coordinates of atoms. shape: nf x (nallx3)
atype_ext: The extended aotm types. shape: nf x nall
nlist: The neighbor list. shape: nf x nloc x nnei
mapping: The index mapping, not required by this descriptor.

Returns

descriptor: The descriptor. shape: nf x nloc x (ng x axis_neuron)
gr: The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3
g2: The rotationally invariant pair-partical representation. this descriptor returns None
h2: The rotationally equivariant pair-partical representation. this descriptor returns None
sw: The smooth switch function.

get_dim_emb() → int[source]: Returns the output dimension.

get_dim_in() → int[source]: Returns the input dimension.

get_dim_out() → int[source]: Returns the output dimension.

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

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

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

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

get_stats() → Dict[str, StatItem][source]: Get the statistics of the descriptor.

mixed_types() → bool[source]

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

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

serialize() → dict[source]: Serialize the obj to dict.

set_stat_mean_and_stddev(mean: Tensor, stddev: Tensor) → None[source]

share_params(base_class, shared_level, resume=False)[source]: Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.

training: bool

classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]

Update the selection and perform neighbor statistics.

Parameters

global_jdatadict: The global data, containing the training section
local_jdatadict: The local data refer to the current class

deepmd.pt.model.descriptor.make_default_type_embedding(ntypes)[source]

deepmd.pt.model.descriptor.prod_env_mat(extended_coord, nlist, atype, mean, stddev, rcut: float, rcut_smth: float, radial_only: bool = False)[source]

Generate smooth environment matrix from atom coordinates and other context.

Args: - extended_coord: Copied atom coordinates with shape [nframes, nall*3]. - atype: Atom types with shape [nframes, nloc]. - mean: Average value of descriptor per element type with shape [len(sec), nnei, 4 or 1]. - stddev: Standard deviation of descriptor per element type with shape [len(sec), nnei, 4 or 1]. - rcut: Cut-off radius. - rcut_smth: Smooth hyper-parameter for pair force & energy. - radial_only: Whether to return a full description or a radial-only descriptor.

Returns

env_mat: Shape is [nframes, natoms[1]*nnei*4].

Submodules

deepmd.pt.model.descriptor.base_descriptor module

deepmd.pt.model.descriptor.descriptor module

class deepmd.pt.model.descriptor.descriptor.DescriptorBlock(*args, **kwargs)[source]

Bases: Module, ABC, PR

The building block of descriptor. Given the input descriptor, provide with the atomic coordinates, atomic types and neighbor list, calculate the new descriptor.

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.
`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_input_stats`(merged[, path])	Compute the input statistics (e.g.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(nlist, extended_coord, extended_atype)	Calculate DescriptorBlock.
`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_emb`()	Returns the embedding dimension.
`get_dim_in`()	Returns the output dimension.
`get_dim_out`()	Returns the output dimension.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_nsel`()	Returns the number of selected atoms in the cut-off radius.
`get_ntypes`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_stats`()	Get the statistics of the descriptor.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`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`(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_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__

compute_input_stats(merged: Union[Callable[[], List[dict]], List[dict]], path: Optional[DPPath] = None)[source]

Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

List[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.
Callable[[], List[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

pathOptional[DPPath]

The path to the stat file.

abstract forward(nlist: Tensor, extended_coord: Tensor, extended_atype: Tensor, extended_atype_embd: Optional[Tensor] = None, mapping: Optional[Tensor] = None)[source]: Calculate DescriptorBlock.

abstract get_dim_emb() → int[source]: Returns the embedding dimension.

abstract get_dim_in() → int[source]: Returns the output dimension.

abstract get_dim_out() → int[source]: Returns the output dimension.

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

abstract get_ntypes() → int[source]: Returns the number of element types.

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

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

get_stats() → Dict[str, StatItem][source]: Get the statistics of the descriptor.

local_cluster = False

share_params(base_class, shared_level, resume=False)[source]: Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.

training: bool

deepmd.pt.model.descriptor.descriptor.make_default_type_embedding(ntypes)[source]

deepmd.pt.model.descriptor.dpa1 module

class deepmd.pt.model.descriptor.dpa1.DescrptDPA1(*args, **kwargs)[source]

Bases: BD, Module

Attributes

dim_emb
dim_out

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.
`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_input_stats`(merged[, path])	Compute the input statistics (e.g.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`deserialize`()	Deserialize from a dict.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(extended_coord, extended_atype, nlist)	Compute the descriptor.
`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_emb`()	Returns the embedding dimension of g2.
`get_dim_out`()	Returns the output dimension.
`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 number of selected atoms in the cut-off radius.
`get_ntypes`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`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 discriptor 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`(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.
`serialize`()	Serialize the obj to dict.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__

compute_input_stats(merged: Union[Callable[[], List[dict]], List[dict]], path: Optional[DPPath] = None)[source]

Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

List[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.
Callable[[], List[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

pathOptional[DPPath]

The path to the stat file.

classmethod deserialize() → DescrptDPA1[source]: Deserialize from a dict.

property dim_emb

property dim_out

forward(extended_coord: Tensor, extended_atype: Tensor, nlist: Tensor, mapping: Optional[Tensor] = None)[source]

Compute the descriptor.

Parameters

coord_ext: The extended coordinates of atoms. shape: nf x (nallx3)
atype_ext: The extended aotm types. shape: nf x nall
nlist: The neighbor list. shape: nf x nloc x nnei
mapping: The index mapping, not required by this descriptor.

Returns

descriptor: The descriptor. shape: nf x nloc x (ng x axis_neuron)
gr: The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3
g2: The rotationally invariant pair-partical representation. shape: nf x nloc x nnei x ng
h2: The rotationally equivariant pair-partical representation. shape: nf x nloc x nnei x 3
sw: The smooth switch function. shape: nf x nloc x nnei

get_dim_emb() → int[source]: Returns the embedding dimension of g2.

get_dim_out() → int[source]: Returns the output dimension.

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

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

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

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

mixed_types() → bool[source]

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

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

serialize() → dict[source]: Serialize the obj to dict.

share_params(base_class, shared_level, resume=False)[source]: Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.

training: bool

classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]

Update the selection and perform neighbor statistics.

Parameters

global_jdatadict: The global data, containing the training section
local_jdatadict: The local data refer to the current class

deepmd.pt.model.descriptor.dpa2 module

class deepmd.pt.model.descriptor.dpa2.DescrptDPA2(*args, **kwargs)[source]

Bases: Module, BD

Attributes

dim_emb: Returns the embedding dimension g2.
dim_out

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.
`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_input_stats`(merged[, path])	Compute the input statistics (e.g.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`deserialize`()	Deserialize from a dict.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(extended_coord, extended_atype, nlist)	Compute the descriptor.
`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_emb`()	Returns the embedding dimension of this descriptor.
`get_dim_out`()	Returns the output dimension of this descriptor.
`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 number of selected atoms in the cut-off radius.
`get_ntypes`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`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 discriptor 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`(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.
`serialize`()	Serialize the obj to dict.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__

compute_input_stats(merged: Union[Callable[[], List[dict]], List[dict]], path: Optional[DPPath] = None)[source]

Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

List[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.
Callable[[], List[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

pathOptional[DPPath]

The path to the stat file.

classmethod deserialize() → DescrptDPA2[source]: Deserialize from a dict.

property dim_emb: Returns the embedding dimension g2.

property dim_out

forward(extended_coord: Tensor, extended_atype: Tensor, nlist: Tensor, mapping: Optional[Tensor] = None)[source]

Compute the descriptor.

Parameters

coord_ext: The extended coordinates of atoms. shape: nf x (nallx3)
atype_ext: The extended aotm types. shape: nf x nall
nlist: The neighbor list. shape: nf x nloc x nnei
mapping: The index mapping, mapps extended region index to local region.

Returns

descriptor: The descriptor. shape: nf x nloc x (ng x axis_neuron)
gr: The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3
g2: The rotationally invariant pair-partical representation. shape: nf x nloc x nnei x ng
h2: The rotationally equivariant pair-partical representation. shape: nf x nloc x nnei x 3
sw: The smooth switch function. shape: nf x nloc x nnei

get_dim_emb() → int[source]: Returns the embedding dimension of this descriptor.

get_dim_out() → int[source]: Returns the output dimension of this descriptor.

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

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

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

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

mixed_types() → bool[source]

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

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

serialize() → dict[source]: Serialize the obj to dict.

share_params(base_class, shared_level, resume=False)[source]: Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.

training: bool

classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]

Update the selection and perform neighbor statistics.

Parameters

global_jdatadict: The global data, containing the training section
local_jdatadict: The local data refer to the current class

deepmd.pt.model.descriptor.env_mat module

deepmd.pt.model.descriptor.env_mat.prod_env_mat(extended_coord, nlist, atype, mean, stddev, rcut: float, rcut_smth: float, radial_only: bool = False)[source]

Generate smooth environment matrix from atom coordinates and other context.

Args: - extended_coord: Copied atom coordinates with shape [nframes, nall*3]. - atype: Atom types with shape [nframes, nloc]. - mean: Average value of descriptor per element type with shape [len(sec), nnei, 4 or 1]. - stddev: Standard deviation of descriptor per element type with shape [len(sec), nnei, 4 or 1]. - rcut: Cut-off radius. - rcut_smth: Smooth hyper-parameter for pair force & energy. - radial_only: Whether to return a full description or a radial-only descriptor.

Returns

env_mat: Shape is [nframes, natoms[1]*nnei*4].

deepmd.pt.model.descriptor.gaussian_lcc module

class deepmd.pt.model.descriptor.gaussian_lcc.DescrptGaussianLcc(*args, **kwargs)[source]

Bases: Module, BD

Attributes

dim_emb: Returns the output dimension of pair representation.
dim_in: Returns the atomic input dimension of this descriptor.
dim_out: Returns the output dimension of atomic representation.

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.
`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_input_stats`(merged[, path])	Update mean and stddev for descriptor elements.
`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.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(extended_coord, nlist, atype, nlist_type)	Calculate the atomic and pair representations of this descriptor.
`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_emb`()	Returns the embedding dimension of g2.
`get_dim_out`()	Returns the output descriptor dimension.
`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`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected neighboring atoms for each type.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`mixed_types`()	Returns if the descriptor requires a neighbor list that distinguish different atomic types or not.
`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`(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.
`serialize`()	Serialize the obj to dict.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__

compute_input_stats(merged: List[dict], path: Optional[DPPath] = None)[source]: Update mean and stddev for descriptor elements.

property dim_emb: Returns the output dimension of pair representation.

property dim_in: Returns the atomic input dimension of this descriptor.

property dim_out: Returns the output dimension of atomic representation.

forward(extended_coord, nlist, atype, nlist_type, nlist_loc=None, atype_tebd=None, nlist_tebd=None, seq_input=None)[source]

Calculate the atomic and pair representations of this descriptor.

Args: - extended_coord: Copied atom coordinates with shape [nframes, nall, 3]. - nlist: Neighbor list with shape [nframes, nloc, nnei]. - atype: Atom type with shape [nframes, nloc]. - nlist_type: Atom type of neighbors with shape [nframes, nloc, nnei]. - nlist_loc: Local index of neighbor list with shape [nframes, nloc, nnei]. - atype_tebd: Atomic type embedding with shape [nframes, nloc, tebd_dim]. - nlist_tebd: Type embeddings of neighbor with shape [nframes, nloc, nnei, tebd_dim]. - seq_input: The sequential input from other descriptor with

shape [nframes, nloc, tebd_dim] or [nframes * nloc, 1 + nnei, tebd_dim]

Returns

result: descriptor with shape [nframes, nloc, self.filter_neuron[-1] * self.axis_neuron].
ret: environment matrix with shape [nframes, nloc, self.neei, out_size]

training: bool

deepmd.pt.model.descriptor.hybrid module

class deepmd.pt.model.descriptor.hybrid.DescrptBlockHybrid(*args, **kwargs)[source]

Bases: DescriptorBlock

Attributes

dim_emb: Returns the output dimension of embedding.
dim_emb_list: Returns the output dimension list of embeddings.
dim_out: Returns the output dimension of this descriptor.

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.
`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_input_stats`(merged[, path])	Compute the input statistics (e.g.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(nlist, extended_coord, extended_atype)	Calculate decoded embedding for each atom.
`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_emb`()	Returns the embedding dimension.
`get_dim_in`()	Returns the input dimension.
`get_dim_out`()	Returns the output dimension.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_nsel`()	Returns the number of selected atoms in the cut-off radius.
`get_ntypes`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_stats`()	Get the statistics of the descriptor.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`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 discriptor 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`(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_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__

compute_input_stats(merged: Union[Callable[[], List[dict]], List[dict]], path: Optional[DPPath] = None)[source]

Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

List[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.
Callable[[], List[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

pathOptional[DPPath]

The path to the stat file.

property dim_emb: Returns the output dimension of embedding.

property dim_emb_list: List[int]: Returns the output dimension list of embeddings.

property dim_out: Returns the output dimension of this descriptor.

forward(nlist: Tensor, extended_coord: Tensor, extended_atype: Tensor, extended_atype_embd: Optional[Tensor] = None, mapping: Optional[Tensor] = None)[source]

Calculate decoded embedding for each atom.

Args: - extended_coord: Tell atom coordinates with shape [nframes, natoms[1]*3]. - nlist: Tell atom types with shape [nframes, natoms[1]]. - atype: Tell atom count and element count. Its shape is [2+self.ntypes]. - nlist_type: Tell simulation box with shape [nframes, 9]. - atype_tebd: Tell simulation box with shape [nframes, 9]. - nlist_tebd: Tell simulation box with shape [nframes, 9].

Returns

result: descriptor with shape [nframes, nloc, self.filter_neuron[-1] * self.axis_neuron].
ret: environment matrix with shape [nframes, nloc, self.neei, out_size]

get_dim_emb()[source]: Returns the embedding dimension.

get_dim_in() → int[source]: Returns the input dimension.

get_dim_out() → int[source]: Returns the output dimension.

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

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

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

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

mixed_types() → bool[source]

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

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

share_params(base_class, shared_level, resume=False)[source]: Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.

training: bool

class deepmd.pt.model.descriptor.hybrid.DescrptHybrid(*args, **kwargs)[source]

Bases: BD, Module

Concate a list of descriptors to form a new descriptor.

Parameters

listlistList[Union[BaseDescriptor, Dict[str, Any]]]: Build a descriptor from the concatenation of the list of descriptors. The descriptor can be either an object or a dictionary.

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.
`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_input_stats`(merged[, path])	Update mean and stddev for descriptor elements.
`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.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(coord_ext, atype_ext, nlist[, mapping])	Compute the descriptor.
`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_emb`()	Returns the output dimension.
`get_dim_out`()	Returns the output dimension.
`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`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`mixed_types`()	Returns if the descriptor requires a neighbor list that distinguish different atomic types or not.
`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`(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.
`serialize`()	Serialize the obj to dict.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__

compute_input_stats(merged: List[dict], path: Optional[DPPath] = None)[source]: Update mean and stddev for descriptor elements.

classmethod deserialize(data: dict) → DescrptHybrid[source]

Deserialize the model.

Parameters

datadict: The serialized data

Returns

BD: The deserialized descriptor

forward(coord_ext: Tensor, atype_ext: Tensor, nlist: Tensor, mapping: Optional[Tensor] = None)[source]

Compute the descriptor.

Parameters

coord_ext: The extended coordinates of atoms. shape: nf x (nallx3)
atype_ext: The extended aotm types. shape: nf x nall
nlist: The neighbor list. shape: nf x nloc x nnei
mapping: The index mapping, not required by this descriptor.

Returns

descriptor: The descriptor. shape: nf x nloc x (ng x axis_neuron)
gr: The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3. This descriptor returns None
g2: The rotationally invariant pair-partical representation. this descriptor returns None
h2: The rotationally equivariant pair-partical representation. this descriptor returns None
sw: The smooth switch function. this descriptor returns None

get_dim_emb() → int[source]: Returns the output dimension.

get_dim_out() → int[source]: Returns the output dimension.

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

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

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

mixed_types()[source]: Returns if the descriptor requires a neighbor list that distinguish different atomic types or not.

nlist_cut_idx: List[torch.Tensor]

serialize() → dict[source]: Serialize the obj to dict.

share_params(base_class, shared_level, resume=False)[source]: Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.

training: bool

classmethod update_sel(global_jdata: dict, local_jdata: dict) → dict[source]

Update the selection and perform neighbor statistics.

Parameters

global_jdatadict: The global data, containing the training section
local_jdatadict: The local data refer to the current class

deepmd.pt.model.descriptor.repformer_layer module

class deepmd.pt.model.descriptor.repformer_layer.Atten2EquiVarApply(ni: int, nh: int)[source]

Bases: Module

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.
`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.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(AA, h2)	Define the computation performed at every call.
`get_buffer`(target)	Return the buffer given by `target` if it exists, otherwise throw an error.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_parameter`(target)	Return the parameter given by `target` if it exists, otherwise throw an error.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`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__

forward(AA: Tensor, h2: Tensor) → Tensor[source]

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.

training: bool

class deepmd.pt.model.descriptor.repformer_layer.Atten2Map(ni: int, nd: int, nh: int, has_gate: bool = False, smooth: bool = True, attnw_shift: float = 20.0)[source]

Bases: Module

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.
`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.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(g2, h2, nlist_mask, sw)	Define the computation performed at every call.
`get_buffer`(target)	Return the buffer given by `target` if it exists, otherwise throw an error.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_parameter`(target)	Return the parameter given by `target` if it exists, otherwise throw an error.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`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__

forward(g2: Tensor, h2: Tensor, nlist_mask: Tensor, sw: Tensor) → Tensor[source]

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.

training: bool

class deepmd.pt.model.descriptor.repformer_layer.Atten2MultiHeadApply(ni: int, nh: int)[source]

Bases: Module

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.
`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.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(AA, g2)	Define the computation performed at every call.
`get_buffer`(target)	Return the buffer given by `target` if it exists, otherwise throw an error.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_parameter`(target)	Return the parameter given by `target` if it exists, otherwise throw an error.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`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__

forward(AA: Tensor, g2: Tensor) → Tensor[source]

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.

training: bool

class deepmd.pt.model.descriptor.repformer_layer.LocalAtten(ni: int, nd: int, nh: int, smooth: bool = True, attnw_shift: float = 20.0)[source]

Bases: Module

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.
`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.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(g1, gg1, nlist_mask, sw)	Define the computation performed at every call.
`get_buffer`(target)	Return the buffer given by `target` if it exists, otherwise throw an error.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_parameter`(target)	Return the parameter given by `target` if it exists, otherwise throw an error.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`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__

forward(g1: Tensor, gg1: Tensor, nlist_mask: Tensor, sw: Tensor) → Tensor[source]

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.

training: bool

class deepmd.pt.model.descriptor.repformer_layer.RepformerLayer(rcut, rcut_smth, sel: int, ntypes: int, g1_dim=128, g2_dim=16, axis_dim: int = 4, update_chnnl_2: bool = True, do_bn_mode: str = 'no', bn_momentum: float = 0.1, update_g1_has_conv: bool = True, update_g1_has_drrd: bool = True, update_g1_has_grrg: bool = True, update_g1_has_attn: bool = True, update_g2_has_g1g1: bool = True, update_g2_has_attn: bool = True, update_h2: bool = False, attn1_hidden: int = 64, attn1_nhead: int = 4, attn2_hidden: int = 16, attn2_nhead: int = 4, attn2_has_gate: bool = False, activation_function: str = 'tanh', update_style: str = 'res_avg', set_davg_zero: bool = True, smooth: bool = True)[source]

Bases: Module

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.
`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.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(g1_ext, g2, h2, nlist, nlist_mask, sw)	Parameters
`get_buffer`(target)	Return the buffer given by `target` if it exists, otherwise throw an error.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_parameter`(target)	Return the parameter given by `target` if it exists, otherwise throw an error.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`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__
cal_1_dim
list_update
list_update_res_avg
list_update_res_incr

cal_1_dim(g1d: int, g2d: int, ax: int) → int[source]

forward(g1_ext: Tensor, g2: Tensor, h2: Tensor, nlist: Tensor, nlist_mask: Tensor, sw: Tensor)[source]

Parameters

g1_extnf x nall x ng1 extended single-atom chanel
g2nf x nloc x nnei x ng2 pair-atom channel, invariant
h2nf x nloc x nnei x 3 pair-atom channel, equivariant
nlistnf x nloc x nnei neighbor list (padded neis are set to 0)
nlist_masknf x nloc x nnei masks of the neighbor list. real nei 1 otherwise 0
swnf x nloc x nnei switch function

Returns

g1: nf x nloc x ng1 updated single-atom chanel
g2: nf x nloc x nnei x ng2 updated pair-atom channel, invariant
h2: nf x nloc x nnei x 3 updated pair-atom channel, equivariant

list_update(update_list: List[Tensor]) → Tensor[source]

list_update_res_avg(update_list: List[Tensor]) → Tensor[source]

list_update_res_incr(update_list: List[Tensor]) → Tensor[source]

training: bool

deepmd.pt.model.descriptor.repformer_layer.torch_linear(*args, **kwargs)[source]

deepmd.pt.model.descriptor.repformers module

class deepmd.pt.model.descriptor.repformers.DescrptBlockRepformers(*args, **kwargs)[source]

Bases: DescriptorBlock

Attributes

dim_emb: Returns the embedding dimension g2.
dim_in: Returns the atomic input dimension of this descriptor.
dim_out: Returns the output dimension of this descriptor.

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.
`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_input_stats`(merged[, path])	Compute the input statistics (e.g.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(nlist, extended_coord, extended_atype)	Calculate DescriptorBlock.
`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_emb`()	Returns the embedding dimension g2.
`get_dim_in`()	Returns the input dimension.
`get_dim_out`()	Returns the output dimension.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_nsel`()	Returns the number of selected atoms in the cut-off radius.
`get_ntypes`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_stats`()	Get the statistics of the descriptor.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`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 discriptor 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`(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_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__

compute_input_stats(merged: Union[Callable[[], List[dict]], List[dict]], path: Optional[DPPath] = None)[source]

Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

List[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.
Callable[[], List[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

pathOptional[DPPath]

The path to the stat file.

property dim_emb: Returns the embedding dimension g2.

property dim_in: Returns the atomic input dimension of this descriptor.

property dim_out: Returns the output dimension of this descriptor.

forward(nlist: Tensor, extended_coord: Tensor, extended_atype: Tensor, extended_atype_embd: Optional[Tensor] = None, mapping: Optional[Tensor] = None)[source]: Calculate DescriptorBlock.

get_dim_emb() → int[source]: Returns the embedding dimension g2.

get_dim_in() → int[source]: Returns the input dimension.

get_dim_out() → int[source]: Returns the output dimension.

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

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

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

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

get_stats() → Dict[str, StatItem][source]: Get the statistics of the descriptor.

mixed_types() → bool[source]

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

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

training: bool

deepmd.pt.model.descriptor.repformers.torch_linear(*args, **kwargs)[source]

deepmd.pt.model.descriptor.se_a module

class deepmd.pt.model.descriptor.se_a.DescrptBlockSeA(*args, **kwargs)[source]

Bases: DescriptorBlock

Attributes

dim_in: Returns the atomic input dimension of this descriptor.
dim_out: Returns the output dimension of this descriptor.

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.
`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_input_stats`(merged[, path])	Compute the input statistics (e.g.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(nlist, extended_coord, extended_atype)	Calculate decoded embedding for each atom.
`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_emb`()	Returns the output dimension.
`get_dim_in`()	Returns the input dimension.
`get_dim_out`()	Returns the output dimension.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_nsel`()	Returns the number of selected atoms in the cut-off radius.
`get_ntypes`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_stats`()	Get the statistics of the descriptor.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`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 discriptor 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`(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_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__
reinit_exclude

compute_input_stats(merged: Union[Callable[[], List[dict]], List[dict]], path: Optional[DPPath] = None)[source]

Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

List[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.
Callable[[], List[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

pathOptional[DPPath]

The path to the stat file.

property dim_in: Returns the atomic input dimension of this descriptor.

property dim_out: Returns the output dimension of this descriptor.

forward(nlist: Tensor, extended_coord: Tensor, extended_atype: Tensor, extended_atype_embd: Optional[Tensor] = None, mapping: Optional[Tensor] = None)[source]

Calculate decoded embedding for each atom.

Args: - coord: Tell atom coordinates with shape [nframes, natoms[1]*3]. - atype: Tell atom types with shape [nframes, natoms[1]]. - natoms: Tell atom count and element count. Its shape is [2+self.ntypes]. - box: Tell simulation box with shape [nframes, 9].

Returns

torch.Tensor: descriptor matrix with shape [nframes, natoms[0]*self.filter_neuron[-1]*self.axis_neuron].

get_dim_emb() → int[source]: Returns the output dimension.

get_dim_in() → int[source]: Returns the input dimension.

get_dim_out() → int[source]: Returns the output dimension.

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

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

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

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

get_stats() → Dict[str, StatItem][source]: Get the statistics of the descriptor.

mixed_types() → bool[source]

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

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

ndescrpt: Final[int]

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

training: bool

class deepmd.pt.model.descriptor.se_a.DescrptSeA(*args, **kwargs)[source]

Bases: BD, Module

Attributes

dim_out: Returns the output dimension of this descriptor.

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.
`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_input_stats`(merged[, path])	Compute the input statistics (e.g.
`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.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(coord_ext, atype_ext, nlist[, mapping])	Compute the descriptor.
`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_emb`()	Returns the output dimension.
`get_dim_out`()	Returns the output dimension.
`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 number of selected atoms in the cut-off radius.
`get_ntypes`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`mixed_types`()	Returns if the descriptor requires a neighbor list that distinguish different atomic types or not.
`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`(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.
`reinit_exclude`([exclude_types])	Update the type exclusions.
`requires_grad_`([requires_grad])	Change if autograd should record operations on parameters in this module.
`serialize`()	Serialize the obj to dict.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__
set_stat_mean_and_stddev

compute_input_stats(merged: Union[Callable[[], List[dict]], List[dict]], path: Optional[DPPath] = None)[source]

Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

List[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.
Callable[[], List[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

pathOptional[DPPath]

The path to the stat file.

classmethod deserialize(data: dict) → DescrptSeA[source]

Deserialize the model.

Parameters

datadict: The serialized data

Returns

BD: The deserialized descriptor

property dim_out: Returns the output dimension of this descriptor.

forward(coord_ext: Tensor, atype_ext: Tensor, nlist: Tensor, mapping: Optional[Tensor] = None)[source]

Compute the descriptor.

Parameters

coord_ext: The extended coordinates of atoms. shape: nf x (nallx3)
atype_ext: The extended aotm types. shape: nf x nall
nlist: The neighbor list. shape: nf x nloc x nnei
mapping: The index mapping, not required by this descriptor.

Returns

descriptor: The descriptor. shape: nf x nloc x (ng x axis_neuron)
gr: The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3
g2: The rotationally invariant pair-partical representation. this descriptor returns None
h2: The rotationally equivariant pair-partical representation. this descriptor returns None
sw: The smooth switch function.

get_dim_emb() → int[source]: Returns the output dimension.

get_dim_out() → int[source]: Returns the output dimension.

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

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

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

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

mixed_types()[source]: Returns if the descriptor requires a neighbor list that distinguish different atomic types or not.

reinit_exclude(exclude_types: List[Tuple[int, int]] = [])[source]: Update the type exclusions.

serialize() → dict[source]: Serialize the obj to dict.

set_stat_mean_and_stddev(mean: Tensor, stddev: Tensor) → None[source]

share_params(base_class, shared_level, resume=False)[source]: Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.

training: bool

classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]

Update the selection and perform neighbor statistics.

Parameters

global_jdatadict: The global data, containing the training section
local_jdatadict: The local data refer to the current class

deepmd.pt.model.descriptor.se_a.analyze_descrpt(matrix, ndescrpt, natoms)[source]: Collect avg, square avg and count of descriptors in a batch.

deepmd.pt.model.descriptor.se_atten module

class deepmd.pt.model.descriptor.se_atten.DescrptBlockSeAtten(*args, **kwargs)[source]

Bases: DescriptorBlock

Attributes

dim_emb: Returns the output dimension of embedding.
dim_in: Returns the atomic input dimension of this descriptor.
dim_out: Returns the output dimension of this descriptor.

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.
`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_input_stats`(merged[, path])	Compute the input statistics (e.g.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(nlist, extended_coord, extended_atype)	Calculate decoded embedding for each atom.
`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_emb`()	Returns the output dimension of embedding.
`get_dim_in`()	Returns the output dimension.
`get_dim_out`()	Returns the output dimension.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_nsel`()	Returns the number of selected atoms in the cut-off radius.
`get_ntypes`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_stats`()	Get the statistics of the descriptor.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`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 discriptor 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`(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_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__

compute_input_stats(merged: Union[Callable[[], List[dict]], List[dict]], path: Optional[DPPath] = None)[source]

Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

List[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.
Callable[[], List[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

pathOptional[DPPath]

The path to the stat file.

property dim_emb: Returns the output dimension of embedding.

property dim_in: Returns the atomic input dimension of this descriptor.

property dim_out: Returns the output dimension of this descriptor.

forward(nlist: Tensor, extended_coord: Tensor, extended_atype: Tensor, extended_atype_embd: Optional[Tensor] = None, mapping: Optional[Tensor] = None) → List[Tensor][source]

Calculate decoded embedding for each atom.

Args: - coord: Tell atom coordinates with shape [nframes, natoms[1]*3]. - atype: Tell atom types with shape [nframes, natoms[1]]. - natoms: Tell atom count and element count. Its shape is [2+self.ntypes]. - box: Tell simulation box with shape [nframes, 9].

Returns

result: descriptor with shape [nframes, nloc, self.filter_neuron[-1] * self.axis_neuron].
ret: environment matrix with shape [nframes, nloc, self.neei, out_size]

get_dim_emb() → int[source]: Returns the output dimension of embedding.

get_dim_in() → int[source]: Returns the output dimension.

get_dim_out() → int[source]: Returns the output dimension.

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

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

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

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

get_stats() → Dict[str, StatItem][source]: Get the statistics of the descriptor.

mixed_types() → bool[source]

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

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

training: bool

deepmd.pt.model.descriptor.se_atten.analyze_descrpt(matrix, ndescrpt, natoms, mixed_types=False, real_atype=None)[source]: Collect avg, square avg and count of descriptors in a batch.

deepmd.pt.model.descriptor.se_r module

class deepmd.pt.model.descriptor.se_r.DescrptSeR(*args, **kwargs)[source]

Bases: BD, Module

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.
`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_input_stats`(merged[, path])	Compute the input statistics (e.g.
`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.
`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.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(coord_ext, atype_ext, nlist[, mapping])	Compute the descriptor.
`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_emb`()	Returns the output dimension.
`get_dim_in`()	Returns the input dimension.
`get_dim_out`()	Returns the output dimension.
`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 number of selected atoms in the cut-off radius.
`get_ntypes`()	Returns the number of element 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`()	Returns the cut-off radius.
`get_sel`()	Returns the number of selected atoms for each type.
`get_stats`()	Get the statistics of the descriptor.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`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 discriptor 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`(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.
`serialize`()	Serialize the obj to dict.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`share_params`(base_class, shared_level[, resume])	Share the parameters of self to the base_class with shared_level during multitask training.
`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__
set_stat_mean_and_stddev

compute_input_stats(merged: Union[Callable[[], List[dict]], List[dict]], path: Optional[DPPath] = None)[source]

Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.

Parameters

mergedUnion[Callable[[], List[dict]], List[dict]]

List[dict]: A list of data samples from various data systems.
Each element, merged[i], is a data dictionary containing keys: torch.Tensor originating from the i-th data system.
Callable[[], List[dict]]: A lazy function that returns data samples in the above format
only when needed. Since the sampling process can be slow and memory-intensive, the lazy function helps by only sampling once.

pathOptional[DPPath]

The path to the stat file.

classmethod deserialize(data: dict) → DescrptSeR[source]

Deserialize the model.

Parameters

datadict: The serialized data

Returns

BD: The deserialized descriptor

forward(coord_ext: Tensor, atype_ext: Tensor, nlist: Tensor, mapping: Optional[Tensor] = None)[source]

Compute the descriptor.

Parameters

coord_ext: The extended coordinates of atoms. shape: nf x (nallx3)
atype_ext: The extended aotm types. shape: nf x nall
nlist: The neighbor list. shape: nf x nloc x nnei
mapping: The index mapping, not required by this descriptor.

Returns

descriptor: The descriptor. shape: nf x nloc x (ng x axis_neuron)
gr: The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3
g2: The rotationally invariant pair-partical representation. this descriptor returns None
h2: The rotationally equivariant pair-partical representation. this descriptor returns None
sw: The smooth switch function.

get_dim_emb() → int[source]: Returns the output dimension.

get_dim_in() → int[source]: Returns the input dimension.

get_dim_out() → int[source]: Returns the output dimension.

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

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

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

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

get_stats() → Dict[str, StatItem][source]: Get the statistics of the descriptor.

mixed_types() → bool[source]

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

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

serialize() → dict[source]: Serialize the obj to dict.

set_stat_mean_and_stddev(mean: Tensor, stddev: Tensor) → None[source]

share_params(base_class, shared_level, resume=False)[source]: Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some seperated parameters (e.g. mean and stddev) will be re-calculated across different classes.

training: bool

classmethod update_sel(global_jdata: dict, local_jdata: dict)[source]

Update the selection and perform neighbor statistics.

Parameters

global_jdatadict: The global data, containing the training section
local_jdatadict: The local data refer to the current class