feat: add dot rbf

Author: iProzd

deepmd/dpmodel/descriptor/dpa3.py

@@ -54,6 +54,8 @@ def __init__(
         edge_attn_hidden: int = 32,
         edge_attn_head: int = 4,
         edge_attn_use_ln: bool = True,
+        edge_rbf_dot_self: bool = False,
+        edge_rbf_dot_message: bool = False,
     ) -> None:
         r"""The constructor for the RepFlowArgs class which defines the parameters of the repflow block in DPA3 descriptor.
 
@@ -157,6 +159,8 @@ def __init__(
         self.edge_attn_hidden = edge_attn_hidden
         self.edge_attn_head = edge_attn_head
         self.edge_attn_use_ln = edge_attn_use_ln
+        self.edge_rbf_dot_self = edge_rbf_dot_self
+        self.edge_rbf_dot_message = edge_rbf_dot_message
 
     def __getitem__(self, key):
         if hasattr(self, key):

deepmd/pt/model/descriptor/dpa3.py

@@ -189,6 +189,8 @@ def init_subclass_params(sub_data, sub_class):
             edge_attn_hidden=self.repflow_args.edge_attn_hidden,
             edge_attn_head=self.repflow_args.edge_attn_head,
             edge_attn_use_ln=self.repflow_args.edge_attn_use_ln,
+            edge_rbf_dot_self=self.repflow_args.edge_rbf_dot_self,
+            edge_rbf_dot_message=self.repflow_args.edge_rbf_dot_message,
             exclude_types=exclude_types,
             env_protection=env_protection,
             precision=precision,

deepmd/pt/model/descriptor/repflow_layer.py

@@ -77,6 +77,9 @@ def __init__(
         edge_attn_hidden: int = 32,
         edge_attn_head: int = 4,
         edge_attn_use_ln: bool = True,
+        edge_rbf_dot_self: bool = False,
+        edge_rbf_dot_message: bool = False,
+        rbf_dim: int = 8,
         activation_function: str = "silu",
         update_style: str = "res_residual",
         update_residual: float = 0.1,
@@ -152,6 +155,30 @@ def __init__(
         self.edge_attn_hidden = edge_attn_hidden
         self.edge_attn_head = edge_attn_head
         self.edge_attn_use_ln = edge_attn_use_ln
+        self.edge_rbf_dot_self = edge_rbf_dot_self
+        self.edge_rbf_dot_message = edge_rbf_dot_message
+        self.rbf_dim = rbf_dim
+
+        if self.edge_rbf_dot_self or self.edge_rbf_dot_message:
+            self.rbf_mlp = MLPLayer(
+                rbf_dim,
+                self.e_dim,
+                precision=precision,
+                seed=child_seed(seed, 30),
+            )
+        else:
+            self.rbf_mlp = None
+
+        if self.edge_rbf_dot_message:
+            self.rbf_mlp_message = MLPLayer(
+                rbf_dim,
+                self.n_dim,
+                precision=precision,
+                seed=child_seed(seed, 31),
+            )
+        else:
+            self.rbf_mlp_message = None
+
         if self.edge_use_attn:
             assert (
                 not self.use_dynamic_sel
@@ -889,6 +916,7 @@ def forward(
         dihedral_index: Optional[torch.Tensor] = None,  # n_dihedral x 2
         dihedral_ebd: Optional[torch.Tensor] = None,  # n_dihedral x d_dim
         d_sw: Optional[torch.Tensor] = None,  # n_dihedral
+        rbf_ebd: Optional[torch.Tensor] = None,  # n_edge x num_b
     ):
         """
         Parameters
@@ -962,6 +990,25 @@ def forward(
             )
         )
 
+        # handle edge rbf
+        if self.edge_rbf_dot_self or self.edge_rbf_dot_message:
+            assert rbf_ebd is not None
+            assert self.rbf_mlp is not None
+            edge_rbf = self.rbf_mlp(rbf_ebd)
+        else:
+            edge_rbf = None
+
+        if self.edge_rbf_dot_message:
+            assert rbf_ebd is not None
+            assert self.rbf_mlp_message is not None
+            edge_rbf_node = self.rbf_mlp_message(rbf_ebd)
+        else:
+            edge_rbf_node = None
+
+        if self.edge_rbf_dot_self:
+            assert edge_rbf is not None
+            edge_ebd = edge_ebd * edge_rbf
+
         n_update_list: list[torch.Tensor] = [node_ebd]
         e_update_list: list[torch.Tensor] = [edge_ebd]
         a_update_list: list[torch.Tensor] = [angle_ebd]
