diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py index 255ea4a385..aa4c9fe79b 100644 --- a/deepmd/dpmodel/infer/deep_eval.py +++ b/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.""" diff --git a/deepmd/dpmodel/model/dos_model.py b/deepmd/dpmodel/model/dos_model.py index 4d854ae007..bc9ee4fcb0 100644 --- a/deepmd/dpmodel/model/dos_model.py +++ b/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, diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 2cf3e4c12b..d907e3db62 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/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() diff --git a/deepmd/infer/deep_dos.py b/deepmd/infer/deep_dos.py index 4e63a156c1..99a03f8e78 100644 --- a/deepmd/infer/deep_dos.py +++ b/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() diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py index 09b5783ba9..d5d769784a 100644 --- a/deepmd/jax/infer/deep_eval.py +++ b/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.""" diff --git a/deepmd/jax/model/hlo.py b/deepmd/jax/model/hlo.py index 8c1e85c59c..0162850eaa 100644 --- a/deepmd/jax/model/hlo.py +++ b/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 diff --git a/deepmd/jax/utils/serialization.py b/deepmd/jax/utils/serialization.py index 59240b41ab..d77fabe6a4 100644 --- a/deepmd/jax/utils/serialization.py +++ b/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(), diff --git a/deepmd/tf/model/dos.py b/deepmd/tf/model/dos.py index 5d53abecf9..5913f42d91 100644 --- a/deepmd/tf/model/dos.py +++ b/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 = {} diff --git a/source/tests/common/dpmodel/test_deep_dos.py b/source/tests/common/dpmodel/test_deep_dos.py new file mode 100644 index 0000000000..7c7a2a6492 --- /dev/null +++ b/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() diff --git a/source/tests/jax/test_deep_dos.py b/source/tests/jax/test_deep_dos.py new file mode 100644 index 0000000000..10d2d4afc1 --- /dev/null +++ b/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() diff --git a/source/tests/tf/test_model_dos.py b/source/tests/tf/test_model_dos.py index a37eae9d75..90fce8bd82 100644 --- a/source/tests/tf/test_model_dos.py +++ b/source/tests/tf/test_model_dos.py @@ -329,6 +329,73 @@ def test_model(self) -> None: ) places = 4 + # o_dos is now [nframes, numb_dos] and o_atom_dos [nframes, natoms, numb_dos] + pred_dos = np.reshape(pred_dos, [-1]) + pred_atom_dos = np.reshape(pred_atom_dos, [natoms, numb_dos]) np.testing.assert_almost_equal(pred_dos, ref_dos, places) np.testing.assert_almost_equal(np.sum(pred_atom_dos, axis=0), ref_dos, places) np.testing.assert_almost_equal(pred_atom_dos[0], ref_ados_1, places) + + def test_multiframe_global_equals_atomic_sum(self) -> None: + # The global DOS must equal the per-frame sum of the atomic DOS. The + # reduction used to reshape/reduce over the wrong axis, mixing frames + # together, so this only held for a single frame. + jdata = j_loader("train_dos.json") + systems = jdata["training"]["systems"] + rcut = jdata["model"]["descriptor"]["rcut"] + data = DataSystem(systems, "set", 1, 1, rcut, run_opt=None) + test_data = data.get_test() + numb_dos = 20 + natoms = test_data["type"].shape[1] + + jdata["model"]["fitting_net"]["numb_dos"] = numb_dos + jdata["model"]["descriptor"]["neuron"] = [5, 5, 5] + jdata["model"]["descriptor"]["axis_neuron"] = 2 + jdata["model"]["descriptor"].pop("type", None) + descrpt = DescrptSeA(**jdata["model"]["descriptor"], uniform_seed=True) + jdata["model"]["fitting_net"].pop("type", None) + jdata["model"]["fitting_net"]["ntypes"] = descrpt.get_ntypes() + jdata["model"]["fitting_net"]["dim_descrpt"] = descrpt.get_dim_out() + fitting = DOSFitting(**jdata["model"]["fitting_net"], uniform_seed=True) + model = DOSModel(descrpt, fitting) + + model._compute_input_stat( + { + "coord": [test_data["coord"]], + "box": [test_data["box"]], + "type": [test_data["type"]], + "natoms_vec": [test_data["natoms_vec"]], + "default_mesh": [test_data["default_mesh"]], + } + ) + + t_coord = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name="i_coord_mf") + t_type = tf.placeholder(tf.int32, [None], name="i_type_mf") + t_natoms = tf.placeholder(tf.int32, [model.ntypes + 2], name="i_natoms_mf") + t_box = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None, 9], name="i_box_mf") + t_mesh = tf.placeholder(tf.int32, [None], name="i_mesh_mf") + + model_pred = model.build( + t_coord, t_type, t_natoms, t_box, t_mesh, None, suffix="dos_mf", reuse=False + ) + dos = model_pred["dos"] + atom_dos = model_pred["atom_dos"] + + nframes = 2 + feed_dict_test = { + t_coord: np.tile(np.reshape(test_data["coord"][:1, :], [-1]), nframes), + t_type: np.tile(np.reshape(test_data["type"][:1, :], [-1]), nframes), + t_natoms: test_data["natoms_vec"], + t_box: np.tile(test_data["box"][:1, :], (nframes, 1)), + t_mesh: test_data["default_mesh"], + } + + with self.cached_session() as sess: + sess.run(tf.global_variables_initializer()) + pred_dos, pred_atom_dos = sess.run( + [dos, atom_dos], feed_dict=feed_dict_test + ) + + pred_dos = np.reshape(pred_dos, [nframes, numb_dos]) + pred_atom_dos = np.reshape(pred_atom_dos, [nframes, natoms, numb_dos]) + np.testing.assert_almost_equal(pred_dos, np.sum(pred_atom_dos, axis=1), 6)