deepmd.tf.descriptor.se_a

Module Contents

Classes

DescrptSeA

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

class deepmd.tf.descriptor.se_a.DescrptSeA(rcut: float, rcut_smth: float, sel: List[int], neuron: List[int] = [24, 48, 96], axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, seed: int | None = None, type_one_side: bool = True, exclude_types: List[List[int]] = [], set_davg_zero: bool = False, activation_function: str = 'tanh', precision: str = 'default', uniform_seed: bool = False, multi_task: bool = False, spin: deepmd.tf.utils.spin.Spin | None = None, stripped_type_embedding: bool = False, env_protection: float = 0.0, **kwargs)

Bases: deepmd.tf.descriptor.se.DescrptSe

DeepPot-SE constructed from all information (both angular and radial) of atomic configurations. The embedding takes the distance between atoms as input.

The descriptor \(\mathcal{D}^i \in \mathcal{R}^{M_1 \times M_2}\) is given by [1]

\[\mathcal{D}^i = (\mathcal{G}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \mathcal{G}^i_<\]

where \(\mathcal{R}^i \in \mathbb{R}^{N \times 4}\) is the coordinate matrix, and each row of \(\mathcal{R}^i\) can be constructed as follows

\[(\mathcal{R}^i)_j = [ \begin{array}{c} s(r_{ji}) & \frac{s(r_{ji})x_{ji}}{r_{ji}} & \frac{s(r_{ji})y_{ji}}{r_{ji}} & \frac{s(r_{ji})z_{ji}}{r_{ji}} \end{array} ]\]

where \(\mathbf{R}_{ji}=\mathbf{R}_j-\mathbf{R}_i = (x_{ji}, y_{ji}, z_{ji})\) is the relative coordinate and \(r_{ji}=\lVert \mathbf{R}_{ji} \lVert\) is its norm. The switching function \(s(r)\) is defined as:

\[\begin{split}s(r)= \begin{cases} \frac{1}{r}, & r<r_s \\ \frac{1}{r} \{ {(\frac{r - r_s}{ r_c - r_s})}^3 (-6 {(\frac{r - r_s}{ r_c - r_s})}^2 +15 \frac{r - r_s}{ r_c - r_s} -10) +1 \}, & r_s \leq r<r_c \\ 0, & r \geq r_c \end{cases}\end{split}\]

Each row of the embedding matrix \(\mathcal{G}^i \in \mathbb{R}^{N \times M_1}\) consists of outputs of a embedding network \(\mathcal{N}\) of \(s(r_{ji})\):

\[(\mathcal{G}^i)_j = \mathcal{N}(s(r_{ji}))\]

\(\mathcal{G}^i_< \in \mathbb{R}^{N \times M_2}\) takes first \(M_2\) columns of \(\mathcal{G}^i\). The equation of embedding network \(\mathcal{N}\) can be found at deepmd.tf.utils.network.embedding_net().

Parameters:

rcut

The cut-off radius \(r_c\)

rcut_smth

From where the environment matrix should be smoothed \(r_s\)

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 \(\mathcal{N}\)

axis_neuron

Number of the axis neuron \(M_2\) (number of columns of the sub-matrix of the embedding matrix)

resnet_dt

Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)

trainable

If the weights of embedding net are trainable.

seed

Random seed for initializing the network parameters.

type_one_side

Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets

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.

set_davg_zero

Set the shift of embedding net input to zero.

activation_function

The activation function in the embedding net. Supported options are “relu”, “tanh”, “none”, “linear”, “softplus”, “sigmoid”, “relu6”, “gelu”, “gelu_tf”.

precision

The precision of the embedding net parameters. Supported options are “float32”, “default”, “float16”, “float64”.

uniform_seed

Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed

multi_task

If the model has multi fitting nets to train.

env_protection: float

Protection parameter to prevent division by zero errors during environment matrix calculations.

References

[1]

Linfeng Zhang, Jiequn Han, Han Wang, Wissam A. Saidi, Roberto Car, and E. Weinan. 2018. End-to-end symmetry preserving inter-atomic potential energy model for finite and extended systems. In Proceedings of the 32nd International Conference on Neural Information Processing Systems (NIPS’18). Curran Associates Inc., Red Hook, NY, USA, 4441-4451.

property explicit_ntypes: bool

Explicit ntypes with type embedding.

get_rcut() → float

Returns the cut-off radius.

get_ntypes() → int

Returns the number of atom types.

get_dim_out() → int

Returns the output dimension of this descriptor.

get_dim_rot_mat_1() → int

Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3.

get_nlist() → Tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, List[int], List[int]]

Returns neighbor information.

Returns:

nlist

