deepmd.pt.model.atomic_model.base_atomic_model#
Attributes#
Classes#
The base of atomic model. |
Module Contents#
- class deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel(type_map: list[str], atom_exclude_types: list[int] = [], pair_exclude_types: list[tuple[int, int]] = [], rcond: float | None = None, preset_out_bias: dict[str, numpy.ndarray] | None = None)[source]#
Bases:
torch.nn.Module
,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
float
,optional
The condition number for the regression of atomic energy.
- preset_out_bias
dict
[str
,list
[Optional
[np.ndarray
]]],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.
- make_atom_mask(atype: torch.Tensor) torch.Tensor [source]#
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:
mask
True for real atoms and False for virtual atoms.
- atomic_output_def() deepmd.dpmodel.output_def.FittingOutputDef [source]#
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.
- forward_common_atomic(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None) dict[str, torch.Tensor] [source]#
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:
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.
- forward(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None) dict[str, torch.Tensor] [source]#
- change_type_map(type_map: list[str], model_with_new_type_stat=None) None [source]#
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.
- classmethod deserialize(data: dict) BaseAtomicModel [source]#
- abstract compute_or_load_stat(merged: Callable[[], list[dict]] | list[dict], stat_file_path: deepmd.utils.path.DPPath | None = None) NoReturn [source]#
Compute the output statistics (e.g. energy bias) for the fitting net from packed data.
- Parameters:
- merged
Union
[Callable
[[],list
[dict
]],list
[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
Optional
[DPPath
] The path to the stat file.
- merged
- compute_or_load_out_stat(merged: Callable[[], list[dict]] | list[dict], stat_file_path: deepmd.utils.path.DPPath | None = None) None [source]#
Compute the output statistics (e.g. energy bias) for the fitting net from packed data.
- Parameters:
- merged
Union
[Callable
[[],list
[dict
]],list
[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
Optional
[DPPath
] The path to the stat file.
- merged
- apply_out_stat(ret: dict[str, torch.Tensor], atype: torch.Tensor)[source]#
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
- change_out_bias(sample_merged, stat_file_path: deepmd.utils.path.DPPath | None = None, bias_adjust_mode='change-by-statistic') None [source]#
Change the output bias according to the input data and the pretrained model.
- Parameters:
- sample_merged
Union
[Callable
[[],list
[dict
]],list
[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
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
Optional
[DPPath
] The path to the stat file.
- sample_merged
- _get_forward_wrapper_func() Callable[Ellipsis, torch.Tensor] [source]#
Get a forward wrapper of the atomic model for output bias calculation.