add NGA edge

Author: iProzd

deepmd/dpmodel/descriptor/dpa3.py

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

@@ -227,6 +227,7 @@ def init_subclass_params(sub_data, sub_class):
             angle_sh_init_lmax=self.repflow_args.angle_sh_init_lmax,
             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,
             exclude_types=exclude_types,
             env_protection=env_protection,
             precision=precision,

deepmd/pt/model/descriptor/repflow_layer.py

@@ -96,6 +96,7 @@ def __init__(
         edge_rbf_cat_message: bool = False,
         edge_message_use_dropout: bool = False,
         angle_message_use_dropout: bool = False,
+        EN_use_NGA: bool = False,
         dropout_rate: float = 0.1,
         activation_function: str = "silu",
         update_style: str = "res_residual",
@@ -146,6 +147,9 @@ def __init__(
         self.sel_reduce_factor = sel_reduce_factor
         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:
+            assert not self.use_dynamic_sel, "NGA does not support dynamic selection!"
 
         self.update_dihedral = update_dihedral
         self.d_dim = d_dim
@@ -375,6 +379,35 @@ def __init__(
                 )
                 residual_idx += 1
 
+        # node edge NGA
+        if self.EN_use_NGA:
+            self.edge_nga_mlp = MLPLayer(
+                e_dim,
+                n_dim,
+                precision=precision,
+                seed=child_seed(seed, 20),
+            )
+            self.node_nga_mlp = MLPLayer(
+                2 * n_dim,
+                n_dim,
+                precision=precision,
+                seed=child_seed(seed, 21),
+            )
+            if self.update_style == "res_residual":
+                self.n_residual.append(
+                    get_residual(
+                        n_dim,
+                        self.update_residual * self.residual_pref[residual_idx],
+                        self.update_residual_init,
+                        precision=precision,
+                        seed=child_seed(seed, 22),
+                    )
+                )
+                residual_idx += 1
+        else:
+            self.edge_nga_mlp = None
+            self.node_nga_mlp = None
+
         # edge self message
         if not self.use_gated_mlp or self.only_angle_gated_mlp:
             self.edge_self_linear = MLPLayer(
@@ -1330,6 +1363,7 @@ def forward(
         else:
             edge_info = None
             edge_info_ffn = None
+            edge_cat_list = None
 
         # edge message use dropout
         if self.edge_message_use_dropout:
@@ -1411,6 +1445,27 @@ def forward(
                     )
                 )
             n_update_list.append(node_edge_update)
+
+        if self.EN_use_NGA:
+            assert self.node_nga_mlp is not None
+            assert self.edge_nga_mlp is not None
+            assert edge_cat_list is not None
+            # nb, nloc, nnei, n_dim
+            attention_weights_nga_i = self.edge_nga_mlp(edge_cat_list[2])
+            attention_weights_nga_i = (attention_weights_nga_i + 20.0) * sw.unsqueeze(
+                -1
+            ) - 20.0
+            attention_weights_nga_i = torch.softmax(attention_weights_nga_i, dim=-2)
+            # nb, nloc, nnei, n_dim
+            attention_value_nga = self.node_nga_mlp(
+                torch.cat(edge_cat_list[:2], dim=-1)
+            )
+
+            # updated value
+            # nb, nloc, n_dim
+            update_node_nga = (attention_weights_nga_i * attention_value_nga).sum(-2)
+            n_update_list.append(update_node_nga)
+
         # update node_ebd
         n_updated = self.list_update(n_update_list, "node")
 

deepmd/pt/model/descriptor/repflows.py

@@ -171,6 +171,7 @@ def __init__(
         angle_sh_init_lmax: int = 3,
         angle_use_fixed_gaussian: bool = False,
         angle_fixed_gaussian_interpolate: bool = False,
+        EN_use_NGA: bool = False,
         seed: Optional[Union[int, list[int]]] = None,
     ) -> None:
         r"""
@@ -316,6 +317,7 @@ def __init__(
             )
         else:
             self.angle_gaussian_encoder = None
+        self.EN_use_NGA = EN_use_NGA
         self.use_env_envelope = use_env_envelope
         self.use_new_sw = use_new_sw
         self.use_force_embedding = use_force_embedding
@@ -582,6 +584,7 @@ def __init__(
                     edge_rbf_cat_message=self.edge_rbf_cat_message,
                     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,
                     dropout_rate=self.dropout_rate,
                     seed=child_seed(child_seed(seed, 1), ii),
                 )

deepmd/utils/argcheck.py

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