deepmd_gnn.mace =============== .. py:module:: deepmd_gnn.mace .. autoapi-nested-parse:: Wrapper for MACE models. Attributes ---------- .. autoapisummary:: deepmd_gnn.mace.ELEMENTS deepmd_gnn.mace.PeriodicTable Classes ------- .. autoapisummary:: deepmd_gnn.mace.MaceModel Functions --------- .. autoapisummary:: deepmd_gnn.mace._load_observed_type_stat_compat Module Contents --------------- .. py:function:: _load_observed_type_stat_compat() -> tuple[Any, Any, Any] .. py:data:: ELEMENTS :value: ['H', 'He', 'Li', 'Be', 'B', 'C', 'N', 'O', 'F', 'Ne', 'Na', 'Mg', 'Al', 'Si', 'P', 'S', 'Cl',... .. py:data:: PeriodicTable .. py:class:: MaceModel(type_map: list[str], sel: int, r_max: float = 5.0, num_radial_basis: int = 8, num_cutoff_basis: int = 5, max_ell: int = 3, interaction: str = 'RealAgnosticResidualInteractionBlock', num_interactions: int = 2, hidden_irreps: str = '128x0e + 128x1o', pair_repulsion: bool = False, distance_transform: str = 'None', correlation: int = 3, gate: str = 'silu', MLP_irreps: str = '16x0e', radial_type: str = 'bessel', radial_MLP: list[int] = [64, 64, 64], std: float = 1, avg_num_neighbors: float | None = None, **kwargs: Any) Bases: :py:obj:`deepmd.pt.model.model.model.BaseModel` Mace 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_radial_basis : int, optional number of radial basis functions num_cutoff_basis : int, optional number of basis functions for smooth cutoff max_ell : int, optional highest ell of spherical harmonics interaction : str, optional name of interaction block num_interactions : int, optional number of interactions hidden_irreps : str, optional hidden irreps pair_repulsion : bool use amsgrad variant of optimizer distance_transform : str, optional distance transform correlation : int correlation order at each layer gate : str, optional non linearity for last readout MLP_irreps : str, optional hidden irreps of the MLP in last readout radial_type : str, optional type of radial basis functions radial_MLP : str, optional width of the radial MLP std : float, optional Standard deviation of force components in the training set .. py:attribute:: mm_types :type: list[int] .. py:attribute:: _observed_type :type: Optional[list[str]] .. py:attribute:: params .. py:attribute:: type_map .. py:attribute:: ntypes .. py:attribute:: rcut :value: 5.0 .. py:attribute:: num_interactions :value: 2 .. py:attribute:: preset_out_bias :type: dict[str, list] .. py:attribute:: sel .. py:attribute:: avg_num_neighbors .. py:attribute:: model .. py:attribute:: atomic_numbers :value: [] .. 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) -> 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. .. py:method:: serialize() -> dict Serialize the model. .. py:method:: deserialize(data: dict) -> MaceModel :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. .. py:method:: get_model(model_params: dict) -> MaceModel :classmethod: Get the model by the parameters. Parameters ---------- model_params : dict The model parameters Returns ------- BaseBaseModel The model