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..0467a1ed3a --- /dev/null +++ b/source/tests/infer/test_dpmodel_deep_eval_spin.py @@ -0,0 +1,98 @@ +# 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)