diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py index a20cc769fd..271198365e 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,15 @@ 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(), + copy_shared=False, + ) # 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 diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index d20626e835..cbfb9955ce 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -30,6 +30,93 @@ 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, + *, + copy_shared: bool = True, +) -> 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. + copy_shared : bool, default: True + Materialize shared inputs as writable arrays. If False, return + read-only broadcast views so compatible backends can slice batches + without allocating the shared parameters for the whole trajectory. + + 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.broadcast_to( + fparam.reshape(1, dim_fparam), (nframes, dim_fparam) + ) + if copy_shared: + fparam = fparam.copy() + 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.broadcast_to( + aparam.reshape(1, natoms, dim_aparam), (nframes, natoms, dim_aparam) + ) + if copy_shared: + aparam = aparam.copy() + elif aparam.size == dim_aparam: + aparam = np.broadcast_to( + aparam.reshape(1, 1, dim_aparam), (nframes, natoms, dim_aparam) + ) + if copy_shared: + aparam = aparam.copy() + 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. @@ -952,28 +1039,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 d2930ff8c7..026cdbb054 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 fd393d901d..f62be9b37c 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, @@ -372,6 +373,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 eab00ce6ba..c4fcbb9651 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/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 ( DeepGlobalPolar, @@ -573,6 +574,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(), + ) if "grid" in kwargs and kwargs["grid"] is not None: grid_input = np.array(kwargs["grid"]) # the directional neighbor list is dense in ngrid x nall, so the @@ -1456,6 +1465,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 7b658196f0..a94389ad78 100644 --- a/deepmd/pt_expt/infer/deep_eval.py +++ b/deepmd/pt_expt/infer/deep_eval.py @@ -50,6 +50,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, @@ -1338,10 +1339,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 @@ -1365,6 +1371,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: @@ -2771,6 +2785,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, @@ -2848,6 +2875,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 7468bdbbd3..3dfeef5a5b 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, @@ -773,6 +774,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, @@ -1099,6 +1108,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 518af01485..442301f91d 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, @@ -301,6 +302,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..f8fa84c58b --- /dev/null +++ b/source/tests/common/test_deep_eval_parameter_shorthand.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for backend-level DeepEval parameter normalization.""" + +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.deep_eval import ( + DeepEval, + _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) + assert actual.flags.writeable + assert actual.flags.c_contiguous + + +@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) + assert actual.flags.writeable + assert actual.flags.c_contiguous + + +@pytest.mark.parametrize("aparam", [APARAM_PER_ATOM, APARAM_ALL_ATOMS]) +def test_shared_parameters_can_be_broadcast_views(aparam: np.ndarray) -> None: + """Opted-in shared inputs retain their storage across a long trajectory.""" + nframes = 10000 + fparam_view, aparam_view = _standardize_fparam_aparam( + FPARAM, + aparam, + nframes, + NATOMS, + DIM_FPARAM, + DIM_APARAM, + copy_shared=False, + ) + assert fparam_view.shape == (nframes, DIM_FPARAM) + assert aparam_view.shape == (nframes, NATOMS, DIM_APARAM) + assert np.shares_memory(fparam_view, FPARAM) + assert np.shares_memory(aparam_view, aparam) + assert fparam_view.strides[0] == aparam_view.strides[0] == 0 + assert not fparam_view.flags.writeable + assert not aparam_view.flags.writeable + expected_aparam = aparam if aparam.ndim == 2 else np.tile(aparam, (NATOMS, 1)) + for frame in (0, nframes - 1): + np.testing.assert_array_equal(fparam_view[frame], FPARAM) + np.testing.assert_array_equal(aparam_view[frame], expected_aparam) + + +@pytest.mark.parametrize("aparam", [APARAM_PER_ATOM, APARAM_ALL_ATOMS]) +def test_dpmodel_batches_shared_parameter_views( + tmp_path, monkeypatch: pytest.MonkeyPatch, aparam: np.ndarray +) -> None: + """The real dpmodel consumes read-only views after slicing their frame axis.""" + model_def = { + "type_map": ["X"], + "descriptor": { + "type": "se_e2_a", + "sel": [4], + "rcut": 3.0, + "rcut_smth": 2.5, + "neuron": [2], + "axis_neuron": 1, + "seed": 1, + }, + "fitting_net": { + "neuron": [2], + "numb_fparam": DIM_FPARAM, + "numb_aparam": DIM_APARAM, + "seed": 1, + }, + } + path = str(tmp_path / "params.dp") + save_dp_model( + path, + { + "model": get_model(model_def).serialize(), + "model_def_script": model_def, + "backend": "dpmodel", + }, + ) + coords = np.tile(np.arange(NATOMS * 3).reshape(1, NATOMS, 3) * 0.1, (NFRAMES, 1, 1)) + cells = np.tile(np.eye(3) * 10.0, (NFRAMES, 1, 1)) + atom_types = np.zeros(NATOMS, dtype=np.int32) + backend = DeepEval(path, auto_batch_size=NATOMS).deep_eval + evaluate_batch = backend._eval_model + batch_sizes = [] + + def check_batch(coords, cells, atom_types, fparam, aparam_batch, request_defs): + batch_sizes.append(coords.shape[0]) + assert np.shares_memory(fparam, FPARAM) + assert np.shares_memory(aparam_batch, aparam) + assert not fparam.flags.writeable + assert not aparam_batch.flags.writeable + return evaluate_batch( + coords, cells, atom_types, fparam, aparam_batch, request_defs + ) + + monkeypatch.setattr(backend, "_eval_model", check_batch) + actual = backend.eval(coords, cells, atom_types, fparam=FPARAM, aparam=aparam) + expected = DeepEval(path, auto_batch_size=False).deep_eval.eval( + coords, + cells, + atom_types, + fparam=np.tile(FPARAM, (NFRAMES, 1)), + aparam=np.tile( + aparam.reshape(-1, DIM_APARAM), + (NFRAMES, NATOMS if aparam.ndim == 1 else 1, 1), + ), + ) + assert len(batch_sizes) > 1 + assert sum(batch_sizes) == NFRAMES + 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( + ("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 913aca4d51..d4717a6de8 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 or 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,108 @@ 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() + shared_params = {} + full_params = {} + frame_params = {} + if dim_fparam: + fparam_shared = np.linspace(0.25, -0.5, dim_fparam) + shared_params["fparam"] = fparam_shared.tolist() + full_params["fparam"] = np.tile(fparam_shared, (nframes, 1)) + frame_params["fparam"] = np.stack((fparam_shared, fparam_shared + 0.375)) + if dim_aparam: + aparam_per_atom = np.linspace(0.1, 0.9, natoms * dim_aparam).reshape( + natoms, dim_aparam + ) + shared_params["aparam"] = aparam_per_atom + full_params["aparam"] = np.tile(aparam_per_atom, (nframes, 1, 1)) + frame_params["aparam"] = np.stack((aparam_per_atom, aparam_per_atom + 0.2)) + all_atom_params = frame_params.copy() + all_atom_full_params = frame_params.copy() + if dim_aparam: + aparam_all_atoms = np.linspace(-0.3, 0.4, dim_aparam) + all_atom_params["aparam"] = aparam_all_atoms + all_atom_full_params["aparam"] = np.tile( + aparam_all_atoms, (nframes, natoms, 1) + ) + + cases = ( + (shared_params, full_params, "shared-frame-and-per-atom"), + (all_atom_params, all_atom_full_params, "shared-all-atoms"), + (frame_params, frame_params, "full-frame-major"), + ) + for params, explicit_params, 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, + **params, + ) + reference_backend = DeepEval(model_file, auto_batch_size=False).deep_eval + expected = reference_backend.eval( + coords, + boxes, + atom_types, + **explicit_params, + ) + self.assertEqual(actual.keys(), expected.keys()) + # Distinct frame inputs must survive batching. Parameter + # sensitivity is checked separately below at fixed geometry. + 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} outputs 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}", + ) + + # Each declared parameter must affect the energy at fixed geometry. + backend = DeepEval(model_file, auto_batch_size=False).deep_eval + baseline = backend.eval(self.coords, self.box, self.atype, **shared_params) + for parameter, value in shared_params.items(): + with self.subTest(parameter=parameter, model_file=model_file): + changed_params = {**shared_params, parameter: np.asarray(value) + 0.5} + changed = backend.eval( + self.coords, self.box, self.atype, **changed_params + ) + self.assertFalse( + np.allclose(baseline[frame_key], changed[frame_key]), + msg=f"backend-direct energy ignores {parameter}", + ) + class TestDeepPot(unittest.TestCase, IOTest): def setUp(self) -> None: @@ -283,6 +390,9 @@ def tearDown(self) -> None: class TestDeepPotFparamAparam(unittest.TestCase, IOTest): + numb_fparam: ClassVar[int] = 2 + numb_aparam: ClassVar[int] = 2 + def setUp(self) -> None: model_def_script = { "type_map": ["O", "H"], @@ -311,8 +421,8 @@ def setUp(self) -> None: "precision": "float64", "atom_ener": [], "seed": 1, - "numb_fparam": 2, - "numb_aparam": 2, + "numb_fparam": self.numb_fparam, + "numb_aparam": self.numb_aparam, }, } model = get_model(copy.deepcopy(model_def_script)) @@ -327,6 +437,14 @@ def tearDown(self) -> None: IOTest.tearDown(self) +class TestDeepPotFparamOnly(TestDeepPotFparamAparam): + numb_aparam = 0 + + +class TestDeepPotAparamOnly(TestDeepPotFparamAparam): + numb_fparam = 0 + + @unittest.skipIf( not DP_TEST_TF2_ONLY, "pair_exclude_types is not supported by the TF v1 backend; it is validated " 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 f49a362cc0..7d21289f98 100644 --- a/source/tests/pt/model/test_embedding.py +++ b/source/tests/pt/model/test_embedding.py @@ -322,6 +322,114 @@ 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)), + ) + + # Each parameter must influence the fitting features at fixed geometry. + backend = DeepPot(path, auto_batch_size=False, no_jit=True).deep_eval + shared_params = {"fparam": fparam_shared, "aparam": aparam_shared} + _, baseline, _ = backend.eval_embedding( + self.coord_np, self.cell_np, self.atype_np, **shared_params + ) + for parameter, value in shared_params.items(): + with self.subTest(parameter=parameter): + changed_params = {**shared_params, parameter: value + 0.5} + _, atomic_feature, _ = backend.eval_embedding( + self.coord_np, self.cell_np, self.atype_np, **changed_params + ) + self.assertFalse( + np.allclose(baseline, atomic_feature), + msg=f"backend-direct atomic feature ignores {parameter}", + ) + 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..6ce73702c2 --- /dev/null +++ b/source/tests/pt_expt/infer/test_parameter_shorthand.py @@ -0,0 +1,163 @@ +# 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 + # Introspection now distinguishes the archive's folded condition from the + # deserialized model's explicit input; this fixture has neither condition. + backend._charge_state_fold = None + backend.get_dim_chg_spin = lambda: 0 + 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)) + )