deepmd.pt.model.descriptor.dpa2
Module Contents
Classes
Base descriptor provides the interfaces of descriptor. |
- class deepmd.pt.model.descriptor.dpa2.DescrptDPA2(ntypes: int, repinit: deepmd.dpmodel.descriptor.dpa2.RepinitArgs | dict, repformer: deepmd.dpmodel.descriptor.dpa2.RepformerArgs | dict, concat_output_tebd: bool = True, precision: str = 'float64', smooth: bool = True, exclude_types: List[Tuple[int, int]] = [], env_protection: float = 0.0, trainable: bool = True, seed: int | None = None, add_tebd_to_repinit_out: bool = False, use_econf_tebd: bool = False, type_map: List[str] | None = None, old_impl: bool = False)[source]
Bases:
deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor
,torch.nn.Module
Base descriptor provides the interfaces of descriptor.
- get_rcut_smth() float [source]
Returns the radius where the neighbor information starts to smoothly decay to 0.
- 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 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)[source]
Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data.
- Parameters:
- merged
Union
[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.
- path
Optional
[DPPath
] The path to the stat file.
- merged
- classmethod deserialize(data: dict) DescrptDPA2 [source]
Deserialize the model.
- Parameters:
- data
dict
The serialized data
- data
- Returns:
BD
The deserialized descriptor
- forward(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, comm_dict: Dict[str, torch.Tensor] | None = None)[source]
Compute the descriptor.
- Parameters:
- extended_coord
The extended coordinates of atoms. shape: nf x (nallx3)
- extended_atype
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.
- 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
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