deepmd.dpmodel.descriptor.dpa2 ============================== .. py:module:: deepmd.dpmodel.descriptor.dpa2 Classes ------- .. autoapisummary:: deepmd.dpmodel.descriptor.dpa2.RepinitArgs deepmd.dpmodel.descriptor.dpa2.RepformerArgs deepmd.dpmodel.descriptor.dpa2.DescrptDPA2 Module Contents --------------- .. py:class:: RepinitArgs(rcut: float, rcut_smth: float, nsel: int, neuron: list[int] = [25, 50, 100], axis_neuron: int = 16, tebd_dim: int = 8, tebd_input_mode: str = 'concat', set_davg_zero: bool = True, activation_function: str = 'tanh', resnet_dt: bool = False, type_one_side: bool = False, use_three_body: bool = False, three_body_neuron: list[int] = [2, 4, 8], three_body_sel: int = 40, three_body_rcut: float = 4.0, three_body_rcut_smth: float = 0.5) .. py:attribute:: rcut .. py:attribute:: rcut_smth .. py:attribute:: nsel .. py:attribute:: neuron :value: [25, 50, 100] .. py:attribute:: axis_neuron :value: 16 .. py:attribute:: tebd_dim :value: 8 .. py:attribute:: tebd_input_mode :value: 'concat' .. py:attribute:: set_davg_zero :value: True .. py:attribute:: activation_function :value: 'tanh' .. py:attribute:: resnet_dt :value: False .. py:attribute:: type_one_side :value: False .. py:attribute:: use_three_body :value: False .. py:attribute:: three_body_neuron :value: [2, 4, 8] .. py:attribute:: three_body_sel :value: 40 .. py:attribute:: three_body_rcut :value: 4.0 .. py:attribute:: three_body_rcut_smth :value: 0.5 .. py:method:: __getitem__(key: str) -> Any .. py:method:: serialize() -> dict .. py:method:: deserialize(data: dict) -> RepinitArgs :classmethod: .. py:class:: RepformerArgs(rcut: float, rcut_smth: float, nsel: int, nlayers: int = 3, g1_dim: int = 128, g2_dim: int = 16, axis_neuron: int = 4, direct_dist: bool = False, 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', update_residual: float = 0.001, update_residual_init: str = 'norm', set_davg_zero: bool = True, trainable_ln: bool = True, use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True, ln_eps: float | None = 1e-05) .. py:attribute:: rcut .. py:attribute:: rcut_smth .. py:attribute:: nsel .. py:attribute:: nlayers :value: 3 .. py:attribute:: g1_dim :value: 128 .. py:attribute:: g2_dim :value: 16 .. py:attribute:: axis_neuron :value: 4 .. py:attribute:: direct_dist :value: False .. py:attribute:: update_g1_has_conv :value: True .. py:attribute:: update_g1_has_drrd :value: True .. py:attribute:: update_g1_has_grrg :value: True .. py:attribute:: update_g1_has_attn :value: True .. py:attribute:: update_g2_has_g1g1 :value: True .. py:attribute:: update_g2_has_attn :value: True .. py:attribute:: update_h2 :value: False .. py:attribute:: attn1_hidden :value: 64 .. py:attribute:: attn1_nhead :value: 4 .. py:attribute:: attn2_hidden :value: 16 .. py:attribute:: attn2_nhead :value: 4 .. py:attribute:: attn2_has_gate :value: False .. py:attribute:: activation_function :value: 'tanh' .. py:attribute:: update_style :value: 'res_avg' .. py:attribute:: update_residual :value: 0.001 .. py:attribute:: update_residual_init :value: 'norm' .. py:attribute:: set_davg_zero :value: True .. py:attribute:: trainable_ln :value: True .. py:attribute:: use_sqrt_nnei :value: True .. py:attribute:: g1_out_conv :value: True .. py:attribute:: g1_out_mlp :value: True .. py:attribute:: ln_eps :value: 1e-05 .. py:method:: __getitem__(key: str) -> Any .. py:method:: serialize() -> dict .. py:method:: deserialize(data: dict) -> RepformerArgs :classmethod: .. py:class:: DescrptDPA2(ntypes: int, repinit: RepinitArgs | dict, repformer: 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 | list[int] | None = None, add_tebd_to_repinit_out: bool = False, use_econf_tebd: bool = False, use_tebd_bias: bool = False, type_map: list[str] | None = None) Bases: :py:obj:`deepmd.dpmodel.NativeOP`, :py:obj:`deepmd.dpmodel.descriptor.base_descriptor.BaseDescriptor` The DPA-2 descriptor[Rf0ed3afb961b-1]_. The DPA-2 descriptor combines a repinit block and a repformer block to extract atomic representations. The overall descriptor is computed as: .. math:: \mathcal{D}^i = \mathrm{Repformer}(\mathrm{Linear}(\mathrm{Repinit}(\mathcal{R}^i, \mathcal{T}^i))), where :math:`\mathcal{R}^i` is the environment matrix and :math:`\mathcal{T}^i` is the type embedding. The repinit block computes initial node and edge representations using attention-based message passing. The repformer block further refines these representations through multiple layers of graph convolution and attention mechanisms. The final output dimension is: .. math:: \dim(\mathcal{D}^i) = \text{g1\_dim} + \text{tebd\_dim} \quad (\text{if concat\_output\_tebd}). :Parameters: **repinit** : :obj:`Union `\[:obj:`RepinitArgs`, :class:`python:dict`] The arguments used to initialize the repinit block, see docstr in `RepinitArgs` for details information. **repformer** : :obj:`Union `\[:obj:`RepformerArgs`, :class:`python:dict`] The arguments used to initialize the repformer block, see docstr in `RepformerArgs` for details information. **concat_output_tebd** : :ref:`bool `, :obj:`optional` Whether to concat type embedding at the output of the descriptor. **precision** : :class:`python:str`, :obj:`optional` The precision of the embedding net parameters. **smooth** : :ref:`bool `, :obj:`optional` Whether to use smoothness in processes such as attention weights calculation. **exclude_types** : :class:`python:list`\[:class:`python:list`\[:class:`python:int`]], :obj:`optional` The excluded pairs of types which have no interaction with each other. For example, `[[0, 1]]` means no interaction between type 0 and type 1. **env_protection** : :class:`python:float`, :obj:`optional` Protection parameter to prevent division by zero errors during environment matrix calculations. For example, when using paddings, there may be zero distances of neighbors, which may make division by zero error during environment matrix calculations without protection. **trainable** : :ref:`bool `, :obj:`optional` If the parameters are trainable. **seed** : :class:`python:int`, :obj:`optional` (Unused yet) Random seed for parameter initialization. **add_tebd_to_repinit_out** : :ref:`bool `, :obj:`optional` Whether to add type embedding to the output representation from repinit before inputting it into repformer. **use_econf_tebd** : :ref:`bool `, :obj:`Optional ` Whether to use electronic configuration type embedding. **use_tebd_bias** : :ref:`bool `, :obj:`Optional ` Whether to use bias in the type embedding layer. **type_map** : :class:`python:list`\[:class:`python:str`], :obj:`Optional ` A list of strings. Give the name to each type of atoms. :Returns: descriptor: :obj:`torch.Tensor` the descriptor of shape nf x nloc x g1_dim. invariant single-atom representation. g2: :obj:`torch.Tensor` invariant pair-atom representation. h2: :obj:`torch.Tensor` equivariant pair-atom representation. rot_mat: :obj:`torch.Tensor` rotation matrix for equivariant fittings sw: :obj:`torch.Tensor` The switch function for decaying inverse distance. .. rubric:: References .. [Rf0ed3afb961b-1] Zhang, D., Liu, X., Zhang, X. et al. DPA-2: a large atomic model as a multi-task learner. npj Comput Mater 10, 293 (2024). https://doi.org/10.1038/s41524-024-01493-2 .. only:: latex [Rf0ed3afb961b-1]_ .. !! processed by numpydoc !! .. py:attribute:: _update_sel_cls .. py:attribute:: repinit_args .. py:attribute:: repformer_args .. 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:: smooth :value: True .. py:attribute:: exclude_types :value: [] .. py:attribute:: env_protection :value: 0.0 .. py:attribute:: rcut_smth .. 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:: ntypes .. py:attribute:: sel .. 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) -> NoReturn :abstractmethod: 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: 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: collections.abc.Callable[[], list[dict]] | list[dict], path: deepmd.utils.path.DPPath | None = 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[deepmd.dpmodel.array_api.Array], stddev: list[deepmd.dpmodel.array_api.Array]) -> None Update mean and stddev for descriptor. .. !! processed by numpydoc !! .. py:method:: get_stat_mean_and_stddev() -> tuple[list[deepmd.dpmodel.array_api.Array], list[deepmd.dpmodel.array_api.Array]] Get mean and stddev for descriptor. .. !! processed by numpydoc !! .. py:method:: call(coord_ext: deepmd.dpmodel.array_api.Array, atype_ext: deepmd.dpmodel.array_api.Array, nlist: deepmd.dpmodel.array_api.Array, mapping: deepmd.dpmodel.array_api.Array | None = None, fparam: deepmd.dpmodel.array_api.Array | None = None) -> tuple[deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array] 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, maps extended region index to local region. :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:: 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:: update_sel(train_data: deepmd.utils.data_system.DeepmdDataSystem, type_map: list[str] | None, local_jdata: dict) -> tuple[deepmd.dpmodel.array_api.Array, deepmd.dpmodel.array_api.Array] :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 !!