deepmd.pt.model.descriptor.se_a

deepmd.pt.model.descriptor.se_a#

Classes#

`DescrptSeA`	Base descriptor provides the interfaces of descriptor.
`DescrptBlockSeA`	The building block of descriptor.

Functions#

tabulate_fusion_se_a(→ list[torch.Tensor])

Module Contents#

deepmd.pt.model.descriptor.se_a.tabulate_fusion_se_a(argument0: torch.Tensor, argument1: torch.Tensor, argument2: torch.Tensor, argument3: torch.Tensor, argument4: int) → list[torch.Tensor][source]#

class deepmd.pt.model.descriptor.se_a.DescrptSeA(rcut: float, rcut_smth: float, sel: list[int] | int, neuron: list[int] = [25, 50, 100], axis_neuron: int = 16, set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = 'float64', resnet_dt: bool = False, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, type_one_side: bool = True, trainable: bool = True, seed: int | list[int] | None = None, ntypes: int | None = None, type_map: list[str] | None = None, spin: Any | None = None)[source]#

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

Base descriptor provides the interfaces of descriptor.

type_map = None[source]#

compress = False[source]#

prec[source]#

sea[source]#

get_dim_chg_spin() → int[source]#: Returns the dimension of charge_spin input (0 if not supported).

has_default_chg_spin() → bool[source]#: Returns whether the descriptor has a default charge_spin value.

get_default_chg_spin() → None[source]#: Returns the default charge_spin value, or None.

get_rcut() → float[source]#: Returns the cut-off radius.

get_rcut_smth() → float[source]#: Returns the radius where the neighbor information starts to smoothly decay to 0.

get_nsel() → int[source]#: Returns the number of selected atoms in the cut-off radius.

get_sel() → list[int][source]#: Returns the number of selected atoms for each type.

get_ntypes() → int[source]#: Returns the number of element types.

get_type_map() → list[str][source]#: Get the name to each type of atoms.

get_dim_out() → int[source]#: Returns the output dimension.

get_dim_emb() → int[source]#: Returns the output dimension.

mixed_types() → bool[source]#: Returns if the descriptor requires a neighbor list that distinguish different atomic types or not.

has_message_passing() → bool[source]#: Returns whether the descriptor has message passing.

need_sorted_nlist_for_lower() → bool[source]#: Returns whether the descriptor needs sorted nlist when using forward_lower.

get_env_protection() → float[source]#: Returns the protection of building environment matrix.

share_params(base_class: Any, shared_level: int, resume: bool = False) → None[source]#: 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.

property dim_out: int[source]#: Returns the output dimension of this descriptor.

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

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

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.

reinit_exclude(exclude_types: list[tuple[int, int]] = []) → None[source]#: Update the type exclusions.

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[source]#

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

Compute the descriptor.

Parameters:

coord_ext: The extended coordinates of atoms. shape: nf x (nallx3)
atype_ext: The extended aotm types. shape: nf x nall
nlist: The neighbor list. shape: nf x nloc x nnei
mapping: The index mapping, not required by this descriptor.
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. this descriptor returns None
h2: The rotationally equivariant pair-partical representation. this descriptor returns None
sw: The smooth switch function.

set_stat_mean_and_stddev(mean: torch.Tensor, stddev: torch.Tensor) → None[source]#: Update mean and stddev for descriptor.

get_stat_mean_and_stddev() → tuple[torch.Tensor, torch.Tensor][source]#: Get mean and stddev for descriptor.

serialize() → dict[source]#: Serialize the obj to dict.

classmethod deserialize(data: dict) → DescrptSeA[source]#

Deserialize the model.

Parameters:

datadict: The serialized data

Returns:

BD: The deserialized descriptor

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

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

class deepmd.pt.model.descriptor.se_a.DescrptBlockSeA(rcut: float, rcut_smth: float, sel: int | list[int], neuron: list[int] = [25, 50, 100], axis_neuron: int = 16, set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = 'float64', resnet_dt: bool = False, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, type_one_side: bool = True, trainable: bool = True, seed: int | list[int] | None = None, **kwargs: Any)[source]#

Bases: deepmd.pt.model.descriptor.DescriptorBlock

The building block of descriptor. Given the input descriptor, provide with the atomic coordinates, atomic types and neighbor list, calculate the new descriptor.

ndescrpt: Final[int][source]#

__constants__: ClassVar[list] = ['ndescrpt'][source]#

rcut[source]#

rcut_smth[source]#

neuron = [25, 50, 100][source]#

filter_neuron = [25, 50, 100][source]#

axis_neuron = 16[source]#

set_davg_zero = False[source]#

activation_function = 'tanh'[source]#

precision = 'float64'[source]#

prec[source]#

resnet_dt = False[source]#

env_protection = 0.0[source]#

ntypes[source]#

type_one_side = True[source]#

seed = None[source]#

sel[source]#

sec[source]#

split_sel[source]#

nnei[source]#

filter_layers[source]#

stats = None[source]#

trainable = True[source]#

compress = False[source]#

compress_info[source]#

compress_data[source]#

get_rcut() → float[source]#: Returns the cut-off radius.

get_rcut_smth() → float[source]#: Returns the radius where the neighbor information starts to smoothly decay to 0.

get_nsel() → int[source]#: Returns the number of selected atoms in the cut-off radius.

get_sel() → list[int][source]#: Returns the number of selected atoms for each type.

get_ntypes() → int[source]#: Returns the number of element types.

get_dim_out() → int[source]#: Returns the output dimension.

get_dim_rot_mat_1() → int[source]#: Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3.

get_dim_emb() → int[source]#: Returns the output dimension.

get_dim_in() → int[source]#: Returns the input dimension.

mixed_types() → bool[source]#

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.

get_env_protection() → float[source]#: Returns the protection of building environment matrix.

property dim_out: int[source]#: Returns the output dimension of this descriptor.

property dim_in: int[source]#: Returns the atomic input dimension of this descriptor.

__setitem__(key: str, value: torch.Tensor) → None[source]#

__getitem__(key: str) → torch.Tensor[source]#

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

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.

get_stats() → dict[str, deepmd.utils.env_mat_stat.StatItem][source]#: Get the statistics of the descriptor.

reinit_exclude(exclude_types: list[tuple[int, int]] = []) → None[source]#

enable_compression(table_data: dict[str, torch.Tensor], table_config: list[int | float], lower: dict[str, int], upper: dict[str, int]) → None[source]#

forward(nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, extended_atype_embd: torch.Tensor | None = None, mapping: torch.Tensor | None = None, type_embedding: torch.Tensor | None = None) → tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None, torch.Tensor | None][source]#

Calculate decoded embedding for each atom.

Args: - coord: Tell atom coordinates with shape [nframes, natoms[1]*3]. - atype: Tell atom types with shape [nframes, natoms[1]]. - natoms: Tell atom count and element count. Its shape is [2+self.ntypes]. - box: Tell simulation box with shape [nframes, 9].

Returns:

torch.Tensor: descriptor matrix with shape [nframes, natoms[0]*self.filter_neuron[-1]*self.axis_neuron].

has_message_passing() → bool[source]#: Returns whether the descriptor block has message passing.

need_sorted_nlist_for_lower() → bool[source]#: Returns whether the descriptor block needs sorted nlist when using forward_lower.