diff --git a/.github/workflows/test_python.yml b/.github/workflows/test_python.yml index 514b552aec..60b5ecf0e0 100644 --- a/.github/workflows/test_python.yml +++ b/.github/workflows/test_python.yml @@ -18,7 +18,7 @@ jobs: - python: 3.8 tf: torch: - - python: "3.12" + - python: "3.11" tf: torch: diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index c004814b60..01bf107c63 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/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"]: diff --git a/deepmd/dpmodel/utils/env_mat.py b/deepmd/dpmodel/utils/env_mat.py index cdd03eb090..0c2ca43c40 100644 --- a/deepmd/dpmodel/utils/env_mat.py +++ b/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)) diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index f46697b990..e17b7c5d54 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/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 diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 641521b3a8..e63a365965 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -388,6 +388,9 @@ 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. @@ -395,7 +398,12 @@ def get_sel_type(self) -> List[int]: 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"]: diff --git a/deepmd/pt/utils/preprocess.py b/deepmd/pt/utils/preprocess.py index 806bacdcd2..ed46292f84 100644 --- a/deepmd/pt/utils/preprocess.py +++ b/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)) diff --git a/source/tests/common/dpmodel/test_fitting_invar_fitting.py b/source/tests/common/dpmodel/test_fitting_invar_fitting.py index a31439d406..87eeb9e06b 100644 --- a/source/tests/common/dpmodel/test_fitting_invar_fitting.py +++ b/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 diff --git a/source/tests/common/dpmodel/test_linear_atomic_model.py b/source/tests/common/dpmodel/test_linear_atomic_model.py index aa56feb3e5..cc08a3b3dd 100644 --- a/source/tests/common/dpmodel/test_linear_atomic_model.py +++ b/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( diff --git a/source/tests/pt/model/test_ener_fitting.py b/source/tests/pt/model/test_ener_fitting.py index a41b4d6b9f..69bd4b42a3 100644 --- a/source/tests/pt/model/test_ener_fitting.py +++ b/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, diff --git a/source/tests/pt/model/test_linear_atomic_model.py b/source/tests/pt/model/test_linear_atomic_model.py index fab6481a6f..e0904097e3 100644 --- a/source/tests/pt/model/test_linear_atomic_model.py +++ b/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(