From f4f5a24e9d1b244bc3990f4913dd5ec1cf9a382f Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 07:06:35 +0800 Subject: [PATCH 1/2] fix(dpmodel): forward spin inputs in DeepEval Normalize spin tensors before automatic batching and pass supported extra model inputs through the dpmodel evaluator. Preserve the magnetic mask in non-atomic DeepPot results and add multi-frame regression coverage. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/dpmodel/infer/deep_eval.py | 36 ++++++- .../infer/test_dpmodel_deep_eval_spin.py | 99 +++++++++++++++++++ 2 files changed, 131 insertions(+), 4 deletions(-) create mode 100644 source/tests/infer/test_dpmodel_deep_eval_spin.py diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py index e86322866b..d735ee3fa4 100644 --- a/deepmd/dpmodel/infer/deep_eval.py +++ b/deepmd/dpmodel/infer/deep_eval.py @@ -240,9 +240,32 @@ def eval( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) + model_kwargs = dict(kwargs) + if self.get_has_spin(): + spin = model_kwargs.get("spin") + if spin is None: + raise ValueError("spin must be provided when evaluating a spin model") + spin = np.asarray(spin) + expected_spin_size = numb_test * natoms * 3 + if spin.size != expected_spin_size: + raise ValueError( + "spin must contain exactly " + f"{expected_spin_size} values for {numb_test} frame(s) and " + f"{natoms} atom(s), but received {spin.size}" + ) + # AutoBatchSize slices only arrays with a frame axis. Normalize a + # flattened public-API input before batching so each model call + # receives the spins belonging to its coordinate frames. + model_kwargs["spin"] = spin.reshape(numb_test, natoms, 3) request_defs = self._get_request_defs(atomic) out = self._eval_func(self._eval_model, numb_test, natoms)( - coords, cells, atom_types, fparam, aparam, request_defs + coords, + cells, + atom_types, + fparam, + aparam, + request_defs, + **model_kwargs, ) # ``AutoBatchSize.execute_all`` unwraps a single-output result out of # its tuple, which would make ``zip`` iterate over the array's frame @@ -287,6 +310,9 @@ def _get_request_defs(self, atomic: bool) -> list[OutputVariableDef]: OutputVariableCategory.DERV_R, OutputVariableCategory.DERV_C_REDU, ) + # DeepPot always returns the magnetic-atom mask for spin + # models, even when per-atom energy was not requested. + or x.name == "mask_mag" ] def _eval_func(self, inner_func: Callable, numb_test: int, natoms: int) -> Callable: @@ -342,6 +368,7 @@ def _eval_model( fparam: Array | None, aparam: Array | None, request_defs: list[OutputVariableDef], + **model_kwargs: Any, ) -> dict[str, Array]: model = self.dp @@ -370,14 +397,15 @@ def _eval_model( do_atomic_virial = any( x.category == OutputVariableCategory.DERV_C_REDU for x in request_defs ) - batch_output = model( - coord_input, - type_input, + # Evaluator-owned arguments take precedence over extra model inputs so + # callers cannot accidentally bypass normalization performed above. + model_kwargs.update( box=box_input, fparam=fparam_input, aparam=aparam_input, do_atomic_virial=do_atomic_virial, ) + batch_output = model(coord_input, type_input, **model_kwargs) if isinstance(batch_output, tuple): batch_output = batch_output[0] diff --git a/source/tests/infer/test_dpmodel_deep_eval_spin.py b/source/tests/infer/test_dpmodel_deep_eval_spin.py new file mode 100644 index 0000000000..46a2ec5ec5 --- /dev/null +++ b/source/tests/infer/test_dpmodel_deep_eval_spin.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Regression tests for spin inputs in the dpmodel DeepEval backend.""" + +from pathlib import ( + Path, +) + +import numpy as np +import pytest + +from deepmd.infer import ( + DeepEval, +) + + +MODEL_FILE = Path(__file__).with_name("deeppot_dpa_spin.yaml") +ATOM_TYPES = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32) +COORD = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 0.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ], + dtype=np.float64, +) +SPIN = np.array( + [ + 0.13, + 0.02, + 0.03, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.14, + 0.10, + 0.12, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + ], + dtype=np.float64, +) +BOX = np.diag([13.0, 13.0, 13.0]).reshape(-1) + + +def test_spin_is_forwarded_and_sliced_by_auto_batch() -> None: + """A flattened multi-frame spin input must follow coordinate batching.""" + coords = np.concatenate([COORD, COORD]) + boxes = np.concatenate([BOX, BOX]) + spins = np.concatenate([SPIN, 2.0 * SPIN]) + + # Six atoms per batch forces the two frames through separate model calls. + batched_eval = DeepEval(MODEL_FILE, auto_batch_size=len(ATOM_TYPES)) + actual = batched_eval.eval(coords, boxes, ATOM_TYPES, spin=spins) + + unbatched_eval = DeepEval(MODEL_FILE, auto_batch_size=False) + expected_by_frame = [ + unbatched_eval.eval(COORD, BOX, ATOM_TYPES, spin=frame_spin) + for frame_spin in (SPIN, 2.0 * SPIN) + ] + expected = tuple( + np.concatenate([frame_result[index] for frame_result in expected_by_frame]) + for index in range(len(actual)) + ) + + assert len(actual) == 5 # energy, force, virial, magnetic force, magnetic mask + for actual_value, expected_value in zip(actual, expected, strict=True): + np.testing.assert_allclose(actual_value, expected_value, equal_nan=True) + # Distinct spin vectors must reach the model instead of being ignored. + assert actual[0][0, 0] != pytest.approx(actual[0][1, 0]) + + +def test_spin_model_requires_spin_input() -> None: + """Report the missing model input at the evaluator boundary.""" + evaluator = DeepEval(MODEL_FILE, auto_batch_size=False) + + with pytest.raises(ValueError, match="spin must be provided"): + evaluator.eval(COORD, BOX, ATOM_TYPES) From c09d167ebb621a65f541e8fd92415fbfadef7c86 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 16 Jul 2026 23:08:14 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/infer/test_dpmodel_deep_eval_spin.py | 1 - 1 file changed, 1 deletion(-) diff --git a/source/tests/infer/test_dpmodel_deep_eval_spin.py b/source/tests/infer/test_dpmodel_deep_eval_spin.py index 46a2ec5ec5..0467a1ed3a 100644 --- a/source/tests/infer/test_dpmodel_deep_eval_spin.py +++ b/source/tests/infer/test_dpmodel_deep_eval_spin.py @@ -12,7 +12,6 @@ DeepEval, ) - MODEL_FILE = Path(__file__).with_name("deeppot_dpa_spin.yaml") ATOM_TYPES = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32) COORD = np.array(