add use_e3nn_angle_conv

Author: iProzd

deepmd/dpmodel/descriptor/dpa3.py

@@ -93,6 +93,8 @@ def __init__(
         use_e3nn_denominator: bool = False,
         e3nn_conv_l_max: int = 3,
         e3nn_use_edge_feat_weights: bool = False,
+        use_e3nn_angle_conv: bool = False,
+        e3nn_angle_conv_l_max: int = 2,
     ) -> None:
         r"""The constructor for the RepFlowArgs class which defines the parameters of the repflow block in DPA3 descriptor.
 
@@ -237,6 +239,8 @@ def __init__(
         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
+        self.use_e3nn_angle_conv = use_e3nn_angle_conv
+        self.e3nn_angle_conv_l_max = e3nn_angle_conv_l_max
 
     def __getitem__(self, key):
         if hasattr(self, key):

deepmd/pt/model/descriptor/dpa3.py

@@ -230,6 +230,8 @@ 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,
+            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,
             e3nn_conv_l_max=self.repflow_args.e3nn_conv_l_max,
             exclude_types=exclude_types,

deepmd/pt/model/descriptor/repflow_layer.py

@@ -44,6 +44,7 @@
 )
 from deepmd.pt.model.network.e3nn_networks import (
     IrrepsBlock,
+    IrrepsAngleBlock,
 )
 
 
@@ -106,6 +107,9 @@ def __init__(
         e3nn_conv_pattern: str = "128x0e+64x1e+32x2e+32x3e",
         e3nn_use_edge_feat_weights: bool = False,
         e3nn_conv_args: dict = {},
+        e3nn_angle_conv_args: dict = {},
+        use_e3nn_angle_conv: bool = False,
+        e3nn_angle_conv_pattern: str = "64x0e+32x1e+32x2e",
         activation_function: str = "silu",
         update_style: str = "res_residual",
         update_residual: float = 0.1,
@@ -563,6 +567,26 @@ def __init__(
                 )
                 residual_idx += 1
 
+            # for edge angle e3nn conv
+            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
+            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":
+                    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, 22),
+                        )
+                    )
+                    residual_idx += 1
+            else:
+                self.e3nn_angle_conv_block = None
+
             # angle self message
             if not self.use_gated_mlp:
                 self.angle_self_linear = MLPLayer(
@@ -1183,6 +1207,8 @@ def forward(
         edge_rbf_ebd: Optional[torch.Tensor] = None,  # n_edge x num_basis
         edge_sph: Optional[torch.Tensor] = None,  # n_edge x num_sph
         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
     ):
         """
         Parameters
@@ -1708,9 +1734,27 @@ def forward(
                 padding_edge_angle_update = self.EAM_rmsnorm(padding_edge_angle_update)
 
             e_update_list.append(padding_edge_angle_update)
+
+            # for edge angle e3nn conv
+            if self.use_e3nn_angle_conv:
+                assert self.e3nn_angle_conv_block is not None
+                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)
+                # 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)
+
             # update edge_ebd
             e_updated = self.list_update(e_update_list, "edge")
 
+            # edge angle e3nn joint update
+            if self.use_e3nn_angle_conv:
+                assert edge_sph_embed is not None
+                edge_sph_embed[:, : self.e_dim] = e_updated
+
             # angle self message
             # nb x nloc x a_nnei x a_nnei x dim_a
             if not self.optim_update:
