Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 34 additions & 7 deletions deepmd/tf/infer/deep_eval.py
Original file line number Diff line number Diff line change
Expand Up @@ -549,15 +549,23 @@ 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
----------
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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -1534,15 +1549,23 @@ 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
----------
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
Expand Down Expand Up @@ -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
Expand Down
10 changes: 8 additions & 2 deletions deepmd/tf/infer/deep_tensor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
35 changes: 34 additions & 1 deletion source/tests/infer/test_models.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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)
29 changes: 29 additions & 0 deletions source/tests/tf/test_deepdipole.py
Original file line number Diff line number Diff line change
Expand Up @@ -1162,3 +1162,32 @@
@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()

Check notice

Code scanning / CodeQL

Should use a 'with' statement Note test

Instance of context-manager class
BaseModifier
is closed in a finally block. Consider using 'with' statement.
Instance of context-manager class
DeepDipole
is closed in a finally block. Consider using 'with' statement.
Instance of context-manager class
DeepDOS
is closed in a finally block. Consider using 'with' statement.
Instance of context-manager class
DeepEval
is closed in a finally block. Consider using 'with' statement.
Instance of context-manager class
DeepPolar
is closed in a finally block. Consider using 'with' statement.
Instance of context-manager class
DeepGlobalPolar
is closed in a finally block. Consider using 'with' statement.
Instance of context-manager class
DeepPopulation
is closed in a finally block. Consider using 'with' statement.
Instance of context-manager class
DeepPot
is closed in a finally block. Consider using 'with' statement.
Instance of context-manager class
DeepProperty
is closed in a finally block. Consider using 'with' statement.
Instance of context-manager class
DeepTensor
is closed in a finally block. Consider using 'with' statement.
Instance of context-manager class
OldDeepTensor
is closed in a finally block. Consider using 'with' statement.
Instance of context-manager class
DeepWFC
is closed in a finally block. Consider using 'with' statement.
Instance of context-manager class
DipoleChargeModifier
is closed in a finally block. Consider using 'with' statement.
Instance of context-manager class
_ConcreteEval
is closed in a finally block. Consider using 'with' statement.
Loading