deepmd.pt_expt.descriptor.se_t

deepmd.pt_expt.descriptor.se_t#

Classes#

DescrptSeT

DeepPot-SE constructed from all information (both angular and radial) of atomic

Module Contents#

class deepmd.pt_expt.descriptor.se_t.DescrptSeT(rcut: float, rcut_smth: float, sel: list[int], neuron: list[int] = [24, 48, 96], resnet_dt: bool = False, set_davg_zero: bool = False, activation_function: str = 'tanh', env_protection: float = 0.0, exclude_types: list[tuple[int, int]] = [], precision: str = DEFAULT_PRECISION, trainable: bool = True, seed: int | list[int] | None = None, type_map: list[str] | None = None, ntypes: int | None = None)[source]#

Bases: deepmd.dpmodel.descriptor.se_t.DescrptSeT

DeepPot-SE constructed from all information (both angular and radial) of atomic configurations.

The embedding takes angles between two neighboring atoms as input.

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

\[\mathcal{D}^i = \sum_{t_j, t_k} \frac{1}{N_{t_j} N_{t_k}} \sum_{j \in t_j, k \in t_k} \tilde{g}_{jk} \, \mathcal{N}_{t_j, t_k}(\tilde{g}_{jk}),\]

where \(\tilde{g}_{jk} = \boldsymbol{rr}_j \cdot \boldsymbol{rr}_k\) is the dot product of the smoothed directional vectors from the environment matrix, \(N_{t_j}\) and \(N_{t_k}\) are the numbers of neighbors of types \(t_j\) and \(t_k\), and \(\mathcal{N}_{t_j, t_k}\) is the embedding network that depends only on the types of neighbor atoms \(j\) and \(k\).

The smoothed directional vector \(\boldsymbol{rr}_j\) is computed as:

\[\boldsymbol{rr}_j = s(r_{ji}) \frac{\boldsymbol{R}_j - \boldsymbol{R}_i}{r_{ji}},\]

where \(s(r)\) is the switching function.

Parameters:

rcutfloat: The cut-off radius
rcut_smthfloat: From where the environment matrix should be smoothed
sellist[int]: sel[i] specifies the maxmum number of type i atoms in the cut-off radius
neuronlist[int]: Number of neurons in each hidden layers of the embedding net
resnet_dtbool: Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
set_davg_zerobool: Set the shift of embedding net input to zero.
activation_functionstr: The activation function in the embedding net. Supported options are “softplus”, “tanh”, “linear”, “gelu_tf”, “silu”, “none”, “relu6”, “gelu”, “silut”, “sigmoid”, “relu”.
env_protectionfloat: Protection parameter to prevent division by zero errors during environment matrix calculations.
exclude_typeslist[list[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.
precisionstr: The precision of the embedding net parameters. Supported options are “bfloat16”, “float16”, “default”, “float64”, “float32”.
trainablebool: If the weights of embedding net are trainable.
seedint, Optional: Random seed for initializing the network parameters.
type_map: list[str], Optional: A list of strings. Give the name to each type of atoms.
ntypesint: Number of element types. Not used in this descriptor, only to be compat with input.

_update_sel_cls[source]#

share_params(base_class: Any, 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 all modules and buffers.

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

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

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
gr: The rotationally equivariant and permutationally invariant single particle representation. This descriptor returns None.
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 custom op.