fix: make global DOS inference work on the dpmodel and JAX backends

deepmd/dpmodel/infer/deep_eval.py

@@ -160,7 +160,7 @@ def get_sel_type(self) -> list[int]:
 
     def get_numb_dos(self) -> int:
         """Get the number of DOS."""
-        return 0
+        return self.dp.get_numb_dos()
 
     def get_has_efield(self) -> bool:
         """Check if the model has efield."""

deepmd/dpmodel/model/dos_model.py

@@ -36,6 +36,10 @@ def __init__(
         DPModelCommon.__init__(self)
         DPDOSModel_.__init__(self, *args, **kwargs)
 
+    def get_numb_dos(self) -> int:
+        """Get the number of DOS for DOSFittingNet."""
+        return self.get_fitting_net().dim_out
+
     def call(
         self,
         coord: Array,

deepmd/dpmodel/model/make_model.py

@@ -275,19 +275,26 @@ def call_common(
             coord
                 The coordinates of the atoms.
                 shape: nf x (nloc x 3)
+
             atype
                 The type of atoms. shape: nf x nloc
+
             box
                 The simulation box. shape: nf x 9
+
             fparam
                 frame parameter. nf x ndf
+
             aparam
                 atomic parameter. nf x nloc x nda
+
             do_atomic_virial
                 If calculate the atomic virial.
+
             coord_corr_for_virial
                 The coordinates correction for virial.
                 shape: nf x (nloc x 3)
+
             neighbor_list
                 Neighbor-list construction strategy for the DENSE-nlist path
                 only.  ``None`` uses the default all-pairs builder; an
@@ -296,6 +303,7 @@ def call_common(
                 is consumed by the dense lower; supplying it forces the dense
                 route (see below) and it is rejected together with an explicit
                 ``neighbor_graph_method``.
+
             neighbor_graph_method
                 Selects the lower the model routes through.  The option strings
                 refer to the neighbor-GRAPH builder, NOT the legacy dense nlist:
@@ -995,6 +1003,10 @@ def get_dim_aparam(self) -> int:
             """Get the number (dimension) of atomic parameters of this atomic model."""
             return self.atomic_model.get_dim_aparam()
 
+        def get_numb_dos(self) -> int:
+            """Get the number of DOS. Zero for models without a DOS output."""
+            return 0
+
         def has_default_fparam(self) -> bool:
             """Check if the model has default frame parameters."""
             return self.atomic_model.has_default_fparam()

deepmd/infer/deep_dos.py

@@ -124,18 +124,20 @@ def eval(
             aparam=aparam,
             **kwargs,
         )
-        # energy = results["dos_redu"].reshape(nframes, self.get_numb_dos())
-        atomic_energy = results["dos"].reshape(nframes, natoms, self.get_numb_dos())
-        # not same as dos_redu... why?
-        energy = np.sum(atomic_energy, axis=1)
-
-        if atomic:
-            return (
-                energy,
-                atomic_energy,
-            )
-        else:
+        # Prefer summing the atomic `dos` output when it is present, preserving
+        # the original global DOS for backends that always return it (TF, PT),
+        # whose reduced output is not necessarily the plain sum of the atomic
+        # DOS. The dpmodel and JAX backends omit the atomic `dos` when
+        # atomic=False, so fall back to the reduced `dos_redu` there (reading
+        # `dos` unconditionally would raise KeyError).
+        if "dos" in results:
+            atomic_energy = results["dos"].reshape(nframes, natoms, self.get_numb_dos())
+            energy = np.sum(atomic_energy, axis=1)
+            if atomic:
+                return (energy, atomic_energy)
             return (energy,)
+        energy = results["dos_redu"].reshape(nframes, self.get_numb_dos())
+        return (energy,)
 
     def get_numb_dos(self) -> int:
         return self.deep_eval.get_numb_dos()

deepmd/jax/infer/deep_eval.py

@@ -177,7 +177,7 @@ def get_sel_type(self) -> list[int]:
 
     def get_numb_dos(self) -> int:
         """Get the number of DOS."""
-        return 0
+        return self.dp.get_numb_dos()
 
     def get_has_efield(self) -> bool:
         """Check if the model has efield."""

deepmd/jax/model/hlo.py

@@ -30,6 +30,13 @@
         r_differentiable=True,
         c_differentiable=True,
     ),
+    "dos": OutputVariableDef(
+        "dos",
+        shape=[-1],
+        reducible=True,
+        r_differentiable=False,
+        c_differentiable=False,
+    ),
     "mask": OutputVariableDef(
         "mask",
         shape=[1],
@@ -61,6 +68,7 @@ def __init__(
         # new in v3.1.1
         has_default_fparam: bool = False,
         default_fparam: list[float] | None = None,
+        numb_dos: int = 0,
     ) -> None:
         self._call_lower = jax_export.deserialize(stablehlo).call
         self._call_lower_atomic_virial = jax_export.deserialize(
@@ -84,6 +92,7 @@ def __init__(
         self.model_def_script = model_def_script
         self._has_default_fparam = has_default_fparam
         self.default_fparam = default_fparam
+        self.numb_dos = numb_dos
 
     def __call__(
         self,
@@ -212,6 +221,10 @@ def get_rcut(self) -> float:
         """Get the cut-off radius."""
         return self.rcut
 
+    def get_numb_dos(self) -> int:
+        """Get the number of DOS."""
+        return self.numb_dos
+
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this atomic model."""
         return self.dim_fparam

deepmd/jax/utils/serialization.py

@@ -310,6 +310,7 @@ def call_lower_with_fixed_do_atomic_virial(
         data["constants"] = {
             "type_map": model.get_type_map(),
             "rcut": model.get_rcut(),
+            "numb_dos": model.get_numb_dos(),
             "dim_fparam": model.get_dim_fparam(),
             "dim_aparam": model.get_dim_aparam(),
             "sel_type": model.get_sel_type(),

deepmd/tf/model/dos.py

@@ -214,11 +214,15 @@ def build(
 
         self.atom_dos = atom_dos
 
-        dos_raw = atom_dos
-
-        dos_raw = tf.reshape(dos_raw, [natoms[0], -1], name="o_atom_dos" + suffix)
+        # Reduce the atomic DOS to the global DOS per frame. Reshaping to
+        # [nframes, nloc, numb_dos] and summing over the atom axis mirrors the
+        # energy model; the previous [nloc, -1] reshape summed over the wrong
+        # axis and mixed frames together for multi-frame inputs.
+        dos_raw = tf.reshape(
+            atom_dos, [-1, natoms[0], self.numb_dos], name="o_atom_dos" + suffix
+        )
         dos = tf.reduce_sum(
-            global_cvt_2_ener_float(dos_raw), axis=0, name="o_dos" + suffix
+            global_cvt_2_ener_float(dos_raw), axis=1, name="o_dos" + suffix
         )
 
         model_dict = {}

source/tests/common/dpmodel/test_deep_dos.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Test global-DOS-only inference on the dpmodel backend.
+
+``DeepDOS.eval`` used to read the atomic ``dos`` output unconditionally and sum
+it to obtain the global DOS.  The dpmodel (and JAX) backends only return the
+atomic ``OUT`` variables when ``atomic=True``; for ``atomic=False`` they return
+the reduced ``dos_redu`` instead, so reading ``results["dos"]`` raised
+``KeyError``.  A global-DOS-only path (e.g. ``dp test`` without atomic DOS
+labels) must use the reduced output.
+"""
+
+import tempfile
+import unittest
+from pathlib import (
+    Path,
+)
+
+import numpy as np
+
+from deepmd.dpmodel.model.model import get_model as get_model_dp
+from deepmd.dpmodel.utils.serialization import (
+    save_dp_model,
+)
+from deepmd.infer.deep_dos import (
+    DeepDOS,
+)
+
+
+def _dos_model_config() -> dict:
+    return {
+        "type_map": ["O", "H"],
+        "descriptor": {
+            "type": "se_e2_a",
+            "sel": [20, 20],
+            "rcut_smth": 1.8,
+            "rcut": 6.0,
+            "neuron": [2, 4, 8],
+            "resnet_dt": False,
+            "axis_neuron": 8,
+            "precision": "float64",
+            "type_one_side": True,
+            "seed": 1,
+        },
+        "fitting_net": {
+            "type": "dos",
+            "numb_dos": 2,
+            "neuron": [4, 4, 4],
+            "resnet_dt": True,
+            "numb_fparam": 0,
+            "precision": "float64",
+            "seed": 1,
+        },
+    }
+
+
+class TestDeepDOSDPModel(unittest.TestCase):
+    def setUp(self) -> None:
+        model = get_model_dp(_dos_model_config())
+        self.tmpdir = tempfile.TemporaryDirectory()
+        model_file = str(Path(self.tmpdir.name) / "dos.dp")
+        save_dp_model(model_file, {"model": model.serialize()})
+        self.dp = DeepDOS(model_file)
+        rng = np.random.default_rng(0)
+        self.coords = rng.random([1, 6, 3]) * 4.0
+        self.cells = (np.eye(3) * 10.0).reshape(1, 9)
+        self.atypes = np.array([[0, 1, 1, 0, 1, 1]], dtype=np.int32)
+
+    def tearDown(self) -> None:
+        self.tmpdir.cleanup()
+
+    def test_global_dos_only(self) -> None:
+        # atomic=False must return the global DOS via the reduced output,
+        # without requiring the atomic `dos` key.
+        (dos,) = self.dp.eval(self.coords, self.cells, self.atypes, atomic=False)
+        self.assertEqual(dos.shape, (1, self.dp.get_numb_dos()))
+
+    def test_global_matches_atomic_sum(self) -> None:
+        # The reduced global DOS must equal the sum of the atomic DOS.
+        (dos,) = self.dp.eval(self.coords, self.cells, self.atypes, atomic=False)
+        _, atomic_dos = self.dp.eval(self.coords, self.cells, self.atypes, atomic=True)
+        np.testing.assert_allclose(dos, np.sum(atomic_dos, axis=1))
+
+
+if __name__ == "__main__":
+    unittest.main()

source/tests/jax/test_deep_dos.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Global-DOS-only inference on the JAX backend (StableHLO export).
+
+Mirrors the dpmodel DOS inference test.  The JAX evaluator wraps an ``HLO``
+object with no live model, so ``numb_dos`` must be carried through the StableHLO
+export for ``DeepDOS.eval`` to reshape the reduced DOS output correctly.
+"""
+
+import tempfile
+import unittest
+from pathlib import (
+    Path,
+)
+
+import numpy as np
+
+from deepmd.dpmodel.model.model import get_model as get_model_dp
+from deepmd.infer.deep_dos import (
+    DeepDOS,
+)
+from deepmd.jax.utils.serialization import (
+    deserialize_to_file,
+)
+
+
+def _dos_model_config() -> dict:
+    return {
+        "type_map": ["O", "H"],
+        "descriptor": {
+            "type": "se_e2_a",
+            "sel": [20, 20],
+            "rcut_smth": 1.8,
+            "rcut": 6.0,
+            "neuron": [2, 4, 8],
+            "resnet_dt": False,
+            "axis_neuron": 8,
+            "precision": "float64",
+            "type_one_side": True,
+            "seed": 1,
+        },
+        "fitting_net": {
+            "type": "dos",
+            "numb_dos": 2,
+            "neuron": [4, 4, 4],
+            "resnet_dt": True,
+            "numb_fparam": 0,
+            "precision": "float64",
+            "seed": 1,
+        },
+    }
+
+
+class TestDeepDOSJAX(unittest.TestCase):
+    def setUp(self) -> None:
+        config = _dos_model_config()
+        model = get_model_dp(config)
+        self.tmpdir = tempfile.TemporaryDirectory()
+        model_file = str(Path(self.tmpdir.name) / "dos.hlo")
+        deserialize_to_file(
+            model_file,
+            {"model": model.serialize(), "model_def_script": {"model": config}},
+        )
+        self.dp = DeepDOS(model_file)
+        rng = np.random.default_rng(0)
+        self.coords = rng.random([1, 6, 3]) * 4.0
+        self.cells = (np.eye(3) * 10.0).reshape(1, 9)
+        self.atypes = np.array([[0, 1, 1, 0, 1, 1]], dtype=np.int32)
+
+    def tearDown(self) -> None:
+        self.tmpdir.cleanup()
+
+    def test_numb_dos_survives_export(self) -> None:
+        self.assertEqual(self.dp.get_numb_dos(), 2)
+
+    def test_global_dos_only(self) -> None:
+        (dos,) = self.dp.eval(self.coords, self.cells, self.atypes, atomic=False)
+        self.assertEqual(dos.shape, (1, 2))
+
+
+if __name__ == "__main__":
+    unittest.main()