add e3nn_conv_use_edge_sh_feat

Author: iProzd

deepmd/dpmodel/descriptor/dpa3.py

@@ -97,6 +97,7 @@ def __init__(
         e3nn_angle_conv_l_max: int = 2,
         e3nn_angle_use_cross: bool = False,
         e3nn_angle_only_single_angle: bool = False,
+        e3nn_conv_use_edge_sh_feat: bool = False,
     ) -> None:
         r"""The constructor for the RepFlowArgs class which defines the parameters of the repflow block in DPA3 descriptor.
 
@@ -245,6 +246,7 @@ def __init__(
         self.e3nn_angle_conv_l_max = e3nn_angle_conv_l_max
         self.e3nn_angle_use_cross = e3nn_angle_use_cross
         self.e3nn_angle_only_single_angle = e3nn_angle_only_single_angle
+        self.e3nn_conv_use_edge_sh_feat = e3nn_conv_use_edge_sh_feat
 
     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_conv_use_edge_sh_feat=self.repflow_args.e3nn_conv_use_edge_sh_feat,
             use_e3nn_angle_conv=self.repflow_args.use_e3nn_angle_conv,
             e3nn_angle_conv_l_max=self.repflow_args.e3nn_angle_conv_l_max,
             e3nn_use_edge_feat_weights=self.repflow_args.e3nn_use_edge_feat_weights,

deepmd/pt/model/descriptor/repflow_layer.py

@@ -106,6 +106,7 @@ def __init__(
         use_e3nn_conv: bool = False,
         e3nn_conv_pattern: str = "128x0e+64x1e+32x2e+32x3e",
         e3nn_use_edge_feat_weights: bool = False,
+        e3nn_conv_use_edge_sh_feat: bool = False,
         e3nn_conv_args: dict = {},
         e3nn_angle_conv_args: dict = {},
         use_e3nn_angle_conv: bool = False,
@@ -394,8 +395,9 @@ def __init__(
         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
+        self.e3nn_conv_use_edge_sh_feat = e3nn_conv_use_edge_sh_feat
         if self.use_e3nn_conv:
-            self.e3nn_conv_block = IrrepsBlock(**self.e3nn_conv_args, weight_layer_act="silu")
+            self.e3nn_conv_block = IrrepsBlock(**self.e3nn_conv_args, e3nn_conv_use_edge_sh_feat=e3nn_conv_use_edge_sh_feat, weight_layer_act="silu")
             if self.update_style == "res_residual":
                 self.n_residual.append(
                     get_residual(
@@ -575,6 +577,10 @@ def __init__(
             self.e3nn_angle_use_cross = e3nn_angle_use_cross
             if self.use_e3nn_angle_conv:
                 self.e3nn_angle_conv_block = IrrepsAngleBlock(**self.e3nn_angle_conv_args, weight_layer_act="silu")
+            else:
+                self.e3nn_angle_conv_block = None
+
+            if self.e3nn_conv_use_edge_sh_feat or self.use_e3nn_angle_conv:
                 if self.update_style == "res_residual":
                     self.e_residual.append(
                         get_residual(
@@ -586,8 +592,7 @@ def __init__(
                         )
                     )
                     residual_idx += 1
-            else:
-                self.e3nn_angle_conv_block = None
+
 
             # angle self message
             if not self.use_gated_mlp:
@@ -1482,7 +1487,11 @@ def forward(
             assert edge_rbf_ebd is not None
             assert edge_index is not None
             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, edge_sph_update = self.e3nn_conv_block(node_sph_embed, edge_sph, edge_weights, edge_index, edge_sph_embed)
+
+            if self.e3nn_conv_use_edge_sh_feat:
+                assert edge_sph_embed is not None
+                edge_sph_embed = edge_sph_embed + 0.1 * edge_sph_update
             # 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)
@@ -1751,14 +1760,17 @@ def forward(
                     angle_weights_input = angle_ebd
                 edge_sph_embed = self.e3nn_angle_conv_block(edge_sph_embed, edge_angle_filter, angle_weights_input, angle_index, a_sw)
                 # node_sph_embed = node_sph_embed
+
+            if self.use_e3nn_angle_conv or self.e3nn_conv_use_edge_sh_feat:
+                assert edge_sph_embed is not None
                 edge_sph_conv_update = edge_sph_embed[:, :self.e_dim].clone()  # avoid following in-place op
                 e_update_list.append(edge_sph_conv_update)
 
             # update edge_ebd
             e_updated = self.list_update(e_update_list, "edge")
 
             # edge angle e3nn joint update
-            if self.use_e3nn_angle_conv:
+            if self.use_e3nn_angle_conv or self.e3nn_conv_use_edge_sh_feat:
                 assert edge_sph_embed is not None
                 edge_sph_embed[:, : self.e_dim] = e_updated
 

deepmd/pt/model/descriptor/repflows.py

@@ -7,7 +7,7 @@
 
 import torch
 
-from e3nn.o3 import Irreps, Linear
+from e3nn.o3 import Irreps, Linear, TensorProduct
 from e3nn.o3 import FullyConnectedTensorProduct as FCTP
 import sevenn.util as util
 from sevenn.nn.edge_embedding import (
@@ -184,6 +184,7 @@ def __init__(
         use_e3nn_conv: bool = False,
         e3nn_conv_pattern: str = "128x0e+64x1e+32x2e+32x3e",
         use_e3nn_denominator: bool = False,
+        e3nn_conv_use_edge_sh_feat: bool = False,
         e3nn_conv_l_max: int = 3,
         e3nn_use_edge_feat_weights: bool = False,
         use_e3nn_angle_conv: bool = False,
@@ -488,6 +489,7 @@ def __init__(
         self.use_e3nn_conv = use_e3nn_conv
         self.e3nn_conv_pattern = e3nn_conv_pattern
         self.use_e3nn_denominator = use_e3nn_denominator
+        self.e3nn_conv_use_edge_sh_feat = e3nn_conv_use_edge_sh_feat
         self.e3nn_conv_l_max = e3nn_conv_l_max
         self.e3nn_use_edge_feat_weights = e3nn_use_edge_feat_weights
         self.use_e3nn_angle_conv = use_e3nn_angle_conv
@@ -564,10 +566,42 @@ def __init__(
             self.edge_rbf_embed = None
             self.edge_env = None
             self.edge_spherical_embd = None
-            self.irreps_filter = "0e"
+            self.irreps_filter = Irreps("0e")
+
+        if self.e3nn_conv_use_edge_sh_feat:
+            irreps_edge = Irreps(f"{self.e_dim}x0e")
+            irreps_filter = self.irreps_filter
+            irreps_out = ((irreps_filter * self.e_dim).sort()[0]).simplify()
+            self.edge_sph_embd_ir = Irreps(self.e3nn_angle_conv_pattern)
+            instructions = []
+            irreps_mid = []
+            weight_numel = 0
+            for i, (mul_x, ir_x) in enumerate(irreps_edge):
+                for j, (_, ir_filter) in enumerate(irreps_filter):
+                    for ir_out in ir_x * ir_filter:
+                        if ir_out in irreps_out:  # here we drop l > lmax
+                            k = len(irreps_mid)
+                            weight_numel += mul_x * 1  # path shape
+                            irreps_mid.append((mul_x, ir_out))
+                            instructions.append((i, j, k, 'uvu', False))
+
+            irreps_mid = Irreps(irreps_mid)
+            self.edge_sph_embd_init_tp = TensorProduct(
+                irreps_in1=irreps_edge,
+                irreps_in2=irreps_filter,
+                irreps_out=irreps_mid,
+                instructions=instructions,
+                shared_weights=False,
+                internal_weights=False,
+            )
+            self.edge_sph_embd_init_linear = Linear(irreps_mid, self.edge_sph_embd_ir)
+        else:
+            self.edge_sph_embd_init_tp = None
+            self.edge_sph_embd_ir = Irreps(f'{self.e_dim}x0e')
+            self.edge_sph_embd_init_linear = None
 
         # for edge angle e3nn conv
-        irreps_edge = Irreps(f'{self.e_dim}x0e')
+        irreps_edge = self.edge_sph_embd_ir
         self.angle_lmax = e3nn_angle_conv_l_max
         self.e3nn_angle_conv_pattern = Irreps("+".join(e3nn_angle_conv_pattern.split("+")[:self.angle_lmax+1]))
         if self.use_e3nn_angle_conv:
@@ -699,6 +733,7 @@ def __init__(
                     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_use_edge_sh_feat=self.e3nn_conv_use_edge_sh_feat,
                     e3nn_conv_args=e3nn_conv_args,
                     use_e3nn_angle_conv=self.use_e3nn_angle_conv,
                     e3nn_angle_conv_args=e3nn_angle_conv_args,
@@ -1249,12 +1284,23 @@ def forward(
             edge_sph = None
             node_sph_embed = None
 
+        if not self.e3nn_conv_use_edge_sh_feat:
+            if self.use_e3nn_angle_conv:
+                edge_sph_embed = edge_ebd
+            else:
+                edge_sph_embed = None
+        else:
+            assert edge_sph is not None
+            assert self.edge_sph_embd_init_tp is not None
+            assert self.edge_sph_embd_init_linear is not None
+            edge_sph_embed = self.edge_sph_embd_init_tp(edge_ebd, edge_sph)
+            edge_sph_embed = self.edge_sph_embd_init_linear(edge_sph_embed)
+
         if self.use_e3nn_angle_conv:
             assert self.use_dynamic_sel, "e3nn conv must use dynamic sel"
             assert self.edge_spherical_embd_for_angle is not None
             assert self.edge_to_angle_filter_prod is not None
             assert self.edge_to_angle_filter_linear is not None
-            edge_sph_embed = edge_ebd
             edge_i_index = angle_index[:, 1]
             edge_j_index = angle_index[:, 2]
             if not self.e3nn_angle_use_cross:
@@ -1277,7 +1323,6 @@ def forward(
                     angle_weights = torch.cat([cosine_ij.unsqueeze(-1), sine_ij.unsqueeze(-1)], dim=-1)
         else:
             edge_angle_filter = None
-            edge_sph_embed = None
             angle_weights = None
 
         for idx, ll in enumerate(self.layers):