fix se_r consistency

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
deepmodeling · Feb 29, 2024 · 42d47f9 · 42d47f9
1 parent 665d716
commit 42d47f9
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 6 deletions.
diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py
@@ -276,9 +276,9 @@ def call(
             gg = self.cal_g(tr, tt)
             gg = np.mean(gg, axis=2)
             # nf x nloc x ng x 1
-            xyz_scatter += gg
+            xyz_scatter += gg * (self.sel[tt] / self.nnei)
 
-        res_rescale = 1.0 / 10.0
+        res_rescale = 1.0 / 5.0
         res = xyz_scatter * res_rescale
         res = res.reshape(nf, nloc, -1).astype(GLOBAL_NP_FLOAT_PRECISION)
         return res, None, None, None, ww

diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
@@ -258,9 +258,9 @@ def forward(
             # nfnl x nt x ng
             gg = ll.forward(ss)
             gg = torch.mean(gg, dim=1).unsqueeze(1)
-            xyz_scatter += gg
+            xyz_scatter += gg * (self.sel[ii] / self.nnei)
 
-        res_rescale = 1.0 / 10.0
+        res_rescale = 1.0 / 5.0
         result = xyz_scatter * res_rescale
         result = result.view(-1, nloc, self.filter_neuron[-1])
         return (

diff --git a/source/tests/consistent/descriptor/test_se_e2_a.py b/source/tests/consistent/descriptor/test_se_e2_a.py
@@ -51,7 +51,7 @@ def data(self) -> dict:
             precision,
         ) = self.param
         return {
-            "sel": [10, 10],
+            "sel": [9, 10],
             "rcut_smth": 5.80,
             "rcut": 6.00,
             "neuron": [6, 12, 24],

diff --git a/source/tests/consistent/descriptor/test_se_r.py b/source/tests/consistent/descriptor/test_se_r.py
@@ -51,7 +51,7 @@ def data(self) -> dict:
             precision,
         ) = self.param
         return {
-            "sel": [10, 10],
+            "sel": [9, 10],
             "rcut_smth": 5.80,
             "rcut": 6.00,
             "neuron": [6, 12, 24],