deepmd.pt.model.atomic_model.linear_atomic_model

Classes

LinearEnergyAtomicModel	Linear model make linear combinations of several existing models.
DPZBLLinearEnergyAtomicModel	Model linearly combine a list of AtomicModels.

Module Contents

class deepmd.pt.model.atomic_model.linear_atomic_model.LinearEnergyAtomicModel(models: list[deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel], type_map: list[str], weights: str | list[float] | None = 'mean', **kwargs)

Bases: deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel

Linear model make linear combinations of several existing models.

Parameters:

modelslist[DPAtomicModel or PairTabAtomicModel]

A list of models to be combined. PairTabAtomicModel must be used together with a DPAtomicModel.

type_maplist[str]

Mapping atom type to the name (str) of the type. For example type_map[1] gives the name of the type 1.

weightsOptional[Union[str,list[float]]]

Weights of the models. If str, must be sum or mean. If list, must be a list of float.

models

mapping_list = []

type_map

mixed_types_list

rcuts

nsels

weights

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.

has_message_passing() → bool

Returns whether the atomic model has message passing.

need_sorted_nlist_for_lower() → bool

Returns whether the atomic model needs sorted nlist when using forward_lower.

get_out_bias() → torch.Tensor

get_rcut() → float

Get the cut-off radius.

get_type_map() → list[str]

Get the type map.

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.

get_model_rcuts() → list[float]

Get the cut-off radius for each individual models.

get_sel() → list[int]

Returns the number of selected atoms for each type.

get_model_nsels() → list[int]

Get the processed sels for each individual models. Not distinguishing types.

get_model_sels() → list[list[int]]

Get the sels for each individual models.

_sort_rcuts_sels() → tuple[list[float], list[int]]

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

forward_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]

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:

result_dict

the result dict, defined by the fitting net output def.

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

static remap_atype(ori_map: list[str], new_map: list[str]) → torch.Tensor

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_maplist[str]

The original type map of an AtomicModel.

new_maplist[str]

The common type map of the DPZBLLinearEnergyAtomicModel, created by the get_type_map method, must be a subset of the ori_map.

Returns:

torch.Tensor

fitting_output_def() → deepmd.dpmodel.FittingOutputDef

Get the output def of developer implemented atomic models.

serialize() → dict

classmethod deserialize(data: dict) → LinearEnergyAtomicModel

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

get_dim_fparam() → int

Get the number (dimension) of frame parameters of this atomic model.

get_dim_aparam() → int

Get the number (dimension) of atomic parameters of this atomic model.

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.

is_aparam_nall() → bool

Check whether the shape of atomic parameters is (nframes, nall, ndim).

If False, the shape is (nframes, nloc, ndim).

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.

compute_or_load_stat(sampled_func, stat_file_path: deepmd.utils.path.DPPath | None = 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.

class deepmd.pt.model.atomic_model.linear_atomic_model.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: float | None = 0.1, **kwargs)

Bases: 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.

sw_rmin

sw_rmax

smin_alpha

zbl_weight

serialize() → dict

classmethod deserialize(data) → DPZBLLinearEnergyAtomicModel

_compute_weight(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlists_: list[torch.Tensor]) → list[torch.Tensor]

ZBL weight.

Returns:

list[torch.Tensor]

the atomic ZBL weight for interpolation. (nframes, nloc, 1)