Fix: dptest label mismatch (#5082)

Currently, when running dptest on atomic polar. The order of atoms for
`coord` and `label` are misaligned.
This is a result of type_sel = None for coord, but not None for label. 
I tried to set it to None for label as well, but it broke TF
compatibility.
So the easiest hack I can think of is to add a check while loading data.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Tests**
* Enhanced output handling in test suite to improve per-atom count
tracking in selection outputs.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: anyangml

deepmd/entrypoints/test.py

@@ -1134,6 +1134,7 @@ def test_polar(
         must=True,
         high_prec=False,
         type_sel=dp.get_sel_type(),
+        output_natoms_for_type_sel=True,
     )
 
     test_data = data.get_test()
@@ -1155,7 +1156,12 @@ def test_polar(
         polar = polar.reshape((polar.shape[0], -1, 9))[:, sel_mask, :].reshape(
             (polar.shape[0], -1)
         )
-        rmse_f = rmse(polar - test_data["atom_polarizability"][:numb_test])
+        label_polar = (
+            test_data["atom_polarizability"][:numb_test]
+            .reshape((numb_test, -1, 9))[:, sel_mask, :]
+            .reshape((numb_test, -1))
+        )
+        rmse_f = rmse(polar - label_polar)
 
     log.info(f"# number of test data : {numb_test:d} ")
     log.info(f"Polarizability  RMSE       : {rmse_f:e}")
@@ -1183,10 +1189,7 @@ def test_polar(
         else:
             pe = np.concatenate(
                 (
-                    np.reshape(
-                        test_data["atom_polarizability"][:numb_test],
-                        [-1, 9 * sel_natoms],
-                    ),
+                    np.reshape(label_polar, [-1, 9 * sel_natoms]),
                     np.reshape(polar, [-1, 9 * sel_natoms]),
                 ),
                 axis=1,
@@ -1275,7 +1278,9 @@ def test_dipole(
         must=True,
         high_prec=False,
         type_sel=dp.get_sel_type(),
+        output_natoms_for_type_sel=True,
     )
+
     test_data = data.get_test()
     dipole, numb_test, atype = run_test(dp, test_data, numb_test, data)
 
@@ -1295,7 +1300,12 @@ def test_dipole(
         dipole = dipole.reshape((dipole.shape[0], -1, 3))[:, sel_mask, :].reshape(
             (dipole.shape[0], -1)
         )
-        rmse_f = rmse(dipole - test_data["atom_dipole"][:numb_test])
+        label_dipole = (
+            test_data["atom_dipole"][:numb_test]
+            .reshape((numb_test, -1, 3))[:, sel_mask, :]
+            .reshape((numb_test, -1))
+        )
+        rmse_f = rmse(dipole - label_dipole)
 
     log.info(f"# number of test data : {numb_test:d}")
     log.info(f"Dipole  RMSE       : {rmse_f:e}")
@@ -1318,9 +1328,7 @@ def test_dipole(
         else:
             pe = np.concatenate(
                 (
-                    np.reshape(
-                        test_data["atom_dipole"][:numb_test], [-1, 3 * sel_natoms]
-                    ),
+                    np.reshape(label_dipole, [-1, 3 * sel_natoms]),
                     np.reshape(dipole, [-1, 3 * sel_natoms]),
                 ),
                 axis=1,

deepmd/utils/data.py

@@ -188,6 +188,9 @@ def add(
         output_natoms_for_type_sel : bool, optional
             if True and type_sel is True, the atomic dimension will be natoms instead of nsel
         """
+        # normalize key: "atomic_" prefix -> "atom_", same convention as _load_set output
+        if key.startswith("atomic_"):
+            key = "atom_" + key[7:]
         self.data_dict[key] = {
             "ndof": ndof,
             "atomic": atomic,
@@ -762,6 +765,15 @@ def _load_set(self, set_name: DPPath) -> dict[str, Any]:
         data = {kk.replace("atomic", "atom"): vv for kk, vv in data.items()}
         return data
 
+    def _get_data_path(self, set_name: "DPPath", key: str) -> "DPPath":
+        """Return the path for a data file, trying both atom_ and atomic_ naming."""
+        path = set_name / (key + ".npy")
+        if not path.is_file() and key.startswith("atom_"):
+            alt = set_name / ("atomic_" + key[5:] + ".npy")
+            if alt.is_file():
+                return alt
+        return path
+
     def _load_data(
         self,
         set_name: str,
@@ -800,7 +812,7 @@ def _load_data(
             dtype = GLOBAL_ENER_FLOAT_PRECISION
         else:
             dtype = GLOBAL_NP_FLOAT_PRECISION
-        path = set_name / (key + ".npy")
+        path = self._get_data_path(set_name, key)
         if path.is_file():
             data = path.load_numpy().astype(dtype)
             try:  # YWolfeee: deal with data shape error
@@ -892,7 +904,7 @@ def _load_single_data(
             The total number of frames in this set (to avoid redundant _get_nframes calls)
         """
         vv = self.data_dict[key]
-        path = set_dir / (key + ".npy")
+        path = self._get_data_path(set_dir, key)
 
         if vv["atomic"]:
             natoms = self.natoms

source/tests/pt/model/test_dipole_fitting.py

@@ -372,6 +372,161 @@ def test_deepdipole_infer(self) -> None:
         load_md.eval_full(coords=coord, atom_types=atype, cells=cell, atomic=True)
         load_md.eval_full(coords=coord, atom_types=atype, cells=cell, atomic=False)
 
+    def test_eval_shuffle_sel_type(self) -> None:
+        # Build a model where only type-0 atoms contribute (exclude types 1 and 2).
+        # This tests that eval() returns per-atom results in the correct input atom
+        # order even when sel_type is a strict subset of all types.
+        ft_sel = DipoleFittingNet(
+            self.nt,
+            self.dd0.dim_out,
+            embedding_width=self.dd0.get_dim_emb(),
+            numb_fparam=0,
+            numb_aparam=0,
+            mixed_types=self.dd0.mixed_types(),
+            exclude_types=[1, 2],
+            seed=GLOBAL_SEED,
+        ).to(env.DEVICE)
+        model_sel = DipoleModel(self.dd0, ft_sel, self.type_mapping)
+        jit_md = torch.jit.script(model_sel)
+        torch.jit.save(jit_md, self.file_path)
+        load_md = DeepDipole(self.file_path)
+
+        atype = to_numpy_array(self.atype)  # [0, 0, 0, 1, 1]
+        coord = to_numpy_array(self.coord.reshape(1, self.natoms, 3))
+        cell = to_numpy_array(self.cell.reshape(1, 9))
+
+        # Reference result with original atom order
+        ref = load_md.eval(coords=coord, atom_types=atype, cells=cell, atomic=True)
+        # ref shape: [nframes, natoms, nout]
+
+        # Shuffle atoms
+        idx_perm = np.array([1, 0, 4, 3, 2], dtype=np.intp)
+        coord_sf = coord.reshape(self.natoms, 3)[idx_perm].reshape(1, -1)
+        atype_sf = atype[idx_perm]
+
+        # Result with shuffled atom order
+        res_sf = load_md.eval(
+            coords=coord_sf, atom_types=atype_sf, cells=cell, atomic=True
+        )
+        # res_sf shape: [nframes, natoms, nout]
+
+        # sel_mask: which atoms in the original order are selected (type 0)
+        sel_mask = np.isin(atype, load_md.get_sel_type())  # [T,T,T,F,F]
+        sel_mask_sf = sel_mask[idx_perm]  # selected atoms in shuffled order
+
+        # Extract selected-atom outputs from each result
+        ref_sel = ref[:, sel_mask]  # [nframes, nsel, nout]
+        at_sf = res_sf[:, sel_mask_sf]  # [nframes, nsel, nout]
+
+        # isel_sf: mapping from shuffled-selected positions to original-selected positions
+        orig_sel_idx = np.where(sel_mask)[0]
+        shuffled_orig = np.array(idx_perm)[sel_mask_sf]
+        isel_sf = np.array(
+            [np.where(orig_sel_idx == x)[0][0] for x in shuffled_orig]
+        )  # [1, 0, 2]
+
+        # Recover original selected order from shuffled selected
+        nat = np.empty_like(at_sf)
+        nat[:, isel_sf] = at_sf
+
+        np.testing.assert_almost_equal(
+            nat.reshape([-1]), ref_sel.reshape([-1]), decimal=10
+        )
+
+    def test_label_order_via_deepmd_data(self) -> None:
+        """Verify that labels loaded via DeepmdData(sort_atoms=False) +
+        output_natoms_for_type_sel=True align with dp.eval() output.
+        Uses a sel_type model (exclude_types=[1,2]) with atype=[0,0,0,1,1]
+        shuffled to [0,0,1,1,0] so selected atoms are non-contiguous.
+        """
+        import shutil
+        import tempfile
+
+        from deepmd.utils.data import (
+            DeepmdData,
+        )
+
+        ft_sel = DipoleFittingNet(
+            self.nt,
+            self.dd0.dim_out,
+            embedding_width=self.dd0.get_dim_emb(),
+            numb_fparam=0,
+            numb_aparam=0,
+            mixed_types=self.dd0.mixed_types(),
+            exclude_types=[1, 2],
+            seed=GLOBAL_SEED,
+        ).to(env.DEVICE)
+        model_sel = DipoleModel(self.dd0, ft_sel, self.type_mapping)
+        jit_md = torch.jit.script(model_sel)
+        torch.jit.save(jit_md, self.file_path)
+        load_md = DeepDipole(self.file_path)
+
+        # Shuffle atoms so selected type-0 atoms are non-contiguous
+        # atype=[0,0,0,1,1] → shuffled idx → atype=[0,0,1,1,0]
+        idx_perm = np.array([1, 0, 4, 3, 2], dtype=np.intp)
+        atype = to_numpy_array(self.atype)  # [0,0,0,1,1]
+        coord = to_numpy_array(self.coord.reshape(1, self.natoms, 3))
+        cell = to_numpy_array(self.cell.reshape(1, 9))
+        atype_sf = atype[idx_perm]
+        coord_sf = coord.reshape(self.natoms, 3)[idx_perm].reshape(1, -1)
+
+        sel_mask_sf = np.isin(atype_sf, load_md.get_sel_type())  # type-0 positions
+
+        # Reference: model output for shuffled atoms, filter to sel atoms
+        ref_sf = load_md.eval(
+            coords=coord_sf, atom_types=atype_sf, cells=cell, atomic=True
+        )  # [1, natoms, nout]
+        ref_sf_sel = ref_sf[:, sel_mask_sf, :]  # [1, nsel, nout]
+
+        tmpdir = tempfile.mkdtemp()
+        try:
+            set_dir = os.path.join(tmpdir, "set.000")
+            os.makedirs(set_dir)
+            np.savetxt(os.path.join(tmpdir, "type.raw"), atype_sf, fmt="%d")
+            np.save(
+                os.path.join(set_dir, "coord.npy"),
+                coord_sf.reshape(1, -1),
+            )
+            np.save(
+                os.path.join(set_dir, "box.npy"),
+                cell.reshape(1, -1),
+            )
+            # Labels: nsel atoms in shuffled atom order (nsel format)
+            np.save(
+                os.path.join(set_dir, "atomic_dipole.npy"),
+                ref_sf_sel.reshape(1, -1),
+            )
+
+            data = DeepmdData(
+                tmpdir,
+                set_prefix="set",
+                shuffle_test=False,
+                type_map=load_md.get_type_map(),
+                sort_atoms=False,
+            )
+            data.add(
+                "atomic_dipole",
+                3,
+                atomic=True,
+                must=True,
+                high_prec=False,
+                type_sel=load_md.get_sel_type(),
+                output_natoms_for_type_sel=True,
+            )
+            test_data = data.get_test()
+
+            # Loaded label shape: [1, natoms*3]. Filter to sel atoms.
+            label_sel = test_data["atom_dipole"].reshape(1, self.natoms, 3)[
+                :, sel_mask_sf, :
+            ]  # [1, nsel, 3]
+
+            # Round-trip: loaded label must match what was written
+            np.testing.assert_almost_equal(
+                label_sel.reshape(-1), ref_sf_sel.reshape(-1), decimal=5
+            )
+        finally:
+            shutil.rmtree(tmpdir)
+
     def tearDown(self) -> None:
         if os.path.exists(self.file_path):
             os.remove(self.file_path)