Neighbor list

rij

The relative distance between the neighbor and the center atom.

sel_a

The number of neighbors with full information

sel_r

The number of neighbors with only radial information

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, **kwargs) → None

Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.

Parameters:

data_coord

The coordinates. Can be generated by deepmd.tf.model.make_stat_input

data_box

The box. Can be generated by deepmd.tf.model.make_stat_input

data_atype

The atom types. Can be generated by deepmd.tf.model.make_stat_input

natoms_vec

The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.tf.model.make_stat_input

mesh

The mesh for neighbor searching. Can be generated by deepmd.tf.model.make_stat_input

input_dict

Dictionary for additional input

**kwargs

Additional keyword arguments.

merge_input_stats(stat_dict)

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

Parameters:

stat_dict

The dict of statisitcs computed from compute_input_stats, including:

sumr

The sum of radial statisitcs.

suma

The sum of relative coord statisitcs.

sumn

The sum of neighbor numbers.

sumr2

The sum of square of radial statisitcs.

suma2

The sum of square of relative coord statisitcs.

enable_compression(min_nbor_dist: float, graph: deepmd.tf.env.tf.Graph, graph_def: deepmd.tf.env.tf.GraphDef, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') → None

Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.

Parameters:

min_nbor_dist

The nearest distance between atoms

graphtf.Graph

The graph of the model

graph_deftf.GraphDef

The graph_def of the model

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

suffixstr, optional

The suffix of the scope

enable_mixed_precision(mixed_prec: dict | None = None) → None

Reveive the mixed precision setting.

Parameters:

mixed_prec

The mixed precision setting used in the embedding net

build(coord_: deepmd.tf.env.tf.Tensor, atype_: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor, box_: deepmd.tf.env.tf.Tensor, mesh: deepmd.tf.env.tf.Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') → deepmd.tf.env.tf.Tensor

Build the computational graph for the descriptor.

Parameters:

coord_

The coordinate of atoms

atype_

The type of atoms

natoms

The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms

box_tf.Tensor

The box of the system

mesh

For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed.

input_dict

Dictionary for additional inputs

reuse

The weights in the networks should be reused when get the variable.

suffix

Name suffix to identify this descriptor

Returns:

descriptor

The output descriptor

get_rot_mat() → deepmd.tf.env.tf.Tensor

Get rotational matrix.

prod_force_virial(atom_ener: deepmd.tf.env.tf.Tensor, natoms: deepmd.tf.env.tf.Tensor) → Tuple[deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor, deepmd.tf.env.tf.Tensor]

Compute force and virial.

Parameters:

atom_ener

The atomic energy

natoms

The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms

Returns:

force

The force on atoms

virial

The total virial

atom_virial

The atomic virial

_pass_filter(inputs, atype, natoms, input_dict, reuse=None, suffix='', trainable=True)

_compute_dstats_sys_smth(data_coord, data_box, data_atype, natoms_vec, mesh)

_compute_std(sumv2, sumv, sumn)

_concat_type_embedding(xyz_scatter, nframes, natoms, type_embedding)

Concatenate type_embedding of neighbors and xyz_scatter. If not self.type_one_side, concatenate type_embedding of center atoms as well.

Parameters:

xyz_scatter:

shape is [nframes*natoms[0]*self.nnei, 1]

nframes:

shape is []

natoms:

shape is [1+1+self.ntypes]

type_embedding:

shape is [self.ntypes, Y] where Y=jdata[‘type_embedding’][‘neuron’][-1]

Returns:

embedding:

environment of each atom represented by embedding.

_filter_lower(type_i, type_input, start_index, incrs_index, inputs, nframes, natoms, type_embedding=None, is_exclude=False, activation_fn=None, bavg=0.0, stddev=1.0, trainable=True, suffix='')

Input env matrix, returns R.G.

_filter(inputs, type_input, natoms, type_embedding=None, activation_fn=tf.nn.tanh, stddev=1.0, bavg=0.0, name='linear', reuse=None, trainable=True)

init_variables(graph: deepmd.tf.env.tf.Graph, graph_def: deepmd.tf.env.tf.GraphDef, suffix: str = '') → None

Init the embedding net variables with the given dict.

Parameters:

graphtf.Graph

The input frozen model graph

graph_deftf.GraphDef

The input frozen model graph_def

suffixstr, optional

The suffix of the scope

classmethod deserialize(data: dict, suffix: str = '')

Deserialize the model.

Parameters:

datadict

The serialized data

Returns:

Model

The deserialized model

serialize(suffix: str = '') → dict

Serialize the model.

Parameters:

suffixstr, optional

The suffix of the scope

Returns:

dict

The serialized data