diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py index e86322866b..a20cc769fd 100644 --- a/deepmd/dpmodel/infer/deep_eval.py +++ b/deepmd/dpmodel/infer/deep_eval.py @@ -240,9 +240,39 @@ def eval( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) + # The public evaluator accepts a superset of backend/model inputs (for + # example ``efield`` is TensorFlow-only). Forward only the dpmodel + # inputs understood here instead of leaking unrelated ``None`` values + # into concrete model ``call`` signatures. + model_kwargs = {} + charge_spin = kwargs.get("charge_spin") + if charge_spin is not None: + model_kwargs["charge_spin"] = charge_spin + if self.get_has_spin(): + spin = 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 +317,12 @@ def _get_request_defs(self, atomic: bool) -> list[OutputVariableDef]: OutputVariableCategory.DERV_R, OutputVariableCategory.DERV_C_REDU, ) + # ``mask_mag`` is exported directly by spin graphs but does not + # fit the category filter. Adding all OUT variables would also + # request atom energy and the general atom mask at + # ``atomic=False``, so keep this low-cost compatibility output + # explicit instead of widening the category set. + or x.name == "mask_mag" ] def _eval_func(self, inner_func: Callable, numb_test: int, natoms: int) -> Callable: @@ -342,6 +378,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 +407,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..21fa41aeb0 --- /dev/null +++ b/source/tests/infer/test_dpmodel_deep_eval_spin.py @@ -0,0 +1,116 @@ +# 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") +PLAIN_MODEL_FILE = Path(__file__).with_name("deeppot_dpa.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) + assert not np.isnan(actual[-1]).any() + # 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) + + +def test_plain_model_ignores_inputs_for_other_backends() -> None: + """The generic ``dp test`` kwarg set must not reach model ``call``.""" + evaluator = DeepEval(PLAIN_MODEL_FILE, auto_batch_size=False) + + result = evaluator.eval( + COORD, + BOX, + ATOM_TYPES, + efield=None, + spin=None, + charge_spin=None, + ) + + assert len(result) == 3