From 5dc1dd1b59fd79ccab8437d402546b4e5c41344d Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 07:15:30 +0800 Subject: [PATCH 1/2] fix(dpmodel): support parameter shorthand in DeepEval Normalize documented frame and atomic parameter shorthand before automatic batching. Add direct-backend coverage for shared inputs, forced batch splits, and invalid parameter sizes. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/dpmodel/infer/deep_eval.py | 34 +++ .../infer/test_dpmodel_deep_eval_params.py | 193 ++++++++++++++++++ 2 files changed, 227 insertions(+) create mode 100644 source/tests/infer/test_dpmodel_deep_eval_params.py diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py index e86322866b..10e83e1ffb 100644 --- a/deepmd/dpmodel/infer/deep_eval.py +++ b/deepmd/dpmodel/infer/deep_eval.py @@ -240,6 +240,40 @@ def eval( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) + # Canonicalize documented shorthand before AutoBatchSize sees it. + # Otherwise a shared per-atom aparam array is sliced on its atom axis, + # while one-dimensional shared inputs cannot be sliced by frame. + if fparam is not None: + fparam = np.asarray(fparam) + dim_fparam = self.get_dim_fparam() + if fparam.size == numb_test * dim_fparam: + fparam = fparam.reshape(numb_test, dim_fparam) + elif fparam.size == dim_fparam: + fparam = np.tile(fparam.reshape(1, dim_fparam), (numb_test, 1)) + else: + raise RuntimeError( + "got wrong size of frame param, should be either " + f"{numb_test} x {dim_fparam} or {dim_fparam}" + ) + if aparam is not None: + aparam = np.asarray(aparam) + dim_aparam = self.get_dim_aparam() + if aparam.size == numb_test * natoms * dim_aparam: + aparam = aparam.reshape(numb_test, natoms, dim_aparam) + elif aparam.size == natoms * dim_aparam: + aparam = np.tile( + aparam.reshape(1, natoms, dim_aparam), (numb_test, 1, 1) + ) + elif aparam.size == dim_aparam: + aparam = np.tile( + aparam.reshape(1, 1, dim_aparam), (numb_test, natoms, 1) + ) + else: + raise RuntimeError( + "got wrong size of atomic param, should be either " + f"{numb_test} x {natoms} x {dim_aparam} or " + f"{natoms} x {dim_aparam} or {dim_aparam}" + ) 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 diff --git a/source/tests/infer/test_dpmodel_deep_eval_params.py b/source/tests/infer/test_dpmodel_deep_eval_params.py new file mode 100644 index 0000000000..9d8801ff9a --- /dev/null +++ b/source/tests/infer/test_dpmodel_deep_eval_params.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Direct-backend tests for dpmodel DeepEval parameter shorthand.""" + +import copy + +import numpy as np +import pytest + +from deepmd.dpmodel.model.model import ( + get_model, +) +from deepmd.dpmodel.utils.serialization import ( + save_dp_model, +) +from deepmd.infer import ( + DeepEval, +) + + +MODEL_CONFIG = { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": [20, 20], + "rcut_smth": 0.5, + "rcut": 6.0, + "neuron": [3, 6], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "type": "ener", + "neuron": [5, 5], + "resnet_dt": True, + "precision": "float64", + "atom_ener": [], + "seed": 1, + "numb_fparam": 2, + "numb_aparam": 2, + }, +} +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, +).reshape(len(ATOM_TYPES), 3) +BOX = np.diag([13.0, 13.0, 13.0]) +NFRAMES = 2 +COORDS = np.tile(COORD, (NFRAMES, 1, 1)) +BOXES = np.tile(BOX, (NFRAMES, 1, 1)) +FPARAM = np.array([0.25, -0.5], dtype=np.float64) +APARAM_PER_ATOM = ( + np.arange(len(ATOM_TYPES) * 2, dtype=np.float64).reshape(len(ATOM_TYPES), 2) / 10.0 +) +APARAM_GLOBAL = np.array([0.3, -0.2], dtype=np.float64) + + +@pytest.fixture(scope="module") +def model_file(tmp_path_factory: pytest.TempPathFactory) -> str: + """Serialize a two-dimensional fparam/aparam model for direct inference.""" + model = get_model(copy.deepcopy(MODEL_CONFIG)) + path = tmp_path_factory.mktemp("dpmodel_params") / "model.dp" + save_dp_model( + str(path), + { + "model": model.serialize(), + "model_def_script": MODEL_CONFIG, + "backend": "dpmodel", + }, + ) + return str(path) + + +def _backend(model_file: str, auto_batch_size: bool | int): + """Return the low-level backend without public input normalization.""" + return DeepEval(model_file, auto_batch_size=auto_batch_size).deep_eval + + +def _assert_outputs_equal(actual: dict, expected: dict) -> None: + assert actual.keys() == expected.keys() + for name in actual: + np.testing.assert_allclose(actual[name], expected[name], equal_nan=True) + + +@pytest.mark.parametrize("auto_batch_size", [False, len(ATOM_TYPES)]) +def test_shared_fparam_is_tiled_before_batching( + model_file: str, auto_batch_size: bool | int +) -> None: + """A single frame-parameter vector applies to every input frame.""" + evaluator = _backend(model_file, auto_batch_size) + full_aparam = np.tile(APARAM_PER_ATOM, (NFRAMES, 1, 1)) + expected = evaluator.eval( + COORDS, + BOXES, + ATOM_TYPES, + fparam=np.tile(FPARAM, (NFRAMES, 1)), + aparam=full_aparam, + ) + actual = evaluator.eval( + COORDS, + BOXES, + ATOM_TYPES, + fparam=FPARAM.tolist(), + aparam=full_aparam, + ) + + _assert_outputs_equal(actual, expected) + + +@pytest.mark.parametrize("auto_batch_size", [False, len(ATOM_TYPES)]) +@pytest.mark.parametrize( + ("shared_aparam", "full_aparam"), + [ + ( + APARAM_PER_ATOM, + np.tile(APARAM_PER_ATOM, (NFRAMES, 1, 1)), + ), + ( + APARAM_GLOBAL, + np.tile(APARAM_GLOBAL, (NFRAMES, len(ATOM_TYPES), 1)), + ), + ], + ids=("per-atom", "all-atoms"), +) +def test_shared_aparam_is_tiled_before_batching( + model_file: str, + auto_batch_size: bool | int, + shared_aparam: np.ndarray, + full_aparam: np.ndarray, +) -> None: + """Both documented atomic-parameter shorthand forms are frame-major.""" + evaluator = _backend(model_file, auto_batch_size) + full_fparam = np.tile(FPARAM, (NFRAMES, 1)) + expected = evaluator.eval( + COORDS, + BOXES, + ATOM_TYPES, + fparam=full_fparam, + aparam=full_aparam, + ) + actual = evaluator.eval( + COORDS, + BOXES, + ATOM_TYPES, + fparam=full_fparam, + aparam=shared_aparam, + ) + + _assert_outputs_equal(actual, expected) + + +@pytest.mark.parametrize( + ("parameter", "value", "message"), + [ + ("fparam", np.zeros(3), "wrong size of frame param"), + ("aparam", np.zeros(3), "wrong size of atomic param"), + ], +) +def test_invalid_parameter_size_has_clear_error( + model_file: str, parameter: str, value: np.ndarray, message: str +) -> None: + """Reject invalid sizes before NumPy emits an opaque reshape error.""" + evaluator = _backend(model_file, auto_batch_size=False) + parameters = { + "fparam": np.tile(FPARAM, (NFRAMES, 1)), + "aparam": np.tile(APARAM_PER_ATOM, (NFRAMES, 1, 1)), + } + parameters[parameter] = value + + with pytest.raises(RuntimeError, match=message): + evaluator.eval(COORDS, BOXES, ATOM_TYPES, **parameters) From 1eb3feb3b351cd45a3befffd7e846fe80966ba30 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:26:28 +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_params.py | 1 - 1 file changed, 1 deletion(-) diff --git a/source/tests/infer/test_dpmodel_deep_eval_params.py b/source/tests/infer/test_dpmodel_deep_eval_params.py index 9d8801ff9a..64f3314280 100644 --- a/source/tests/infer/test_dpmodel_deep_eval_params.py +++ b/source/tests/infer/test_dpmodel_deep_eval_params.py @@ -16,7 +16,6 @@ DeepEval, ) - MODEL_CONFIG = { "type_map": ["O", "H"], "descriptor": {