Skip to content

[Code scan] Preserve virtual atom types when expanding SpinModel inputs #5663

Description

@njzjz

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

Problem

SpinModel expands spin inputs before the backbone atomic-model wrapper clamps virtual atom types. Negative placeholder types can therefore be converted into valid real spin type ids.

The atomic-model contract treats atype < 0 as virtual placeholders:

def make_atom_mask(
self,
atype: t_tensor,
) -> t_tensor:
"""The atoms with type < 0 are treated as virtual atoms,
which serves as place-holders for multi-frame calculations

The normal atomic wrapper masks virtual atoms and passes clamped atom types into forward_atomic():

ext_atom_mask = self.make_atom_mask(extended_atype)
ret_dict = self.forward_atomic(
extended_coord,
xp.where(ext_atom_mask, extended_atype, 0),
nlist,
mapping=mapping,
fparam=fparam,
aparam=aparam,
comm_dict=comm_dict,
charge_spin=charge_spin,
)
atom_mask = xp_take_first_n(ext_atom_mask, 1, nloc)
return self._finalize_atomic_ret(ret_dict, atom_mask, atype)

SpinModel.process_spin_input() runs before that wrapper. It concatenates atype + self.ntypes_real for every atom:

xp = array_api_compat.array_namespace(coord)
nframes, nloc = coord.shape[:-1]
atype_spin = xp.concat([atype, atype + self.ntypes_real], axis=-1)
vsm = self._to_xp(self.virtual_scale_mask, xp, coord)
spin_dist = spin * xp.reshape(vsm[atype], (nframes, nloc, 1))
virtual_coord = coord + spin_dist

The lower-interface spin path does the same for extended atom types:

xp = array_api_compat.array_namespace(extended_coord)
nframes, nall = extended_coord.shape[:2]
nloc = nlist.shape[1]
vsm = self._to_xp(self.virtual_scale_mask, xp, extended_coord)
extended_spin_dist = extended_spin * xp.reshape(
vsm[extended_atype], (nframes, nall, 1)
)
virtual_extended_coord = extended_coord + extended_spin_dist
virtual_extended_atype = extended_atype + self.ntypes_real

For a placeholder atom with type -1, the generated virtual spin type becomes ntypes_real - 1, which is a valid type id. The same code also indexes virtual_scale_mask with the negative type id before clamping.

Impact

Padded or virtual atoms in spin batches can leak into the backbone model as valid spin atoms instead of remaining masked placeholders. This can corrupt spin-model outputs for mixed-size or placeholder-padded frames.

Suggested fix

Preserve negative types through spin expansion, or clamp/mask them before creating virtual spin atoms. The virtual partner of a placeholder atom should also stay negative and should not index spin masks as a real type. Add dense and lower-interface spin tests with 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