add e3nn_angle_use_cross

Author: iProzd

deepmd/dpmodel/descriptor/dpa3.py

@@ -95,6 +95,7 @@ def __init__(
         e3nn_use_edge_feat_weights: bool = False,
         use_e3nn_angle_conv: bool = False,
         e3nn_angle_conv_l_max: int = 2,
+        e3nn_angle_use_cross: bool = False,
     ) -> None:
         r"""The constructor for the RepFlowArgs class which defines the parameters of the repflow block in DPA3 descriptor.
 
@@ -241,6 +242,7 @@ def __init__(
         self.e3nn_use_edge_feat_weights = e3nn_use_edge_feat_weights
         self.use_e3nn_angle_conv = use_e3nn_angle_conv
         self.e3nn_angle_conv_l_max = e3nn_angle_conv_l_max
+        self.e3nn_angle_use_cross = e3nn_angle_use_cross
 
     def __getitem__(self, key):
         if hasattr(self, key):

deepmd/pt/model/descriptor/dpa3.py

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

deepmd/pt/model/descriptor/repflow_layer.py

@@ -110,6 +110,7 @@ def __init__(
         e3nn_angle_conv_args: dict = {},
         use_e3nn_angle_conv: bool = False,
         e3nn_angle_conv_pattern: str = "64x0e+32x1e+32x2e",
+        e3nn_angle_use_cross: bool = False,
         activation_function: str = "silu",
         update_style: str = "res_residual",
         update_residual: float = 0.1,
@@ -571,6 +572,7 @@ def __init__(
             self.use_e3nn_angle_conv = use_e3nn_angle_conv
             self.e3nn_angle_conv_pattern = e3nn_angle_conv_pattern
             self.e3nn_angle_conv_args = e3nn_angle_conv_args
+            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=self.activation_function)
                 if self.update_style == "res_residual":
@@ -1209,6 +1211,7 @@ def forward(
         node_sph_embed: Optional[torch.Tensor] = None,  # nf x nloc x num_sph_node
         edge_angle_filter: Optional[torch.Tensor] = None,  # n_angle x num_sph
         edge_sph_embed: Optional[torch.Tensor] = None,  # n_edge x num_sph
+        angle_weights: Optional[torch.Tensor] = None,  # n_angle x 8
     ):
         """
         Parameters
@@ -1741,8 +1744,12 @@ def forward(
                 assert edge_sph_embed is not None
                 assert edge_angle_filter is not None
                 assert angle_index is not None
-                angle_weights = angle_ebd
-                edge_sph_embed = self.e3nn_angle_conv_block(edge_sph_embed, edge_angle_filter, angle_weights, angle_index, a_sw)
+                if self.e3nn_angle_use_cross:
+                    assert angle_weights is not None
+                    angle_weights_input = angle_weights
+                else:
+                    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
                 edge_sph_conv_update = edge_sph_embed[:, :self.e_dim].clone()  # avoid following in-place op
                 e_update_list.append(edge_sph_conv_update)

deepmd/pt/model/descriptor/repflows.py

@@ -189,6 +189,7 @@ def __init__(
         use_e3nn_angle_conv: bool = False,
         e3nn_angle_conv_l_max: int = 2,
         e3nn_angle_conv_pattern: str = "64x0e+32x1e+32x2e",
+        e3nn_angle_use_cross: bool = False,
         seed: Optional[Union[int, list[int]]] = None,
     ) -> None:
         r"""
@@ -491,6 +492,7 @@ def __init__(
         self.use_e3nn_angle_conv = use_e3nn_angle_conv
         self.e3nn_angle_conv_l_max = e3nn_angle_conv_l_max
         self.e3nn_angle_conv_pattern = e3nn_angle_conv_pattern
+        self.e3nn_angle_use_cross = e3nn_angle_use_cross
 
         if not self.edge_use_esen_rbf:
             self.edge_embd = MLPLayer(
@@ -607,7 +609,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],
+                "weight_layer_input_to_hidden": [8, 64, 64] if not self.e3nn_use_edge_feat_weights else [self.e_dim, 64, 64],
             }
             irreps_x = irreps_out
 
@@ -627,7 +629,7 @@ def __init__(
                 "irreps_out": irreps_edge_out,
                 "denominator": 1.0 if not self.use_e3nn_denominator else self.dynamic_a_sel / 4,
                 "train_denominator": True,
-                "weight_layer_input_to_hidden": [self.a_dim],
+                "weight_layer_input_to_hidden": [8, 64, 64] if self.e3nn_angle_use_cross else [self.a_dim],
             }
             irreps_edge = irreps_edge_out
 
@@ -699,6 +701,7 @@ def __init__(
                     use_e3nn_angle_conv=self.use_e3nn_angle_conv,
                     e3nn_angle_conv_args=e3nn_angle_conv_args,
                     e3nn_angle_conv_pattern=self.e3nn_angle_conv_pattern,
+                    e3nn_angle_use_cross=self.e3nn_angle_use_cross,
                     seed=child_seed(child_seed(seed, 1), ii),
                 )
             )
@@ -1249,15 +1252,28 @@ def forward(
             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_for_angle = self.edge_spherical_embd_for_angle(diff)
+            edge_sph_embed = edge_ebd
             edge_i_index = angle_index[:, 1]
             edge_j_index = angle_index[:, 2]
-            edge_angle_filter_tp = self.edge_to_angle_filter_prod(edge_sph_for_angle[edge_i_index], edge_sph_for_angle[edge_j_index])
-            edge_angle_filter = self.edge_to_angle_filter_linear(edge_angle_filter_tp)
-            edge_sph_embed = edge_ebd
+            if not self.e3nn_angle_use_cross:
+                edge_sph_for_angle = self.edge_spherical_embd_for_angle(diff)
+                edge_angle_filter_tp = self.edge_to_angle_filter_prod(edge_sph_for_angle[edge_i_index], edge_sph_for_angle[edge_j_index])
+                edge_angle_filter = self.edge_to_angle_filter_linear(edge_angle_filter_tp)
+                angle_weights = None
+            else:
+                angle_cross_vec = torch.cross(diff[edge_i_index], diff[edge_j_index], dim=-1)
+                edge_angle_filter = self.edge_spherical_embd_for_angle(angle_cross_vec)
+                # angle basis as weights
+                # 1 - 1e-6 for torch.acos stability
+                cosine_ij = cosine_ij[a_nlist_mask]
+                sine_ij = torch.sqrt(1 - cosine_ij ** 2)
+                theta = torch.acos(cosine_ij).unsqueeze(-1)
+                theta_list = torch.cat([theta*2, theta*4, theta*8], dim=-1)
+                angle_weights = torch.cat([cosine_ij.unsqueeze(-1), torch.cos(theta_list), sine_ij.unsqueeze(-1), torch.sin(theta_list)], dim=-1)
         else:
             edge_angle_filter = None
             edge_sph_embed = None
+            angle_weights = None
 
         for idx, ll in enumerate(self.layers):
             # node_ebd:     nb x nloc x n_dim
@@ -1379,6 +1395,7 @@ def forward(
                         node_sph_embed=node_sph_embed,
                         edge_angle_filter=edge_angle_filter,
                         edge_sph_embed=edge_sph_embed,
+                        angle_weights=angle_weights,
                     )
                 # may cause jit slow, todo fix
                 elif not self.rk_update_diff_layer:

deepmd/utils/argcheck.py

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