add e3nn_use_edge_feat_weights

Author: iProzd

deepmd/dpmodel/descriptor/dpa3.py

@@ -92,6 +92,7 @@ def __init__(
         e3nn_conv_pattern: str = "128x0e+64x1e+32x2e+32x3e",
         use_e3nn_denominator: bool = False,
         e3nn_conv_l_max: int = 3,
+        e3nn_use_edge_feat_weights: bool = False,
     ) -> None:
         r"""The constructor for the RepFlowArgs class which defines the parameters of the repflow block in DPA3 descriptor.
 
@@ -235,6 +236,7 @@ def __init__(
         self.e3nn_conv_pattern = e3nn_conv_pattern
         self.use_e3nn_denominator = use_e3nn_denominator
         self.e3nn_conv_l_max = e3nn_conv_l_max
+        self.e3nn_use_edge_feat_weights = e3nn_use_edge_feat_weights
 
     def __getitem__(self, key):
         if hasattr(self, key):

deepmd/pt/model/descriptor/dpa3.py

@@ -230,6 +230,7 @@ def init_subclass_params(sub_data, sub_class):
             use_e3nn_conv=self.repflow_args.use_e3nn_conv,
             e3nn_conv_pattern=self.repflow_args.e3nn_conv_pattern,
             use_e3nn_denominator=self.repflow_args.use_e3nn_denominator,
+            e3nn_use_edge_feat_weights=self.repflow_args.e3nn_use_edge_feat_weights,
             e3nn_conv_l_max=self.repflow_args.e3nn_conv_l_max,
             exclude_types=exclude_types,
             env_protection=env_protection,

deepmd/pt/model/descriptor/repflow_layer.py

@@ -104,6 +104,7 @@ def __init__(
         dropout_rate: float = 0.1,
         use_e3nn_conv: bool = False,
         e3nn_conv_pattern: str = "128x0e+64x1e+32x2e+32x3e",
+        e3nn_use_edge_feat_weights: bool = False,
         e3nn_conv_args: dict = {},
         activation_function: str = "silu",
         update_style: str = "res_residual",
@@ -387,6 +388,7 @@ def __init__(
         self.use_e3nn_conv = use_e3nn_conv
         self.e3nn_conv_pattern = e3nn_conv_pattern
         self.e3nn_conv_args = e3nn_conv_args
+        self.e3nn_use_edge_feat_weights = e3nn_use_edge_feat_weights
         if self.use_e3nn_conv:
             self.e3nn_conv_block = IrrepsBlock(**self.e3nn_conv_args, weight_layer_act=self.activation_function)
             if self.update_style == "res_residual":
@@ -1450,7 +1452,8 @@ def forward(
             assert edge_sph is not None
             assert edge_rbf_ebd is not None
             assert edge_index is not None
-            node_sph_embed = self.e3nn_conv_block(node_sph_embed, edge_sph, edge_rbf_ebd, edge_index)
+            edge_weights = edge_rbf_ebd if not self.e3nn_use_edge_feat_weights else edge_ebd
+            node_sph_embed = self.e3nn_conv_block(node_sph_embed, edge_sph, edge_weights, edge_index)
             # node_sph_embed = node_sph_embed
             sph_conv_update = node_sph_embed[:, :, :self.n_dim].clone()  # avoid following in-place op
             n_update_list.append(sph_conv_update)

deepmd/pt/model/descriptor/repflows.py

@@ -184,6 +184,7 @@ def __init__(
         e3nn_conv_pattern: str = "128x0e+64x1e+32x2e+32x3e",
         use_e3nn_denominator: bool = False,
         e3nn_conv_l_max: int = 3,
+        e3nn_use_edge_feat_weights: bool = False,
         seed: Optional[Union[int, list[int]]] = None,
     ) -> None:
         r"""
@@ -482,6 +483,7 @@ def __init__(
         self.e3nn_conv_pattern = e3nn_conv_pattern
         self.use_e3nn_denominator = use_e3nn_denominator
         self.e3nn_conv_l_max = e3nn_conv_l_max
+        self.e3nn_use_edge_feat_weights = e3nn_use_edge_feat_weights
 
         if not self.edge_use_esen_rbf:
             self.edge_embd = MLPLayer(
@@ -568,6 +570,7 @@ def __init__(
                 "irreps_out": irreps_out,
                 "denominator": 1.0 if not self.use_e3nn_denominator else self.dynamic_e_sel / 4,
                 "train_denominator": True,
+                "weight_layer_input_to_hidden": [8, 64, 64] if not self.e3nn_use_edge_feat_weights else [self.e_dim],
             }
             irreps_x = irreps_out
             layers.append(
@@ -633,6 +636,7 @@ def __init__(
                     dropout_rate=self.dropout_rate,
                     use_e3nn_conv=self.use_e3nn_conv,
                     e3nn_conv_pattern=self.e3nn_conv_pattern,
+                    e3nn_use_edge_feat_weights=self.e3nn_use_edge_feat_weights,
                     e3nn_conv_args=e3nn_conv_args,
                     seed=child_seed(child_seed(seed, 1), ii),
                 )

deepmd/pt/model/network/e3nn_networks.py

@@ -245,6 +245,7 @@ def __init__(
             weight_layer_act="tanh",
             denominator: float = 1.0,
             train_denominator: bool = False,
+            weight_layer_input_to_hidden: list[int] = [8, 64, 64],
     ) -> None:
         super().__init__()
 
@@ -256,6 +257,7 @@ def __init__(
             weight_layer_act=weight_layer_act,
             denominator=denominator,
             train_denominator=train_denominator,
+            weight_layer_input_to_hidden=weight_layer_input_to_hidden,
         )
 
         # 3. gate

deepmd/utils/argcheck.py

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