[Code scan] Map TensorFlow spin virtual types by use_spin

[njzjz]

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

Problem

Spin represents virtual spin types only for real atom types where use_spin is true:

deepmd-kit/deepmd/utils/spin.py

Lines 41 to 50 in 73de44b

	self.ntypes_real = len(use_spin)
	self.ntypes_spin = use_spin.count(True)
	self.use_spin = np.array(use_spin)
	self.spin_mask = self.use_spin.astype(np.int64)
	self.ntypes_real_and_spin = self.ntypes_real + self.ntypes_spin
	self.ntypes_placeholder = self.ntypes_real - self.ntypes_spin
	self.ntypes_input = 2 * self.ntypes_real # with placeholder for input types
	self.real_type = np.arange(self.ntypes_real, dtype=type_dtype)
	self.spin_type = self.real_type[self.use_spin] + self.ntypes_real
	self.real_and_spin_type = np.concatenate([self.real_type, self.spin_type])

For example, use_spin=[False, True] produces one virtual type for real type 1, not a dense virtual block for all real types.

The legacy TensorFlow SE-A descriptor extends neighbor selections by taking the first ntypes_spin real selections:

deepmd-kit/deepmd/tf/descriptor/se_a.py

Lines 234 to 239 in 73de44b

	# extend sel_a for spin system
	if self.spin is not None:
	self.ntypes_spin = self.spin.get_ntypes_spin()
	self.sel_a_spin = self.sel_a[: self.ntypes_spin]
	self.sel_a.extend(self.sel_a_spin)
	else:

This duplicates type 0's selection for use_spin=[False, True], instead of duplicating the selected spin type's selection. Related coordinate and force logic then assumes the virtual type for real type i is located at i + len(use_spin):

deepmd-kit/deepmd/tf/descriptor/se_a.py

Lines 1407 to 1423 in 73de44b

	def natoms_match(self, coord: tf.Tensor, natoms: tf.Tensor) -> tf.Tensor:
	natoms_index = tf.concat([[0], tf.cumsum(natoms[2:])], axis=0)
	diff_coord_loc = []
	for i in range(self.ntypes):
	if i + self.ntypes_spin >= self.ntypes:
	diff_coord_loc.append(
	tf.slice(
	coord,
	[0, natoms_index[i] * 3],
	[-1, natoms[2 + i] * 3],
	)
	- tf.slice(
	coord,
	[0, natoms_index[i - len(self.spin.use_spin)] * 3],
	[-1, natoms[2 + i - len(self.spin.use_spin)] * 3],
	)
	)

deepmd-kit/deepmd/tf/model/ener.py

Lines 500 to 535 in 73de44b

	use_spin = self.spin.use_spin
	virtual_len = self.spin.virtual_len
	spin_norm = self.spin.spin_norm
	natoms_index = tf.concat([[0], tf.cumsum(natoms[2:])], axis=0)
	force_real_list = []
	for idx, use in enumerate(use_spin):
	if use is True:
	force_real_list.append(
	tf.slice(
	force, [0, natoms_index[idx] * 3], [-1, natoms[idx + 2] * 3]
	)
	+ tf.slice(
	force,
	[0, natoms_index[idx + len(use_spin)] * 3],
	[-1, natoms[idx + 2 + len(use_spin)] * 3],
	)
	)
	else:
	force_real_list.append(
	tf.slice(
	force, [0, natoms_index[idx] * 3], [-1, natoms[idx + 2] * 3]
	)
	)
	force_mag_list = []
	for idx, use in enumerate(use_spin):
	if use is True:
	force_mag_list.append(
	tf.slice(
	force,
	[0, natoms_index[idx + len(use_spin)] * 3],
	[-1, natoms[idx + 2 + len(use_spin)] * 3],
	)
	)
	force_mag_list[idx] *= virtual_len[idx] / spin_norm[idx]
	force_real = tf.concat(force_real_list, axis=1)

Bias adjustment uses the same dense-offset assumption:

deepmd-kit/deepmd/tf/fit/ener.py

Lines 535 to 553 in 73de44b

	ntypes_atom = self.ntypes - self.ntypes_spin
	if self.spin is not None:
	for type_i in range(ntypes_atom):
	if self.bias_atom_e.shape[0] != self.ntypes:
	self.bias_atom_e = np.pad(
	self.bias_atom_e,
	(0, self.ntypes_spin),
	"constant",
	constant_values=(0, 0),
	)
	bias_atom_e = self.bias_atom_e
	if self.spin.use_spin[type_i]:
	self.bias_atom_e[type_i] = (
	self.bias_atom_e[type_i]
	+ self.bias_atom_e[type_i + ntypes_atom]
	)
	else:
	self.bias_atom_e[type_i] = self.bias_atom_e[type_i]
	self.bias_atom_e = self.bias_atom_e[:ntypes_atom]

Impact

TensorFlow spin models only work reliably when all spin-enabled real types form a prefix of the type map. For non-prefix layouts such as use_spin=[False, True], neighbor selection, coordinate splitting, force splitting, and bias merging can read the wrong real/virtual type ranges.

Suggested fix

Drive the TensorFlow spin layout from Spin.spin_type / Spin.use_spin instead of assuming a dense virtual block. When extending sel_a, append selections for the actual spin-enabled real types, and use an explicit mapping from real type to virtual type in coordinate, force, and bias logic.