From f75ab0b8f2e5efaef369bb4f048f7425e690e40e Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 08:34:24 +0800 Subject: [PATCH 1/3] fix(tf): preserve non-PBC ASE neighbor semantics Keep the original open-boundary decision separate from the identity box required by TensorFlow feeds, and let both ASE builders handle a missing cell without creating periodic ghosts. Add collected DeepPotential and DeepTensor regressions for non-periodic external neighbor-list inference. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/tf/infer/deep_eval.py | 41 +++++++++++++++++++++++++----- deepmd/tf/infer/deep_tensor.py | 10 ++++++-- source/tests/infer/test_models.py | 35 ++++++++++++++++++++++++- source/tests/tf/test_deepdipole.py | 29 +++++++++++++++++++++ 4 files changed, 105 insertions(+), 10 deletions(-) diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 30107cb693..4688dd949c 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -549,7 +549,14 @@ def build_neighbor_list( atype: np.ndarray, imap: np.ndarray, neighbor_list: "ase.neighborlist.NeighborList | None", - ) -> tuple[np.ndarray, np.ndarray]: + ) -> tuple[ + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, + ]: """Make the mesh with neighbor list for a single frame. Parameters @@ -557,7 +564,8 @@ def build_neighbor_list( coords : np.ndarray The coordinates of atoms. Should be of shape [natoms, 3] cell : Optional[np.ndarray] - The cell of the system. Should be of shape [3, 3] + The cell of the system. Should be of shape [3, 3]. None denotes + open boundary conditions. atype : np.ndarray The type of atoms. Should be of shape [natoms] imap : np.ndarray @@ -587,7 +595,11 @@ def build_neighbor_list( The index map of ghost atoms. Should be of shape [nghost] """ pbc = np.repeat(cell is not None, 3) - cell = cell.reshape(3, 3) + # ASE still requires a 3x3 cell for non-periodic systems, but the cell + # must not be used to infer periodicity or create ghost atoms. + cell = ( + np.zeros((3, 3), dtype=coords.dtype) if cell is None else cell.reshape(3, 3) + ) positions = coords.reshape(-1, 3) neighbor_list.bothways = True neighbor_list.self_interaction = False @@ -814,6 +826,9 @@ def _prepare_feed_dict( else: pbc = True cells = np.array(cells).reshape([nframes, 9]) + # Keep the original boundary semantics separate from the identity box + # used only to satisfy TensorFlow's non-optional box placeholder. + neighbor_cell = cells if pbc else None if self.has_fparam: assert fparam is not None @@ -884,7 +899,7 @@ def _prepare_feed_dict( ghost_map, ) = self.build_neighbor_list( coords, - cells if cells is not None else None, + neighbor_cell, atom_types, imap, self.neighbor_list, @@ -1534,7 +1549,14 @@ def build_neighbor_list( atype: np.ndarray, imap: np.ndarray, neighbor_list: "ase.neighborlist.NeighborList | None", - ) -> tuple[np.ndarray, np.ndarray]: + ) -> tuple[ + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, + np.ndarray, + ]: """Make the mesh with neighbor list for a single frame. Parameters @@ -1542,7 +1564,8 @@ def build_neighbor_list( coords : np.ndarray The coordinates of atoms. Should be of shape [natoms, 3] cell : Optional[np.ndarray] - The cell of the system. Should be of shape [3, 3] + The cell of the system. Should be of shape [3, 3]. None denotes + open boundary conditions. atype : np.ndarray The type of atoms. Should be of shape [natoms] imap : np.ndarray @@ -1572,7 +1595,11 @@ def build_neighbor_list( The index map of ghost atoms. Should be of shape [nghost] """ pbc = np.repeat(cell is not None, 3) - cell = cell.reshape(3, 3) + # ASE still requires a 3x3 cell for non-periodic systems, but the cell + # must not be used to infer periodicity or create ghost atoms. + cell = ( + np.zeros((3, 3), dtype=coords.dtype) if cell is None else cell.reshape(3, 3) + ) positions = coords.reshape(-1, 3) neighbor_list.bothways = True neighbor_list.self_interaction = False diff --git a/deepmd/tf/infer/deep_tensor.py b/deepmd/tf/infer/deep_tensor.py index dbb3266da4..7256157513 100644 --- a/deepmd/tf/infer/deep_tensor.py +++ b/deepmd/tf/infer/deep_tensor.py @@ -202,6 +202,9 @@ def eval( else: pbc = True cells = np.array(cells).reshape([nframes, 9]) + # Keep the original boundary semantics separate from the identity box + # used only to satisfy TensorFlow's non-optional box placeholder. + neighbor_cell = cells if pbc else None # sort inputs coords, atom_types, imap, sel_at, sel_imap = self.sort_input( @@ -227,7 +230,7 @@ def eval( _, ) = self.build_neighbor_list( coords, - cells if cells is not None else None, + neighbor_cell, atom_types, imap, self.neighbor_list, @@ -346,6 +349,9 @@ def eval_full( else: pbc = True cells = np.array(cells).reshape([nframes, 9]) + # Keep the original boundary semantics separate from the identity box + # used only to satisfy TensorFlow's non-optional box placeholder. + neighbor_cell = cells if pbc else None nout = self.output_dim # sort inputs @@ -373,7 +379,7 @@ def eval_full( ghost_map, ) = self.build_neighbor_list( coords, - cells if cells is not None else None, + neighbor_cell, atom_types, imap, self.neighbor_list, diff --git a/source/tests/infer/test_models.py b/source/tests/infer/test_models.py index 2e6295afaf..2f8e99e0c7 100644 --- a/source/tests/infer/test_models.py +++ b/source/tests/infer/test_models.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import unittest -import ase +import ase.neighborlist import dpdata import numpy as np @@ -421,3 +421,36 @@ def test_2frame_atm(self) -> None: @unittest.skip("Zero atoms not supported") def test_zero_input(self) -> None: pass + + +def test_deep_pot_neighbor_list_nopbc() -> None: + """The ASE path must preserve a testcase's open-boundary semantics.""" + # This is intentionally standalone: the parameterized TestDeepPot symbol is + # replaced by ``object``, so its neighbor-list subclass inherits no tests. + case = get_cases()["se_e2_a"] + result = next(result for result in case.results if result.box is None) + with DeepEval( + case.get_model(".pb"), + neighbor_list=ase.neighborlist.NewPrimitiveNeighborList( + cutoffs=case.rcut, + bothways=True, + ), + ) as dp: + ee, ff, vv, ae, av = dp.eval( + result.coord, + None, + result.atype, + atomic=True, + fparam=result.fparam, + aparam=result.aparam, + )[:5] + + np.testing.assert_almost_equal(ff.ravel(), result.force.ravel(), STRICT_PLACES) + np.testing.assert_almost_equal( + ae.ravel(), result.atomic_energy.ravel(), STRICT_PLACES + ) + np.testing.assert_almost_equal( + av.ravel(), result.atomic_virial.ravel(), STRICT_PLACES + ) + np.testing.assert_almost_equal(ee.ravel(), result.energy, STRICT_PLACES) + np.testing.assert_almost_equal(vv.ravel(), result.virial, STRICT_PLACES) diff --git a/source/tests/tf/test_deepdipole.py b/source/tests/tf/test_deepdipole.py index 33da480e12..78aa024350 100644 --- a/source/tests/tf/test_deepdipole.py +++ b/source/tests/tf/test_deepdipole.py @@ -1162,3 +1162,32 @@ def test_2frame_full_atm(self) -> None: @unittest.skip("multiple frames not supported") def test_2frame_old_atm(self) -> None: pass + + def test_nopbc_matches_native_neighbor_building(self) -> None: + """ASE and native tensor inference must agree for an open system.""" + native = DeepDipole("deepdipole_new.pb") + try: + actual_tensor = self.dp.eval(self.coords, None, self.atype, atomic=True) + expected_tensor = native.eval(self.coords, None, self.atype, atomic=True) + np.testing.assert_almost_equal( + actual_tensor, + expected_tensor, + default_places, + ) + + actual_full = self.dp.eval_full( + self.coords, + None, + self.atype, + atomic=True, + ) + expected_full = native.eval_full( + self.coords, + None, + self.atype, + atomic=True, + ) + for actual, expected in zip(actual_full, expected_full, strict=True): + np.testing.assert_almost_equal(actual, expected, default_places) + finally: + native.close() From 8d56cb1487fcc94a486b7a445e8b7177555043c0 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 23 Jul 2026 20:38:49 +0800 Subject: [PATCH 2/3] test(tf): scope native dipole evaluator Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- source/tests/tf/test_deepdipole.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/source/tests/tf/test_deepdipole.py b/source/tests/tf/test_deepdipole.py index 78aa024350..b0a91ab1de 100644 --- a/source/tests/tf/test_deepdipole.py +++ b/source/tests/tf/test_deepdipole.py @@ -1165,8 +1165,7 @@ def test_2frame_old_atm(self) -> None: def test_nopbc_matches_native_neighbor_building(self) -> None: """ASE and native tensor inference must agree for an open system.""" - native = DeepDipole("deepdipole_new.pb") - try: + with DeepDipole("deepdipole_new.pb") as native: actual_tensor = self.dp.eval(self.coords, None, self.atype, atomic=True) expected_tensor = native.eval(self.coords, None, self.atype, atomic=True) np.testing.assert_almost_equal( @@ -1189,5 +1188,3 @@ def test_nopbc_matches_native_neighbor_building(self) -> None: ) for actual, expected in zip(actual_full, expected_full, strict=True): np.testing.assert_almost_equal(actual, expected, default_places) - finally: - native.close() From 9dd42102db0fcffd648114ed9910d77578224b2b Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Tue, 28 Jul 2026 13:47:39 +0800 Subject: [PATCH 3/3] fix(tf): strengthen non-PBC neighbor coverage Restore the parameterized ASE neighbor-list test inheritance, exercise the legacy TF DeepTensor path directly, and align open-boundary fallback cell precision with the reference implementation. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/tf/infer/deep_eval.py | 4 +- source/tests/infer/test_models.py | 69 ++++++++++++++---------------- source/tests/tf/test_deepdipole.py | 27 ++++++++++-- 3 files changed, 59 insertions(+), 41 deletions(-) diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index 4688dd949c..7468bdbbd3 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -598,7 +598,7 @@ def build_neighbor_list( # ASE still requires a 3x3 cell for non-periodic systems, but the cell # must not be used to infer periodicity or create ghost atoms. cell = ( - np.zeros((3, 3), dtype=coords.dtype) if cell is None else cell.reshape(3, 3) + np.zeros((3, 3), dtype=np.float64) if cell is None else cell.reshape(3, 3) ) positions = coords.reshape(-1, 3) neighbor_list.bothways = True @@ -1598,7 +1598,7 @@ def build_neighbor_list( # ASE still requires a 3x3 cell for non-periodic systems, but the cell # must not be used to infer periodicity or create ghost atoms. cell = ( - np.zeros((3, 3), dtype=coords.dtype) if cell is None else cell.reshape(3, 3) + np.zeros((3, 3), dtype=np.float64) if cell is None else cell.reshape(3, 3) ) positions = coords.reshape(-1, 3) neighbor_list.bothways = True diff --git a/source/tests/infer/test_models.py b/source/tests/infer/test_models.py index 2f8e99e0c7..b5785535e0 100644 --- a/source/tests/infer/test_models.py +++ b/source/tests/infer/test_models.py @@ -32,15 +32,9 @@ STRICT_KEYS = frozenset(("se_e2_a", "se_e2_r")) -@parameterized( - ( - "se_e2_a", - "se_e2_r", - "fparam_aparam", - ), # key - (".pb", ".pth", ".pte", ".pt2"), # model extension -) -class TestDeepPot(unittest.TestCase): +class DeepPotTestMixin: + """Shared DeepPotential checks for native and external neighbor lists.""" + # moved from tests/tf/test_deeppot_a.py @classmethod @@ -396,19 +390,33 @@ def test_model_script_def(self) -> None: ) +@parameterized( + ( + "se_e2_a", + "se_e2_r", + "fparam_aparam", + ), # key + (".pb", ".pth", ".pte", ".pt2"), # model extension +) +class TestDeepPot(DeepPotTestMixin, unittest.TestCase): + """Run the common inference checks with native neighbor construction.""" + + @parameterized( ("se_e2_a",), # key (".pb",), # model extension ) -class TestDeepPotNeighborList(TestDeepPot): +class TestDeepPotNeighborList(DeepPotTestMixin, unittest.TestCase): + """Run the common inference checks with an external ASE neighbor list.""" + @classmethod def setUpClass(cls) -> None: key, extension = cls.param cls.places = STRICT_PLACES if key in STRICT_KEYS else default_places cls.case = get_cases()[key] - model_name = cls.case.get_model(extension) + cls.model_name = cls.case.get_model(extension) cls.dp = DeepEval( - model_name, + cls.model_name, neighbor_list=ase.neighborlist.NewPrimitiveNeighborList( cutoffs=cls.case.rcut, bothways=True ), @@ -422,21 +430,10 @@ def test_2frame_atm(self) -> None: def test_zero_input(self) -> None: pass - -def test_deep_pot_neighbor_list_nopbc() -> None: - """The ASE path must preserve a testcase's open-boundary semantics.""" - # This is intentionally standalone: the parameterized TestDeepPot symbol is - # replaced by ``object``, so its neighbor-list subclass inherits no tests. - case = get_cases()["se_e2_a"] - result = next(result for result in case.results if result.box is None) - with DeepEval( - case.get_model(".pb"), - neighbor_list=ase.neighborlist.NewPrimitiveNeighborList( - cutoffs=case.rcut, - bothways=True, - ), - ) as dp: - ee, ff, vv, ae, av = dp.eval( + def test_nopbc_matches_reference(self) -> None: + """The ASE path must preserve a testcase's open-boundary semantics.""" + result = next(result for result in self.case.results if result.box is None) + ee, ff, vv, ae, av = self.dp.eval( result.coord, None, result.atype, @@ -445,12 +442,12 @@ def test_deep_pot_neighbor_list_nopbc() -> None: aparam=result.aparam, )[:5] - np.testing.assert_almost_equal(ff.ravel(), result.force.ravel(), STRICT_PLACES) - np.testing.assert_almost_equal( - ae.ravel(), result.atomic_energy.ravel(), STRICT_PLACES - ) - np.testing.assert_almost_equal( - av.ravel(), result.atomic_virial.ravel(), STRICT_PLACES - ) - np.testing.assert_almost_equal(ee.ravel(), result.energy, STRICT_PLACES) - np.testing.assert_almost_equal(vv.ravel(), result.virial, STRICT_PLACES) + np.testing.assert_almost_equal(ff.ravel(), result.force.ravel(), STRICT_PLACES) + np.testing.assert_almost_equal( + ae.ravel(), result.atomic_energy.ravel(), STRICT_PLACES + ) + np.testing.assert_almost_equal( + av.ravel(), result.atomic_virial.ravel(), STRICT_PLACES + ) + np.testing.assert_almost_equal(ee.ravel(), result.energy, STRICT_PLACES) + np.testing.assert_almost_equal(vv.ravel(), result.virial, STRICT_PLACES) diff --git a/source/tests/tf/test_deepdipole.py b/source/tests/tf/test_deepdipole.py index b0a91ab1de..863a696e11 100644 --- a/source/tests/tf/test_deepdipole.py +++ b/source/tests/tf/test_deepdipole.py @@ -12,6 +12,9 @@ from deepmd.tf.infer import ( DeepDipole, ) +from deepmd.tf.infer.deep_dipole import ( + DeepDipoleOld, +) from deepmd.tf.utils.convert import ( convert_pbtxt_to_pb, ) @@ -1165,8 +1168,9 @@ def test_2frame_old_atm(self) -> None: def test_nopbc_matches_native_neighbor_building(self) -> None: """ASE and native tensor inference must agree for an open system.""" - with DeepDipole("deepdipole_new.pb") as native: - actual_tensor = self.dp.eval(self.coords, None, self.atype, atomic=True) + + def assert_evaluators_match(external, native) -> None: + actual_tensor = external.eval(self.coords, None, self.atype, atomic=True) expected_tensor = native.eval(self.coords, None, self.atype, atomic=True) np.testing.assert_almost_equal( actual_tensor, @@ -1174,7 +1178,7 @@ def test_nopbc_matches_native_neighbor_building(self) -> None: default_places, ) - actual_full = self.dp.eval_full( + actual_full = external.eval_full( self.coords, None, self.atype, @@ -1188,3 +1192,20 @@ def test_nopbc_matches_native_neighbor_building(self) -> None: ) for actual, expected in zip(actual_full, expected_full, strict=True): np.testing.assert_almost_equal(actual, expected, default_places) + + # The public wrapper reaches deepmd/tf/infer/deep_eval.py. + with DeepDipole("deepdipole_new.pb") as native: + assert_evaluators_match(self.dp, native) + + # Exercise the legacy TF-specific deep_tensor.py path separately; the + # public DeepDipole wrapper does not instantiate this implementation. + with ( + DeepDipoleOld( + "deepdipole_new.pb", + neighbor_list=ase.neighborlist.NewPrimitiveNeighborList( + cutoffs=6, bothways=True + ), + ) as external_old, + DeepDipoleOld("deepdipole_new.pb") as native_old, + ): + assert_evaluators_match(external_old, native_old)