fixup

Author: OutisLi

deepmd/pt_expt/infer/deep_eval.py

@@ -138,14 +138,10 @@ def _init_from_model_json(self, model_json_str: str) -> None:
             self._dpmodel = BaseModel.deserialize(model_data)
             self._is_spin = False
 
-        self.rcut = self._dpmodel.get_rcut()
-        self.type_map = self._dpmodel.get_type_map()
-        # Hoist sel / mixed_types to plain attributes so the inference hot
-        # path (`_build_nlist_*`) does not need the dpmodel instance.  This
-        # matches what `_init_from_metadata` sets and keeps both code paths
-        # numerically identical.
-        self.sel = list(self._dpmodel.get_sel())
-        self.mixed_types = bool(self._dpmodel.mixed_types())
+        self._rcut = self._dpmodel.get_rcut()
+        self._type_map = self._dpmodel.get_type_map()
+        self._sel = list(self._dpmodel.get_sel())
+        self._mixed_types = bool(self._dpmodel.mixed_types())
         if self._is_spin:
             self._model_output_def = ModelOutputDef(
                 FittingOutputDef(
@@ -177,15 +173,15 @@ def _init_from_metadata(self) -> None:
         ``self._dpmodel`` is left as ``None`` to signal the metadata-only
         mode.  Inference does not need it: it runs through
         ``aoti_load_package`` / the exported module and uses plain
-        attributes (``self.rcut``, ``self.sel``, ``self.mixed_types``,
+        attributes (``self._rcut``, ``self._sel``, ``self._mixed_types``,
         ``self._model_output_def``) for all metadata-level queries.
         """
         self._dpmodel = None
         self._is_spin = bool(self.metadata.get("is_spin", False))
-        self.rcut = float(self.metadata["rcut"])
-        self.type_map = list(self.metadata["type_map"])
-        self.sel = [int(s) for s in self.metadata["sel"]]
-        self.mixed_types = bool(self.metadata["mixed_types"])
+        self._rcut = float(self.metadata["rcut"])
+        self._type_map = list(self.metadata["type_map"])
+        self._sel = [int(s) for s in self.metadata["sel"]]
+        self._mixed_types = bool(self.metadata["mixed_types"])
 
         fitting_defs = []
         for vdef in self.metadata["fitting_output_defs"]:
@@ -294,15 +290,15 @@ def _load_pt2(self, model_file: str) -> None:
 
     def get_rcut(self) -> float:
         """Get the cutoff radius of this model."""
-        return self.rcut
+        return self._rcut
 
     def get_ntypes(self) -> int:
         """Get the number of atom types of this model."""
-        return len(self.type_map)
+        return len(self._type_map)
 
     def get_type_map(self) -> list[str]:
         """Get the type map (element name of the atom types) of this model."""
-        return self.type_map
+        return self._type_map
 
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this DP."""
@@ -318,7 +314,7 @@ def get_dim_aparam(self) -> int:
 
     @property
     def model_type(self) -> type["DeepEvalWrapper"]:
-        """The the evaluator of the model type."""
+        """The evaluator of the model type."""
         if self._dpmodel is not None:
             model_output_type = self._dpmodel.model_output_type()
         else:
@@ -366,11 +362,11 @@ def get_has_efield(self) -> bool:
 
     def get_has_spin(self) -> bool:
         """Check if the model has spin atom types."""
-        return getattr(self, "_is_spin", False)
+        return self._is_spin
 
     def get_use_spin(self) -> list[bool]:
         """Get the per-type spin usage of this model."""
-        if not getattr(self, "_is_spin", False):
+        if not self._is_spin:
             return []
         if self._dpmodel is not None:
             return self._dpmodel.spin.use_spin.tolist()
@@ -528,12 +524,9 @@ def _build_nlist_native(
         """
         nframes = coords.shape[0]
         natoms = coords.shape[1]
-        rcut = self.rcut
-        # ``self.sel`` / ``self.mixed_types`` are populated in both
-        # :meth:`_init_from_model_json` and :meth:`_init_from_metadata`,
-        # so this works whether or not ``model.json`` was available.
-        sel = self.sel
-        mixed_types = self.mixed_types
+        rcut = self._rcut
+        sel = self._sel
+        mixed_types = self._mixed_types
 
         if cells is not None:
             box_input = cells.reshape(nframes, 3, 3)
@@ -644,8 +637,8 @@ def _build_nlist_ase_single(
         nlist : np.ndarray, shape (nloc, nsel)
         mapping : np.ndarray, shape (nall,)
         """
-        sel = self.sel
-        mixed_types = self.mixed_types
+        sel = self._sel
+        mixed_types = self._mixed_types
         nsel = sum(sel)
 
         natoms = positions.shape[0]
@@ -688,7 +681,7 @@ def _build_nlist_ase_single(
         ghost_remap[out_mask] = np.arange(nloc, nloc + nghost, dtype=np.int64)
 
         # Build nlist: vectorized CSR-to-dense conversion
-        rcut = self.rcut
+        rcut = self._rcut
         counts = np.diff(first_neigh)
         max_nn = int(counts.max()) if counts.size > 0 else 0
 

source/tests/pt_expt/infer/test_deep_eval.py

@@ -551,13 +551,26 @@ def setUpClass(cls) -> None:
         finally:
             torch.set_default_device("cuda:9999999")
 
+        cls.meta_tmpfile = tempfile.NamedTemporaryFile(suffix=".pt2", delete=False)
+        cls.meta_tmpfile.close()
+        with (
+            zipfile.ZipFile(cls.tmpfile.name, "r") as zin,
+            zipfile.ZipFile(cls.meta_tmpfile.name, "w") as zout,
+        ):
+            for info in zin.infolist():
+                if info.filename == "model/extra/model.json":
+                    continue
+                zout.writestr(info, zin.read(info.filename))
+
         # Also save to .pte for cross-format comparison
         cls.pte_tmpfile = tempfile.NamedTemporaryFile(suffix=".pte", delete=False)
         cls.pte_tmpfile.close()
         deserialize_to_file(cls.pte_tmpfile.name, cls.model_data)
 
         # Create DeepPot for .pt2
         cls.dp = DeepPot(cls.tmpfile.name)
+        # Create DeepPot for metadata-only .pt2
+        cls.dp_meta = DeepPot(cls.meta_tmpfile.name)
         # Create DeepPot for .pte reference
         cls.dp_pte = DeepPot(cls.pte_tmpfile.name)
 
@@ -566,6 +579,7 @@ def tearDownClass(cls) -> None:
         import os
 
         os.unlink(cls.tmpfile.name)
+        os.unlink(cls.meta_tmpfile.name)
         os.unlink(cls.pte_tmpfile.name)
 
     def test_get_rcut(self) -> None:
@@ -647,6 +661,44 @@ def test_pt2_has_metadata(self) -> None:
             self.assertNotIn("model/extra/output_keys.json", names)
             self.assertNotIn("model/extra/model_params.json", names)
 
+    def test_metadata_only_pt2_has_no_model_json(self) -> None:
+        """The metadata-only .pt2 keeps metadata but drops model.json."""
+        with zipfile.ZipFile(self.meta_tmpfile.name, "r") as zf:
+            names = zf.namelist()
+            self.assertIn("model/extra/metadata.json", names)
+            self.assertNotIn("model/extra/model.json", names)
+
+    def test_metadata_only_pt2_accessors_match(self) -> None:
+        """Metadata-only .pt2 archives expose the same metadata API."""
+        full = self.dp.deep_eval
+        meta = self.dp_meta.deep_eval
+        self.assertIsNotNone(full._dpmodel)
+        self.assertIsNone(meta._dpmodel)
+        self.assertEqual(full.get_rcut(), meta.get_rcut())
+        self.assertEqual(full.get_ntypes(), meta.get_ntypes())
+        self.assertEqual(full.get_type_map(), meta.get_type_map())
+        self.assertEqual(full.get_dim_fparam(), meta.get_dim_fparam())
+        self.assertEqual(full.get_dim_aparam(), meta.get_dim_aparam())
+        self.assertEqual(full.get_sel_type(), meta.get_sel_type())
+        self.assertEqual(full.get_has_spin(), meta.get_has_spin())
+        self.assertEqual(full.get_use_spin(), meta.get_use_spin())
+        self.assertIs(full.model_type, meta.model_type)
+
+    def test_metadata_only_pt2_eval_parity(self) -> None:
+        """Metadata-only .pt2 inference matches the full archive exactly."""
+        rng = np.random.default_rng(GLOBAL_SEED + 29)
+        natoms = 5
+        coords = rng.random((1, natoms, 3)) * 8.0
+        cells = np.eye(3).reshape(1, 9) * 10.0
+        atom_types = np.array([i % self.nt for i in range(natoms)], dtype=np.int32)
+
+        full_out = self.dp.eval(coords, cells, atom_types, atomic=True)
+        meta_out = self.dp_meta.eval(coords, cells, atom_types, atomic=True)
+
+        self.assertEqual(len(full_out), len(meta_out))
+        for ref, test in zip(full_out, meta_out, strict=True):
+            np.testing.assert_array_equal(test, ref)
+
     def test_eval_consistency(self) -> None:
         """Test that DeepPot.eval gives same results as direct model forward."""
         rng = np.random.default_rng(GLOBAL_SEED)

source/tests/pt_expt/infer/test_deep_eval_metadata_only.py

@@ -12,8 +12,8 @@
    ``.pte`` (the fast path; ``.pt2`` AOTInductor compilation is too
    heavy for a routine unit test).
 2. Read back that ``.pte`` and record the reference outputs.
-3. Rewrite the archive byte-for-byte except for the ``extra/model.json``
-   entry, producing a metadata-only variant.
+3. Copy all archive entries except ``extra/model.json`` into a
+   metadata-only variant.
 4. Load the metadata-only archive via ``DeepPot`` and assert that the
    metadata-level accessors and the numeric ``eval`` result are
    **bitwise identical** to the reference.
@@ -151,10 +151,10 @@ def test_internal_attributes_match(self) -> None:
         """The hot-path attributes hoisted in both init paths must agree."""
         full = self.dp_full.deep_eval
         meta = self.dp_meta.deep_eval
-        self.assertEqual(list(full.sel), list(meta.sel))
-        self.assertEqual(bool(full.mixed_types), bool(meta.mixed_types))
-        self.assertEqual(full.rcut, meta.rcut)
-        self.assertEqual(list(full.type_map), list(meta.type_map))
+        self.assertEqual(list(full._sel), list(meta._sel))
+        self.assertEqual(bool(full._mixed_types), bool(meta._mixed_types))
+        self.assertEqual(full._rcut, meta._rcut)
+        self.assertEqual(list(full._type_map), list(meta._type_map))
 
     def test_dpmodel_presence(self) -> None:
         """``_dpmodel`` is the single signal that separates the two modes."""

source/tests/pt_expt/infer/test_deep_eval_spin.py

@@ -8,6 +8,7 @@
 import copy
 import os
 import tempfile
+import zipfile
 
 import numpy as np
 import pytest
@@ -105,6 +106,37 @@
 BOX = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], dtype=np.float64)
 
 
+def _strip_extra_model_json(src: str, dst: str) -> None:
+    """Copy ``src`` to ``dst`` dropping any ``extra/model.json`` entry."""
+    with zipfile.ZipFile(src, "r") as zin, zipfile.ZipFile(dst, "w") as zout:
+        for info in zin.infolist():
+            if info.filename.endswith("extra/model.json"):
+                continue
+            zout.writestr(info, zin.read(info.filename))
+
+
+def _assert_fitting_output_defs_match(full_eval, meta_eval) -> None:
+    """Assert that metadata rebuilds the same fitting output definitions."""
+    full_defs = full_eval._model_output_def.def_outp.get_data()
+    meta_defs = meta_eval._model_output_def.def_outp.get_data()
+    assert full_defs.keys() == meta_defs.keys()
+    attrs = (
+        "shape",
+        "reducible",
+        "r_differentiable",
+        "c_differentiable",
+        "atomic",
+        "category",
+        "r_hessian",
+        "magnetic",
+        "intensive",
+    )
+    for name, full_def in full_defs.items():
+        meta_def = meta_defs[name]
+        for attr in attrs:
+            assert getattr(meta_def, attr) == getattr(full_def, attr)
+
+
 def _build_reference():
     """Build pt_expt model and run eager reference inference.
 
@@ -163,6 +195,9 @@ def spin_model_files():
         finally:
             torch.set_default_device(prev)
         files[ext] = path
+    meta_path = os.path.join(tmpdir, "spin_test_metadata_only.pte")
+    _strip_extra_model_json(files[".pte"], meta_path)
+    files[".pte.meta"] = meta_path
     yield files, ref_pbc, ref_nopbc
     for path in files.values():
         if os.path.exists(path):
@@ -341,6 +376,32 @@ def test_eval_nopbc_nonatomic(self, spin_model_files, ext) -> None:
         )
 
 
+class TestSpinMetadataOnly:
+    """Test metadata-only spin model inference through DeepPot."""
+
+    def test_metadata_only_spin_pte_parity(self, spin_model_files) -> None:
+        """Metadata-only spin .pte matches full archive metadata and outputs."""
+        from deepmd.infer import (
+            DeepPot,
+        )
+
+        files, _, _ = spin_model_files
+        full_dp = DeepPot(files[".pte"])
+        meta_dp = DeepPot(files[".pte.meta"])
+
+        assert meta_dp.has_spin == full_dp.has_spin
+        assert meta_dp.use_spin == full_dp.use_spin
+        assert meta_dp.get_ntypes_spin() == full_dp.get_ntypes_spin()
+        _assert_fitting_output_defs_match(full_dp.deep_eval, meta_dp.deep_eval)
+
+        full_out = full_dp.eval(COORD, BOX, ATYPE, atomic=True, spin=SPIN)
+        meta_out = meta_dp.eval(COORD, BOX, ATYPE, atomic=True, spin=SPIN)
+
+        assert len(full_out) == len(meta_out)
+        for ref, test in zip(full_out, meta_out, strict=True):
+            np.testing.assert_array_equal(test, ref)
+
+
 SPIN_FPARAM_CONFIG = copy.deepcopy(SPIN_CONFIG)
 SPIN_FPARAM_CONFIG["fitting_net"]["numb_fparam"] = 1
 SPIN_FPARAM_CONFIG["fitting_net"]["default_fparam"] = [0.5]