deepmd.pt.model.descriptor.dpa2 =============================== .. py:module:: deepmd.pt.model.descriptor.dpa2 Classes ------- .. autoapisummary:: deepmd.pt.model.descriptor.dpa2.DescrptDPA2 Module Contents --------------- .. py:class:: DescrptDPA2(ntypes: int, repinit: Union[deepmd.dpmodel.descriptor.dpa2.RepinitArgs, dict], repformer: Union[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: Optional[Union[int, list[int]]] = None, add_tebd_to_repinit_out: bool = False, use_econf_tebd: bool = False, use_tebd_bias: bool = False, type_map: Optional[list[str]] = None) Bases: :py:obj:`deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor`, :py:obj:`torch.nn.Module` Base descriptor provides the interfaces of descriptor. .. !! processed by numpydoc !! .. py:attribute:: repinit_args .. py:attribute:: repformer_args .. py:attribute:: tebd_input_mode :value: 'concat' .. py:attribute:: repinit .. py:attribute:: use_three_body :value: False .. py:attribute:: repformers .. py:attribute:: rcsl_list .. py:attribute:: rcut_list .. py:attribute:: nsel_list .. py:attribute:: use_econf_tebd :value: False .. py:attribute:: use_tebd_bias :value: False .. py:attribute:: type_map :value: None .. py:attribute:: type_embedding .. py:attribute:: concat_output_tebd :value: True .. py:attribute:: precision :value: 'float64' .. py:attribute:: prec .. py:attribute:: smooth :value: True .. py:attribute:: exclude_types :value: [] .. py:attribute:: env_protection :value: 0.0 .. py:attribute:: trainable :value: True .. py:attribute:: add_tebd_to_repinit_out :value: False .. py:attribute:: repinit_out_dim .. py:attribute:: tebd_transform :value: None .. py:attribute:: tebd_dim :value: 8 .. py:attribute:: rcut .. py:attribute:: rcut_smth .. py:attribute:: ntypes .. py:attribute:: sel .. py:attribute:: compress :value: False .. py:method:: get_rcut() -> float Returns the cut-off radius. .. !! processed by numpydoc !! .. py:method:: get_rcut_smth() -> float Returns the radius where the neighbor information starts to smoothly decay to 0. .. !! processed by numpydoc !! .. py:method:: get_nsel() -> int Returns the number of selected atoms in the cut-off radius. .. !! processed by numpydoc !! .. py:method:: get_sel() -> list[int] Returns the number of selected atoms for each type. .. !! processed by numpydoc !! .. py:method:: get_ntypes() -> int Returns the number of element types. .. !! processed by numpydoc !! .. py:method:: get_type_map() -> list[str] Get the name to each type of atoms. .. !! processed by numpydoc !! .. py:method:: get_dim_out() -> int Returns the output dimension of this descriptor. .. !! processed by numpydoc !! .. py:method:: get_dim_emb() -> int Returns the embedding dimension of this descriptor. .. !! processed by numpydoc !! .. py:method:: mixed_types() -> bool If true, the descriptor 1. assumes total number of atoms aligned across frames; 2. requires a neighbor list that does not distinguish different atomic types. If false, the descriptor 1. assumes total number of atoms of each atom type aligned across frames; 2. requires a neighbor list that distinguishes different atomic types. .. !! processed by numpydoc !! .. py:method:: has_message_passing() -> bool Returns whether the descriptor has message passing. .. !! processed by numpydoc !! .. py:method:: need_sorted_nlist_for_lower() -> bool Returns whether the descriptor needs sorted nlist when using `forward_lower`. .. !! processed by numpydoc !! .. py:method:: get_env_protection() -> float Returns the protection of building environment matrix. .. !! processed by numpydoc !! .. py:method:: share_params(base_class: Any, shared_level: int, resume: bool = False) -> None Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes. .. !! processed by numpydoc !! .. py:method:: change_type_map(type_map: list[str], model_with_new_type_stat: Optional[Any] = None) -> None Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. .. !! processed by numpydoc !! .. py:property:: dim_out :type: int .. py:property:: dim_emb :type: int Returns the embedding dimension g2. .. !! processed by numpydoc !! .. py:method:: compute_input_stats(merged: Union[Callable[[], list[dict]], list[dict]], path: Optional[deepmd.utils.path.DPPath] = None) -> None Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. :Parameters: **merged** : :obj:`Union `\[:obj:`Callable `\[[], :class:`python:list`\[:class:`python:dict`]], :class:`python:list`\[:class:`python: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** : :obj:`Optional `\[:obj:`DPPath`] The path to the stat file. .. !! processed by numpydoc !! .. py:method:: set_stat_mean_and_stddev(mean: list[torch.Tensor], stddev: list[torch.Tensor]) -> None Update mean and stddev for descriptor. .. !! processed by numpydoc !! .. py:method:: get_stat_mean_and_stddev() -> tuple[list[torch.Tensor], list[torch.Tensor]] Get mean and stddev for descriptor. .. !! processed by numpydoc !! .. py:method:: serialize() -> dict Serialize the obj to dict. .. !! processed by numpydoc !! .. py:method:: deserialize(data: dict) -> DescrptDPA2 :classmethod: Deserialize the model. :Parameters: **data** : :class:`python:dict` The serialized data :Returns: :obj:`BD` The deserialized descriptor .. !! processed by numpydoc !! .. py:method:: forward(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: Optional[torch.Tensor] = None, comm_dict: Optional[dict[str, torch.Tensor]] = None) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]] 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: :obj:`descriptor` The descriptor. shape: nf x nloc x (ng x axis_neuron) :obj:`gr` The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3 :obj:`g2` The rotationally invariant pair-partical representation. shape: nf x nloc x nnei x ng :obj:`h2` The rotationally equivariant pair-partical representation. shape: nf x nloc x nnei x 3 :obj:`sw` The smooth switch function. shape: nf x nloc x nnei .. !! processed by numpydoc !! .. py:method:: update_sel(train_data: deepmd.utils.data_system.DeepmdDataSystem, type_map: Optional[list[str]], local_jdata: dict) -> tuple[dict, Optional[float]] :classmethod: Update the selection and perform neighbor statistics. :Parameters: **train_data** : :obj:`DeepmdDataSystem` data used to do neighbor statistics **type_map** : :class:`python:list`\[:class:`python:str`], :obj:`optional` The name of each type of atoms **local_jdata** : :class:`python:dict` The local data refer to the current class :Returns: :class:`python:dict` The updated local data :class:`python:float` The minimum distance between two atoms .. !! processed by numpydoc !! .. py:method:: enable_compression(min_nbor_dist: float, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1) -> None Receive the statistics (distance, max_nbor_size and env_mat_range) of the training data. :Parameters: **min_nbor_dist** The nearest distance between atoms **table_extrapolate** The scale of model extrapolation **table_stride_1** The uniform stride of the first table **table_stride_2** The uniform stride of the second table **check_frequency** The overflow check frequency .. !! processed by numpydoc !!