Fix: PairTab Data

deepmd/pt/model/atomic_model/linear_atomic_model.py

@@ -501,8 +501,13 @@ def _compute_weight(
             extended_coord, masked_nlist
         )
         numerator = torch.sum(
-            pairwise_rr * torch.exp(-pairwise_rr / self.smin_alpha), dim=-1
-        )  # masked nnei will be zero, no need to handle
+            torch.where(
+                nlist_larger != -1,
+                pairwise_rr * torch.exp(-pairwise_rr / self.smin_alpha),
+                torch.zeros_like(nlist_larger),
+            ),
+            dim=-1,
+        )
         denominator = torch.sum(
             torch.where(
                 nlist_larger != -1,
@@ -522,5 +527,8 @@ def _compute_weight(
         smooth = -6 * u**5 + 15 * u**4 - 10 * u**3 + 1
         coef[mid_mask] = smooth[mid_mask]
         coef[right_mask] = 0
+
+        # to handle masked atoms
+        coef = torch.where(sigma != 0, coef, torch.zeros_like(coef))
         self.zbl_weight = coef  # nframes, nloc
         return [1 - coef.unsqueeze(-1), coef.unsqueeze(-1)]  # to match the model order.

source/tests/pt/model/test_autodiff.py

@@ -55,7 +55,7 @@ class ForceTest:
     def test(
         self,
     ):
-        places = 8
+        places = 6
         delta = 1e-5
         natoms = 5
         cell = torch.rand([3, 3], dtype=dtype, device="cpu")
@@ -126,7 +126,7 @@ class VirialTest:
     def test(
         self,
     ):
-        places = 8
+        places = 6
         delta = 1e-4
         natoms = 5
         cell = torch.rand([3, 3], dtype=dtype, device="cpu")

source/tests/pt/model/test_permutation.py

@@ -66,7 +66,7 @@
     "use_srtab": f"{CUR_DIR}/water/data/zbl_tab_potential/H2O_tab_potential.txt",
     "smin_alpha": 0.1,
     "sw_rmin": 0.2,
-    "sw_rmax": 1.0,
+    "sw_rmax": 4.0,
     "descriptor": {
         "type": "se_atten",
         "sel": 40,