deepmd.pt.model.descriptor.repflows

Classes

DescrptBlockRepflows

The repflow descriptor block.

Functions

border_op(→ torch.Tensor)

Module Contents

deepmd.pt.model.descriptor.repflows.border_op(argument0, argument1, argument2, argument3, argument4, argument5, argument6, argument7, argument8) → torch.Tensor

class deepmd.pt.model.descriptor.repflows.DescrptBlockRepflows(e_rcut, e_rcut_smth, e_sel: int, a_rcut, a_rcut_smth, a_sel: int, ntypes: int, nlayers: int = 6, n_dim: int = 128, e_dim: int = 64, a_dim: int = 64, a_compress_rate: int = 0, a_compress_e_rate: int = 1, a_compress_use_split: bool = False, n_multi_edge_message: int = 1, axis_neuron: int = 4, update_angle: bool = True, activation_function: str = 'silu', update_style: str = 'res_residual', update_residual: float = 0.1, update_residual_init: str = 'const', set_davg_zero: bool = True, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, precision: str = 'float64', fix_stat_std: float = 0.3, smooth_edge_update: bool = False, edge_init_use_dist: bool = False, use_exp_switch: bool = False, use_dynamic_sel: bool = False, sel_reduce_factor: float = 10.0, use_loc_mapping: bool = True, optim_update: bool = True, seed: int | list[int] | None = None)

Bases: deepmd.pt.model.descriptor.descriptor.DescriptorBlock

The repflow descriptor block.

Parameters:

n_dimint, optional

The dimension of node representation.

e_dimint, optional

The dimension of edge representation.

a_dimint, optional

The dimension of angle representation.

nlayersint, optional

Number of repflow layers.

e_rcutfloat, optional

The edge cut-off radius.

e_rcut_smthfloat, optional

Where to start smoothing for edge. For example the 1/r term is smoothed from rcut to rcut_smth.

e_selint, optional

Maximally possible number of selected edge neighbors.

a_rcutfloat, optional

The angle cut-off radius.

a_rcut_smthfloat, optional

Where to start smoothing for angle. For example the 1/r term is smoothed from rcut to rcut_smth.

a_selint, optional

Maximally possible number of selected angle neighbors.

a_compress_rateint, optional

The compression rate for angular messages. The default value is 0, indicating no compression. If a non-zero integer c is provided, the node and edge dimensions will be compressed to a_dim/c and a_dim/2c, respectively, within the angular message.

a_compress_e_rateint, optional

The extra compression rate for edge in angular message compression. The default value is 1. When using angular message compression with a_compress_rate c and a_compress_e_rate c_e, the edge dimension will be compressed to (c_e * a_dim / 2c) within the angular message.

a_compress_use_splitbool, optional

Whether to split first sub-vectors instead of linear mapping during angular message compression. The default value is False.

n_multi_edge_messageint, optional

The head number of multiple edge messages to update node feature. Default is 1, indicating one head edge message.

axis_neuronint, optional

The number of dimension of submatrix in the symmetrization ops.

update_anglebool, optional

Where to update the angle rep. If not, only node and edge rep will be used.

update_stylestr, optional

Style to update a representation. Supported options are: -‘res_avg’: Updates a rep u with: u = 1/\sqrt{n+1} (u + u_1 + u_2 + … + u_n) -‘res_incr’: Updates a rep u with: u = u + 1/\sqrt{n} (u_1 + u_2 + … + u_n) -‘res_residual’: Updates a rep u with: u = u + (r1*u_1 + r2*u_2 + … + r3*u_n) where r1, r2 … r3 are residual weights defined by update_residual and update_residual_init.

update_residualfloat, optional

When update using residual mode, the initial std of residual vector weights.

update_residual_initstr, optional

When update using residual mode, the initialization mode of residual vector weights.

fix_stat_stdfloat, optional

If non-zero (default is 0.3), use this constant as the normalization standard deviation instead of computing it from data statistics.

smooth_edge_updatebool, optional

Whether to make edge update smooth. If True, the edge update from angle message will not use self as padding.

edge_init_use_distbool, optional

