Skip to content

[Code scan] Filter virtual atoms in ASE neighbor-graph construction #5664

Description

@njzjz

This issue comes from a Codex global scan of deepmodeling/deepmd-kit at commit 73de44b1f94471b2e3bdb6b11f57b34d7bc791bb.

Problem

The ASE neighbor-graph builder documents atype as unused and runs ASE neighbor search over every coordinate row, so virtual atoms with atype < 0 can become graph nodes and edges.

The ASE builder states that atype is unused:

Parameters
----------
coord
(nf, nloc, 3) local coordinates.
atype
(nf, nloc) local atom types (unused for the search; carried for API parity).
box

It then constructs ASE Atoms from all positions in each frame:

coord_np = _to_cpu_numpy(coord)
nf, nloc = coord_np.shape[:2]
coord_np = coord_np.reshape(nf, nloc, 3)
box_np = _to_cpu_numpy(box).reshape(nf, 3, 3) if box is not None else None
periodic = box is not None
i_parts = []
j_parts = []
S_parts = []
nframe_parts = []
for f in range(nf):
atoms = Atoms(
positions=coord_np[f],
cell=(box_np[f] if periodic else None),
pbc=periodic,
)
ii, jj, SS = neighbor_list("ijS", atoms, rcut)

The sparse (i, j, S) converter has no atype argument, so it cannot filter virtual atoms later:

def neighbor_graph_from_ijs(
i: Array,
j: Array,
S: Array,
coord: Array,
box: Array | None,
nframe_id: Array,
nloc: int,
layout: GraphLayout | None = None,
) -> NeighborGraph:

This differs from the dense graph builder, which explicitly treats type < 0 as virtual and excludes virtual centers and neighbors:

Parameters
----------
coord
(nf, nloc, 3) or (nf, nloc*3) local coordinates.
atype
(nf, nloc) local atom types; ``type < 0`` marks a virtual atom (excluded
as both a center and a neighbor).

# keep neighbors within rcut, dropping: the self extended atom (i==j; a periodic
# IMAGE of i has j!=i and is kept), virtual neighbors, and virtual centers.
not_self = j_grid != i_grid
vir_nei = xp.broadcast_to((extended_atype < 0)[:, None, :], (nf, nloc, nall))
vir_cen = xp.broadcast_to((atype < 0)[:, :, None], (nf, nloc, nall))
keep_mask = (
(dist <= rcut) & not_self & xp.logical_not(vir_nei) & xp.logical_not(vir_cen)
)

Impact

neighbor_graph_method="ase" can create edges to or from placeholder atoms in mixed-size batches, diverging from the dense graph path and potentially feeding negative atom types into graph descriptors.

Suggested fix

Filter virtual atoms before running the ASE search, or pass atype through to the sparse converter and drop edges whose center or neighbor has atype < 0. Add a parity regression comparing neighbor_graph_method="dense" and "ase" on a frame containing atype == -1 placeholders.

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    Status
    In Progress

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions