fix: make global DOS inference work on the dpmodel and JAX backends (deepmodeling#5722)

## Problem

Fixes deepmodeling#5674. Three related defects in global DOS inference, uncovered in
sequence:

1. `DeepDOS.eval` unconditionally read the atomic `dos` output and
summed it, even for `atomic=False`. The dpmodel and JAX backends only
return the atomic output when `atomic=True`; on the global-DOS-only path
(e.g. `dp test` without atomic DOS labels) `results["dos"]` raised
`KeyError`. TF and PyTorch masked this because they always include the
atomic output.

2. Fixing the `KeyError` exposed that dpmodel and JAX DOS inference was
broken more deeply: both `DeepEval.get_numb_dos` implementations
hard-returned `0`, so the DOS reshape target was `(nframes, 0)` and
inference failed on every path, not just the missing-key case.

3. On TF, the global DOS did not equal the sum of the atomic DOS for
multi-frame inputs — even though, by construction of the model, it must.
`deepmd/tf/model/dos.py` reduced the atomic DOS with
`reshape([natoms[0], -1])` + `reduce_sum(axis=0)`, which sums across the
wrong axis and mixes atoms from different frames together. Single-frame
inputs happened to give the right answer, so no test caught it.

## Fix

Backend-agnostic (`deep_dos.py`): prefer the atomic `dos` output and sum
it whenever the backend returns it (this is the exact global DOS on
TF/PT, whose reduced output is not necessarily the plain sum), and fall
back to the reduced `dos_redu` only when the atomic output is absent
(dpmodel/JAX at `atomic=False`). Reading `dos` unconditionally is what
raised the original `KeyError`.

dpmodel: add `get_numb_dos` to the dpmodel `DOSModel` (mirroring the
PyTorch model), add a default `get_numb_dos` returning 0 on the shared
base model so non-DOS models can still be serialized, and delegate
`dpmodel/infer/deep_eval.py:get_numb_dos` to the model.

JAX: the evaluator wraps a deserialized `HLO` object with no live model,
so `numb_dos` is now persisted into the StableHLO export constants and
exposed via `HLO.get_numb_dos`; the `dos` output is registered in the
HLO `OUTPUT_DEFS` table; and `jax/infer/deep_eval.py:get_numb_dos`
delegates to the model. With these, JAX DOS inference works end to end.

TF: reduce the atomic DOS per frame — `reshape([-1, natoms[0],
numb_dos])` + `reduce_sum(axis=1)`, mirroring the energy model — so the
global DOS equals the atomic sum for multi-frame inputs.

## Test

- `source/tests/common/dpmodel/test_deep_dos.py`: builds a dpmodel DOS
model and evaluates it — `atomic=False` returns the global DOS
(`KeyError` on master), and the global DOS equals the sum of the atomic
DOS (guarding the `dos_redu == sum(dos)` invariant relied on by all
backends).
- `source/tests/jax/test_deep_dos.py`: exports a DOS model to `.hlo`,
checks `numb_dos` survives the round trip, and evaluates the global DOS.
- `source/tests/tf/test_model_dos.py`: adds
`test_multiframe_global_equals_atomic_sum`, which builds a two-frame DOS
graph and asserts the global DOS equals the per-frame atomic sum — this
fails on the old axis-0 reduction and passes after the per-frame fix.
The existing single-frame assertions were updated to the corrected
output shapes.

dpmodel and JAX DOS inference previously had no test, and the TF path
had only single-frame coverage; DOS was effectively exercised only where
the bugs were masked.


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

## Summary by CodeRabbit

* **Bug Fixes**
* Corrected DOS inference so global DOS is reported properly instead of
using a fixed default.
* Fixed multi-frame DOS aggregation to keep results separated by frame
and sum across atoms correctly.
* Improved consistency when using the model in different backends and
after export, so DOS output counts are preserved.

* **New Features**
* Added support for exposing DOS output counts in model inference and
export workflows.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

Author: wanghan-iapcm

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:
@@ -1000,6 +1008,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,14 @@ 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)
-
+        energy = results["dos_redu"].reshape(nframes, self.get_numb_dos())
         if atomic:
+            atomic_energy = results["dos"].reshape(nframes, natoms, self.get_numb_dos())
             return (
                 energy,
                 atomic_energy,
             )
-        else:
-            return (energy,)
+        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()