deepmd.pt.model.descriptor.hybrid

Module Contents

Classes

DescrptHybrid	Concate a list of descriptors to form a new descriptor.
DescrptBlockHybrid	The building block of descriptor.

class deepmd.pt.model.descriptor.hybrid.DescrptHybrid(list: List[deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor | Dict[str, Any]], **kwargs)

Bases: deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor, torch.nn.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.

get_rcut() → float

Returns the cut-off radius.

get_sel() → List[int]

Returns the number of selected atoms for each type.

get_ntypes() → int

Returns the number of element types.

get_dim_out() → int

Returns the output dimension.

get_dim_emb() → int

Returns the output dimension.

mixed_types()

Returns if the descriptor requires a neighbor list that distinguish different atomic types or not.

share_params(base_class, shared_level, resume=False)

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.

compute_input_stats(merged: List[dict], path: deepmd.utils.path.DPPath | None = None)

Update mean and stddev for descriptor elements.

forward(coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, comm_dict: Dict[str, torch.Tensor] | None = None)

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.

comm_dict

The data needed for communication for parallel inference.

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

classmethod update_sel(global_jdata: dict, local_jdata: dict) → dict

Update the selection and perform neighbor statistics.

Parameters:

global_jdatadict

The global data, containing the training section

local_jdatadict

The local data refer to the current class

serialize() → dict

Serialize the obj to dict.

classmethod deserialize(data: dict) → DescrptHybrid

Deserialize the model.

Parameters:

datadict

The serialized data

Returns:

BD

The deserialized descriptor

class deepmd.pt.model.descriptor.hybrid.DescrptBlockHybrid(list, ntypes: int, tebd_dim: int = 8, tebd_input_mode: str = 'concat', hybrid_mode: str = 'concat', **kwargs)

Bases: deepmd.pt.model.descriptor.DescriptorBlock

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

property dim_out

Returns the output dimension of this descriptor.

property dim_emb_list: List[int]

Returns the output dimension list of embeddings.

property dim_emb

Returns the output dimension of embedding.

get_rcut() → float

Returns the cut-off radius.

get_nsel() → int

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

get_sel() → List[int]

Returns the number of selected atoms for each type.

get_ntypes() → int

Returns the number of element types.

get_dim_out() → int

Returns the output dimension.

get_dim_in() → int

Returns the input dimension.

get_dim_emb()

Returns the embedding dimension.

mixed_types() → bool

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)

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.

compute_input_stats(merged: Callable[[], List[dict]] | List[dict], path: deepmd.utils.path.DPPath | None = None)

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.

forward(nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, extended_atype_embd: torch.Tensor | None = None, mapping: torch.Tensor | None = None, comm_dict: Dict[str, torch.Tensor] | None = None)

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]