From f75ab0b8f2e5efaef369bb4f048f7425e690e40e Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 08:34:24 +0800 Subject: [PATCH] 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()