diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py index e86322866b..bfa4a80d96 100644 --- a/deepmd/dpmodel/infer/deep_eval.py +++ b/deepmd/dpmodel/infer/deep_eval.py @@ -40,6 +40,7 @@ from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper from deepmd.infer.deep_eval import ( DeepEvalBackend, + _standardize_fparam_aparam, ) from deepmd.infer.deep_polar import ( DeepPolar, @@ -240,6 +241,14 @@ def eval( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + numb_test, + natoms, + self.get_dim_fparam(), + self.get_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/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 915d216672..2621607cb7 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -30,6 +30,75 @@ import ase.neighborlist +def _standardize_fparam_aparam( + fparam: np.ndarray | list | None, + aparam: np.ndarray | list | None, + nframes: int, + natoms: int, + dim_fparam: int, + dim_aparam: int, +) -> tuple[np.ndarray | None, np.ndarray | None]: + """Normalize documented parameter shorthand to frame-major arrays. + + This normalization must happen before automatic batching. In particular, + an ``(natoms, dim_aparam)`` shared atomic parameter has an atom axis first; + a batcher would otherwise mistake that axis for frames and slice it. + + Parameters + ---------- + fparam : np.ndarray or list or None + Frame parameters in full ``(nframes, dim_fparam)`` form or shared + ``(dim_fparam,)`` shorthand. + aparam : np.ndarray or list or None + Atomic parameters in full ``(nframes, natoms, dim_aparam)`` form, + shared-per-atom ``(natoms, dim_aparam)`` shorthand, or shared-for-all + ``(dim_aparam,)`` shorthand. + nframes : int + Number of input frames. + natoms : int + Number of atoms in each frame. + dim_fparam : int + Number of frame-parameter components. + dim_aparam : int + Number of atomic-parameter components. + + Returns + ------- + normalized_fparam : np.ndarray or None + Frame parameters with shape ``(nframes, dim_fparam)`` when provided. + normalized_aparam : np.ndarray or None + Atomic parameters with shape ``(nframes, natoms, dim_aparam)`` when + provided. The public wrapper may subsequently flatten the last two + axes to preserve its historical backend ABI. + """ + if fparam is not None: + fparam = np.asarray(fparam) + if fparam.size == nframes * dim_fparam: + fparam = fparam.reshape(nframes, dim_fparam) + elif fparam.size == dim_fparam: + fparam = np.tile(fparam.reshape(1, dim_fparam), (nframes, 1)) + else: + raise RuntimeError( + "got wrong size of frame param, should be either " + f"{nframes} x {dim_fparam} or {dim_fparam}" + ) + if aparam is not None: + aparam = np.asarray(aparam) + if aparam.size == nframes * natoms * dim_aparam: + aparam = aparam.reshape(nframes, natoms, dim_aparam) + elif aparam.size == natoms * dim_aparam: + aparam = np.tile(aparam.reshape(1, natoms, dim_aparam), (nframes, 1, 1)) + elif aparam.size == dim_aparam: + aparam = np.tile(aparam.reshape(1, 1, dim_aparam), (nframes, natoms, 1)) + else: + raise RuntimeError( + "got wrong size of atomic param, should be either " + f"{nframes} x {natoms} x {dim_aparam} or " + f"{natoms} x {dim_aparam} or {dim_aparam}" + ) + return fparam, aparam + + class DeepEvalBackend(ABC): """Low-level Deep Evaluator interface. @@ -948,28 +1017,18 @@ def _standard_input( coords = coords.reshape(nframes, natoms, 3) if cells is not None: cells = cells.reshape(nframes, 3, 3) - if fparam is not None: - fdim = self.get_dim_fparam() - if fparam.size == nframes * fdim: - fparam = np.reshape(fparam, [nframes, fdim]) - elif fparam.size == fdim: - fparam = np.tile(fparam.reshape([-1]), [nframes, 1]) - else: - raise RuntimeError( - f"got wrong size of frame param, should be either {nframes} x {fdim} or {fdim}" - ) + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + nframes, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) if aparam is not None: - fdim = self.get_dim_aparam() - if aparam.size == nframes * natoms * fdim: - aparam = np.reshape(aparam, [nframes, natoms * fdim]) - elif aparam.size == natoms * fdim: - aparam = np.tile(aparam.reshape([-1]), [nframes, 1]) - elif aparam.size == fdim: - aparam = np.tile(aparam.reshape([-1]), [nframes, natoms]) - else: - raise RuntimeError( - f"got wrong size of frame param, should be either {nframes} x {natoms} x {fdim} or {natoms} x {fdim} or {fdim}" - ) + # Preserve the historical flattened backend ABI used by the public + # wrapper; backend adapters normalize it back to frame-major 3-D. + aparam = aparam.reshape(nframes, natoms * self.get_dim_aparam()) return coords, cells, atom_types, fparam, aparam, nframes, natoms def get_sel_type(self) -> list[int]: diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py index 0e6c11ede6..9c4a151e81 100644 --- a/deepmd/jax/infer/deep_eval.py +++ b/deepmd/jax/infer/deep_eval.py @@ -35,6 +35,7 @@ from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper from deepmd.infer.deep_eval import ( DeepEvalBackend, + _standardize_fparam_aparam, ) from deepmd.infer.deep_polar import ( DeepPolar, @@ -278,6 +279,14 @@ def eval( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + numb_test, + natoms, + self.get_dim_fparam(), + self.get_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, charge_spin, request_defs diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py index c8bb113495..5658a0e73c 100644 --- a/deepmd/pd/infer/deep_eval.py +++ b/deepmd/pd/infer/deep_eval.py @@ -28,6 +28,7 @@ from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper from deepmd.infer.deep_eval import ( DeepEvalBackend, + _standardize_fparam_aparam, ) from deepmd.infer.deep_polar import ( DeepGlobalPolar, @@ -376,6 +377,14 @@ def eval( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + numb_test, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) request_defs = self._get_request_defs(atomic) if "spin" not in kwargs or kwargs["spin"] is None: out = self._eval_func(self._eval_model, numb_test, natoms)( diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index dffdbee5e1..4f30448747 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -29,6 +29,7 @@ from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper from deepmd.infer.deep_eval import ( DeepEvalBackend, + _standardize_fparam_aparam, ) from deepmd.infer.deep_polar import ( DeepGlobalPolar, @@ -544,6 +545,14 @@ def eval( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + numb_test, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) request_defs = self._get_request_defs(atomic) if "spin" not in kwargs or kwargs["spin"] is None: out = self._eval_func(self._eval_model, numb_test, natoms)( @@ -1312,6 +1321,17 @@ def eval_embedding( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) + # Normalize shared parameter shorthand before auto batching. Otherwise + # a one-dimensional fparam/aparam is passed unchanged to every split, + # and _eval_embedding cannot reshape it to the split frame count. + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + numb_test, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) return self._eval_func(self._eval_embedding, numb_test, natoms)( coords, cells, atom_types, fparam, aparam, charge_spin, dtype ) diff --git a/deepmd/pt_expt/infer/deep_eval.py b/deepmd/pt_expt/infer/deep_eval.py index 1b97ef7f50..3daf55b825 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -41,6 +41,7 @@ from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper from deepmd.infer.deep_eval import ( DeepEvalBackend, + _standardize_fparam_aparam, ) from deepmd.infer.deep_polar import ( DeepPolar, @@ -848,10 +849,15 @@ def eval( Calculate the atomic energy and virial fparam The frame parameter. - The array should be of size nframes x dim_fparam. + The array can be of size : + - nframes x dim_fparam. + - dim_fparam. Then all frames are assumed to be provided with the same fparam. aparam - The atomic parameter. - The array should be of size nframes x natoms x dim_aparam. + The atomic parameter + The array can be of size : + - nframes x natoms x dim_aparam. + - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. + - dim_aparam. Then all frames and atoms are provided with the same aparam. charge_spin The charge and spin values for each frame. The array should be reshape-compatible with nframes x 2, where the first @@ -875,6 +881,14 @@ def eval( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + numb_test, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) request_defs = self._get_request_defs(atomic) spins = kwargs.get("spin") if self._is_spin and spins is None: @@ -2203,6 +2217,19 @@ def eval_descriptor( "eval_descriptor is not supported for this model type " f"({type(self._dpmodel).__name__})." ) + nframes = coords.shape[0] + natoms = len(atom_types) if len(atom_types.shape) == 1 else atom_types.shape[1] + # Canonicalize shared parameter shorthands before _prepare_nlist_inputs + # reshapes them; otherwise a shared per-atom array can be mistaken for + # a batch of frames. + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + nframes, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) ( ext_coord_t, ext_atype_t, @@ -2274,6 +2301,19 @@ def eval_fitting_last_layer( "eval_fitting_last_layer is not supported for this model type " f"({type(self._dpmodel).__name__})." ) + nframes = coords.shape[0] + natoms = len(atom_types) if len(atom_types.shape) == 1 else atom_types.shape[1] + # Canonicalize shared parameter shorthands before _prepare_nlist_inputs + # reshapes them; otherwise a shared per-atom array can be mistaken for + # a batch of frames. + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + nframes, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) ( ext_coord_t, ext_atype_t, diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 30107cb693..0209bc694f 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -32,6 +32,7 @@ ) from deepmd.infer.deep_eval import ( DeepEvalBackend, + _standardize_fparam_aparam, ) from deepmd.infer.deep_polar import ( DeepGlobalPolar, @@ -761,6 +762,14 @@ def eval( coords, atom_types, ) + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + numb_test, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) output = self._eval_func(self._eval_inner, numb_test, natoms)( coords, cells, @@ -1084,6 +1093,17 @@ def eval_descriptor( coords, atom_types, ) + # Canonicalize shared parameter shorthands before AutoBatchSize slices + # the frame axis; otherwise a shared per-atom array can be mistaken for + # a batch of frames. + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + numb_test, + natoms, + self.get_dim_fparam(), + self.get_dim_aparam(), + ) descriptor = self._eval_func(self._eval_descriptor_inner, numb_test, natoms)( coords, cells, diff --git a/deepmd/tf2/infer/deep_eval.py b/deepmd/tf2/infer/deep_eval.py index af5037e2e6..35f78957af 100644 --- a/deepmd/tf2/infer/deep_eval.py +++ b/deepmd/tf2/infer/deep_eval.py @@ -32,6 +32,7 @@ from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper from deepmd.infer.deep_eval import ( DeepEvalBackend, + _standardize_fparam_aparam, ) from deepmd.infer.deep_polar import ( DeepPolar, @@ -292,6 +293,14 @@ def eval( natoms, numb_test = self._get_natoms_and_nframes( coords, atom_types, len(atom_types.shape) > 1 ) + fparam, aparam = _standardize_fparam_aparam( + fparam, + aparam, + numb_test, + natoms, + self.get_dim_fparam(), + self.get_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/common/test_deep_eval_parameter_shorthand.py b/source/tests/common/test_deep_eval_parameter_shorthand.py new file mode 100644 index 0000000000..4aa18fe93a --- /dev/null +++ b/source/tests/common/test_deep_eval_parameter_shorthand.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for backend-level DeepEval parameter normalization.""" + +import numpy as np +import pytest + +from deepmd.infer.deep_eval import ( + _standardize_fparam_aparam, +) + +NFRAMES = 3 +NATOMS = 4 +DIM_FPARAM = 2 +DIM_APARAM = 2 +FPARAM = np.array([0.25, -0.5], dtype=np.float64) +APARAM_PER_ATOM = np.arange(NATOMS * DIM_APARAM, dtype=np.float64).reshape( + NATOMS, DIM_APARAM +) +APARAM_ALL_ATOMS = np.array([0.3, -0.2], dtype=np.float64) + + +@pytest.mark.parametrize( + ("fparam", "expected"), + [ + (FPARAM.tolist(), np.tile(FPARAM, (NFRAMES, 1))), + ( + np.arange(NFRAMES * DIM_FPARAM).reshape(NFRAMES, DIM_FPARAM), + np.arange(NFRAMES * DIM_FPARAM).reshape(NFRAMES, DIM_FPARAM), + ), + ], + ids=("shared", "per-frame"), +) +def test_standardize_fparam(fparam, expected) -> None: + """Frame parameters become a canonical frame-major matrix.""" + actual, _ = _standardize_fparam_aparam( + fparam, + None, + NFRAMES, + NATOMS, + DIM_FPARAM, + DIM_APARAM, + ) + + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize( + ("aparam", "expected"), + [ + ( + APARAM_PER_ATOM, + np.tile(APARAM_PER_ATOM, (NFRAMES, 1, 1)), + ), + ( + APARAM_ALL_ATOMS.tolist(), + np.tile(APARAM_ALL_ATOMS, (NFRAMES, NATOMS, 1)), + ), + ( + np.arange(NFRAMES * NATOMS * DIM_APARAM).reshape( + NFRAMES, NATOMS, DIM_APARAM + ), + np.arange(NFRAMES * NATOMS * DIM_APARAM).reshape( + NFRAMES, NATOMS, DIM_APARAM + ), + ), + ], + ids=("shared-per-atom", "shared-all-atoms", "per-frame"), +) +def test_standardize_aparam(aparam, expected) -> None: + """Atomic shorthand is expanded before a batcher can slice its atom axis.""" + _, actual = _standardize_fparam_aparam( + None, + aparam, + NFRAMES, + NATOMS, + DIM_FPARAM, + DIM_APARAM, + ) + + np.testing.assert_array_equal(actual, expected) + + +@pytest.mark.parametrize( + ("fparam", "aparam", "message"), + [ + (np.zeros(3), None, "wrong size of frame param"), + (None, np.zeros(3), "wrong size of atomic param"), + ], +) +def test_invalid_parameter_size_is_rejected(fparam, aparam, message) -> None: + """Report the documented contract instead of a backend reshape failure.""" + with pytest.raises(RuntimeError, match=message): + _standardize_fparam_aparam( + fparam, + aparam, + NFRAMES, + NATOMS, + DIM_FPARAM, + DIM_APARAM, + ) diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index 2a34e2bbe5..662f9243cb 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -181,6 +181,11 @@ def test_deep_eval(self) -> None: aparam = np.ones((nframes, natoms, deep_eval.get_dim_aparam())) else: aparam = None + # Paddle is absent from the loop above because its frozen .json + # reader rejects fparam/aparam. A dedicated adapter test executes + # its normalization path without relying on that frozen format. + if deep_eval.get_dim_fparam() > 0 and deep_eval.get_dim_aparam() > 0: + self._assert_backend_parameter_shorthand(model_file, deep_eval) ret = deep_eval.eval( self.coords, self.box, @@ -238,6 +243,102 @@ def test_deep_eval(self) -> None: err_msg=f"backend {idx + 1} for rets_idx {rets_idx}", ) + def _assert_backend_parameter_shorthand( + self, model_file: str, deep_eval: DeepEval + ) -> None: + """Compare backend-direct shorthand with explicit frame-major inputs. + + Calling ``deep_eval.deep_eval`` deliberately bypasses the public + ``_standard_input`` normalization. A one-frame auto-batch size also + proves that shared per-atom parameters are expanded before the batcher + can mistake their atom axis for a frame axis. + """ + natoms = self.atype.shape[1] + nframes = 2 + coords = np.concatenate((self.coords, self.coords + 0.125), axis=0) + boxes = np.repeat(self.box, nframes, axis=0) + boxes[1, 0] += 0.25 + atom_types = np.repeat(self.atype, nframes, axis=0) + dim_fparam = deep_eval.get_dim_fparam() + dim_aparam = deep_eval.get_dim_aparam() + fparam_shared = np.linspace(0.25, -0.5, dim_fparam) + aparam_per_atom = np.linspace(0.1, 0.9, natoms * dim_aparam).reshape( + natoms, dim_aparam + ) + fparam_full = np.tile(fparam_shared, (nframes, 1)) + aparam_full = np.tile(aparam_per_atom, (nframes, 1, 1)) + fparam_per_frame = np.stack((fparam_shared, fparam_shared + 0.375)) + aparam_per_frame = np.stack((aparam_per_atom, aparam_per_atom + 0.2)) + aparam_all_atoms = np.linspace(-0.3, 0.4, dim_aparam) + + cases = ( + ( + fparam_shared.tolist(), + aparam_per_atom, + fparam_full, + aparam_full, + "shared-frame-and-per-atom", + ), + ( + fparam_per_frame, + aparam_all_atoms, + fparam_per_frame, + np.tile(aparam_all_atoms, (nframes, natoms, 1)), + "shared-all-atoms", + ), + ( + fparam_per_frame, + aparam_per_frame, + fparam_per_frame, + aparam_per_frame, + "full-frame-major", + ), + ) + for fparam, aparam, full_fparam, full_aparam, case_name in cases: + # Run the shorthand/full input first on a fresh auto-batcher so a + # GPU runner cannot grow the batch size on a preceding reference + # call and silently stop exercising the split path. + backend = DeepEval(model_file, auto_batch_size=natoms).deep_eval + actual = backend.eval( + coords, + boxes, + atom_types, + fparam=fparam, + aparam=aparam, + ) + reference_backend = DeepEval(model_file, auto_batch_size=False).deep_eval + expected = reference_backend.eval( + coords, + boxes, + atom_types, + fparam=full_fparam, + aparam=full_aparam, + ) + self.assertEqual(actual.keys(), expected.keys()) + # The frame-major expansion must preserve per-frame parameter + # variation. Without this, a degenerate expansion that reuses one + # frame's parameters for every frame would also make the reference + # degenerate, so the allclose comparison alone could not catch it. + frame_key = next( + (key for key in ("energy_redu", "energy") if key in actual), None + ) + self.assertIsNotNone(frame_key) + frame_values = actual[frame_key] + self.assertEqual(frame_values.shape[0], 2) + self.assertFalse( + np.allclose(frame_values[0], frame_values[1]), + msg=f"backend-direct {case_name} params collapsed across frames", + ) + for name in actual: + np.testing.assert_allclose( + actual[name], + expected[name], + rtol=1e-10, + atol=1e-10, + equal_nan=True, + err_msg=f"backend-direct {case_name} output {name}", + ) + class TestDeepPot(unittest.TestCase, IOTest): def setUp(self) -> None: 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..321089caa6 --- /dev/null +++ b/source/tests/infer/test_dpmodel_deep_eval_params.py @@ -0,0 +1,226 @@ +# 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.stack((COORD, COORD + 0.125)) +BOXES = np.tile(BOX, (NFRAMES, 1, 1)) +BOXES[1, 0, 0] += 0.25 +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], + rtol=1e-10, + atol=1e-10, + 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.""" + full_aparam = np.tile(APARAM_PER_ATOM, (NFRAMES, 1, 1)) + # Evaluate shorthand first on a fresh batcher. On GPU, a successful + # reference call would otherwise grow the batch size and remove the split + # this regression is intended to exercise. + actual = _backend(model_file, auto_batch_size).eval( + COORDS, + BOXES, + ATOM_TYPES, + fparam=FPARAM.tolist(), + aparam=full_aparam, + ) + expected = _backend(model_file, auto_batch_size=False).eval( + COORDS, + BOXES, + ATOM_TYPES, + fparam=np.tile(FPARAM, (NFRAMES, 1)), + 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.""" + full_fparam = np.tile(FPARAM, (NFRAMES, 1)) + actual = _backend(model_file, auto_batch_size).eval( + COORDS, + BOXES, + ATOM_TYPES, + fparam=full_fparam, + aparam=shared_aparam, + ) + expected = _backend(model_file, auto_batch_size=False).eval( + COORDS, + BOXES, + ATOM_TYPES, + fparam=full_fparam, + aparam=full_aparam, + ) + + _assert_outputs_equal(actual, expected) + + +def test_distinct_full_parameters_preserve_frame_order_when_split( + model_file: str, +) -> None: + """Split evaluation preserves distinct frame-major parameter rows.""" + fparam = np.stack((FPARAM, FPARAM + np.array([0.4, -0.2]))) + aparam = np.stack((APARAM_PER_ATOM, APARAM_PER_ATOM + 0.35)) + + split = _backend(model_file, auto_batch_size=len(ATOM_TYPES)).eval( + COORDS, + BOXES, + ATOM_TYPES, + fparam=fparam, + aparam=aparam, + ) + unsplit = _backend(model_file, auto_batch_size=False).eval( + COORDS, + BOXES, + ATOM_TYPES, + fparam=fparam, + aparam=aparam, + ) + + _assert_outputs_equal(split, unsplit) + assert not np.allclose(split["energy_redu"][0], split["energy_redu"][1]) + + +@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) diff --git a/source/tests/pd/test_deep_eval_parameter_shorthand.py b/source/tests/pd/test_deep_eval_parameter_shorthand.py new file mode 100644 index 0000000000..b3c065ae97 --- /dev/null +++ b/source/tests/pd/test_deep_eval_parameter_shorthand.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Exercise parameter shorthand through the Paddle DeepEval adapter.""" + +from types import ( + SimpleNamespace, +) +from unittest.mock import ( + MagicMock, +) + +import numpy as np +import pytest + +pytest.importorskip("paddle") + +from deepmd.pd.infer.deep_eval import ( + DeepEval, +) + + +def test_eval_standardizes_parameter_shorthand_before_dispatch() -> None: + """The Paddle adapter must pass frame-major parameters to its evaluator.""" + abstract_methods = getattr(DeepEval, "__abstractmethods__", frozenset()) + try: + DeepEval.__abstractmethods__ = frozenset() + backend = object.__new__(DeepEval) + finally: + DeepEval.__abstractmethods__ = abstract_methods + + nframes = 2 + natoms = 3 + fparam = np.array([0.25, -0.5]) + aparam = np.arange(natoms, dtype=np.float64)[:, None] + backend.get_dim_fparam = lambda: 2 + backend.get_dim_aparam = lambda: 1 + backend._get_request_defs = lambda atomic: [SimpleNamespace(name="energy")] + backend._eval_func = lambda inner, numb_test, numb_atoms: inner + backend._eval_model = MagicMock(return_value=(np.zeros((nframes, 1)),)) + + result = backend.eval( + np.zeros((nframes, natoms, 3)), + None, + np.zeros(natoms, dtype=np.int32), + fparam=fparam, + aparam=aparam, + ) + + np.testing.assert_array_equal( + backend._eval_model.call_args.args[3], np.tile(fparam, (nframes, 1)) + ) + np.testing.assert_array_equal( + backend._eval_model.call_args.args[4], np.tile(aparam, (nframes, 1, 1)) + ) + assert result.keys() == {"energy"} diff --git a/source/tests/pt/model/test_embedding.py b/source/tests/pt/model/test_embedding.py index 2617a0afba..afc9a827d0 100644 --- a/source/tests/pt/model/test_embedding.py +++ b/source/tests/pt/model/test_embedding.py @@ -314,6 +314,97 @@ def test_eval_embedding_dtype_fp64(self) -> None: np.float64, ) + def test_backend_embedding_normalizes_parameter_shorthand_before_batching( + self, + ) -> None: + """Shared fparam/aparam forms must survive one-frame auto batching.""" + params = _se_e2_a_params() + params["fitting_net"]["numb_fparam"] = 2 + params["fitting_net"]["numb_aparam"] = 1 + model = get_model(params) + _randomize(model) + path = self._save_checkpoint(model, params, "se_e2_a_params.pt") + + natoms = int(self.atype_np.shape[0]) + coords = np.concatenate((self.coord_np, self.coord_np + 0.125), axis=0) + cells = np.repeat(self.cell_np, 2, axis=0) + cells[1, 0] += 0.25 + fparam_shared = np.array([0.25, -0.5], dtype=np.float64) + fparam_full = np.repeat(fparam_shared[None, :], 2, axis=0) + aparam_shared = np.linspace(0.1, 0.7, natoms, dtype=np.float64)[:, None] + aparam_full = np.repeat(aparam_shared[None, :, :], 2, axis=0) + + def assert_method_matches_full( + method_name: str, + fparam: np.ndarray, + aparam: np.ndarray, + full_fparam: np.ndarray, + full_aparam: np.ndarray, + ) -> None: + # A fresh batch budget of exactly natoms forces the shorthand call + # to split the two frames even on GPU, where successful calls grow + # the auto-batch size. The reference is intentionally unbatched. + backend = DeepPot(path, auto_batch_size=natoms, no_jit=True).deep_eval + self.assertIsInstance(backend, PTDeepEval) + self.assertEqual(backend.auto_batch_size.current_batch_size, natoms) + actual = getattr(backend, method_name)( + coords, + cells, + self.atype_np, + fparam=fparam, + aparam=aparam, + ) + + reference_backend = DeepPot( + path, auto_batch_size=False, no_jit=True + ).deep_eval + expected = getattr(reference_backend, method_name)( + coords, + cells, + self.atype_np, + fparam=full_fparam, + aparam=full_aparam, + ) + actual_values = actual if isinstance(actual, tuple) else (actual,) + expected_values = expected if isinstance(expected, tuple) else (expected,) + for actual_value, expected_value in zip( + actual_values, expected_values, strict=True + ): + np.testing.assert_allclose(actual_value, expected_value) + + for method_name in ( + "eval_embedding", + "eval_descriptor", + "eval_fitting_last_layer", + ): + assert_method_matches_full( + method_name, + fparam_shared, + aparam_shared, + fparam_full, + aparam_full, + ) + + scalar_aparam = np.array([0.35], dtype=np.float64) + scalar_full = np.full((2, natoms, 1), scalar_aparam.item()) + assert_method_matches_full( + "eval_embedding", + fparam_shared, + scalar_aparam, + fparam_full, + scalar_full, + ) + + # Distinct frame-major values ensure batching preserves the frame and + # atom axes instead of merely producing a shape-compatible result. + assert_method_matches_full( + "eval_embedding", + np.stack((fparam_shared, fparam_shared + 0.375)), + np.stack((aparam_shared, aparam_shared + 0.2)), + np.stack((fparam_shared, fparam_shared + 0.375)), + np.stack((aparam_shared, aparam_shared + 0.2)), + ) + def test_legacy_frozen_model_uses_baked_in_hook(self) -> None: # Frozen ``.pth`` files predating ``forward_embedding`` still carry the # descriptor / fitting hooks baked into the TorchScript module. The diff --git a/source/tests/pt_expt/infer/test_parameter_shorthand.py b/source/tests/pt_expt/infer/test_parameter_shorthand.py new file mode 100644 index 0000000000..8be4168cb1 --- /dev/null +++ b/source/tests/pt_expt/infer/test_parameter_shorthand.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Exercise parameter shorthand through the pt_expt DeepEval adapter.""" + +from types import ( + SimpleNamespace, +) +from unittest.mock import ( + MagicMock, +) + +import numpy as np + +from deepmd.pt_expt.infer.deep_eval import ( + DeepEval, +) + + +def test_eval_standardizes_parameter_shorthand_before_dispatch() -> None: + """The pt_expt adapter must normalize parameters before auto batching.""" + abstract_methods = getattr(DeepEval, "__abstractmethods__", frozenset()) + try: + DeepEval.__abstractmethods__ = frozenset() + backend = object.__new__(DeepEval) + finally: + DeepEval.__abstractmethods__ = abstract_methods + + nframes = 2 + natoms = 3 + fparam = np.array([0.25, -0.5]) + aparam = np.arange(natoms, dtype=np.float64)[:, None] + backend._is_spin = False + backend.get_dim_fparam = lambda: 2 + backend.get_dim_aparam = lambda: 1 + backend._get_request_defs = lambda atomic: [SimpleNamespace(name="energy")] + backend._eval_func = lambda inner, numb_test, numb_atoms: inner + backend._eval_model = MagicMock(return_value=(np.zeros((nframes, 1)),)) + + result = backend.eval( + np.zeros((nframes, natoms, 3)), + None, + np.zeros(natoms, dtype=np.int32), + fparam=fparam, + aparam=aparam, + ) + + np.testing.assert_array_equal( + backend._eval_model.call_args.args[3], np.tile(fparam, (nframes, 1)) + ) + np.testing.assert_array_equal( + backend._eval_model.call_args.args[4], np.tile(aparam, (nframes, 1, 1)) + ) + assert result.keys() == {"energy"} + + +def _new_backend() -> DeepEval: + """Build an uninitialized pt_expt DeepEval for adapter mocking.""" + abstract_methods = getattr(DeepEval, "__abstractmethods__", frozenset()) + try: + DeepEval.__abstractmethods__ = frozenset() + backend = object.__new__(DeepEval) + finally: + DeepEval.__abstractmethods__ = abstract_methods + backend.get_dim_fparam = lambda: 2 + backend.get_dim_aparam = lambda: 1 + backend._require_dpmodel = lambda name: None + backend._is_spin_model = lambda: False + return backend + + +def test_eval_descriptor_standardizes_parameter_shorthand() -> None: + """The descriptor route must normalize parameters before reshaping.""" + nframes = 2 + natoms = 3 + fparam = np.array([0.25, -0.5]) + aparam = np.arange(natoms, dtype=np.float64)[:, None] + backend = _new_backend() + dp_am = MagicMock() + backend._dpmodel = MagicMock() + backend._dpmodel.get_dp_atomic_model.return_value = dp_am + descriptor = MagicMock() + descriptor.detach.return_value = descriptor + descriptor.cpu.return_value = descriptor + descriptor.numpy.return_value = np.zeros((nframes, natoms, 4)) + dp_am.descriptor.return_value = (descriptor,) + backend._prepare_nlist_inputs = MagicMock( + return_value=(None, None, None, None, None, None, None, nframes, natoms) + ) + + backend.eval_descriptor( + np.zeros((nframes, natoms, 3)), + None, + np.zeros(natoms, dtype=np.int32), + fparam=fparam, + aparam=aparam, + ) + + np.testing.assert_array_equal( + backend._prepare_nlist_inputs.call_args.args[3], + np.tile(fparam, (nframes, 1)), + ) + np.testing.assert_array_equal( + backend._prepare_nlist_inputs.call_args.args[4], + np.tile(aparam, (nframes, 1, 1)), + ) + + +def test_eval_fitting_last_layer_standardizes_parameter_shorthand() -> None: + """The fitting-last-layer route must normalize parameters before reshaping.""" + nframes = 2 + natoms = 3 + fparam = np.array([0.25, -0.5]) + aparam = np.arange(natoms, dtype=np.float64)[:, None] + backend = _new_backend() + dp_am = MagicMock() + backend._dpmodel = MagicMock() + backend._dpmodel.get_dp_atomic_model.return_value = dp_am + dp_am.descriptor.return_value = ( + np.zeros((nframes, natoms, 4)), + np.zeros((nframes, natoms, 3, 3)), + np.zeros((nframes, natoms, 3, 3)), + np.zeros((nframes, natoms, 3, 3)), + np.zeros((nframes, natoms, 3)), + ) + out = MagicMock() + out.detach.return_value = out + out.cpu.return_value = out + out.numpy.return_value = np.zeros((nframes, natoms, 4)) + dp_am.fitting_net = MagicMock(return_value={"middle_output": out}) + ext_atype_t = np.zeros((nframes, natoms), dtype=np.int64) + backend._prepare_nlist_inputs = MagicMock( + return_value=( + None, + ext_atype_t, + None, + None, + None, + None, + None, + nframes, + natoms, + ) + ) + + backend.eval_fitting_last_layer( + np.zeros((nframes, natoms, 3)), + None, + np.zeros(natoms, dtype=np.int32), + fparam=fparam, + aparam=aparam, + ) + + np.testing.assert_array_equal( + backend._prepare_nlist_inputs.call_args.args[3], + np.tile(fparam, (nframes, 1)), + ) + np.testing.assert_array_equal( + backend._prepare_nlist_inputs.call_args.args[4], + np.tile(aparam, (nframes, 1, 1)), + ) diff --git a/source/tests/tf/test_deep_eval_parameter_shorthand.py b/source/tests/tf/test_deep_eval_parameter_shorthand.py new file mode 100644 index 0000000000..1b359f5949 --- /dev/null +++ b/source/tests/tf/test_deep_eval_parameter_shorthand.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Regression tests for TensorFlow DeepEval parameter shorthand handling.""" + +import numpy as np + +from deepmd.tf.infer.deep_eval import ( + DeepEval, +) + + +class _CapturingAutoBatchSize: + """Capture inputs at the auto-batching boundary without loading a model.""" + + def __init__(self) -> None: + self.kwargs: dict | None = None + + def execute_all( + self, + inner_func, + nframes: int, + natoms: int, + *args, + **kwargs, + ) -> np.ndarray: + """Record canonical inputs that would be sliced into frame batches.""" + self.kwargs = kwargs + return np.zeros((nframes, natoms, 1), dtype=np.float64) + + +def test_eval_descriptor_normalizes_shared_parameters_before_batching() -> None: + """Shared parameters must gain a frame axis before auto batching.""" + auto_batch_size = _CapturingAutoBatchSize() + backend = object.__new__(DeepEval) + backend.dfparam = 2 + backend.daparam = 1 + backend.auto_batch_size = auto_batch_size + + fparam = np.array([0.25, -0.5]) + aparam = np.array([[0.1], [0.2]]) + descriptor = backend.eval_descriptor( + np.zeros((2, 2, 3)), + None, + np.array([0, 1]), + fparam=fparam, + aparam=aparam, + ) + + assert descriptor.shape == (2, 2, 1) + assert auto_batch_size.kwargs is not None + np.testing.assert_array_equal( + auto_batch_size.kwargs["fparam"], np.tile(fparam, (2, 1)) + ) + np.testing.assert_array_equal( + auto_batch_size.kwargs["aparam"], np.tile(aparam[None, :, :], (2, 1, 1)) + )