convert exclude_types to sel_type

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/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"]:

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/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,