Skip to content

Commit 5cc4c57

Browse files
fix(dpmodel): forward spin inputs in DeepEval (#5852)
Closes #5661. ## Summary - normalize and validate spin inputs before automatic batching so flattened multi-frame inputs are sliced with their coordinate frames - forward extra model inputs through the dpmodel `DeepEval` adapter while keeping evaluator-owned inputs canonical - retain the magnetic-atom mask in non-atomic `DeepPot` results - add public-API regressions for forced one-frame batching and missing spin inputs ## Why existing tests missed this Existing dpmodel spin tests exercised direct model calls and serialization parity, not the `DeepEval` adapter. Public spin inference coverage targeted the PyTorch and PT2 backends, which already have dedicated spin forwarding paths. Those fixtures were also primarily single-frame, so they would not detect a flat spin tensor being left unsliced when automatic batching splits a multi-frame evaluation. ## Validation - `pytest source/tests/infer/test_dpmodel_deep_eval_spin.py -q` (2 passed) - `source/tests/common/test_auto_batch_size.py::TestAutoBatchSize::test_execute_all` (passed) - dpmodel non-spin `DeepEval` smoke evaluation (passed) - `ruff format .` (1664 files unchanged) - `ruff check .` (passed) - `git diff --check` (passed) Coding agent: Codex Codex version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning effort: xhigh <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved inference handling for spin-enabled models, including validation and consistent batching of spin data. * Preserved optional charge and spin inputs during evaluation. * Ensured non-atomic evaluations include magnetic mask output when applicable. * Prevented unsupported inputs from being passed to models that do not use them. * **Tests** * Added regression coverage for spin input batching, required spin data, and irrelevant input handling. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: njzjz-bot <njzjz.bot@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: njzjz-bot <njzjz-bot@users.noreply.github.com>
1 parent f573ca0 commit 5cc4c57

2 files changed

Lines changed: 158 additions & 4 deletions

File tree

deepmd/dpmodel/infer/deep_eval.py

Lines changed: 42 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -240,9 +240,39 @@ def eval(
240240
natoms, numb_test = self._get_natoms_and_nframes(
241241
coords, atom_types, len(atom_types.shape) > 1
242242
)
243+
# The public evaluator accepts a superset of backend/model inputs (for
244+
# example ``efield`` is TensorFlow-only). Forward only the dpmodel
245+
# inputs understood here instead of leaking unrelated ``None`` values
246+
# into concrete model ``call`` signatures.
247+
model_kwargs = {}
248+
charge_spin = kwargs.get("charge_spin")
249+
if charge_spin is not None:
250+
model_kwargs["charge_spin"] = charge_spin
251+
if self.get_has_spin():
252+
spin = kwargs.get("spin")
253+
if spin is None:
254+
raise ValueError("spin must be provided when evaluating a spin model")
255+
spin = np.asarray(spin)
256+
expected_spin_size = numb_test * natoms * 3
257+
if spin.size != expected_spin_size:
258+
raise ValueError(
259+
"spin must contain exactly "
260+
f"{expected_spin_size} values for {numb_test} frame(s) and "
261+
f"{natoms} atom(s), but received {spin.size}"
262+
)
263+
# AutoBatchSize slices only arrays with a frame axis. Normalize a
264+
# flattened public-API input before batching so each model call
265+
# receives the spins belonging to its coordinate frames.
266+
model_kwargs["spin"] = spin.reshape(numb_test, natoms, 3)
243267
request_defs = self._get_request_defs(atomic)
244268
out = self._eval_func(self._eval_model, numb_test, natoms)(
245-
coords, cells, atom_types, fparam, aparam, request_defs
269+
coords,
270+
cells,
271+
atom_types,
272+
fparam,
273+
aparam,
274+
request_defs,
275+
**model_kwargs,
246276
)
247277
# ``AutoBatchSize.execute_all`` unwraps a single-output result out of
248278
# 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]:
287317
OutputVariableCategory.DERV_R,
288318
OutputVariableCategory.DERV_C_REDU,
289319
)
320+
# ``mask_mag`` is exported directly by spin graphs but does not
321+
# fit the category filter. Adding all OUT variables would also
322+
# request atom energy and the general atom mask at
323+
# ``atomic=False``, so keep this low-cost compatibility output
324+
# explicit instead of widening the category set.
325+
or x.name == "mask_mag"
290326
]
291327

292328
def _eval_func(self, inner_func: Callable, numb_test: int, natoms: int) -> Callable:
@@ -342,6 +378,7 @@ def _eval_model(
342378
fparam: Array | None,
343379
aparam: Array | None,
344380
request_defs: list[OutputVariableDef],
381+
**model_kwargs: Any,
345382
) -> dict[str, Array]:
346383
model = self.dp
347384

