[Code scan] Ignore virtual atoms in dpmodel neighbor statistics #5630

Found during a Codex global scan of deepmodeling/deepmd-kit at commit 73de44b1f94471b2e3bdb6b11f57b34d7bc791bb.

Problem

NeighborStatOP.call() computes the minimum pair distance and per-center neighbor counts without consistently masking virtual atoms (atype < 0).

Evidence:

• dpmodel neighbor-list code documents negative types as virtual atoms:

  deepmd-kit/deepmd/dpmodel/utils/nlist.py

  Lines 89 to 117 in 73de44b

  	atype : Array
  	extended atomic types of shape [batch_size, nall]
  	type < 0 the atom is treat as virtual atoms.
  	nloc : int
  	number of local atoms.
  	rcut : float
  	cut-off radius
  	sel : int or list[int]
  	maximal number of neighbors (of each type).
  	if distinguish_types==True, nsel should be list and
  	the length of nsel should be equal to number of
  	types.
  	distinguish_types : bool
  	distinguish different types.
  	Returns
  	-------
  	neighbor_list : Array
  	Neighbor list of shape [batch_size, nloc, nsel], the neighbors
  	are stored in an ascending order. If the number of
  	neighbors is less than nsel, the positions are masked
  	with -1. The neighbor list of an atom looks like
  	|------ nsel ------|
  	xx xx xx xx -1 -1 -1
  	if distinguish_types==True and we have two types
  	|---- nsel[0] -----| |---- nsel[1] -----|
  	xx xx xx xx -1 -1 -1 xx xx xx -1 -1 -1 -1
  	For virtual atoms all neighboring positions are filled with -1.

• extend_coord_with_ghosts preserves negative atom types in the extended atom-type array:

  deepmd-kit/deepmd/dpmodel/utils/nlist.py

  Lines 438 to 500 in 73de44b

  	xp = array_api_compat.array_namespace(coord, atype)
  	nf, nloc = atype.shape
  	# int64 for index
  	aidx = xp.tile(
  	xp.arange(nloc, dtype=xp.int64, device=array_api_compat.device(atype))[
  	xp.newaxis, :
  	],
  	(nf, 1),
  	)
  	if cell is None:
  	nall = nloc
  	extend_coord = coord
  	extend_atype = atype
  	extend_aidx = aidx
  	else:
  	coord = xp.reshape(coord, (nf, nloc, 3))
  	cell = xp.reshape(cell, (nf, 3, 3))
  	to_face = to_face_distance(cell)
  	nbuff = xp.astype(xp.ceil(rcut / to_face), xp.int64)
  	nbuff = xp.max(nbuff, axis=0)
  	device = array_api_compat.device(coord)
  	xi = _arange_nbuff(nbuff, 0, xp, device)
  	yi = _arange_nbuff(nbuff, 1, xp, device)
  	zi = _arange_nbuff(nbuff, 2, xp, device)
  	xyz = xp.linalg.outer(
  	xi, xp.asarray([1, 0, 0], device=array_api_compat.device(xi))
  	)[:, xp.newaxis, xp.newaxis, :]
  	xyz = (
  	xyz
  	+ xp.linalg.outer(
  	yi, xp.asarray([0, 1, 0], device=array_api_compat.device(yi))
  	)[xp.newaxis, :, xp.newaxis, :]
  	)
  	xyz = (
  	xyz
  	+ xp.linalg.outer(
  	zi, xp.asarray([0, 0, 1], device=array_api_compat.device(zi))
  	)[xp.newaxis, xp.newaxis, :, :]
  	)
  	xyz = xp.reshape(xyz, (-1, 3))
  	xyz = xp.astype(xyz, coord.dtype)
  	shift_idx = xp.take(xyz, xp.argsort(xp.linalg.vector_norm(xyz, axis=1)), axis=0)
  	ns, _ = shift_idx.shape
  	nall = ns * nloc
  	if array_api_compat.is_jax_namespace(xp):
  	# Avoid JAX internal errors in tensordot.
  	shift_vec = xp.sum(
  	shift_idx[xp.newaxis, :, :, xp.newaxis] * cell[:, xp.newaxis, :, :],
  	axis=2,
  	)
  	else:
  	# shift_vec = xp.einsum("sd,fdk->fsk", shift_idx, cell)
  	shift_vec = xp.tensordot(shift_idx, cell, axes=([1], [1]))
  	shift_vec = xp.permute_dims(shift_vec, (1, 0, 2))
  	extend_coord = coord[:, None, :, :] + shift_vec[:, :, None, :]
  	extend_atype = xp.tile(atype[:, :, xp.newaxis], (1, ns, 1))
  	extend_aidx = xp.tile(aidx[:, :, xp.newaxis], (1, ns, 1))
  	return (
  	xp.reshape(extend_coord, (nf, nall * 3)),
  	xp.reshape(extend_atype, (nf, nall)),
  	xp.reshape(extend_aidx, (nf, nall)),
  	)

• NeighborStatOP.call() removes only self pairs before computing min_rr2; it does not mask pairs where the center or neighbor atom is virtual:

  deepmd-kit/deepmd/dpmodel/utils/neighbor_stat.py

  Lines 81 to 95 in 73de44b

  	coord1 = xp.reshape(extend_coord, (nframes, -1))
  	nall = coord1.shape[1] // 3
  	coord0 = coord1[:, : nloc * 3]
  	diff = (
  	xp.reshape(coord1, (nframes, -1, 3))[:, None, :, :]
  	- xp.reshape(coord0, (nframes, -1, 3))[:, :, None, :]
  	)
  	assert list(diff.shape) == [nframes, nloc, nall, 3]
  	# remove the diagonal elements
  	mask = xp.eye(nloc, nall, dtype=xp.bool, device=array_api_compat.device(diff))
  	mask = xp.tile(mask[None, :, :, None], (nframes, 1, 1, 3))
  	diff = xp.where(mask, xp.full_like(diff, xp.inf), diff)
  	rr2 = xp.sum(xp.square(diff), axis=-1)
  	min_rr2 = xp.min(rr2, axis=-1)
  	# count the number of neighbors

• The mixed-types count masks negative neighbors, but it still counts real neighbors around virtual centers; the non-mixed branch also lacks a center mask:

  deepmd-kit/deepmd/dpmodel/utils/neighbor_stat.py

  Lines 96 to 108 in 73de44b

  	if not self.mixed_types:
  	mask = rr2 < self.rcut**2
  	nneis = []
  	for ii in range(self.ntypes):
  	nneis.append(xp.sum(mask & (extend_atype == ii)[:, None, :], axis=-1))
  	nnei = xp.stack(nneis, axis=-1)
  	else:
  	mask = rr2 < self.rcut**2
  	# virtual type (<0) are not counted
  	nnei = xp.sum(mask & (extend_atype >= 0)[:, None, :], axis=-1)
  	nnei = xp.reshape(nnei, (nframes, nloc, 1))
  	max_nnei = xp.max(nnei, axis=1)
  	return min_rr2, max_nnei

A minimal external run with one real atom and one overlapping virtual atom returned min_rr2 == 0 and nonzero neighbor counts, even though the virtual atom should not contribute to statistics.

Impact

Virtual atoms can drive the reported minimum distance to zero and can inflate max_nnei through virtual center rows. These statistics feed automatic neighbor selection and environment-matrix range checks, so they can produce overly conservative selections or false close-contact diagnostics.

Suggested Fix

Apply a real-center and real-neighbor mask before both min_rr2 and neighbor counting. Add regression tests with overlapping real/virtual atoms and with a virtual center near real atoms, for both mixed_types=True and False.