fix(dpmodel): forward spin inputs in DeepEval

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]
 

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)