deepmd.pt.model.atomic_model ============================ .. py:module:: deepmd.pt.model.atomic_model .. autoapi-nested-parse:: The atomic model provides the prediction of some property on each atom. All the atomic models are not supposed to be directly accessed by users, but it provides a convenient interface for the implementation of models. Taking the energy models for example, the developeres only needs to implement the atomic energy prediction via an atomic model, and the model can be automatically made by the `deepmd.dpmodel.make_model` method. The `DPModel` is made by ``` DPModel = make_model(DPAtomicModel) ``` .. !! processed by numpydoc !! Submodules ---------- .. toctree:: :maxdepth: 1 /autoapi/deepmd/pt/model/atomic_model/base_atomic_model/index /autoapi/deepmd/pt/model/atomic_model/dipole_atomic_model/index /autoapi/deepmd/pt/model/atomic_model/dos_atomic_model/index /autoapi/deepmd/pt/model/atomic_model/dp_atomic_model/index /autoapi/deepmd/pt/model/atomic_model/energy_atomic_model/index /autoapi/deepmd/pt/model/atomic_model/linear_atomic_model/index /autoapi/deepmd/pt/model/atomic_model/pairtab_atomic_model/index /autoapi/deepmd/pt/model/atomic_model/polar_atomic_model/index /autoapi/deepmd/pt/model/atomic_model/property_atomic_model/index Classes ------- .. autoapisummary:: deepmd.pt.model.atomic_model.BaseAtomicModel deepmd.pt.model.atomic_model.DPDipoleAtomicModel deepmd.pt.model.atomic_model.DPDOSAtomicModel deepmd.pt.model.atomic_model.DPAtomicModel deepmd.pt.model.atomic_model.DPEnergyAtomicModel deepmd.pt.model.atomic_model.DPZBLLinearEnergyAtomicModel deepmd.pt.model.atomic_model.LinearEnergyAtomicModel deepmd.pt.model.atomic_model.PairTabAtomicModel deepmd.pt.model.atomic_model.DPPolarAtomicModel deepmd.pt.model.atomic_model.DPPropertyAtomicModel Package Contents ---------------- .. py:class:: BaseAtomicModel(type_map: list[str], atom_exclude_types: list[int] = [], pair_exclude_types: list[tuple[int, int]] = [], rcond: Optional[float] = None, preset_out_bias: Optional[dict[str, numpy.ndarray]] = None) Bases: :py:obj:`torch.nn.Module`, :py:obj:`BaseAtomicModel_` The base of atomic model. :Parameters: **type_map** Mapping atom type to the name (str) of the type. For example `type_map[1]` gives the name of the type 1. **atom_exclude_types** Exclude the atomic contribution of the given types **pair_exclude_types** Exclude the pair of atoms of the given types from computing the output of the atomic model. Implemented by removing the pairs from the nlist. **rcond** : :class:`python:float`, :obj:`optional` The condition number for the regression of atomic energy. **preset_out_bias** : :class:`python:dict`\[:class:`python:str`, :class:`python:list`\[:obj:`Optional `\[:obj:`np.ndarray `]]], :obj:`optional` Specifying atomic energy contribution in vacuum. Given by key:value pairs. The value is a list specifying the bias. the elements can be None or np.ndarray of output shape. For example: [None, [2.]] means type 0 is not set, type 1 is set to [2.] The `set_davg_zero` key in the descriptor should be set. .. !! processed by numpydoc !! .. py:attribute:: type_map .. py:attribute:: rcond .. py:attribute:: preset_out_bias .. py:method:: init_out_stat() -> None Initialize the output bias. .. !! processed by numpydoc !! .. py:method:: set_out_bias(out_bias: torch.Tensor) -> None .. py:method:: __setitem__(key, value) -> None .. py:method:: __getitem__(key) .. py:method:: get_type_map() -> list[str] Get the type map. .. !! processed by numpydoc !! .. py:method:: reinit_atom_exclude(exclude_types: list[int] = []) -> None .. py:method:: reinit_pair_exclude(exclude_types: list[tuple[int, int]] = []) -> None .. py:method:: make_atom_mask(atype: torch.Tensor) -> torch.Tensor The atoms with type < 0 are treated as virtual atoms, which serves as place-holders for multi-frame calculations with different number of atoms in different frames. :Parameters: **atype** Atom types. >= 0 for real atoms <0 for virtual atoms. :Returns: :obj:`mask` True for real atoms and False for virtual atoms. .. !! processed by numpydoc !! .. py:method:: atomic_output_def() -> deepmd.dpmodel.output_def.FittingOutputDef Get the output def of the atomic model. By default it is the same as FittingOutputDef, but it allows model level wrapper of the output defined by the developer. .. !! processed by numpydoc !! .. py:method:: forward_common_atomic(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, comm_dict: Optional[dict[str, torch.Tensor]] = None) -> dict[str, torch.Tensor] Common interface for atomic inference. This method accept extended coordinates, extended atom typs, neighbor list, and predict the atomic contribution of the fit property. :Parameters: **extended_coord** extended coordinates, shape: nf x (nall x 3) **extended_atype** extended atom typs, shape: nf x nall for a type < 0 indicating the atomic is virtual. **nlist** neighbor list, shape: nf x nloc x nsel **mapping** extended to local index mapping, shape: nf x nall **fparam** frame parameters, shape: nf x dim_fparam **aparam** atomic parameter, shape: nf x nloc x dim_aparam **comm_dict** The data needed for communication for parallel inference. :Returns: :obj:`ret_dict` dict of output atomic properties. should implement the definition of `fitting_output_def`. ret_dict["mask"] of shape nf x nloc will be provided. ret_dict["mask"][ff,ii] == 1 indicating the ii-th atom of the ff-th frame is real. ret_dict["mask"][ff,ii] == 0 indicating the ii-th atom of the ff-th frame is virtual. .. !! processed by numpydoc !! .. py:method:: forward(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, comm_dict: Optional[dict[str, torch.Tensor]] = None) -> dict[str, torch.Tensor] .. py:method:: change_type_map(type_map: list[str], model_with_new_type_stat=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:method:: serialize() -> dict .. py:method:: deserialize(data: dict) -> BaseAtomicModel :classmethod: .. py:method:: compute_or_load_stat(merged: Union[Callable[[], list[dict]], list[dict]], stat_file_path: Optional[deepmd.utils.path.DPPath] = None) -> NoReturn :abstractmethod: Compute the output statistics (e.g. energy bias) for the fitting net 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. **stat_file_path** : :obj:`Optional `\[:obj:`DPPath`] The path to the stat file. .. !! processed by numpydoc !! .. py:method:: compute_or_load_out_stat(merged: Union[Callable[[], list[dict]], list[dict]], stat_file_path: Optional[deepmd.utils.path.DPPath] = None) -> None Compute the output statistics (e.g. energy bias) for the fitting net 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. **stat_file_path** : :obj:`Optional `\[:obj:`DPPath`] The path to the stat file. .. !! processed by numpydoc !! .. py:method:: apply_out_stat(ret: dict[str, torch.Tensor], atype: torch.Tensor) Apply the stat to each atomic output. The developer may override the method to define how the bias is applied to the atomic output of the model. :Parameters: **ret** The returned dict by the forward_atomic method **atype** The atom types. nf x nloc .. !! processed by numpydoc !! .. py:method:: change_out_bias(sample_merged, stat_file_path: Optional[deepmd.utils.path.DPPath] = None, bias_adjust_mode='change-by-statistic') -> None Change the output bias according to the input data and the pretrained model. :Parameters: **sample_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. **bias_adjust_mode** : :class:`python:str` The mode for changing output bias : ['change-by-statistic', 'set-by-statistic'] 'change-by-statistic' : perform predictions on labels of target dataset, and do least square on the errors to obtain the target shift as bias. 'set-by-statistic' : directly use the statistic output bias in the target dataset. **stat_file_path** : :obj:`Optional `\[:obj:`DPPath`] The path to the stat file. .. !! processed by numpydoc !! .. py:method:: _get_forward_wrapper_func() -> Callable[Ellipsis, torch.Tensor] Get a forward wrapper of the atomic model for output bias calculation. .. !! processed by numpydoc !! .. py:method:: _default_bias() .. py:method:: _default_std() .. py:method:: _varsize(shape: list[int]) -> int .. py:method:: _get_bias_index(kk: str) -> int .. py:method:: _store_out_stat(out_bias: dict[str, torch.Tensor], out_std: dict[str, torch.Tensor], add: bool = False) -> None .. py:method:: _fetch_out_stat(keys: list[str]) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]] .. py:class:: DPDipoleAtomicModel(descriptor, fitting, type_map, **kwargs) Bases: :py:obj:`deepmd.pt.model.atomic_model.dp_atomic_model.DPAtomicModel` Model give atomic prediction of some physical property. :Parameters: **descriptor** Descriptor **fitting_net** Fitting net **type_map** Mapping atom type to the name (str) of the type. For example `type_map[1]` gives the name of the type 1. .. !! processed by numpydoc !! .. py:method:: apply_out_stat(ret: dict[str, torch.Tensor], atype: torch.Tensor) Apply the stat to each atomic output. The developer may override the method to define how the bias is applied to the atomic output of the model. :Parameters: **ret** The returned dict by the forward_atomic method **atype** The atom types. nf x nloc .. !! processed by numpydoc !! .. py:class:: DPDOSAtomicModel(descriptor, fitting, type_map, **kwargs) Bases: :py:obj:`deepmd.pt.model.atomic_model.dp_atomic_model.DPAtomicModel` Model give atomic prediction of some physical property. :Parameters: **descriptor** Descriptor **fitting_net** Fitting net **type_map** Mapping atom type to the name (str) of the type. For example `type_map[1]` gives the name of the type 1. .. !! processed by numpydoc !! .. py:class:: DPAtomicModel(descriptor, fitting, type_map: list[str], **kwargs) Bases: :py:obj:`deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel` Model give atomic prediction of some physical property. :Parameters: **descriptor** Descriptor **fitting_net** Fitting net **type_map** Mapping atom type to the name (str) of the type. For example `type_map[1]` gives the name of the type 1. .. !! processed by numpydoc !! .. py:attribute:: type_map .. py:attribute:: ntypes .. py:attribute:: descriptor .. py:attribute:: rcut .. py:attribute:: sel .. py:attribute:: fitting_net .. py:attribute:: enable_eval_descriptor_hook :value: False .. py:attribute:: eval_descriptor_list :value: [] .. py:method:: set_eval_descriptor_hook(enable: bool) -> None Set the hook for evaluating descriptor and clear the cache for descriptor list. .. !! processed by numpydoc !! .. py:method:: eval_descriptor() -> torch.Tensor Evaluate the descriptor. .. !! processed by numpydoc !! .. py:method:: fitting_output_def() -> deepmd.dpmodel.FittingOutputDef Get the output def of the fitting net. .. !! processed by numpydoc !! .. py:method:: get_rcut() -> float Get the cut-off radius. .. !! processed by numpydoc !! .. py:method:: get_sel() -> list[int] Get the neighbor selection. .. !! processed by numpydoc !! .. py:method:: mixed_types() -> bool 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. .. !! processed by numpydoc !! .. py:method:: change_type_map(type_map: list[str], model_with_new_type_stat=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:method:: has_message_passing() -> bool Returns whether the atomic model has message passing. .. !! processed by numpydoc !! .. py:method:: need_sorted_nlist_for_lower() -> bool Returns whether the atomic model needs sorted nlist when using `forward_lower`. .. !! processed by numpydoc !! .. py:method:: serialize() -> dict .. py:method:: deserialize(data) -> DPAtomicModel :classmethod: .. 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 Call descriptor enable_compression(). :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 !! .. py:method:: forward_atomic(extended_coord, extended_atype, nlist, mapping: Optional[torch.Tensor] = None, fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, comm_dict: Optional[dict[str, torch.Tensor]] = None) -> dict[str, torch.Tensor] Return atomic prediction. :Parameters: **extended_coord** coordinates in extended region **extended_atype** atomic type in extended region **nlist** neighbor list. nf x nloc x nsel **mapping** mapps the extended indices to local indices **fparam** frame parameter. nf x ndf **aparam** atomic parameter. nf x nloc x nda :Returns: :obj:`result_dict` the result dict, defined by the `FittingOutputDef`. .. !! processed by numpydoc !! .. py:method:: get_out_bias() -> torch.Tensor .. py:method:: compute_or_load_stat(sampled_func, stat_file_path: Optional[deepmd.utils.path.DPPath] = None) -> None Compute or load the statistics parameters of the model, such as 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 lazy sampled function to get data frames from different data systems. **stat_file_path** The dictionary of paths to the statistics files. .. !! processed by numpydoc !! .. py:method:: get_dim_fparam() -> int Get the number (dimension) of frame parameters of this atomic model. .. !! processed by numpydoc !! .. py:method:: get_dim_aparam() -> int Get the number (dimension) of atomic parameters of this atomic model. .. !! processed by numpydoc !! .. 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. .. !! processed by numpydoc !! .. 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). .. !! processed by numpydoc !! .. py:class:: DPEnergyAtomicModel(descriptor, fitting, type_map, **kwargs) Bases: :py:obj:`deepmd.pt.model.atomic_model.dp_atomic_model.DPAtomicModel` Model give atomic prediction of some physical property. :Parameters: **descriptor** Descriptor **fitting_net** Fitting net **type_map** Mapping atom type to the name (str) of the type. For example `type_map[1]` gives the name of the type 1. .. !! processed by numpydoc !! .. py:class:: DPZBLLinearEnergyAtomicModel(dp_model: deepmd.pt.model.atomic_model.dp_atomic_model.DPAtomicModel, zbl_model: deepmd.pt.model.atomic_model.pairtab_atomic_model.PairTabAtomicModel, sw_rmin: float, sw_rmax: float, type_map: list[str], smin_alpha: Optional[float] = 0.1, **kwargs) Bases: :py:obj:`LinearEnergyAtomicModel` Model linearly combine a list of AtomicModels. :Parameters: **dp_model** The DPAtomicModel being combined. **zbl_model** The PairTable model being combined. **sw_rmin** The lower boundary of the interpolation between short-range tabulated interaction and DP. **sw_rmax** The upper boundary of the interpolation between short-range tabulated interaction and DP. **type_map** Mapping atom type to the name (str) of the type. For example `type_map[1]` gives the name of the type 1. **smin_alpha** The short-range tabulated interaction will be switched according to the distance of the nearest neighbor. This distance is calculated by softmin. .. !! processed by numpydoc !! .. py:attribute:: sw_rmin .. py:attribute:: sw_rmax .. py:attribute:: smin_alpha .. py:attribute:: zbl_weight .. py:method:: serialize() -> dict .. py:method:: deserialize(data) -> DPZBLLinearEnergyAtomicModel :classmethod: .. py:method:: _compute_weight(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlists_: list[torch.Tensor]) -> list[torch.Tensor] ZBL weight. :Returns: :class:`python:list`\[:obj:`torch.Tensor`] the atomic ZBL weight for interpolation. (nframes, nloc, 1) .. !! processed by numpydoc !! .. py:class:: LinearEnergyAtomicModel(models: list[deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel], type_map: list[str], weights: Optional[Union[str, list[float]]] = 'mean', **kwargs) Bases: :py:obj:`deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel` Linear model make linear combinations of several existing models. :Parameters: **models** : :class:`python:list`\[:obj:`DPAtomicModel` or :obj:`PairTabAtomicModel`] A list of models to be combined. PairTabAtomicModel must be used together with a DPAtomicModel. **type_map** : :class:`python:list`\[:class:`python:str`] Mapping atom type to the name (str) of the type. For example `type_map[1]` gives the name of the type 1. **weights** : :obj:`Optional `\[:obj:`Union `\[str,list[:class:`python:float`]]] Weights of the models. If str, must be `sum` or `mean`. If list, must be a list of float. .. !! processed by numpydoc !! .. py:attribute:: models .. py:attribute:: mapping_list :value: [] .. py:attribute:: type_map .. py:attribute:: mixed_types_list .. py:attribute:: rcuts .. py:attribute:: nsels .. py:attribute:: weights .. py:method:: mixed_types() -> bool 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. .. !! processed by numpydoc !! .. py:method:: has_message_passing() -> bool Returns whether the atomic model has message passing. .. !! processed by numpydoc !! .. py:method:: need_sorted_nlist_for_lower() -> bool Returns whether the atomic model needs sorted nlist when using `forward_lower`. .. !! processed by numpydoc !! .. py:method:: get_out_bias() -> torch.Tensor .. py:method:: get_rcut() -> float Get the cut-off radius. .. !! processed by numpydoc !! .. py:method:: get_type_map() -> list[str] Get the type map. .. !! processed by numpydoc !! .. py:method:: change_type_map(type_map: list[str], model_with_new_type_stat=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:method:: get_model_rcuts() -> list[float] Get the cut-off radius for each individual models. .. !! processed by numpydoc !! .. py:method:: get_sel() -> list[int] Returns the number of selected atoms for each type. .. !! processed by numpydoc !! .. py:method:: get_model_nsels() -> list[int] Get the processed sels for each individual models. Not distinguishing types. .. !! processed by numpydoc !! .. py:method:: get_model_sels() -> list[list[int]] Get the sels for each individual models. .. !! processed by numpydoc !! .. py:method:: _sort_rcuts_sels() -> tuple[list[float], list[int]] .. 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 Compress model. :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 !! .. py:method:: forward_atomic(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, comm_dict: Optional[dict[str, torch.Tensor]] = None) -> dict[str, torch.Tensor] Return atomic prediction. :Parameters: **extended_coord** coordinates in extended region, (nframes, nall * 3) **extended_atype** atomic type in extended region, (nframes, nall) **nlist** neighbor list, (nframes, nloc, nsel). **mapping** mapps the extended indices to local indices. **fparam** frame parameter. (nframes, ndf) **aparam** atomic parameter. (nframes, nloc, nda) :Returns: :obj:`result_dict` the result dict, defined by the fitting net output def. .. !! processed by numpydoc !! .. py:method:: apply_out_stat(ret: dict[str, torch.Tensor], atype: torch.Tensor) Apply the stat to each atomic output. The developer may override the method to define how the bias is applied to the atomic output of the model. :Parameters: **ret** The returned dict by the forward_atomic method **atype** The atom types. nf x nloc .. !! processed by numpydoc !! .. py:method:: remap_atype(ori_map: list[str], new_map: list[str]) -> torch.Tensor :staticmethod: This method is used to map the atype from the common type_map to the original type_map of indivial AtomicModels. It creates a index mapping for the conversion. :Parameters: **ori_map** : :class:`python:list`\[:class:`python:str`] The original type map of an AtomicModel. **new_map** : :class:`python:list`\[:class:`python:str`] The common type map of the DPZBLLinearEnergyAtomicModel, created by the `get_type_map` method, must be a subset of the ori_map. :Returns: :obj:`torch.Tensor` .. .. !! processed by numpydoc !! .. py:method:: fitting_output_def() -> deepmd.dpmodel.FittingOutputDef Get the output def of developer implemented atomic models. .. !! processed by numpydoc !! .. py:method:: serialize() -> dict .. py:method:: deserialize(data: dict) -> LinearEnergyAtomicModel :classmethod: .. py:method:: _compute_weight(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlists_: list[torch.Tensor]) -> list[torch.Tensor] This should be a list of user defined weights that matches the number of models to be combined. .. !! processed by numpydoc !! .. py:method:: get_dim_fparam() -> int Get the number (dimension) of frame parameters of this atomic model. .. !! processed by numpydoc !! .. py:method:: get_dim_aparam() -> int Get the number (dimension) of atomic parameters of this atomic model. .. !! processed by numpydoc !! .. 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. .. !! processed by numpydoc !! .. 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). .. !! processed by numpydoc !! .. py:method:: compute_or_load_out_stat(merged: Union[Callable[[], list[dict]], list[dict]], stat_file_path: Optional[deepmd.utils.path.DPPath] = None) -> None Compute the output statistics (e.g. energy bias) for the fitting net 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. **stat_file_path** : :obj:`Optional `\[:obj:`DPPath`] The path to the stat file. .. !! processed by numpydoc !! .. py:method:: compute_or_load_stat(sampled_func, stat_file_path: Optional[deepmd.utils.path.DPPath] = None) -> None Compute or load the statistics parameters of the model, such as 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 lazy sampled function to get data frames from different data systems. **stat_file_path** The dictionary of paths to the statistics files. .. !! processed by numpydoc !! .. py:class:: PairTabAtomicModel(tab_file: str, rcut: float, sel: Union[int, list[int]], type_map: list[str], **kwargs) Bases: :py:obj:`deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel` Pairwise tabulation energy model. This model can be used to tabulate the pairwise energy between atoms for either short-range or long-range interactions, such as D3, LJ, ZBL, etc. It should not be used alone, but rather as one submodel of a linear (sum) model, such as DP+D3. Do not put the model on the first model of a linear model, since the linear model fetches the type map from the first model. At this moment, the model does not smooth the energy at the cutoff radius, so one needs to make sure the energy has been smoothed to zero. :Parameters: **tab_file** : :class:`python:str` The path to the tabulation file. **rcut** : :class:`python:float` The cutoff radius. **sel** : :class:`python:int` or :class:`python:list`\[:class:`python:int`] The maxmum number of atoms in the cut-off radius. **type_map** : :class:`python:list`\[:class:`python:str`] Mapping atom type to the name (str) of the type. For example `type_map[1]` gives the name of the type 1. **rcond** : :class:`python:float`, :obj:`optional` The condition number for the regression of atomic energy. **atom_ener** Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descriptor should be set. .. !! processed by numpydoc !! .. py:attribute:: tab_file .. py:attribute:: rcut .. py:attribute:: tab .. py:attribute:: type_map .. py:attribute:: ntypes .. py:attribute:: bias_atom_e .. py:method:: _set_pairtab(tab_file: str, rcut: float) -> deepmd.utils.pair_tab.PairTab .. py:method:: fitting_output_def() -> deepmd.dpmodel.FittingOutputDef Get the output def of developer implemented atomic models. .. !! processed by numpydoc !! .. py:method:: get_out_bias() -> torch.Tensor .. py:method:: get_rcut() -> float Get the cut-off radius. .. !! processed by numpydoc !! .. py:method:: get_type_map() -> list[str] Get the type map. .. !! processed by numpydoc !! .. py:method:: get_sel() -> list[int] Returns the number of selected atoms for each type. .. !! processed by numpydoc !! .. py:method:: get_nsel() -> int Returns the total number of selected neighboring atoms in the cut-off radius. .. !! processed by numpydoc !! .. py:method:: mixed_types() -> bool 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. .. !! processed by numpydoc !! .. py:method:: has_message_passing() -> bool Returns whether the atomic model has message passing. .. !! processed by numpydoc !! .. py:method:: need_sorted_nlist_for_lower() -> bool Returns whether the atomic model needs sorted nlist when using `forward_lower`. .. !! processed by numpydoc !! .. py:method:: change_type_map(type_map: list[str], model_with_new_type_stat=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:method:: serialize() -> dict .. py:method:: deserialize(data) -> PairTabAtomicModel :classmethod: .. py:method:: compute_or_load_stat(merged: Union[Callable[[], list[dict]], list[dict]], stat_file_path: Optional[deepmd.utils.path.DPPath] = None) -> None Compute the output statistics (e.g. energy bias) for the fitting net 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. **stat_file_path** : :obj:`Optional `\[:obj:`DPPath`] The path to the stat file. .. !! processed by numpydoc !! .. py:method:: forward_atomic(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] .. py:method:: _pair_tabulated_inter(nlist: torch.Tensor, i_type: torch.Tensor, j_type: torch.Tensor, rr: torch.Tensor) -> torch.Tensor Pairwise tabulated energy. :Parameters: **nlist** : :obj:`torch.Tensor` The unmasked neighbour list. (nframes, nloc) **i_type** : :obj:`torch.Tensor` The integer representation of atom type for all local atoms for all frames. (nframes, nloc) **j_type** : :obj:`torch.Tensor` The integer representation of atom type for all neighbour atoms of all local atoms for all frames. (nframes, nloc, nnei) **rr** : :obj:`torch.Tensor` The salar distance vector between two atoms. (nframes, nloc, nnei) :Returns: :obj:`torch.Tensor` The masked atomic energy for all local atoms for all frames. (nframes, nloc, nnei) :Raises: :obj:`Exception` If the distance is beyond the table. .. rubric:: Notes This function is used to calculate the pairwise energy between two atoms. It uses a table containing cubic spline coefficients calculated in PairTab. .. !! processed by numpydoc !! .. py:method:: _get_pairwise_dist(coords: torch.Tensor, nlist: torch.Tensor) -> torch.Tensor :staticmethod: Get pairwise distance `dr`. :Parameters: **coords** : :obj:`torch.Tensor` The coordinate of the atoms, shape of (nframes, nall, 3). **nlist** The masked nlist, shape of (nframes, nloc, nnei) :Returns: :obj:`torch.Tensor` The pairwise distance between the atoms (nframes, nloc, nnei). .. !! processed by numpydoc !! .. py:method:: _extract_spline_coefficient(i_type: torch.Tensor, j_type: torch.Tensor, idx: torch.Tensor, tab_data: torch.Tensor, nspline: int) -> torch.Tensor :staticmethod: Extract the spline coefficient from the table. :Parameters: **i_type** : :obj:`torch.Tensor` The integer representation of atom type for all local atoms for all frames. (nframes, nloc) **j_type** : :obj:`torch.Tensor` The integer representation of atom type for all neighbour atoms of all local atoms for all frames. (nframes, nloc, nnei) **idx** : :obj:`torch.Tensor` The index of the spline coefficient. (nframes, nloc, nnei) **tab_data** : :obj:`torch.Tensor` The table storing all the spline coefficient. (ntype, ntype, nspline, 4) **nspline** : :class:`python:int` The number of splines in the table. :Returns: :obj:`torch.Tensor` The spline coefficient. (nframes, nloc, nnei, 4), shape may be squeezed. .. !! processed by numpydoc !! .. py:method:: _calculate_ener(coef: torch.Tensor, uu: torch.Tensor) -> torch.Tensor :staticmethod: Calculate energy using spline coeeficients. :Parameters: **coef** : :obj:`torch.Tensor` The spline coefficients. (nframes, nloc, nnei, 4) **uu** : :obj:`torch.Tensor` The atom displancemnt used in interpolation and extrapolation (nframes, nloc, nnei) :Returns: :obj:`torch.Tensor` The atomic energy for all local atoms for all frames. (nframes, nloc, nnei) .. !! processed by numpydoc !! .. py:method:: get_dim_fparam() -> int Get the number (dimension) of frame parameters of this atomic model. .. !! processed by numpydoc !! .. py:method:: get_dim_aparam() -> int Get the number (dimension) of atomic parameters of this atomic model. .. !! processed by numpydoc !! .. 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. .. !! processed by numpydoc !! .. 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). .. !! processed by numpydoc !! .. py:class:: DPPolarAtomicModel(descriptor, fitting, type_map, **kwargs) Bases: :py:obj:`deepmd.pt.model.atomic_model.dp_atomic_model.DPAtomicModel` Model give atomic prediction of some physical property. :Parameters: **descriptor** Descriptor **fitting_net** Fitting net **type_map** Mapping atom type to the name (str) of the type. For example `type_map[1]` gives the name of the type 1. .. !! processed by numpydoc !! .. py:method:: apply_out_stat(ret: dict[str, torch.Tensor], atype: torch.Tensor) Apply the stat to each atomic output. :Parameters: **ret** The returned dict by the forward_atomic method **atype** The atom types. nf x nloc .. !! processed by numpydoc !! .. py:class:: DPPropertyAtomicModel(descriptor, fitting, type_map, **kwargs) Bases: :py:obj:`deepmd.pt.model.atomic_model.dp_atomic_model.DPAtomicModel` Model give atomic prediction of some physical property. :Parameters: **descriptor** Descriptor **fitting_net** Fitting net **type_map** Mapping atom type to the name (str) of the type. For example `type_map[1]` gives the name of the type 1. .. !! processed by numpydoc !! .. py:method:: apply_out_stat(ret: dict[str, torch.Tensor], atype: torch.Tensor) Apply the stat to each atomic output. This function defines how the bias is applied to the atomic output of the model. :Parameters: **ret** The returned dict by the forward_atomic method **atype** The atom types. nf x nloc .. !! processed by numpydoc !!