deepmd.pt.model.descriptor.descriptor

Module Contents

Classes

DescriptorBlock

The building block of descriptor.

Functions

make_default_type_embedding(ntypes)

Attributes

log

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

Bases: torch.nn.Module, abc.ABC, make_plugin_registry('DescriptorBlock')

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

local_cluster = False[source]
abstract get_rcut() → float[source]

Returns the cut-off radius.

abstract get_nsel() → int[source]

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

abstract get_sel() → List[int][source]

Returns the number of selected atoms for each type.

abstract get_ntypes() → int[source]

Returns the number of element types.

abstract get_dim_out() → int[source]

Returns the output dimension.

abstract get_dim_in() → int[source]

Returns the output dimension.

abstract get_dim_emb() → int[source]

Returns the embedding dimension.

abstract compute_input_stats(merged: Callable[[], List[dict]] | List[dict], path: deepmd.utils.path.DPPath | None = 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 get_stats() → Dict[str, deepmd.utils.env_mat_stat.StatItem][source]

Get the statistics of the descriptor.

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.

abstract forward(nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, extended_atype_embd: torch.Tensor | None = None, mapping: torch.Tensor | None = None)[source]

Calculate DescriptorBlock.

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