deepmd.pt.model.atomic_model.base_atomic_model

Module Contents

Classes

BaseAtomicModel

The base of atomic model.

Attributes

log
dtype
device
BaseAtomicModel_

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, torch.Tensor] | None = None)

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[torch.Tensor]]], 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.array 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 descrptor should be set.

init_out_stat()

Initialize the output bias.

__setitem__(key, value)

__getitem__(key)

get_type_map() → List[str]

Get the type map.

reinit_atom_exclude(exclude_types: List[int] = [])

reinit_pair_exclude(exclude_types: List[Tuple[int, int]] = [])

make_atom_mask(atype: torch.Tensor) → torch.Tensor

The atoms with type < 0 are treated as virutal 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 virutal 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 coodinates, 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.

serialize() → dict

classmethod deserialize(data: dict) → BaseAtomicModel

abstract compute_or_load_stat(merged: Callable[[], List[dict]] | List[dict], stat_file_path: deepmd.utils.path.DPPath | 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.

compute_or_load_out_stat(merged: Callable[[], List[dict]] | List[dict], stat_file_path: deepmd.utils.path.DPPath | 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)

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

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()

_default_std()

_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)

_fetch_out_stat(keys: List[str]) → Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]