pt(feat): add rbf concat mess

Author: iProzd

deepmd/dpmodel/descriptor/dpa3.py

@@ -57,6 +57,7 @@ def __init__(
         edge_attn_use_ln: bool = True,
         edge_rbf_dot_self: bool = False,
         edge_rbf_dot_message: bool = False,
+        edge_rbf_cat_message: bool = False,
         edge_use_esen_rbf: bool = False,
         edge_use_esen_atom_ebd: bool = False,
         edge_use_esen_env: bool = False,
@@ -208,6 +209,7 @@ def __init__(
         self.angle_self_attention = angle_self_attention
         self.angle_self_attention_gate = angle_self_attention_gate
         self.rmsnorm_mode = rmsnorm_mode
+        self.edge_rbf_cat_message = edge_rbf_cat_message
         assert (
             fix_stat_std == 0.3
         ), "fix_stat_std is not implemented in this version, please use skip_stat instead."

deepmd/pt/model/descriptor/dpa3.py

@@ -217,6 +217,7 @@ def init_subclass_params(sub_data, sub_class):
             angle_self_attention=self.repflow_args.angle_self_attention,
             angle_self_attention_gate=self.repflow_args.angle_self_attention_gate,
             rmsnorm_mode=self.repflow_args.rmsnorm_mode,
+            edge_rbf_cat_message=self.repflow_args.edge_rbf_cat_message,
             exclude_types=exclude_types,
             env_protection=env_protection,
             precision=precision,

deepmd/pt/model/descriptor/repflow_layer.py

@@ -93,6 +93,7 @@ def __init__(
         angle_self_attention: bool = False,
         angle_self_attention_gate: str = "none",
         rmsnorm_mode: str = "none",
+        edge_rbf_cat_message: bool = False,
         activation_function: str = "silu",
         update_style: str = "res_residual",
         update_residual: float = 0.1,
@@ -168,6 +169,7 @@ def __init__(
         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.edge_rbf_cat_message = edge_rbf_cat_message
         self.rbf_dim = rbf_dim
         self.residual_pref = residual_pref
         self.residual_pref += [1.0] * 10
@@ -228,7 +230,11 @@ def __init__(
         self.angle_self_attention = angle_self_attention
         self.angle_self_attention_gate = angle_self_attention_gate
 
-        if self.edge_rbf_dot_self or self.edge_rbf_dot_message:
+        if (
+            self.edge_rbf_dot_self
+            or self.edge_rbf_dot_message
+            or self.edge_rbf_cat_message
+        ):
             self.rbf_mlp = MLPLayer(
                 rbf_dim,
                 self.e_dim,
@@ -264,7 +270,11 @@ def __init__(
         self.e_residual = []
         self.a_residual = []
         self.d_residual = []
-        self.edge_info_dim = self.n_dim * 2 + self.e_dim
+        self.edge_info_dim = (
+            self.n_dim * 2 + self.e_dim
+            if not self.edge_rbf_cat_message
+            else self.n_dim * 2 + self.e_dim * 2
+        )
 
         # node self mlp
         self.node_self_mlp = MLPLayer(
@@ -1198,7 +1208,11 @@ def forward(
         )
 
         # handle edge rbf
-        if self.edge_rbf_dot_self or self.edge_rbf_dot_message:
+        if (
+            self.edge_rbf_dot_self
+            or self.edge_rbf_dot_message
+            or self.edge_rbf_cat_message
+        ):
             assert rbf_ebd is not None
             assert self.rbf_mlp is not None
             edge_rbf = self.rbf_mlp(rbf_ebd)
@@ -1272,26 +1286,26 @@ def forward(
         if not self.optim_update:
             if not self.use_dynamic_sel:
                 # nb x nloc x nnei x (n_dim * 2 + e_dim)
-                edge_info = torch.cat(
-                    [
-                        torch.tile(node_ebd.unsqueeze(-2), [1, 1, self.nnei, 1]),
-                        nei_node_ebd,
-                        edge_ebd,
-                    ],
-                    dim=-1,
-                )
+                edge_cat_list = [
+                    torch.tile(node_ebd.unsqueeze(-2), [1, 1, self.nnei, 1]),
+                    nei_node_ebd,
+                    edge_ebd,
+                ]
+                if self.edge_rbf_cat_message:
+                    assert edge_rbf is not None
+                    edge_cat_list += [edge_rbf]
+                edge_info = torch.cat(edge_cat_list, dim=-1)
             else:
                 # n_edge x (n_dim * 2 + e_dim)
-                edge_info = torch.cat(
-                    [
-                        torch.index_select(
-                            node_ebd.reshape(-1, self.n_dim), 0, n2e_index
-                        ),
-                        nei_node_ebd,
-                        edge_ebd,
-                    ],
-                    dim=-1,
-                )
+                edge_cat_list = [
+                    torch.index_select(node_ebd.reshape(-1, self.n_dim), 0, n2e_index),
+                    nei_node_ebd,
+                    edge_ebd,
+                ]
+                if self.edge_rbf_cat_message:
+                    assert edge_rbf is not None
+                    edge_cat_list += [edge_rbf]
+                edge_info = torch.cat(edge_cat_list, dim=-1)
             if self.use_ffn_node_edge_message or self.use_ffn_edge_edge_message:
                 assert self.edge_message_ffn1 is not None
                 edge_info_ffn = self.act(self.edge_message_ffn1(edge_info))

deepmd/pt/model/descriptor/repflows.py

@@ -161,6 +161,7 @@ def __init__(
         angle_self_attention: bool = False,
         angle_self_attention_gate: str = "none",
         rmsnorm_mode: str = "none",
+        edge_rbf_cat_message: bool = False,
         seed: Optional[Union[int, list[int]]] = None,
     ) -> None:
         r"""
@@ -314,6 +315,7 @@ def __init__(
         self.edge_use_esen_rbf = edge_use_esen_rbf
         self.edge_use_esen_atom_ebd = edge_use_esen_atom_ebd
         self.edge_use_esen_env = edge_use_esen_env
+        self.edge_rbf_cat_message = edge_rbf_cat_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
@@ -333,6 +335,9 @@ def __init__(
         elif self.edge_use_rbf:
             self.rbf = BesselBasis(self.e_rcut)
             self.edge_embed_input_dim = self.rbf.num_basis
+        elif self.edge_rbf_cat_message:
+            # edge can use dist itself
+            self.rbf = BesselBasis(self.e_rcut)
         else:
             self.rbf = None
 
@@ -379,6 +384,11 @@ def __init__(
                 not self.optim_update
             ), "optim_update must be False when angle_use_node is False"
 
+        if self.edge_rbf_cat_message:
+            assert (
+                not self.optim_update
+            ), "optim_update must be False when edge_rbf_cat_message is True"
+
         if self.edge_use_esen_atom_ebd:
             self.source_embedding = torch.nn.Embedding(self.ntypes, self.e_dim)
             self.target_embedding = torch.nn.Embedding(self.ntypes, self.e_dim)
@@ -508,7 +518,9 @@ def __init__(
                     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,
+                    rbf_dim=self.edge_embed_input_dim
+                    if not self.edge_rbf_cat_message
+                    else self.rbf.num_basis,
                     residual_pref=self.residual_pref,
                     message_use_self_concat=self.message_use_self_concat,
                     use_slim_message=self.use_slim_message,
@@ -520,6 +532,7 @@ def __init__(
                     angle_self_attention=self.angle_self_attention,
                     angle_self_attention_gate=self.angle_self_attention_gate,
                     rmsnorm_mode=self.rmsnorm_mode,
+                    edge_rbf_cat_message=self.edge_rbf_cat_message,
                     seed=child_seed(child_seed(seed, 1), ii),
                 )
             )
@@ -934,7 +947,11 @@ def forward(
             edge_ebd = self.edge_embd(rbf_input)
         elif self.edge_use_dist:
             edge_ebd = self.edge_embd(edge_input)
-            rbf_ebd = None
+            if not self.edge_rbf_cat_message:
+                rbf_ebd = None
+            else:
+                assert self.rbf is not None
+                rbf_ebd = self.rbf(edge_input)
         elif self.edge_use_concat_rbf:
             assert self.rbf is not None
             rbf_ebd = torch.cat([dmatrix[..., :1], self.rbf(edge_input)], dim=-1)

deepmd/utils/argcheck.py

@@ -1749,6 +1749,12 @@ def dpa3_repflow_args():
             optional=True,
             default=False,
         ),
+        Argument(
+            "edge_rbf_cat_message",
+            bool,
+            optional=True,
+            default=False,
+        ),
         Argument(
             "edge_use_esen_rbf",
             bool,