deepmd.pt_expt.descriptor.se_t_tebd

deepmd.pt_expt.descriptor.se_t_tebd#

Classes#

DescrptSeTTebd

Construct an embedding net that takes angles between two neighboring atoms and type embeddings as input.

Module Contents#

class deepmd.pt_expt.descriptor.se_t_tebd.DescrptSeTTebd(rcut: float, rcut_smth: float, sel: list[int] | int, ntypes: int, neuron: list = [2, 4, 8], tebd_dim: int = 8, tebd_input_mode: str = 'concat', resnet_dt: bool = False, set_davg_zero: bool = True, activation_function: str = 'tanh', env_protection: float = 0.0, exclude_types: list[tuple[int, int]] = [], precision: str = 'float64', trainable: bool = True, seed: int | list[int] | None = None, type_map: list[str] | None = None, concat_output_tebd: bool = True, use_econf_tebd: bool = False, use_tebd_bias: bool = False, smooth: bool = True)[source]#

Bases: deepmd.dpmodel.descriptor.se_t_tebd.DescrptSeTTebd

Construct an embedding net that takes angles between two neighboring atoms and type embeddings as input.

The descriptor \(\mathcal{D}^i \in \mathbb{R}^{M}\) is given by

\[\mathcal{D}^i = \frac{1}{N_c^2} \sum_{j,k} \mathcal{N}(\cos\theta_{jik}, \mathcal{T}_j, \mathcal{T}_k),\]

where \(\theta_{jik}\) is the angle between neighbors \(j\) and \(k\) around the central atom \(i\), \(\mathcal{T}_j\) and \(\mathcal{T}_k\) are the type embeddings of atoms \(j\) and \(k\), and \(\mathcal{N}\) is the embedding network.

The cosine of the angle is computed from the normalized relative coordinates:

\[\cos\theta_{jik} = \frac{\boldsymbol{r}_{ij} \cdot \boldsymbol{r}_{ik}}{|\boldsymbol{r}_{ij}| |\boldsymbol{r}_{ik}|}.\]

The type embedding can be incorporated in two modes:

“concat”: Concatenate \([\cos\theta_{jik}, \mathcal{T}_j, \mathcal{T}_k]\) as input to the embedding network.
“strip”: Use separate embedding networks for \(\cos\theta_{jik}\) and \([\mathcal{T}_j, \mathcal{T}_k]\), then combine their outputs multiplicatively.

Parameters:

rcut: The cut-off radius
rcut_smth: From where the environment matrix should be smoothed
selUnion[list[int], int]: list[int]: sel[i] specifies the maxmum number of type i atoms in the cut-off radius int: the total maxmum number of atoms in the cut-off radius
ntypesint: Number of element types
neuronlist[int]: Number of neurons in each hidden layers of the embedding net
tebd_dimint: Dimension of the type embedding
tebd_input_modestr: The input mode of the type embedding. Supported modes are [“concat”, “strip”]. - “concat”: Concatenate the type embedding with the smoothed angular information as the union input for the embedding network. - “strip”: Use a separated embedding network for the type embedding and combine the output with the angular embedding network output.
resnet_dt: Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
set_davg_zero: Set the shift of embedding net input to zero.
activation_function: The activation function in the embedding net. Supported options are “silut”, “relu”, “tanh”, “sigmoid”, “silu”, “linear”, “softplus”, “none”, “gelu”, “relu6”, “gelu_tf”.
env_protection: float: Protection parameter to prevent division by zero errors during environment matrix calculations.
exclude_typeslist[tuple[int, int]]: 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.
precision: The precision of the embedding net parameters. Supported options are “float32”, “default”, “float16”, “float64”, “bfloat16”.
trainable: If the weights of embedding net are trainable.
seed: Random seed for initializing the network parameters.
type_map: list[str], Optional: A list of strings. Give the name to each type of atoms.
concat_output_tebd: bool: Whether to concat type embedding at the output of the descriptor.
use_econf_tebd: bool, Optional: Whether to use electronic configuration type embedding.
use_tebd_biasbool, Optional: Whether to use bias in the type embedding layer.
smooth: bool: Whether to use smooth process in calculation.

_update_sel_cls[source]#

share_params(base_class: DescrptSeTTebd, shared_level: int, model_prob: float = 1.0, resume: bool = False) → None[source]#

Share parameters with base_class for multi-task training.

Level 0: share type_embedding and se_ttebd. Level 1: share type_embedding only.

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

Enable compression for the SE_T_TEBD descriptor.

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

_store_compress_data() → None[source]#: Store tabulated data as buffers for the compressed forward path.

_store_type_embd_data() → None[source]#: Precompute type embedding pairs and store as a buffer.

call(coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, comm_dict: dict | None = None, charge_spin: torch.Tensor | None = None) → Any[source]#

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 from extended to local region. not used by this descriptor.

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.

_call_compressed(coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor) → Any[source]#: Compressed forward using tabulate_fusion_se_t_tebd custom op.