breaking: change DeepTensor output dim from nsel_atoms to natoms (#3390)

Signed-off-by: Jinzhe Zeng <[email protected]>

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}")

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

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,

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, :]

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,14 +152,15 @@ def setUp(self):
                 1.667798310054391e00,
             ]
         )
+        self.sel_mask = np.isin(self.atype, self.dp.get_sel_type())
 
     @classmethod
     def tearDownClass(cls):
         os.remove("deepdipole.pb")
         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)

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,14 +174,15 @@ def setUp(self):
                 4.382376148484938e-01,
             ]
         )
+        self.sel_mask = np.isin(self.atype, self.dp.get_sel_type())
 
     @classmethod
     def tearDownClass(cls):
         os.remove("deeppolar.pb")
         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)