deepmd.pt.model.descriptor.dpa2

Classes

DescrptDPA2

Base descriptor provides the interfaces of descriptor.

Module Contents

class deepmd.pt.model.descriptor.dpa2.DescrptDPA2(ntypes: int, repinit: deepmd.dpmodel.descriptor.dpa2.RepinitArgs | dict, repformer: deepmd.dpmodel.descriptor.dpa2.RepformerArgs | dict, concat_output_tebd: bool = True, precision: str = 'float64', smooth: bool = True, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable: bool = True, seed: int | list[int] | None = None, add_tebd_to_repinit_out: bool = False, use_econf_tebd: bool = False, use_tebd_bias: bool = False, type_map: list[str] | None = None)

Bases: deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor, torch.nn.Module

Base descriptor provides the interfaces of descriptor.

repinit_args

repformer_args

tebd_input_mode = 'concat'

repinit

use_three_body = False

repformers

rcsl_list

rcut_list

nsel_list

use_econf_tebd = False

use_tebd_bias = False

type_map = None

type_embedding

concat_output_tebd = True

precision = 'float64'

prec

smooth = True

exclude_types = []

env_protection = 0.0

trainable = True

add_tebd_to_repinit_out = False

repinit_out_dim

tebd_transform = None

tebd_dim = 8

rcut

rcut_smth

ntypes

sel

compress = False

get_rcut() → float

Returns the cut-off radius.

get_rcut_smth() → float

Returns the radius where the neighbor information starts to smoothly decay to 0.

get_nsel() → int

Returns the number of selected atoms in the cut-off radius.

get_sel() → list[int]

Returns the number of selected atoms for each type.

get_ntypes() → int

Returns the number of element types.

get_type_map() → list[str]

Get the name to each type of atoms.

get_dim_out() → int

Returns the output dimension of this descriptor.

get_dim_emb() → int

Returns the embedding dimension of this descriptor.

mixed_types() → bool

If true, the descriptor 1. assumes total number of atoms aligned across frames; 2. requires a neighbor list that does not distinguish different atomic types.

If false, the descriptor 1. assumes total number of atoms of each atom type aligned across frames; 2. requires a neighbor list that distinguishes different atomic types.

has_message_passing() → bool

Returns whether the descriptor has message passing.

need_sorted_nlist_for_lower() → bool

Returns whether the descriptor needs sorted nlist when using forward_lower.

get_env_protection() → float

Returns the protection of building environment matrix.

share_params(base_class: Any, shared_level: int, resume: bool = False) → None

Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), some separated parameters (e.g. mean and stddev) will be re-calculated across different classes.

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

property dim_out: int

property dim_emb: int

Returns the embedding dimension g2.

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

Compute the input statistics (e.g. mean and stddev) for the descriptors 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.

pathOptional[DPPath]

The path to the stat file.

set_stat_mean_and_stddev(mean: list[torch.Tensor], stddev: list[torch.Tensor]) → None

Update mean and stddev for descriptor.

get_stat_mean_and_stddev() → tuple[list[torch.Tensor], list[torch.Tensor]]

Get mean and stddev for descriptor.

serialize() → dict

Serialize the obj to dict.

classmethod deserialize(data: dict) → DescrptDPA2

Deserialize the model.

Parameters:

datadict

The serialized data

Returns:

BD

The deserialized descriptor

forward(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, comm_dict: dict[str, torch.Tensor] | None = None) → tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]

Compute the descriptor.

Parameters:

extended_coord

The extended coordinates of atoms. shape: nf x (nallx3)

extended_atype

The extended aotm types. shape: nf x nall

nlist

The neighbor list. shape: nf x nloc x nnei

mapping

The index mapping, mapps extended region index to local region.

comm_dict

The data needed for communication for parallel inference.

Returns:

descriptor

The descriptor. shape: nf x nloc x (ng x axis_neuron)

gr

The rotationally equivariant and permutationally invariant single particle representation. shape: nf x nloc x ng x 3

g2

The rotationally invariant pair-partical representation. shape: nf x nloc x nnei x ng

h2

The rotationally equivariant pair-partical representation. shape: nf x nloc x nnei x 3

sw

The smooth switch function. shape: nf x nloc x nnei

classmethod update_sel(train_data: deepmd.utils.data_system.DeepmdDataSystem, type_map: list[str] | None, local_jdata: dict) → tuple[dict, float | None]

Update the selection and perform neighbor statistics.

Parameters:

train_dataDeepmdDataSystem

data used to do neighbor statistics

type_maplist[str], optional

The name of each type of atoms

local_jdatadict

The local data refer to the current class

Returns:

dict

The updated local data

float

The minimum distance between two atoms

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

Receive the statistics (distance, max_nbor_size and env_mat_range) of the training data.

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