@@ -370,14 +407,15 @@ def _eval_model(
370407
do_atomic_virial = any(
371408
x.category == OutputVariableCategory.DERV_C_REDU for x in request_defs
372409
)
373-
batch_output = model(
374-
coord_input,
375-
type_input,
410+
# Evaluator-owned arguments take precedence over extra model inputs so
411+
# callers cannot accidentally bypass normalization performed above.
412+
model_kwargs.update(
376413
box=box_input,
377414
fparam=fparam_input,
378415
aparam=aparam_input,
379416
do_atomic_virial=do_atomic_virial,
380417
)
418+
batch_output = model(coord_input, type_input, **model_kwargs)
381419
if isinstance(batch_output, tuple):
382420
batch_output = batch_output[0]
383421

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
# SPDX-License-Identifier: LGPL-3.0-or-later
2+
"""Regression tests for spin inputs in the dpmodel DeepEval backend."""
3+
4+
from pathlib import (
5+
Path,
6+
)
7+
8+
import numpy as np
9+
import pytest
10+
11+
from deepmd.infer import (
12+
DeepEval,
13+
)
14+
15+
MODEL_FILE = Path(__file__).with_name("deeppot_dpa_spin.yaml")
16+
PLAIN_MODEL_FILE = Path(__file__).with_name("deeppot_dpa.yaml")
17+
ATOM_TYPES = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32)
18+
COORD = np.array(
19+
[
20+
12.83,
21+
2.56,
22+
2.18,
23+
12.09,
24+
2.87,
25+
2.74,
26+
0.25,
27+
3.32,
28+
1.68,
29+
3.36,
30+
3.00,
31+
1.81,
32+
3.51,
33+
2.51,
34+
2.60,
35+
4.27,
36+
3.22,
37+
1.56,
38+
],
39+
dtype=np.float64,
40+
)
41+
SPIN = np.array(
42+
[
43+
0.13,
44+
0.02,
45+
0.03,
46+
0.0,
47+
0.0,
48+
0.0,
49+
0.0,
50+
0.0,
51+
0.0,
52+
0.14,
53+
0.10,
54+
0.12,
55+
0.0,
56+
0.0,
57+
0.0,
58+
0.0,
59+
0.0,
60+
0.0,
61+
],
62+
dtype=np.float64,
63+
)
64+
BOX = np.diag([13.0, 13.0, 13.0]).reshape(-1)
65+
66+
67+
def test_spin_is_forwarded_and_sliced_by_auto_batch() -> None:
68+
"""A flattened multi-frame spin input must follow coordinate batching."""
69+
coords = np.concatenate([COORD, COORD])
70+
boxes = np.concatenate([BOX, BOX])
71+
spins = np.concatenate([SPIN, 2.0 * SPIN])
72+
73+
# Six atoms per batch forces the two frames through separate model calls.
74+
batched_eval = DeepEval(MODEL_FILE, auto_batch_size=len(ATOM_TYPES))
75+
actual = batched_eval.eval(coords, boxes, ATOM_TYPES, spin=spins)
76+
77+
unbatched_eval = DeepEval(MODEL_FILE, auto_batch_size=False)
78+
expected_by_frame = [
79+
unbatched_eval.eval(COORD, BOX, ATOM_TYPES, spin=frame_spin)
80+
for frame_spin in (SPIN, 2.0 * SPIN)
81+
]
82+
expected = tuple(
83+
np.concatenate([frame_result[index] for frame_result in expected_by_frame])
84+
for index in range(len(actual))
85+
)
86+
87+
assert len(actual) == 5 # energy, force, virial, magnetic force, magnetic mask
88+
for actual_value, expected_value in zip(actual, expected, strict=True):
89+
np.testing.assert_allclose(actual_value, expected_value, equal_nan=True)
90+
assert not np.isnan(actual[-1]).any()
91+
# Distinct spin vectors must reach the model instead of being ignored.
92+
assert actual[0][0, 0] != pytest.approx(actual[0][1, 0])
93+
94+
95+
def test_spin_model_requires_spin_input() -> None:
96+
"""Report the missing model input at the evaluator boundary."""
97+
evaluator = DeepEval(MODEL_FILE, auto_batch_size=False)
98+
99+
with pytest.raises(ValueError, match="spin must be provided"):
100+
evaluator.eval(COORD, BOX, ATOM_TYPES)
101+
102+
103+
def test_plain_model_ignores_inputs_for_other_backends() -> None:
104+
"""The generic ``dp test`` kwarg set must not reach model ``call``."""
105+
evaluator = DeepEval(PLAIN_MODEL_FILE, auto_batch_size=False)
106+
107+
result = evaluator.eval(
108+
COORD,
109+
BOX,
110+
ATOM_TYPES,
111+
efield=None,
112+
spin=None,
113+
charge_spin=None,
114+
)
115+
116+
assert len(result) == 3

0 commit comments

Comments
 (0)