deepmd.pt.model.atomic_model.base_atomic_model

Attributes

log
dtype
device
BaseAtomicModel_

Classes

BaseAtomicModel

The base of atomic model.

Module Contents

deepmd.pt.model.atomic_model.base_atomic_model.log

deepmd.pt.model.atomic_model.base_atomic_model.dtype

deepmd.pt.model.atomic_model.base_atomic_model.device

deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel_

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, data_stat_protect: float = 0.01)

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.

rcondfloat, optional

The condition number for the regression of atomic energy.

preset_out_biasdict[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.

type_map

rcond = None

preset_out_bias = None

data_stat_protect = 0.01

init_out_stat() → None

Initialize the output bias.

set_out_bias(out_bias: torch.Tensor) → None

__setitem__(key: str, value: torch.Tensor) → None

__getitem__(key: str) → torch.Tensor

get_type_map() → list[str]

Get the type map.

get_compute_stats_distinguish_types() → bool

Get whether the fitting net computes stats which are not distinguished between different types of atoms.

get_intensive() → bool

Whether the fitting property is intensive.

has_default_fparam() → bool

Check if the model has default frame parameters.

reinit_atom_exclude(exclude_types: list[int] = []) → None

reinit_pair_exclude(exclude_types: list[tuple[int, int]] = []) → None

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:

mask

True for real atoms and False for virtual atoms.

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.

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]

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]

change_type_map(type_map: list[str], model_with_new_type_stat: BaseAtomicModel | None = 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.

serialize() → dict

classmethod deserialize(data: dict) → BaseAtomicModel

abstractmethod compute_or_load_stat(merged: Callable[[], list[dict]] | list[dict], stat_file_path: deepmd.utils.path.DPPath | None = None, compute_or_load_out_stat: bool = True) → NoReturn

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:

merged

The lazy sampled function to get data frames from different data systems.

stat_file_path

The dictionary of paths to the statistics files.

compute_or_load_out_statbool

Whether to compute the output statistics. If False, it will only compute the input statistics (e.g. mean and standard deviation of descriptors).

compute_or_load_out_stat(merged: Callable[[], list[dict]] | list[dict], stat_file_path: deepmd.utils.path.DPPath | None = None) → None

Compute the output statistics (e.g. energy bias) for the fitting net from packed data.

Parameters:

mergedUnion[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_pathOptional[DPPath]

The path to the stat file.

apply_out_stat(ret: dict[str, torch.Tensor], atype: torch.Tensor) → dict[str, 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

change_out_bias(sample_merged: Callable[[], list[dict]] | list[dict], stat_file_path: deepmd.utils.path.DPPath | None = None, bias_adjust_mode: str = 'change-by-statistic') → None

Change the output bias according to the input data and the pretrained model.

Parameters:

sample_mergedUnion[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_modestr

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_pathOptional[DPPath]

The path to the stat file.

_get_forward_wrapper_func() → Callable[Ellipsis, torch.Tensor]

Get a forward wrapper of the atomic model for output bias calculation.

_default_bias() → torch.Tensor

_default_std() → torch.Tensor

_varsize(shape: list[int]) → int

_get_bias_index(kk: str) → int

_store_out_stat(out_bias: dict[str, torch.Tensor], out_std: dict[str, torch.Tensor], add: bool = False) → None

_fetch_out_stat(keys: list[str]) → tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]