@@ -1927,7 +1971,7 @@ def forward(
 
         # update angle_ebd
         a_updated = self.list_update(a_update_list, "angle")
-        return n_updated, e_updated, a_updated, d_updated, node_sph_embed
+        return n_updated, e_updated, a_updated, d_updated, node_sph_embed, edge_sph_embed
 
     @torch.jit.export
     def list_update_res_avg(

deepmd/pt/model/descriptor/repflows.py

@@ -7,7 +7,8 @@
 
 import torch
 
-from e3nn.o3 import Irreps
+from e3nn.o3 import Irreps, Linear
+from e3nn.o3 import FullyConnectedTensorProduct as FCTP
 import sevenn.util as util
 from sevenn.nn.edge_embedding import (
     SphericalEncoding,
@@ -185,6 +186,9 @@ def __init__(
         use_e3nn_denominator: bool = False,
         e3nn_conv_l_max: int = 3,
         e3nn_use_edge_feat_weights: bool = False,
+        use_e3nn_angle_conv: bool = False,
+        e3nn_angle_conv_l_max: int = 2,
+        e3nn_angle_conv_pattern: str = "64x0e+32x1e+32x2e",
         seed: Optional[Union[int, list[int]]] = None,
     ) -> None:
         r"""
@@ -484,6 +488,9 @@ def __init__(
         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
+        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
 
         if not self.edge_use_esen_rbf:
             self.edge_embd = MLPLayer(
@@ -540,6 +547,7 @@ def __init__(
             self.force_embedding_linear = None
 
         layers = []
+        # for node edge e3nn conv
         irreps_x = Irreps(f'{self.n_dim}x0e')
         self.lmax = e3nn_conv_l_max
         self.e3nn_conv_pattern = Irreps("+".join(e3nn_conv_pattern.split("+")[:self.lmax+1]))
@@ -554,7 +562,36 @@ def __init__(
             self.edge_spherical_embd = None
             self.irreps_filter = "0e"
 
+        # for edge angle e3nn conv
+        irreps_edge = Irreps(f'{self.e_dim}x0e')
+        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:
+            self.edge_spherical_embd_for_angle = SphericalEncoding(self.angle_lmax, parity=1, normalize=True)
+            self.irreps_angle_filter = self.edge_spherical_embd_for_angle.irreps_out
+            self.angle_edge_filter_out = util.infer_irreps_out(
+                self.irreps_angle_filter,  # type: ignore
+                self.irreps_angle_filter,
+                self.angle_lmax,  # type: ignore
+                'full',
+                False,
+            )
+            # edge_i x edge_j --> linear --> angle_filter
+            self.edge_to_angle_filter_prod = FCTP(self.irreps_angle_filter, self.irreps_angle_filter, self.angle_edge_filter_out)
+            self.edge_to_angle_filter_linear = Linear(
+                irreps_in=self.angle_edge_filter_out,
+                irreps_out=self.irreps_angle_filter,
+                biases=False,
+            )
+        else:
+            self.edge_spherical_embd_for_angle = None
+            self.irreps_angle_filter = "0e"
+            self.edge_to_angle_filter_prod = None
+            self.edge_to_angle_filter_linear = None
+
+
         for ii in range(nlayers):
+            # for node edge e3nn conv
             irreps_out = Irreps(self.e3nn_conv_pattern)
             irreps_out_tp = util.infer_irreps_out(
                 irreps_x,  # type: ignore
@@ -573,6 +610,27 @@ def __init__(
                 "weight_layer_input_to_hidden": [8, 64, 64] if not self.e3nn_use_edge_feat_weights else [self.e_dim],
             }
             irreps_x = irreps_out
+
+            # for edge angle e3nn conv
+            irreps_edge_out = Irreps(self.e3nn_angle_conv_pattern)
+            irreps_out_tp = util.infer_irreps_out(
+                irreps_edge,  # type: ignore
+                self.irreps_angle_filter,
+                irreps_edge_out.lmax,  # type: ignore
+                'full',
+                False,
+            )
+            e3nn_angle_conv_args = {
+                "irreps_x": irreps_edge,
+                "irreps_filter": self.irreps_angle_filter,
+                "irreps_out_tp": irreps_out_tp,
+                "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],
+            }
+            irreps_edge = irreps_edge_out
+
             layers.append(
                 RepFlowLayer(
                     e_rcut=self.e_rcut,
@@ -638,6 +696,9 @@ def __init__(
                     e3nn_conv_pattern=self.e3nn_conv_pattern,
                     e3nn_use_edge_feat_weights=self.e3nn_use_edge_feat_weights,
                     e3nn_conv_args=e3nn_conv_args,
+                    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,
                     seed=child_seed(child_seed(seed, 1), ii),
                 )
             )
@@ -1183,6 +1244,21 @@ def forward(
             edge_sph = None
             node_sph_embed = None
 
+        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_for_angle = self.edge_spherical_embd_for_angle(diff)
+            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
+        else:
+            edge_angle_filter = None
+            edge_sph_embed = None
+
         for idx, ll in enumerate(self.layers):
             # node_ebd:     nb x nloc x n_dim
             # node_ebd_ext: nb x nall x n_dim [OR] nb x nloc x n_dim when not parrallel_mode
@@ -1256,7 +1332,7 @@ def forward(
                 # for jit
                 assert not self.use_rk_update
                 assert dihedral_ebd is not None
-                node_ebd, edge_ebd, angle_ebd, dihedral_ebd, ___ = ll.forward(
+                node_ebd, edge_ebd, angle_ebd, dihedral_ebd, ___, ___, = ll.forward(
                     node_ebd_ext,
                     edge_ebd,
                     h2,
@@ -1279,7 +1355,7 @@ def forward(
             else:
                 assert dihedral_ebd is None
                 if not self.use_rk_update:
-                    node_ebd, edge_ebd, angle_ebd, ___, node_sph_embed = ll.forward(
+                    node_ebd, edge_ebd, angle_ebd, ___, node_sph_embed, edge_sph_embed = ll.forward(
                         node_ebd_ext,
                         edge_ebd,
                         h2,
@@ -1301,6 +1377,8 @@ def forward(
                         edge_rbf_ebd=edge_rbf_ebd,
                         edge_sph=edge_sph,
                         node_sph_embed=node_sph_embed,
+                        edge_angle_filter=edge_angle_filter,
+                        edge_sph_embed=edge_sph_embed,
                     )
                 # may cause jit slow, todo fix
                 elif not self.rk_update_diff_layer:
@@ -1321,6 +1399,7 @@ def forward(
                                 angle_ebd_k1,
                                 ___,
                                 ___,
+                                ___,
                             ) = ll.forward(
                                 node_ebd_k_in,
                                 edge_ebd_k_in,