deepmd_gnn.nequip ================= .. py:module:: deepmd_gnn.nequip .. autoapi-nested-parse:: Nequip model. Classes ------- .. autoapisummary:: deepmd_gnn.nequip.NequipModel Functions --------- .. autoapisummary:: deepmd_gnn.nequip._load_observed_type_stat_compat Module Contents --------------- .. py:function:: _load_observed_type_stat_compat() -> tuple[Any, Any, Any] .. py:class:: NequipModel(type_map: list[str], sel: int, r_max: float = 6.0, num_layers: int = 4, l_max: int = 2, num_features: int = 32, nonlinearity_type: str = 'gate', parity: bool = True, num_basis: int = 8, BesselBasis_trainable: bool = True, PolynomialCutoff_p: int = 6, invariant_layers: int = 2, invariant_neurons: int = 64, use_sc: bool = True, irreps_edge_sh: str = '0e + 1e', feature_irreps_hidden: str = '32x0o + 32x0e + 32x1o + 32x1e', chemical_embedding_irreps_out: str = '32x0e', conv_to_output_hidden_irreps_out: str = '16x0e', precision: str = 'float32', **kwargs: Any) Bases: :py:obj:`deepmd.pt.model.model.model.BaseModel` Nequip model. Parameters ---------- type_map : list[str] The name of each type of atoms sel : int Maximum number of neighbor atoms r_max : float, optional distance cutoff (in Ang) num_layers : int number of interaction blocks, we find 3-5 to work best l_max : int the maximum irrep order (rotation order) for the network's features, l=1 is a good default, l=2 is more accurate but slower num_features : int the multiplicity of the features, 32 is a good default for accurate network, if you want to be more accurate, go larger, if you want to be faster, go lower nonlinearity_type : str may be 'gate' or 'norm', 'gate' is recommended parity : bool whether to include features with odd mirror parityy; often turning parity off gives equally good results but faster networks, so do consider this num_basis : int number of basis functions used in the radial basis, 8 usually works best BesselBasis_trainable : bool set true to train the bessel weights PolynomialCutoff_p : int p-exponent used in polynomial cutoff function, smaller p corresponds to stronger decay with distance invariant_layers : int number of radial layers, usually 1-3 works best, smaller is faster invariant_neurons : int number of hidden neurons in radial function, smaller is faster use_sc : bool use self-connection or not, usually gives big improvement irreps_edge_sh : str irreps for the chemical embedding of species feature_irreps_hidden : str irreps used for hidden features, here we go up to lmax=1, with even and odd parities; for more accurate but slower networks, use l=2 or higher, smaller number of features is faster chemical_embedding_irreps_out : str irreps of the spherical harmonics used for edges. If a single integer, indicates the full SH up to L_max=that_integer conv_to_output_hidden_irreps_out : str irreps used in hidden layer of output block .. py:attribute:: mm_types :type: list[int] .. py:attribute:: e0 :type: torch.Tensor .. py:attribute:: _observed_type :type: Optional[list[str]] .. py:attribute:: params .. py:attribute:: type_map .. py:attribute:: ntypes .. py:attribute:: preset_out_bias :type: dict[str, list] .. py:attribute:: sel .. py:attribute:: num_layers :value: 4 .. py:attribute:: rcut :value: 6.0 .. py:attribute:: model .. py:property:: atomic_model :type: Any Provide a compatibility view matching wrapped deepmd-kit models. .. py:property:: observed_type :type: Optional[list[str]] Observed element types collected during statistics. .. py:method:: get_observed_type_list() -> list[str] Get observed element types collected during statistics. .. py:method:: compute_or_load_stat(sampled_func, stat_file_path: Optional[deepmd.utils.path.DPPath] = None, preset_observed_type: Optional[list[str]] = None) -> None Compute or load the statistics parameters of the model. For example, mean and standard deviation of descriptors or the energy bias of the fitting net. When `sampled` is provided, all the statistics parameters will be calculated (or re-calculated for update), and saved in the `stat_file_path`(s). When `sampled` is not provided, it will check the existence of `stat_file_path`(s) and load the calculated statistics parameters. Parameters ---------- sampled_func The sampled data frames from different data systems. stat_file_path The path to the statistics files. preset_observed_type Optional observed element types to seed or override ``self._observed_type``. This compatibility parameter is accepted for newer deepmd-kit versions; when provided, it is used directly instead of restoring or collecting observed types from statistics data. .. py:method:: fitting_output_def() -> deepmd.dpmodel.output_def.FittingOutputDef Get the output def of developer implemented atomic models. .. py:method:: get_rcut() -> float Get the cut-off radius. .. py:method:: get_type_map() -> list[str] Get the type map. .. py:method:: get_sel() -> list[int] Return the number of selected atoms for each type. .. py:method:: get_dim_fparam() -> int Get the number (dimension) of frame parameters of this atomic model. .. py:method:: get_dim_aparam() -> int Get the number (dimension) of atomic parameters of this atomic model. .. py:method:: get_sel_type() -> list[int] Get the selected atom types of this model. Only atoms with selected atom types have atomic contribution to the result of the model. If returning an empty list, all atom types are selected. .. py:method:: is_aparam_nall() -> bool Check whether the shape of atomic parameters is (nframes, nall, ndim). If False, the shape is (nframes, nloc, ndim). .. py:method:: mixed_types() -> bool Return whether the model is in mixed-types mode. If true, the model 1. assumes total number of atoms aligned across frames; 2. uses a neighbor list that does not distinguish different atomic types. If false, the model 1. assumes total number of atoms of each atom type aligned across frames; 2. uses a neighbor list that distinguishes different atomic types. .. py:method:: has_message_passing() -> bool Return whether the descriptor has message passing. .. py:method:: forward(coord: torch.Tensor, atype: torch.Tensor, box: Optional[torch.Tensor] = None, fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, do_atomic_virial: bool = False) -> dict[str, torch.Tensor] Forward pass of the model. Parameters ---------- coord : torch.Tensor The coordinates of atoms. atype : torch.Tensor The atomic types of atoms. box : torch.Tensor, optional The box tensor. fparam : torch.Tensor, optional The frame parameters. aparam : torch.Tensor, optional The atomic parameters. do_atomic_virial : bool, optional Whether to compute atomic virial. .. py:method:: forward_lower(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: Optional[torch.Tensor] = None, fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, do_atomic_virial: bool = False, comm_dict: Optional[dict[str, torch.Tensor]] = None) -> dict[str, torch.Tensor] Forward lower pass of the model. Parameters ---------- extended_coord : torch.Tensor The extended coordinates of atoms. extended_atype : torch.Tensor The extended atomic types of atoms. nlist : torch.Tensor The neighbor list. mapping : torch.Tensor, optional The mapping tensor. fparam : torch.Tensor, optional The frame parameters. aparam : torch.Tensor, optional The atomic parameters. do_atomic_virial : bool, optional Whether to compute atomic virial. comm_dict : dict[str, torch.Tensor], optional The communication dictionary. .. py:method:: forward_lower_common(nloc: int, extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: Optional[torch.Tensor] = None, fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, do_atomic_virial: bool = False, comm_dict: Optional[dict[str, torch.Tensor]] = None, box: Optional[torch.Tensor] = None) -> dict[str, torch.Tensor] Forward lower common pass of the model. Parameters ---------- extended_coord : torch.Tensor The extended coordinates of atoms. extended_atype : torch.Tensor The extended atomic types of atoms. nlist : torch.Tensor The neighbor list. mapping : torch.Tensor, optional The mapping tensor. fparam : torch.Tensor, optional The frame parameters. aparam : torch.Tensor, optional The atomic parameters. do_atomic_virial : bool, optional Whether to compute atomic virial. comm_dict : dict[str, torch.Tensor], optional The communication dictionary. box : torch.Tensor, optional The box tensor. .. py:method:: serialize() -> dict Serialize the model. .. py:method:: deserialize(data: dict) -> NequipModel :classmethod: Deserialize the model. .. py:method:: get_nnei() -> int Return the total number of selected neighboring atoms in cut-off radius. .. py:method:: get_nsel() -> int Return the total number of selected neighboring atoms in cut-off radius. .. 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 : DeepmdDataSystem data used to do neighbor statictics type_map : list[str], optional The name of each type of atoms local_jdata : dict The local data refer to the current class Returns ------- dict The updated local data float The minimum distance between two atoms .. py:method:: model_output_type() -> list[str] Get the output type for the model. .. py:method:: translated_output_def() -> dict[str, Any] Get the translated output def for the model. .. py:method:: model_output_def() -> deepmd.dpmodel.output_def.ModelOutputDef Get the output def for the model.