add angle NGAi

Author: iProzd

deepmd/dpmodel/descriptor/dpa3.py

@@ -89,6 +89,7 @@ def __init__(
         angle_use_fixed_gaussian: bool = False,
         angle_fixed_gaussian_interpolate: bool = False,
         EN_use_NGA: bool = False,
+        AE_use_NGA: bool = False,
     ) -> None:
         r"""The constructor for the RepFlowArgs class which defines the parameters of the repflow block in DPA3 descriptor.
 
@@ -226,6 +227,7 @@ def __init__(
         self.angle_use_fixed_gaussian = angle_use_fixed_gaussian
         self.angle_fixed_gaussian_interpolate = angle_fixed_gaussian_interpolate
         self.EN_use_NGA = EN_use_NGA
+        self.AE_use_NGA = AE_use_NGA
         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

@@ -228,6 +228,7 @@ def init_subclass_params(sub_data, sub_class):
             angle_use_fixed_gaussian=self.repflow_args.angle_use_fixed_gaussian,
             angle_fixed_gaussian_interpolate=self.repflow_args.angle_fixed_gaussian_interpolate,
             EN_use_NGA=self.repflow_args.EN_use_NGA,
+            AE_use_NGA=self.repflow_args.AE_use_NGA,
             exclude_types=exclude_types,
             env_protection=env_protection,
             precision=precision,

deepmd/pt/model/descriptor/repflow_layer.py

@@ -97,6 +97,7 @@ def __init__(
         edge_message_use_dropout: bool = False,
         angle_message_use_dropout: bool = False,
         EN_use_NGA: bool = False,
+        AE_use_NGA: bool = False,
         dropout_rate: float = 0.1,
         activation_function: str = "silu",
         update_style: str = "res_residual",
@@ -148,7 +149,8 @@ def __init__(
         self.dynamic_e_sel = self.nnei / self.sel_reduce_factor
         self.dynamic_a_sel = self.a_sel / self.sel_reduce_factor
         self.EN_use_NGA = EN_use_NGA
-        if self.EN_use_NGA:
+        self.AE_use_NGA = AE_use_NGA
+        if self.EN_use_NGA or self.AE_use_NGA:
             assert (
                 not self.use_dynamic_sel and not self.optim_update
             ), "NGA does not support dynamic selection or optim update!"
@@ -575,6 +577,42 @@ def __init__(
                 )
                 residual_idx += 1
 
+            # edge angle NGA
+            if self.AE_use_NGA:
+                self.AE_angle_nga_mlp = MLPLayer(
+                    a_dim,
+                    e_dim,
+                    precision=precision,
+                    seed=child_seed(seed, 24),
+                )
+                self.AE_edge_nga_mlp = MLPLayer(
+                    2 * self.e_a_compress_dim,
+                    e_dim,
+                    precision=precision,
+                    seed=child_seed(seed, 25),
+                )
+                self.AE_edge_nga_mlp_out = MLPLayer(
+                    e_dim,
+                    e_dim,
+                    precision=precision,
+                    seed=child_seed(seed, 26),
+                )
+                if self.update_style == "res_residual":
+                    self.e_residual.append(
+                        get_residual(
+                            e_dim,
+                            self.update_residual * self.residual_pref[residual_idx],
+                            self.update_residual_init,
+                            precision=precision,
+                            seed=child_seed(seed, 27),
+                        )
+                    )
+                    residual_idx += 1
+            else:
+                self.AE_angle_nga_mlp = None
+                self.AE_edge_nga_mlp = None
+                self.AE_edge_nga_mlp_out = None
+
             # angle self message
             if not self.use_gated_mlp:
                 self.angle_self_linear = MLPLayer(
@@ -1455,6 +1493,7 @@ def forward(
                 )
             n_update_list.append(node_edge_update)
 
+        # node edge nga
         if self.EN_use_NGA:
             assert self.node_nga_mlp is not None
             assert self.edge_nga_mlp is not None
@@ -1613,6 +1652,7 @@ def forward(
             else:
                 angle_info = None
                 angle_info_ffn = None
+                angle_info_list = None
 
             # angle message use dropout
             if self.angle_message_use_dropout:
@@ -1725,6 +1765,44 @@ def forward(
                 padding_edge_angle_update = self.EAM_rmsnorm(padding_edge_angle_update)
 
             e_update_list.append(padding_edge_angle_update)
+
+            # edge angle NGA
+            if self.AE_use_NGA:
+                assert self.AE_angle_nga_mlp is not None
+                assert self.AE_edge_nga_mlp is not None
+                assert self.AE_edge_nga_mlp_out is not None
+                assert angle_info_list is not None
+
+                # nb, nloc, a_nnei, a_nnei, e_dim
+                attention_weights_nga_i = self.AE_angle_nga_mlp(angle_info_list[0])
+                attention_weights_nga_i = (
+                    attention_weights_nga_i + 20.0
+                ) * a_sw.unsqueeze(-1).unsqueeze(-1) - 20.0
+                attention_weights_nga_i = torch.softmax(attention_weights_nga_i, dim=-2)
+                # nb, nloc, a_nnei, a_nnei, e_dim
+                attention_value_nga = self.act(self.AE_edge_nga_mlp(angle_info_list[2]))
+
+                # updated value
+                # nb, nloc, a_nnei, e_dim
+                reduce_edge_nga = (attention_weights_nga_i * attention_value_nga).sum(
+                    -2
+                )
+                # nb x nloc x nnei x e_dim
+                padding_edge_nga = torch.concat(
+                    [
+                        reduce_edge_nga,
+                        torch.zeros(
+                            [nb, nloc, self.nnei - self.a_sel, self.e_dim],
+                            dtype=edge_ebd.dtype,
+                            device=edge_ebd.device,
+                        ),
+                    ],
+                    dim=2,
+                )
+                # nb x nloc x nnei x e_dim
+                update_edge_nga = self.act(self.AE_edge_nga_mlp_out(padding_edge_nga))
+                e_update_list.append(update_edge_nga)
+
             # update edge_ebd
             e_updated = self.list_update(e_update_list, "edge")
 

deepmd/pt/model/descriptor/repflows.py

@@ -172,6 +172,7 @@ def __init__(
         angle_use_fixed_gaussian: bool = False,
         angle_fixed_gaussian_interpolate: bool = False,
         EN_use_NGA: bool = False,
+        AE_use_NGA: bool = False,
         seed: Optional[Union[int, list[int]]] = None,
     ) -> None:
         r"""
@@ -318,6 +319,7 @@ def __init__(
         else:
             self.angle_gaussian_encoder = None
         self.EN_use_NGA = EN_use_NGA
+        self.AE_use_NGA = AE_use_NGA
         self.use_env_envelope = use_env_envelope
         self.use_new_sw = use_new_sw
         self.use_force_embedding = use_force_embedding
@@ -585,6 +587,7 @@ def __init__(
                     edge_message_use_dropout=self.edge_message_use_dropout,
                     angle_message_use_dropout=self.angle_message_use_dropout,
                     EN_use_NGA=self.EN_use_NGA,
+                    AE_use_NGA=self.AE_use_NGA,
                     dropout_rate=self.dropout_rate,
                     seed=child_seed(child_seed(seed, 1), ii),
                 )

deepmd/utils/argcheck.py

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