@@ -1079,6 +1126,9 @@ def forward(
                     "node",
                 )
             ) * sw.unsqueeze(-1)
+        if self.edge_rbf_dot_message:
+            assert edge_rbf_node is not None
+            node_edge_update = node_edge_update * edge_rbf_node
         node_edge_update = (
             (torch.sum(node_edge_update, dim=-2) / self.nnei)
             if not self.use_dynamic_sel
@@ -1132,6 +1182,9 @@ def forward(
                     "edge",
                 )
             )
+        if self.edge_rbf_dot_message:
+            assert edge_rbf is not None
+            edge_self_update = edge_self_update * edge_rbf
         e_update_list.append(edge_self_update)
 
         # edge attention message

deepmd/pt/model/descriptor/repflows.py

@@ -130,6 +130,8 @@ def __init__(
         edge_attn_hidden: int = 32,
         edge_attn_head: int = 4,
         edge_attn_use_ln: bool = True,
+        edge_rbf_dot_self: bool = False,
+        edge_rbf_dot_message: bool = False,
         optim_update: bool = True,
         seed: Optional[Union[int, list[int]]] = None,
     ) -> None:
@@ -275,6 +277,10 @@ def __init__(
         self.edge_attn_hidden = edge_attn_hidden
         self.edge_attn_head = edge_attn_head
         self.edge_attn_use_ln = edge_attn_use_ln
+        self.edge_rbf_dot_self = edge_rbf_dot_self
+        self.edge_rbf_dot_message = edge_rbf_dot_message
+        if self.edge_rbf_dot_self or self.edge_rbf_dot_message:
+            assert self.edge_use_rbf or self.edge_use_concat_rbf, "rbf is not used"
         self.edge_embed_input_dim = 1
         if self.edge_use_concat_rbf:
             self.rbf = BesselBasis(self.e_rcut)
@@ -370,6 +376,9 @@ def __init__(
                     edge_attn_hidden=self.edge_attn_hidden,
                     edge_attn_head=self.edge_attn_head,
                     edge_attn_use_ln=self.edge_attn_use_ln,
+                    edge_rbf_dot_self=self.edge_rbf_dot_self,
+                    edge_rbf_dot_message=self.edge_rbf_dot_message,
+                    rbf_dim=self.edge_embed_input_dim,
                     seed=child_seed(child_seed(seed, 1), ii),
                 )
             )
@@ -687,15 +696,17 @@ def forward(
         # nb x nloc x nnei x e_dim [OR] n_edge x e_dim
         if self.edge_use_dist:
             edge_ebd = self.edge_embd(edge_input)
+            rbf_ebd = None
         elif self.edge_use_concat_rbf:
             assert self.rbf is not None
-            edge_ebd = self.edge_embd(
-                torch.cat([dmatrix[..., :1], self.rbf(edge_input)], dim=-1)
-            )
+            rbf_ebd = torch.cat([dmatrix[..., :1], self.rbf(edge_input)], dim=-1)
+            edge_ebd = self.edge_embd(rbf_ebd)
         elif self.edge_use_rbf:
             assert self.rbf is not None
-            edge_ebd = self.edge_embd(self.rbf(edge_input))
+            rbf_ebd = self.rbf(edge_input)
+            edge_ebd = self.edge_embd(rbf_ebd)
         else:
+            rbf_ebd = None
             edge_ebd = self.act(self.edge_embd(edge_input))
 
         # nf x nloc x a_nnei x a_nnei x a_dim [OR] n_angle x a_dim
@@ -794,6 +805,7 @@ def forward(
                 dihedral_index=dihedral_index,
                 dihedral_ebd=dihedral_ebd,
                 d_sw=d_sw,
+                rbf_ebd=rbf_ebd,
             )
 
         # nb x nloc x 3 x e_dim

deepmd/utils/argcheck.py

@@ -1724,6 +1724,18 @@ def dpa3_repflow_args():
             optional=True,
             default=True,
         ),
+        Argument(
+            "edge_rbf_dot_self",
+            bool,
+            optional=True,
+            default=False,
+        ),
+        Argument(
+            "edge_rbf_dot_message",
+            bool,
+            optional=True,
+            default=False,
+        ),
     ]