fix(pt): unsmooth update of edge update in DPA3

Author: iProzd

deepmd/dpmodel/descriptor/dpa3.py

@@ -54,6 +54,9 @@ def __init__(
         auto_batchsize: int = 0,
         optim_update: bool = True,
         no_sym: bool = False,
+        smooth_angle_init: bool = False,
+        angle_init_use_sin: bool = False,
+        smooth_edge_update: bool = False,
     ) -> None:
         r"""The constructor for the RepFlowArgs class which defines the parameters of the repflow block in DPA3 descriptor.
 
@@ -153,6 +156,9 @@ def __init__(
         self.auto_batchsize = auto_batchsize
         self.optim_update = optim_update
         self.no_sym = no_sym
+        self.smooth_angle_init = smooth_angle_init
+        self.angle_init_use_sin = angle_init_use_sin
+        self.smooth_edge_update = smooth_edge_update
 
     def __getitem__(self, key):
         if hasattr(self, key):

deepmd/pt/model/descriptor/dpa3.py

@@ -192,6 +192,9 @@ def init_subclass_params(sub_data, sub_class):
             optim_update=self.repflow_args.optim_update,
             skip_stat=self.repflow_args.skip_stat,
             no_sym=self.repflow_args.no_sym,
+            smooth_angle_init=self.repflow_args.smooth_angle_init,
+            angle_init_use_sin=self.repflow_args.angle_init_use_sin,
+            smooth_edge_update=self.repflow_args.smooth_edge_update,
             exclude_types=exclude_types,
             env_protection=env_protection,
             precision=precision,

deepmd/pt/model/descriptor/repflow_layer.py

@@ -78,6 +78,7 @@ def __init__(
         bn_moment: float = 0.1,
         optim_update: bool = True,
         no_sym: bool = False,
+        smooth_edge_update: bool = False,
         activation_function: str = "silu",
         update_style: str = "res_residual",
         update_residual: float = 0.1,
@@ -145,6 +146,7 @@ def __init__(
         self.n_update_has_a_first_sum = n_update_has_a_first_sum
         self.optim_update = optim_update
         self.no_sym = no_sym
+        self.smooth_edge_update = smooth_edge_update
 
         assert update_residual_init in [
             "norm",
@@ -1146,20 +1148,22 @@ def forward(
                 ],
                 dim=2,
             )
-            full_mask = torch.concat(
-                [
-                    a_nlist_mask,
-                    torch.zeros(
-                        [nb, nloc, self.nnei - self.a_sel],
-                        dtype=a_nlist_mask.dtype,
-                        device=a_nlist_mask.device,
-                    ),
-                ],
-                dim=-1,
-            )
-            padding_edge_angle_update = torch.where(
-                full_mask.unsqueeze(-1), padding_edge_angle_update, edge_ebd
-            )
+            if not self.smooth_edge_update:
+                # will be deprecated in the future
+                full_mask = torch.concat(
+                    [
+                        a_nlist_mask,
+                        torch.zeros(
+                            [nb, nloc, self.nnei - self.a_sel],
+                            dtype=a_nlist_mask.dtype,
+                            device=a_nlist_mask.device,
+                        ),
+                    ],
+                    dim=-1,
+                )
+                padding_edge_angle_update = torch.where(
+                    full_mask.unsqueeze(-1), padding_edge_angle_update, edge_ebd
+                )
             e_update_list.append(
                 self.act(self.edge_angle_linear2(padding_edge_angle_update))
             )

deepmd/pt/model/descriptor/repflows.py

@@ -112,6 +112,9 @@ def __init__(
         precision: str = "float64",
         skip_stat: bool = True,
         no_sym: bool = False,
+        smooth_angle_init: bool = False,
+        angle_init_use_sin: bool = False,
+        smooth_edge_update: bool = False,
         pre_ln: bool = False,
         only_e_ln: bool = False,
         pre_bn: bool = False,
@@ -239,6 +242,9 @@ def __init__(
         self.auto_batchsize = auto_batchsize
         self.optim_update = optim_update
         self.no_sym = no_sym
+        self.smooth_angle_init = smooth_angle_init
+        self.angle_init_use_sin = angle_init_use_sin
+        self.smooth_edge_update = smooth_edge_update
 
         self.n_dim = n_dim
         self.e_dim = e_dim
@@ -299,7 +305,11 @@ def __init__(
             1, self.e_dim, precision=precision, seed=child_seed(seed, 0)
         )
         self.angle_embd = MLPLayer(
-            1, self.a_dim, precision=precision, bias=False, seed=child_seed(seed, 1)
+            1 if not self.angle_init_use_sin else 2,
+            self.a_dim,
+            precision=precision,
+            bias=False,
+            seed=child_seed(seed, 1),
         )
         self.has_h1 = self.update_n_has_h1 or self.update_e_has_h1
         if self.has_h1:
@@ -452,6 +462,7 @@ def __init__(
                     bn_moment=self.bn_moment,
                     optim_update=self.optim_update,
                     no_sym=self.no_sym,
+                    smooth_edge_update=self.smooth_edge_update,
                     seed=child_seed(child_seed(seed, 1), ii),
                 )
             )
@@ -632,10 +643,21 @@ def forward(
         # nf x nloc x a_nnei x a_nnei
         # 1 - 1e-6 for torch.acos stability
         cosine_ij = torch.matmul(normalized_diff_i, normalized_diff_j) * (1 - 1e-6)
-        # nf x nloc x a_nnei x a_nnei x 1
-        cosine_ij = cosine_ij.unsqueeze(-1) / (torch.pi**0.5)
+        sine_ij = torch.sqrt(1 - cosine_ij**2)
+        if self.smooth_angle_init:
+            cosine_ij = cosine_ij * a_sw.unsqueeze(-1) * a_sw.unsqueeze(-2)
+            sine_ij = sine_ij * a_sw.unsqueeze(-1) * a_sw.unsqueeze(-2)
+
+        if not self.angle_init_use_sin:
+            # nf x nloc x a_nnei x a_nnei x 1
+            angle_input = cosine_ij.unsqueeze(-1) / (torch.pi**0.5)
+        else:
+            angle_input = torch.cat(
+                [cosine_ij.unsqueeze(-1), sine_ij.unsqueeze(-1)], dim=-1
+            ) / (torch.pi**0.5)
+
         # nf x nloc x a_nnei x a_nnei x a_dim
-        angle_ebd = self.angle_embd(cosine_ij).reshape(
+        angle_ebd = self.angle_embd(angle_input).reshape(
             nframes, nloc, self.a_sel, self.a_sel, self.a_dim
         )
         if self.has_h1:

deepmd/utils/argcheck.py

@@ -1727,6 +1727,24 @@ def dpa3_repflow_args():
             optional=True,
             default=False,
         ),
+        Argument(
+            "smooth_angle_init",
+            bool,
+            optional=True,
+            default=False,
+        ),
+        Argument(
+            "angle_init_use_sin",
+            bool,
+            optional=True,
+            default=False,
+        ),
+        Argument(
+            "smooth_edge_update",
+            bool,
+            optional=True,
+            default=False,
+        ),
     ]
 
 

source/tests/pt/model/test_permutation.py

@@ -167,6 +167,46 @@
     },
 }
 
+model_dpa3 = {
+    "type_map": ["O", "H", "B"],
+    "descriptor": {
+        "type": "dpa3",
+        "repflow": {
+            "n_dim": 30,
+            "e_dim": 20,
+            "a_dim": 10,
+            "nlayers": 3,
+            "e_rcut": 6.0,
+            "e_rcut_smth": 3.5,
+            "e_sel": 20,
+            "a_rcut": 4.0,
+            "a_rcut_smth": 3.5,
+            "a_sel": 12,
+            "axis_neuron": 4,
+            "skip_stat": True,
+            "a_compress_rate": 1,
+            "a_compress_e_rate": 2,
+            "a_compress_use_split": True,
+            "update_angle": True,
+            "smooth_edge_update": True,
+            "update_style": "res_residual",
+            "update_residual": 0.1,
+            "update_residual_init": "const",
+        },
+        "activation_function": "custom_silu:10.0",
+        "use_tebd_bias": False,
+        "precision": "float32",
+        "concat_output_tebd": False,
+    },
+    "fitting_net": {
+        "neuron": [24, 24],
+        "activation_function": "custom_silu:10.0",
+        "resnet_dt": True,
+        "precision": "float32",
+        "seed": 1,
+    },
+}
+
 model_dpa2tebd = {
     "type_map": ["O", "H", "B"],
     "descriptor": {

source/tests/pt/model/test_smooth.py

@@ -21,6 +21,7 @@
     model_dos,
     model_dpa1,
     model_dpa2,
+    model_dpa3,
     model_hybrid,
     model_se_e2_a,
     model_spin,
@@ -53,6 +54,12 @@ def test(
                 0.0,
                 0.0,
                 0.0,
+                6.0 - 0.5 * epsilon,
+                0.0,
+                0.0,
+                0.0,
+                6.0 - 0.5 * epsilon,
+                0.0,
                 4.0 - 0.5 * epsilon,
                 0.0,
                 0.0,
@@ -77,11 +84,15 @@ def test(
         coord0 = torch.clone(coord)
         coord1 = torch.clone(coord)
         coord1[1][0] += epsilon
+        coord1[3][0] += epsilon
         coord2 = torch.clone(coord)
         coord2[2][1] += epsilon
+        coord2[4][1] += epsilon
         coord3 = torch.clone(coord)
         coord3[1][0] += epsilon
+        coord1[3][0] += epsilon
         coord3[2][1] += epsilon
+        coord2[4][1] += epsilon
         test_spin = getattr(self, "test_spin", False)
         if not test_spin:
             test_keys = ["energy", "force", "virial"]
@@ -226,6 +237,17 @@ def setUp(self) -> None:
         self.epsilon, self.aprec = None, None
 
 
+class TestEnergyModelDPA3(unittest.TestCase, SmoothTest):
+    def setUp(self) -> None:
+        model_params = copy.deepcopy(model_dpa3)
+        self.type_split = True
+        self.model = get_model(model_params).to(env.DEVICE)
+        # less degree of smoothness,
+        # error can be systematically removed by reducing epsilon
+        self.epsilon = 1e-5
+        self.aprec = 1e-5
+
+
 class TestEnergyModelHybrid(unittest.TestCase, SmoothTest):
     def setUp(self) -> None:
         model_params = copy.deepcopy(model_hybrid)