feat(pt): add angle_use_node

Author: iProzd

deepmd/dpmodel/descriptor/dpa3.py

@@ -74,6 +74,7 @@ def __init__(
         use_rk_update: bool = False,
         rk_order: int = 4,
         rk_update_diff_layer: bool = False,
+        angle_use_node: bool = True,
     ) -> None:
         r"""The constructor for the RepFlowArgs class which defines the parameters of the repflow block in DPA3 descriptor.
 
@@ -197,6 +198,7 @@ def __init__(
         self.use_rk_update = use_rk_update
         self.rk_order = rk_order
         self.rk_update_diff_layer = rk_update_diff_layer
+        self.angle_use_node = angle_use_node
 
     def __getitem__(self, key):
         if hasattr(self, key):

deepmd/pt/model/descriptor/dpa3.py

@@ -211,6 +211,7 @@ def init_subclass_params(sub_data, sub_class):
             use_rk_update=self.repflow_args.use_rk_update,
             rk_order=self.repflow_args.rk_order,
             rk_update_diff_layer=self.repflow_args.rk_update_diff_layer,
+            angle_use_node=self.repflow_args.angle_use_node,
             use_loc_mapping=use_loc_mapping,
             exclude_types=exclude_types,
             env_protection=env_protection,

deepmd/pt/model/descriptor/repflow_layer.py

@@ -88,6 +88,7 @@ def __init__(
         use_gated_mlp: bool = False,
         gated_mlp_norm: str = "none",
         node_use_rmsnorm: bool = False,
+        angle_use_node: bool = True,
         activation_function: str = "silu",
         update_style: str = "res_residual",
         update_residual: float = 0.1,
@@ -184,6 +185,8 @@ def __init__(
         else:
             self.node_rmsnorm = None
 
+        self.angle_use_node = angle_use_node
+
         if self.edge_rbf_dot_self or self.edge_rbf_dot_message:
             self.rbf_mlp = MLPLayer(
                 rbf_dim,
@@ -380,16 +383,22 @@ def __init__(
             self.angle_dim = self.a_dim
             if self.a_compress_rate == 0:
                 # angle + node + edge * 2
-                self.angle_dim += self.n_dim + 2 * self.e_dim
+                self.angle_dim += (
+                    self.n_dim + 2 * self.e_dim
+                    if self.angle_use_node
+                    else 2 * self.e_dim
+                )
                 self.a_compress_n_linear = None
                 self.a_compress_e_linear = None
                 self.e_a_compress_dim = e_dim
                 self.n_a_compress_dim = n_dim
             else:
                 # angle + a_dim/c + a_dim/2c * 2 * e_rate
-                self.angle_dim += (1 + self.a_compress_e_rate) * (
-                    self.a_dim // self.a_compress_rate
-                )
+                self.angle_dim += (
+                    (1 + self.a_compress_e_rate)
+                    if self.angle_use_node
+                    else self.a_compress_e_rate
+                ) * (self.a_dim // self.a_compress_rate)
                 self.e_a_compress_dim = (
                     self.a_dim // (2 * self.a_compress_rate) * self.a_compress_e_rate
                 )
@@ -1383,20 +1392,6 @@ def forward(
                     a_nlist_mask.unsqueeze(-1), edge_ebd_for_angle, 0.0
                 )
             if not self.optim_update:
-                # nb x nloc x a_nnei x a_nnei x n_dim [OR] n_angle x n_dim
-                node_for_angle_info = (
-                    torch.tile(
-                        node_ebd_for_angle.unsqueeze(2).unsqueeze(2),
-                        (1, 1, self.a_sel, self.a_sel, 1),
-                    )
-                    if not self.use_dynamic_sel
-                    else torch.index_select(
-                        node_ebd_for_angle.reshape(-1, self.n_a_compress_dim),
-                        0,
-                        n2a_index,
-                    )
-                )
-
                 # nb x nloc x (a_nnei) x a_nnei x e_dim [OR] n_angle x e_dim
                 edge_for_angle_k = (
                     torch.tile(
@@ -1418,7 +1413,21 @@ def forward(
                     [edge_for_angle_k, edge_for_angle_j], dim=-1
                 )
                 angle_info_list = [angle_ebd]
-                angle_info_list.append(node_for_angle_info)
+                if self.angle_use_node:
+                    # nb x nloc x a_nnei x a_nnei x n_dim [OR] n_angle x n_dim
+                    node_for_angle_info = (
+                        torch.tile(
+                            node_ebd_for_angle.unsqueeze(2).unsqueeze(2),
+                            (1, 1, self.a_sel, self.a_sel, 1),
+                        )
+                        if not self.use_dynamic_sel
+                        else torch.index_select(
+                            node_ebd_for_angle.reshape(-1, self.n_a_compress_dim),
+                            0,
+                            n2a_index,
+                        )
+                    )
+                    angle_info_list.append(node_for_angle_info)
                 angle_info_list.append(edge_for_angle_info)
                 # nb x nloc x a_nnei x a_nnei x (a + n_dim + e_dim*2) or (a + a/c + a/c)
                 # [OR]

deepmd/pt/model/descriptor/repflows.py

@@ -155,6 +155,7 @@ def __init__(
         use_rk_update: bool = False,
         rk_order: int = 4,
         rk_update_diff_layer: bool = False,
+        angle_use_node: bool = True,
         optim_update: bool = True,
         seed: Optional[Union[int, list[int]]] = None,
     ) -> None:
@@ -367,6 +368,12 @@ def __init__(
 
         self.node_use_rmsnorm = node_use_rmsnorm
 
+        self.angle_use_node = angle_use_node
+        if not self.angle_use_node:
+            assert (
+                not self.optim_update
+            ), "optim_update must be False when angle_use_node is False"
+
         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)
@@ -490,6 +497,7 @@ def __init__(
                     use_gated_mlp=self.use_gated_mlp,
                     gated_mlp_norm=self.gated_mlp_norm,
                     node_use_rmsnorm=self.node_use_rmsnorm,
+                    angle_use_node=self.angle_use_node,
                     seed=child_seed(child_seed(seed, 1), ii),
                 )
             )

deepmd/utils/argcheck.py

@@ -1857,6 +1857,12 @@ def dpa3_repflow_args():
             optional=True,
             default=False,
         ),
+        Argument(
+            "angle_use_node",
+            bool,
+            optional=True,
+            default=True,
+        ),
     ]