Merge branch 'devel' into spin_rf

.github/workflows/test_python.yml

@@ -18,7 +18,7 @@ jobs:
           - python: 3.8
             tf:
             torch:
-          - python: "3.12"
+          - python: "3.11"
             tf:
             torch:
 

deepmd/dpmodel/fitting/general_fitting.py

@@ -179,7 +179,7 @@ def get_sel_type(self) -> List[int]:
         to the result of the model.
         If returning an empty list, all atom types are selected.
         """
-        return []
+        return [ii for ii in range(self.ntypes) if ii not in self.exclude_types]
 
     def __setitem__(self, key, value):
         if key in ["bias_atom_e"]:

deepmd/dpmodel/utils/env_mat.py

@@ -17,6 +17,8 @@ def compute_smooth_weight(
     rmax: float,
 ):
     """Compute smooth weight for descriptor elements."""
+    if rmin >= rmax:
+        raise ValueError("rmin should be less than rmax.")
     min_mask = distance <= rmin
     max_mask = distance >= rmax
     mid_mask = np.logical_not(np.logical_or(min_mask, max_mask))

deepmd/pt/model/descriptor/se_a.py

@@ -347,9 +347,8 @@ def __init__(
         self.reinit_exclude(exclude_types)
 
         self.sel = sel
-        self.sec = torch.tensor(
-            np.append([0], np.cumsum(self.sel)), dtype=int, device=env.DEVICE
-        )
+        # should be on CPU to avoid D2H, as it is used as slice index
+        self.sec = [0, *np.cumsum(self.sel).tolist()]
         self.split_sel = self.sel
         self.nnei = sum(sel)
         self.ndescrpt = self.nnei * 4

deepmd/pt/model/task/fitting.py

@@ -388,14 +388,22 @@ def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this atomic model."""
         return self.numb_aparam
 
+    # make jit happy
+    exclude_types: List[int]
+
     def get_sel_type(self) -> List[int]:
         """Get the selected atom types of this model.
 
         Only atoms with selected atom types have atomic contribution
         to the result of the model.
         If returning an empty list, all atom types are selected.
         """
-        return []
+        # make jit happy
+        sel_type: List[int] = []
+        for ii in range(self.ntypes):
+            if ii not in self.exclude_types:
+                sel_type.append(ii)
+        return sel_type
 
     def __setitem__(self, key, value):
         if key in ["bias_atom_e"]:

deepmd/pt/utils/preprocess.py

@@ -228,6 +228,8 @@ def build_neighbor_list(
 
 def compute_smooth_weight(distance, rmin: float, rmax: float):
     """Compute smooth weight for descriptor elements."""
+    if rmin >= rmax:
+        raise ValueError("rmin should be less than rmax.")
     min_mask = distance <= rmin
     max_mask = distance >= rmax
     mid_mask = torch.logical_not(torch.logical_or(min_mask, max_mask))

source/tests/common/dpmodel/test_fitting_invar_fitting.py

@@ -64,6 +64,10 @@ def test_self_consistency(
             ret0 = ifn0(dd[0], atype, fparam=ifp, aparam=iap)
             ret1 = ifn1(dd[0], atype, fparam=ifp, aparam=iap)
             np.testing.assert_allclose(ret0["energy"], ret1["energy"])
+            sel_set = set(ifn0.get_sel_type())
+            exclude_set = set(et)
+            self.assertEqual(sel_set | exclude_set, set(range(self.nt)))
+            self.assertEqual(sel_set & exclude_set, set())
 
     def test_mask(self):
         nf, nloc, nnei = self.nlist.shape

source/tests/common/dpmodel/test_linear_atomic_model.py

@@ -40,8 +40,8 @@ def test_pairwise(self, mock_loadtxt):
         nlist = np.array([[[1], [-1]]])
 
         ds = DescrptSeA(
-            rcut=0.3,
-            rcut_smth=0.4,
+            rcut_smth=0.3,
+            rcut=0.4,
             sel=[3],
         )
         ft = InvarFitting(
@@ -122,8 +122,8 @@ def setUp(self, mock_loadtxt):
             ],
             dtype=int,
         ).reshape([1, self.nloc, sum(self.sel)])
-        self.rcut = 0.4
-        self.rcut_smth = 2.2
+        self.rcut_smth = 0.4
+        self.rcut = 2.2
 
         file_path = "dummy_path"
         mock_loadtxt.return_value = np.array(

source/tests/pt/model/test_ener_fitting.py

@@ -95,6 +95,7 @@ def test_consistency(
                 to_numpy_array(ret0["foo"]),
                 to_numpy_array(ret2["foo"]),
             )
+            self.assertEqual(ft0.get_sel_type(), ft1.get_sel_type())
 
     def test_new_old(
         self,

source/tests/pt/model/test_linear_atomic_model.py

@@ -56,8 +56,8 @@ def test_pairwise(self, mock_loadtxt):
         nlist = torch.tensor([[[1], [-1]]], device=env.DEVICE)
 
         ds = DescrptSeA(
-            rcut=0.3,
-            rcut_smth=0.4,
+            rcut_smth=0.3,
+            rcut=0.4,
             sel=[3],
         ).to(env.DEVICE)
         ft = InvarFitting(