From 8e0d3964018e98bdf4b3067c59c57676aad22fc5 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sat, 2 Mar 2024 04:27:54 -0500 Subject: [PATCH 1/2] breaking: change DeepTensor output dim from nsel_atoms to natoms Signed-off-by: Jinzhe Zeng --- deepmd/infer/deep_tensor.py | 10 ++-------- deepmd/tf/infer/deep_eval.py | 12 ++++++++++++ source/tests/tf/test_deepdipole.py | 21 +++++++++++++++------ source/tests/tf/test_deeppolar.py | 18 ++++++++++++------ 4 files changed, 41 insertions(+), 20 deletions(-) diff --git a/deepmd/infer/deep_tensor.py b/deepmd/infer/deep_tensor.py index 1bdc459920..14e13e7f84 100644 --- a/deepmd/infer/deep_tensor.py +++ b/deepmd/infer/deep_tensor.py @@ -104,11 +104,8 @@ def eval( aparam=aparam, **kwargs, ) - sel_natoms = self._get_sel_natoms(atom_types[0]) - if sel_natoms == 0: - sel_natoms = atom_types.shape[-1] # set to natoms if atomic: - return results[self.output_tensor_name].reshape(nframes, sel_natoms, -1) + return results[self.output_tensor_name].reshape(nframes, natoms, -1) else: return results[f"{self.output_tensor_name}_redu"].reshape(nframes, -1) @@ -187,9 +184,6 @@ def eval_full( **kwargs, ) - sel_natoms = self._get_sel_natoms(atom_types[0]) - if sel_natoms == 0: - sel_natoms = atom_types.shape[-1] # set to natoms energy = results[f"{self.output_tensor_name}_redu"].reshape(nframes, -1) force = results[f"{self.output_tensor_name}_derv_r"].reshape( nframes, -1, natoms, 3 @@ -199,7 +193,7 @@ def eval_full( ) if atomic: atomic_energy = results[self.output_tensor_name].reshape( - nframes, sel_natoms, -1 + nframes, natoms, -1 ) atomic_virial = results[f"{self.output_tensor_name}_derv_c"].reshape( nframes, -1, natoms, 9 diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 27152b9f12..45eda3392f 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -967,6 +967,18 @@ def _eval_inner( v_out[ii] = self.reverse_map( np.reshape(v_out[ii], odef_shape), sel_imap[:natoms_real] ) + if nloc_sel < nloc: + # convert shape from nsel to nloc + # sel_atoms was applied before sort; see sort_input + # do not consider mixed_types here (as it is never supported) + sel_mask = np.isin(atom_types[0], self.sel_type) + out_nsel = v_out[ii] + out_nloc = np.zeros( + (nframes, nloc, *out_nsel.shape[2:]), dtype=out_nsel.dtype + ) + out_nloc[:, sel_mask] = out_nsel + v_out[ii] = out_nloc + odef_shape = self._get_output_shape(odef, nframes, nloc) v_out[ii] = np.reshape(v_out[ii], odef_shape) elif odef.category in ( OutputVariableCategory.REDU, diff --git a/source/tests/tf/test_deepdipole.py b/source/tests/tf/test_deepdipole.py index 2c8ec7cc66..1e2f6dd45a 100644 --- a/source/tests/tf/test_deepdipole.py +++ b/source/tests/tf/test_deepdipole.py @@ -72,6 +72,7 @@ def setUp(self): 1.667785136187720063e00, ] ) + self.sel_mask = np.isin(self.atype, self.dp.get_sel_type()) @classmethod def tearDownClass(cls): @@ -85,7 +86,7 @@ def test_attrs(self): self.assertEqual(self.dp.get_sel_type(), [0]) def test_1frame_atm(self): - dd = self.dp.eval(self.coords, self.box, self.atype) + dd = self.dp.eval(self.coords, self.box, self.atype)[:, self.sel_mask] # check shape of the returns nframes = 1 natoms = len(self.atype) @@ -97,7 +98,7 @@ def test_1frame_atm(self): def test_2frame_atm(self): coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) - dd = self.dp.eval(coords2, box2, self.atype) + dd = self.dp.eval(coords2, box2, self.atype)[:, self.sel_mask] # check shape of the returns nframes = 2 natoms = len(self.atype) @@ -151,6 +152,7 @@ def setUp(self): 1.667798310054391e00, ] ) + self.sel_mask = np.isin(self.atype, self.dp.get_sel_type()) @classmethod def tearDownClass(cls): @@ -158,7 +160,7 @@ def tearDownClass(cls): cls.dp = None def test_1frame_atm(self): - dd = self.dp.eval(self.coords, None, self.atype) + dd = self.dp.eval(self.coords, None, self.atype)[:, self.sel_mask] # check shape of the returns nframes = 1 natoms = len(self.atype) @@ -168,7 +170,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(dd.ravel(), self.expected_d, default_places) def test_1frame_atm_large_box(self): - dd = self.dp.eval(self.coords, self.box, self.atype) + dd = self.dp.eval(self.coords, self.box, self.atype)[:, self.sel_mask] # check shape of the returns nframes = 1 natoms = len(self.atype) @@ -455,6 +457,7 @@ def setUp(self): self.expected_gv = ( self.expected_v.reshape(1, self.nout, 6, 9).sum(-2).reshape(-1) ) + self.sel_mask = np.isin(self.atype, self.dp.get_sel_type()) @classmethod def tearDownClass(cls): @@ -476,7 +479,7 @@ def test_1frame_old(self): np.testing.assert_almost_equal(gt.ravel(), self.expected_gt, default_places) def test_1frame_old_atm(self): - at = self.dp.eval(self.coords, self.box, self.atype) + at = self.dp.eval(self.coords, self.box, self.atype)[:, self.sel_mask] # check shape of the returns nframes = 1 natoms = len(self.atype) @@ -488,7 +491,7 @@ def test_1frame_old_atm(self): def test_2frame_old_atm(self): coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) - at = self.dp.eval(coords2, box2, self.atype) + at = self.dp.eval(coords2, box2, self.atype)[:, self.sel_mask] # check shape of the returns nframes = 2 natoms = len(self.atype) @@ -515,6 +518,7 @@ def test_1frame_full_atm(self): gt, ff, vv, at, av = self.dp.eval_full( self.coords, self.box, self.atype, atomic=True ) + at = at[:, self.sel_mask] # check shape of the returns nframes = 1 natoms = len(self.atype) @@ -550,6 +554,7 @@ def test_1frame_full_atm_shuffle(self): self.atype[i_sf], atomic=True, ) + at = at[:, self.sel_mask[i_sf]] # check shape of the returns nframes = 1 natoms = len(self.atype) @@ -617,6 +622,7 @@ def test_2frame_full_atm(self): coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) gt, ff, vv, at, av = self.dp.eval_full(coords2, box2, self.atype, atomic=True) + at = at[:, self.sel_mask] # check shape of the returns nframes = 2 natoms = len(self.atype) @@ -949,6 +955,7 @@ def setUp(self): ) fake_target = fake_target - 13 * np.rint(fake_target / 13) self.target_t = fake_target.reshape(-1) + self.sel_mask = np.isin(self.atype, self.dp.get_sel_type()) @classmethod def tearDownClass(cls): @@ -966,6 +973,7 @@ def test_1frame_full_atm(self): gt, ff, vv, at, av = self.dp.eval_full( self.coords, self.box, self.atype, atomic=True ) + at = at[:, self.sel_mask] # check shape of the returns nframes = 1 natoms = len(self.atype) @@ -1001,6 +1009,7 @@ def test_1frame_full_atm_shuffle(self): self.atype[i_sf], atomic=True, ) + at = at[:, self.sel_mask[i_sf]] # check shape of the returns nframes = 1 natoms = len(self.atype) diff --git a/source/tests/tf/test_deeppolar.py b/source/tests/tf/test_deeppolar.py index cfa115c59f..b4f3fe7d0d 100644 --- a/source/tests/tf/test_deeppolar.py +++ b/source/tests/tf/test_deeppolar.py @@ -82,6 +82,7 @@ def setUp(self): 4.448255365635306879e-01, ] ) + self.sel_mask = np.isin(self.atype, self.dp.get_sel_type()) @classmethod def tearDownClass(cls): @@ -95,7 +96,7 @@ def test_attrs(self): self.assertEqual(self.dp.get_sel_type(), [0]) def test_1frame_atm(self): - dd = self.dp.eval(self.coords, self.box, self.atype) + dd = self.dp.eval(self.coords, self.box, self.atype)[:, self.sel_mask] # check shape of the returns nframes = 1 natoms = len(self.atype) @@ -107,7 +108,7 @@ def test_1frame_atm(self): def test_2frame_atm(self): coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) - dd = self.dp.eval(coords2, box2, self.atype) + dd = self.dp.eval(coords2, box2, self.atype)[:, self.sel_mask] # check shape of the returns nframes = 2 natoms = len(self.atype) @@ -173,6 +174,7 @@ def setUp(self): 4.382376148484938e-01, ] ) + self.sel_mask = np.isin(self.atype, self.dp.get_sel_type()) @classmethod def tearDownClass(cls): @@ -180,7 +182,7 @@ def tearDownClass(cls): cls.dp = None def test_1frame_atm(self): - dd = self.dp.eval(self.coords, None, self.atype) + dd = self.dp.eval(self.coords, None, self.atype)[:, self.sel_mask] # check shape of the returns nframes = 1 natoms = len(self.atype) @@ -190,7 +192,7 @@ def test_1frame_atm(self): np.testing.assert_almost_equal(dd.ravel(), self.expected_d, default_places) def test_1frame_atm_large_box(self): - dd = self.dp.eval(self.coords, self.box, self.atype) + dd = self.dp.eval(self.coords, self.box, self.atype)[:, self.sel_mask] # check shape of the returns nframes = 1 natoms = len(self.atype) @@ -921,6 +923,7 @@ def setUp(self): self.expected_gv = ( self.expected_v.reshape(1, self.nout, 6, 9).sum(-2).reshape(-1) ) + self.sel_mask = np.isin(self.atype, self.dp.get_sel_type()) @classmethod def tearDownClass(cls): @@ -942,7 +945,7 @@ def test_1frame_old(self): np.testing.assert_almost_equal(gt.ravel(), self.expected_gt, default_places) def test_1frame_old_atm(self): - at = self.dp.eval(self.coords, self.box, self.atype) + at = self.dp.eval(self.coords, self.box, self.atype)[:, self.sel_mask] # check shape of the returns nframes = 1 natoms = len(self.atype) @@ -954,7 +957,7 @@ def test_1frame_old_atm(self): def test_2frame_old_atm(self): coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) - at = self.dp.eval(coords2, box2, self.atype) + at = self.dp.eval(coords2, box2, self.atype)[:, self.sel_mask] # check shape of the returns nframes = 2 natoms = len(self.atype) @@ -981,6 +984,7 @@ def test_1frame_full_atm(self): gt, ff, vv, at, av = self.dp.eval_full( self.coords, self.box, self.atype, atomic=True ) + at = at[:, self.sel_mask] # check shape of the returns nframes = 1 @@ -1017,6 +1021,7 @@ def test_1frame_full_atm_shuffle(self): self.atype[i_sf], atomic=True, ) + at = at[:, self.sel_mask[i_sf]] # check shape of the returns nframes = 1 natoms = len(self.atype) @@ -1054,6 +1059,7 @@ def test_2frame_full_atm(self): coords2 = np.concatenate((self.coords, self.coords)) box2 = np.concatenate((self.box, self.box)) gt, ff, vv, at, av = self.dp.eval_full(coords2, box2, self.atype, atomic=True) + at = at[:, self.sel_mask] # check shape of the returns nframes = 2 natoms = len(self.atype) From e3756ce243d3cd97a61ee5f0b06fb09d1567a02d Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sat, 2 Mar 2024 05:25:19 -0500 Subject: [PATCH 2/2] fix dp test and other tests Signed-off-by: Jinzhe Zeng --- deepmd/entrypoints/test.py | 8 ++++++++ source/tests/tf/test_data_modifier_shuffle.py | 6 ++++-- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py index b9421bb198..efc75e31a7 100644 --- a/deepmd/entrypoints/test.py +++ b/deepmd/entrypoints/test.py @@ -860,6 +860,10 @@ def test_polar( rmse_fs = rmse_f / np.sqrt(sel_natoms) rmse_fa = rmse_f / sel_natoms else: + sel_mask = np.isin(atype, sel_type) + polar = polar.reshape((polar.shape[0], -1, 9))[:, sel_mask, :].reshape( + (polar.shape[0], -1) + ) rmse_f = rmse(polar - test_data["atomic_polarizability"][:numb_test]) log.info(f"# number of test data : {numb_test:d} ") @@ -996,6 +1000,10 @@ def test_dipole( rmse_fs = rmse_f / np.sqrt(sel_natoms) rmse_fa = rmse_f / sel_natoms else: + sel_mask = np.isin(atype, sel_type) + dipole = dipole.reshape((dipole.shape[0], -1, 3))[:, sel_mask, :].reshape( + (dipole.shape[0], -1) + ) rmse_f = rmse(dipole - test_data["atomic_dipole"][:numb_test]) log.info(f"# number of test data : {numb_test:d}") diff --git a/source/tests/tf/test_data_modifier_shuffle.py b/source/tests/tf/test_data_modifier_shuffle.py index fb9da948f1..41086cc775 100644 --- a/source/tests/tf/test_data_modifier_shuffle.py +++ b/source/tests/tf/test_data_modifier_shuffle.py @@ -127,6 +127,8 @@ def _setUp_data(self): self.coords1 = np.reshape(self.coords1, [self.nframes, self.natoms * 3]) self.dipoles1 = self.dipoles0[:, self.sel_idx_map] self.box1 = self.box0 + self.sel_mask0 = np.isin(self.atom_types0, self.sel_type) + self.sel_mask1 = np.isin(self.atom_types1, self.sel_type) def _write_sys_data(self, dirname, atom_types, coords, dipoles, box): os.makedirs(dirname, exist_ok=True) @@ -185,8 +187,8 @@ def _setUp_jdata(self): def test_z_dipole(self): dd = DeepDipole(os.path.join(modifier_datapath, "dipole.pb")) - dv0 = dd.eval(self.coords0, self.box0, self.atom_types0) - dv1 = dd.eval(self.coords1, self.box1, self.atom_types1) + dv0 = dd.eval(self.coords0, self.box0, self.atom_types0)[:, self.sel_mask0] + dv1 = dd.eval(self.coords1, self.box1, self.atom_types1)[:, self.sel_mask1] dv01 = dv0.reshape([self.nframes, -1, 3]) dv01 = dv01[:, self.sel_idx_map, :]