Skip to content

Commit 1f0317e

Browse files
njzjz-botnjzjz-bot
authored
fix(tf): preserve non-PBC ASE neighbor semantics (#5859)
Closes #5668. ## Summary - preserve the original cells=None decision separately from the identity box required by TensorFlow placeholders - pass a missing cell to ASE for open systems and make both TF1 neighbor-list builders handle that documented input without creating periodic ghosts - correct the builders return annotations and document their open-boundary behavior - add collected regressions for DeepPotential and for DeepTensor eval/eval_full ## Why existing tests missed this - the original external ASE neighbor-list tests were designed around periodic systems and never exercised cells=None - the parameterized TestDeepPot symbol is replaced by object, so TestDeepPotNeighborList did not inherit the base evaluation tests; collection contained only two locally declared skipped tests - although the DeepPotential testcase YAML contains a box: null result, it therefore never ran through the ASE-backed evaluator - DeepTensor neighbor-list tests inherited PBC fixtures and had no native-versus-ASE open-boundary comparison ## Validation - focused DeepPotential no-PBC regression: passed - focused DeepTensor eval/eval_full no-PBC regression: passed - full TestDeepDipoleNewPBCNeighborList class: 10 passed, 2 skipped - TF1 DeepPotential periodic native-versus-ASE parity smoke: passed - pytest collect-only confirms both new regressions are collected - ruff format . - ruff check . - git diff --check Coding agent: Codex Codex version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning effort: xhigh <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Bug Fixes** - Fixed neighbor-list construction for open-boundary (no PBC) systems to preserve correct boundary semantics. - Ensured consistent neighbor-list handling across TensorFlow inference paths (including energy/tensor and dipole evaluation), so results no longer depend on placeholder cell inputs. - **Tests** - Added new no-PBC neighbor-list verification against reference outputs (forces, energies, and virials). - Added dipole coverage confirming neighbor-list-based results match native neighbor-list behavior, including legacy implementation checks. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: njzjz-bot <njzjz.bot@gmail.com> Co-authored-by: njzjz-bot <njzjz-bot@users.noreply.github.com>
1 parent e37cc00 commit 1f0317e

4 files changed

Lines changed: 132 additions & 22 deletions

File tree

‎deepmd/tf/infer/deep_eval.py‎

Lines changed: 34 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -549,15 +549,23 @@ def build_neighbor_list(
549549
atype: np.ndarray,
550550
imap: np.ndarray,
551551
neighbor_list: "ase.neighborlist.NeighborList | None",
552-
) -> tuple[np.ndarray, np.ndarray]:
552+
) -> tuple[
553+
np.ndarray,
554+
np.ndarray,
555+
np.ndarray,
556+
np.ndarray,
557+
np.ndarray,
558+
np.ndarray,
559+
]:
553560
"""Make the mesh with neighbor list for a single frame.
554561
555562
Parameters
556563
----------
557564
coords : np.ndarray
558565
The coordinates of atoms. Should be of shape [natoms, 3]
559566
cell : Optional[np.ndarray]
560-
The cell of the system. Should be of shape [3, 3]
567+
The cell of the system. Should be of shape [3, 3]. None denotes
568+
open boundary conditions.
561569
atype : np.ndarray
562570
The type of atoms. Should be of shape [natoms]
563571
imap : np.ndarray
@@ -587,7 +595,11 @@ def build_neighbor_list(
587595
The index map of ghost atoms. Should be of shape [nghost]
588596
"""
589597
pbc = np.repeat(cell is not None, 3)
590-
cell = cell.reshape(3, 3)
598+
# ASE still requires a 3x3 cell for non-periodic systems, but the cell
599+
# must not be used to infer periodicity or create ghost atoms.
600+
cell = (
601+
np.zeros((3, 3), dtype=np.float64) if cell is None else cell.reshape(3, 3)
602+
)
591603
positions = coords.reshape(-1, 3)
592604
neighbor_list.bothways = True
593605
neighbor_list.self_interaction = False
@@ -814,6 +826,9 @@ def _prepare_feed_dict(
814826
else:
815827
pbc = True
816828
cells = np.array(cells).reshape([nframes, 9])
829+
# Keep the original boundary semantics separate from the identity box
830+
# used only to satisfy TensorFlow's non-optional box placeholder.
831+
neighbor_cell = cells if pbc else None
817832

818833
if self.has_fparam:
819834
assert fparam is not None
@@ -884,7 +899,7 @@ def _prepare_feed_dict(
884899
ghost_map,
885900
) = self.build_neighbor_list(
886901
coords,
887-
cells if cells is not None else None,
902+
neighbor_cell,
888903
atom_types,
889904
imap,
890905
self.neighbor_list,
@@ -1534,15 +1549,23 @@ def build_neighbor_list(
15341549
atype: np.ndarray,
15351550
imap: np.ndarray,
15361551
neighbor_list: "ase.neighborlist.NeighborList | None",
1537-
) -> tuple[np.ndarray, np.ndarray]:
1552+
) -> tuple[
1553+
np.ndarray,
1554+
np.ndarray,
1555+
np.ndarray,
1556+
np.ndarray,
1557+
np.ndarray,
1558+
np.ndarray,
1559+
]:
15381560
"""Make the mesh with neighbor list for a single frame.
15391561
15401562
Parameters
15411563
----------
15421564
coords : np.ndarray
15431565
The coordinates of atoms. Should be of shape [natoms, 3]
15441566
cell : Optional[np.ndarray]
1545-
The cell of the system. Should be of shape [3, 3]
1567+
The cell of the system. Should be of shape [3, 3]. None denotes
1568+
open boundary conditions.
15461569
atype : np.ndarray
15471570
The type of atoms. Should be of shape [natoms]
15481571
imap : np.ndarray
@@ -1572,7 +1595,11 @@ def build_neighbor_list(
15721595
The index map of ghost atoms. Should be of shape [nghost]
15731596
"""
15741597
pbc = np.repeat(cell is not None, 3)
1575-
cell = cell.reshape(3, 3)
1598+
# ASE still requires a 3x3 cell for non-periodic systems, but the cell
1599+
# must not be used to infer periodicity or create ghost atoms.
1600+
cell = (
1601+
np.zeros((3, 3), dtype=np.float64) if cell is None else cell.reshape(3, 3)
1602+
)
15761603
positions = coords.reshape(-1, 3)
15771604
neighbor_list.bothways = True
15781605
neighbor_list.self_interaction = False

‎deepmd/tf/infer/deep_tensor.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,9 @@ def eval(
202202
else:
203203
pbc = True
204204
cells = np.array(cells).reshape([nframes, 9])
205+
# Keep the original boundary semantics separate from the identity box
206+
# used only to satisfy TensorFlow's non-optional box placeholder.
207+
neighbor_cell = cells if pbc else None
205208

206209
# sort inputs
207210
coords, atom_types, imap, sel_at, sel_imap = self.sort_input(
@@ -227,7 +230,7 @@ def eval(
227230
_,
228231
) = self.build_neighbor_list(
229232
coords,
230-
cells if cells is not None else None,
233+
neighbor_cell,
231234
atom_types,
232235
imap,
233236
self.neighbor_list,
@@ -346,6 +349,9 @@ def eval_full(
346349
else:
347350
pbc = True
348351
cells = np.array(cells).reshape([nframes, 9])
352+
# Keep the original boundary semantics separate from the identity box
353+
# used only to satisfy TensorFlow's non-optional box placeholder.
354+
neighbor_cell = cells if pbc else None
349355
nout = self.output_dim
350356

351357
# sort inputs
@@ -373,7 +379,7 @@ def eval_full(
373379
ghost_map,
374380
) = self.build_neighbor_list(
375381
coords,
376-
cells if cells is not None else None,
382+
neighbor_cell,
377383
atom_types,
378384
imap,
379385
self.neighbor_list,

‎source/tests/infer/test_models.py‎

Lines changed: 43 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
# SPDX-License-Identifier: LGPL-3.0-or-later
22
import unittest
33

4-
import ase
4+
import ase.neighborlist
55
import dpdata
66
import numpy as np
77

@@ -32,15 +32,9 @@
3232
STRICT_KEYS = frozenset(("se_e2_a", "se_e2_r"))
3333

3434

35-
@parameterized(
36-
(
37-
"se_e2_a",
38-
"se_e2_r",
39-
"fparam_aparam",
40-
), # key
41-
(".pb", ".pth", ".pte", ".pt2"), # model extension
42-
)
43-
class TestDeepPot(unittest.TestCase):
35+
class DeepPotTestMixin:
36+
"""Shared DeepPotential checks for native and external neighbor lists."""
37+
4438
# moved from tests/tf/test_deeppot_a.py
4539

4640
@classmethod
@@ -396,19 +390,33 @@ def test_model_script_def(self) -> None:
396390
)
397391

398392

393+
@parameterized(
394+
(
395+
"se_e2_a",
396+
"se_e2_r",
397+
"fparam_aparam",
398+
), # key
399+
(".pb", ".pth", ".pte", ".pt2"), # model extension
400+
)
401+
class TestDeepPot(DeepPotTestMixin, unittest.TestCase):
402+
"""Run the common inference checks with native neighbor construction."""
403+
404+
399405
@parameterized(
400406
("se_e2_a",), # key
401407
(".pb",), # model extension
402408
)
403-
class TestDeepPotNeighborList(TestDeepPot):
409+
class TestDeepPotNeighborList(DeepPotTestMixin, unittest.TestCase):
410+
"""Run the common inference checks with an external ASE neighbor list."""
411+
404412
@classmethod
405413
def setUpClass(cls) -> None:
406414
key, extension = cls.param
407415
cls.places = STRICT_PLACES if key in STRICT_KEYS else default_places
408416
cls.case = get_cases()[key]
409-
model_name = cls.case.get_model(extension)
417+
cls.model_name = cls.case.get_model(extension)
410418
cls.dp = DeepEval(
411-
model_name,
419+
cls.model_name,
412420
neighbor_list=ase.neighborlist.NewPrimitiveNeighborList(
413421
cutoffs=cls.case.rcut, bothways=True
414422
),
@@ -421,3 +429,25 @@ def test_2frame_atm(self) -> None:
421429
@unittest.skip("Zero atoms not supported")
422430
def test_zero_input(self) -> None:
423431
pass
432+
433+
def test_nopbc_matches_reference(self) -> None:
434+
"""The ASE path must preserve a testcase's open-boundary semantics."""
435+
result = next(result for result in self.case.results if result.box is None)
436+
ee, ff, vv, ae, av = self.dp.eval(
437+
result.coord,
438+
None,
439+
result.atype,
440+
atomic=True,
441+
fparam=result.fparam,
442+
aparam=result.aparam,
443+
)[:5]
444+
445+
np.testing.assert_almost_equal(ff.ravel(), result.force.ravel(), STRICT_PLACES)
446+
np.testing.assert_almost_equal(
447+
ae.ravel(), result.atomic_energy.ravel(), STRICT_PLACES
448+
)
449+
np.testing.assert_almost_equal(
450+
av.ravel(), result.atomic_virial.ravel(), STRICT_PLACES
451+
)
452+
np.testing.assert_almost_equal(ee.ravel(), result.energy, STRICT_PLACES)
453+
np.testing.assert_almost_equal(vv.ravel(), result.virial, STRICT_PLACES)

‎source/tests/tf/test_deepdipole.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@
1212
from deepmd.tf.infer import (
1313
DeepDipole,
1414
)
15+
from deepmd.tf.infer.deep_dipole import (
16+
DeepDipoleOld,
17+
)
1518
from deepmd.tf.utils.convert import (
1619
convert_pbtxt_to_pb,
1720
)
@@ -1162,3 +1165,47 @@ def test_2frame_full_atm(self) -> None:
11621165
@unittest.skip("multiple frames not supported")
11631166
def test_2frame_old_atm(self) -> None:
11641167
pass
1168+
1169+
def test_nopbc_matches_native_neighbor_building(self) -> None:
1170+
"""ASE and native tensor inference must agree for an open system."""
1171+
1172+
def assert_evaluators_match(external, native) -> None:
1173+
actual_tensor = external.eval(self.coords, None, self.atype, atomic=True)
1174+
expected_tensor = native.eval(self.coords, None, self.atype, atomic=True)
1175+
np.testing.assert_almost_equal(
1176+
actual_tensor,
1177+
expected_tensor,
1178+
default_places,
1179+
)
1180+
1181+
actual_full = external.eval_full(
1182+
self.coords,
1183+
None,
1184+
self.atype,
1185+
atomic=True,
1186+
)
1187+
expected_full = native.eval_full(
1188+
self.coords,
1189+
None,
1190+
self.atype,
1191+
atomic=True,
1192+
)
1193+
for actual, expected in zip(actual_full, expected_full, strict=True):
1194+
np.testing.assert_almost_equal(actual, expected, default_places)
1195+
1196+
# The public wrapper reaches deepmd/tf/infer/deep_eval.py.
1197+
with DeepDipole("deepdipole_new.pb") as native:
1198+
assert_evaluators_match(self.dp, native)
1199+
1200+
# Exercise the legacy TF-specific deep_tensor.py path separately; the
1201+
# public DeepDipole wrapper does not instantiate this implementation.
1202+
with (
1203+
DeepDipoleOld(
1204+
"deepdipole_new.pb",
1205+
neighbor_list=ase.neighborlist.NewPrimitiveNeighborList(
1206+
cutoffs=6, bothways=True
1207+
),
1208+
) as external_old,
1209+
DeepDipoleOld("deepdipole_new.pb") as native_old,
1210+
):
1211+
assert_evaluators_match(external_old, native_old)

0 commit comments

Comments
 (0)