diff --git a/deepmd/descriptor/__init__.py b/deepmd/descriptor/__init__.py index 3e2c0e59b0..4a8356b46a 100644 --- a/deepmd/descriptor/__init__.py +++ b/deepmd/descriptor/__init__.py @@ -17,6 +17,9 @@ DescrptSeAEf, DescrptSeAEfLower, ) +from .se_a_mask import ( + DescrptSeAMask, +) from .se_atten import ( DescrptSeAtten, ) diff --git a/deepmd/descriptor/se_a_mask.py b/deepmd/descriptor/se_a_mask.py new file mode 100644 index 0000000000..6e5d686c19 --- /dev/null +++ b/deepmd/descriptor/se_a_mask.py @@ -0,0 +1,433 @@ +import math +import warnings +from typing import ( + Any, + Dict, + List, + Optional, + Tuple, +) + +import numpy as np + +from deepmd.common import ( + get_activation_func, + get_precision, +) +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, + GLOBAL_TF_FLOAT_PRECISION, + default_tf_session_config, + op_module, + tf, +) +from deepmd.utils.argcheck import ( + list_to_doc, +) +from deepmd.utils.graph import ( + get_tensor_by_name_from_graph, + load_graph_def, +) +from deepmd.utils.network import ( + embedding_net, + embedding_net_rand_seed_shift, +) +from deepmd.utils.sess import ( + run_sess, +) +from deepmd.utils.tabulate import ( + DPTabulate, +) +from deepmd.utils.type_embed import ( + embed_atom_type, +) + +from .descriptor import ( + Descriptor, +) +from .se_a import ( + DescrptSeA, +) + + +@Descriptor.register("se_a_mask") +class DescrptSeAMask(DescrptSeA): + r"""DeepPot-SE constructed from all information (both angular and radial) of + atomic configurations. The embedding takes the distance between atoms as input. + + The descriptor :math:`\mathcal{D}^i \in \mathcal{R}^{M_1 \times M_2}` is given by [1]_ + + .. math:: + \mathcal{D}^i = (\mathcal{G}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \mathcal{G}^i_< + + where :math:`\mathcal{R}^i \in \mathbb{R}^{N \times 4}` is the coordinate + matrix, and each row of :math:`\mathcal{R}^i` can be constructed as follows + + .. math:: + (\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 :math:`\mathbf{R}_{ji}=\mathbf{R}_j-\mathbf{R}_i = (x_{ji}, y_{ji}, z_{ji})` is + the relative coordinate and :math:`r_{ji}=\lVert \mathbf{R}_{ji} \lVert` is its norm. + The switching function :math:`s(r)` is defined as: + + .. math:: + s(r)= + \begin{cases} + \frac{1}{r}, & r