update angle embedding

iProzd · Nov 27, 2024 · af3bfd9 · af3bfd9
1 parent d1712c9
commit af3bfd9
Show file tree

Hide file tree

Showing 6 changed files with 758 additions and 354 deletions.
diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py
@@ -199,6 +199,17 @@ def __init__(
         update_g2_has_g1g1: bool = True,
         update_g2_has_attn: bool = True,
         update_h2: bool = False,
+        has_angle: bool = True,
+        update_a_has_g1: bool = True,
+        update_a_has_g2: bool = True,
+        update_g2_has_a: bool = True,
+        a_dim: int = 64,
+        num_a: int = 9,
+        a_rcut: float = 4.0,
+        a_sel: int = 40,
+        angle_use_self_g2_padding: bool = True,
+        update_g1_has_edge: bool = False,
+        update_g2_has_edge: bool = False,
         attn1_hidden: int = 64,
         attn1_nhead: int = 4,
         attn2_hidden: int = 16,
@@ -300,6 +311,17 @@ def __init__(
         self.update_g1_has_attn = update_g1_has_attn
         self.update_g2_has_g1g1 = update_g2_has_g1g1
         self.update_g2_has_attn = update_g2_has_attn
+        self.has_angle = has_angle
+        self.update_a_has_g1 = update_a_has_g1
+        self.update_a_has_g2 = update_a_has_g2
+        self.update_g1_has_edge = update_g1_has_edge
+        self.update_g2_has_edge = update_g2_has_edge
+        self.update_g2_has_a = update_g2_has_a
+        self.a_dim = a_dim
+        self.num_a = num_a
+        self.a_rcut = a_rcut
+        self.a_sel = a_sel
+        self.angle_use_self_g2_padding = angle_use_self_g2_padding
         self.update_h2 = update_h2
         self.attn1_hidden = attn1_hidden
         self.attn1_nhead = attn1_nhead

diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py
@@ -220,6 +220,17 @@ def init_subclass_params(sub_data, sub_class):
             update_g1_has_attn=self.repformer_args.update_g1_has_attn,
             update_g2_has_g1g1=self.repformer_args.update_g2_has_g1g1,
             update_g2_has_attn=self.repformer_args.update_g2_has_attn,
+            has_angle=self.repformer_args.has_angle,
+            update_a_has_g1=self.repformer_args.update_a_has_g1,
+            update_a_has_g2=self.repformer_args.update_a_has_g2,
+            update_g2_has_a=self.repformer_args.update_g2_has_a,
+            update_g1_has_edge=self.repformer_args.update_g1_has_edge,
+            update_g2_has_edge=self.repformer_args.update_g2_has_edge,
+            a_dim=self.repformer_args.a_dim,
+            num_a=self.repformer_args.num_a,
+            a_rcut=self.repformer_args.a_rcut,
+            a_sel=self.repformer_args.a_sel,
+            angle_use_self_g2_padding=self.repformer_args.angle_use_self_g2_padding,
             update_h2=self.repformer_args.update_h2,
             attn1_hidden=self.repformer_args.attn1_hidden,
             attn1_nhead=self.repformer_args.attn1_nhead,