Whether to use direct distance r to initialize the edge features instead of 1/r. Note that when using this option, the activation function will not be used when initializing edge features.

use_exp_switchbool, optional

Whether to use an exponential switch function instead of a polynomial one in the neighbor update. The exponential switch function ensures neighbor contributions smoothly diminish as the interatomic distance r approaches the cutoff radius rcut. Specifically, the function is defined as: s(r) = \exp(-\exp(20 * (r - rcut_smth) / rcut_smth)) for 0 < r \leq rcut, and s(r) = 0 for r > rcut. Here, rcut_smth is an adjustable smoothing factor and rcut_smth should be chosen carefully according to rcut, ensuring s(r) approaches zero smoothly at the cutoff. Typical recommended values are rcut_smth = 5.3 for rcut = 6.0, and 3.5 for rcut = 4.0.

use_dynamic_selbool, optional

Whether to dynamically select neighbors within the cutoff radius. If True, the exact number of neighbors within the cutoff radius is used without padding to a fixed selection numbers. When enabled, users can safely set larger values for e_sel or a_sel (e.g., 1200 or 300, respectively) to guarantee capturing all neighbors within the cutoff radius. Note that when using dynamic selection, the smooth_edge_update must be True.

sel_reduce_factorfloat, optional

Reduction factor applied to neighbor-scale normalization when use_dynamic_sel is True. In the dynamic selection case, neighbor-scale normalization will use e_sel / sel_reduce_factor or a_sel / sel_reduce_factor instead of the raw e_sel or a_sel values, accommodating larger selection numbers.

use_loc_mappingbool, Optional

Whether to use local atom index mapping in training or non-parallel inference. When True, local indexing and mapping are applied to neighbor lists and embeddings during descriptor computation.

optim_updatebool, optional

Whether to enable the optimized update method. Uses a more efficient process when enabled. Defaults to True

ntypesint

Number of element types

activation_functionstr, optional

The activation function in the embedding net.

set_davg_zerobool, optional

Set the normalization average to zero.

precisionstr, optional

The precision of the embedding net parameters.

exclude_typeslist[list[int]], optional

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

env_protectionfloat, optional

Protection parameter to prevent division by zero errors during environment matrix calculations. For example, when using paddings, there may be zero distances of neighbors, which may make division by zero error during environment matrix calculations without protection.

seedint, optional

Random seed for parameter initialization.

e_rcut

e_rcut_smth

e_sel

a_rcut

a_rcut_smth

a_sel

ntypes

nlayers = 6

nnei

ndescrpt

sel

rcut

rcut_smth

sec

split_sel

a_compress_rate = 0

a_compress_e_rate = 1

n_multi_edge_message = 1

axis_neuron = 4

set_davg_zero = True

fix_stat_std = 0.3

set_stddev_constant

a_compress_use_split = False

use_loc_mapping = True

optim_update = True

smooth_edge_update = False

edge_init_use_dist = False

use_exp_switch = False

use_dynamic_sel = False

sel_reduce_factor = 10.0

n_dim = 128

e_dim = 64

a_dim = 64

update_angle = True

activation_function = 'silu'

update_style = 'res_residual'

update_residual = 0.1

update_residual_init = 'const'

act

prec

env_protection = 0.0

precision = 'float64'

epsilon = 0.0001

seed = None

edge_embd

angle_embd

layers

stats = None

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_dim_out() → int

Returns the output dimension.

get_dim_in() → int

Returns the input dimension.

get_dim_emb() → int

Returns the embedding dimension e_dim.

__setitem__(key, value) → None

__getitem__(key)

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.

get_env_protection() → float

Returns the protection of building environment matrix.

property dim_out

Returns the output dimension of this descriptor.

property dim_in

Returns the atomic input dimension of this descriptor.

property dim_emb

Returns the embedding dimension e_dim.

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

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

Calculate DescriptorBlock.

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.

get_stats() → dict[str, deepmd.utils.env_mat_stat.StatItem]

Get the statistics of the descriptor.

has_message_passing() → bool

Returns whether the descriptor block has message passing.

need_sorted_nlist_for_lower() → bool

Returns whether the descriptor block needs sorted nlist when using forward_lower.