diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index 0d9da07355..e91598d121 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -81,6 +81,20 @@ def _to_xp(self, arr: Any, xp: Any, ref_arr: Any) -> Any: """Convert a numpy array to the same namespace as ref_arr.""" return xp.asarray(arr, device=array_api_compat.device(ref_arr)) + def _lookup_type_values(self, values: Any, atype: Array, ref_arr: Array) -> Array: + """Gather per-type values while mapping virtual atom types to zero. + + Negative atom types are padding placeholders, not Python-style indices + from the end of the type table. Their spin scale and mask must remain + zero until the backbone model applies its normal virtual-atom mask. + """ + xp = array_api_compat.array_namespace(ref_arr) + values = self._to_xp(values, xp, ref_arr) + real_atom = atype >= 0 + safe_atype = xp.where(real_atom, atype, xp.zeros_like(atype)) + gathered = values[safe_atype] + return xp.where(real_atom, gathered, xp.zeros_like(gathered)) + def process_spin_input( self, coord: Array, atype: Array, spin: Array ) -> tuple[Array, Array, Array]: @@ -97,9 +111,12 @@ def process_spin_input( """ 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_atype = xp.where(atype >= 0, atype + self.ntypes_real, atype) + atype_spin = xp.concat([atype, virtual_atype], axis=-1) + spin_dist = spin * xp.reshape( + self._lookup_type_values(self.virtual_scale_mask, atype, coord), + (nframes, nloc, 1), + ) virtual_coord = coord + spin_dist coord_spin = xp.concat([coord, virtual_coord], axis=-2) # for spin virial correction @@ -151,12 +168,18 @@ def process_spin_input_lower( 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) + self._lookup_type_values( + self.virtual_scale_mask, extended_atype, extended_coord + ), + (nframes, nall, 1), ) virtual_extended_coord = extended_coord + extended_spin_dist - virtual_extended_atype = extended_atype + self.ntypes_real + virtual_extended_atype = xp.where( + extended_atype >= 0, + extended_atype + self.ntypes_real, + extended_atype, + ) extended_coord_updated = self.concat_switch_virtual( extended_coord, virtual_extended_coord, nloc ) @@ -223,7 +246,10 @@ def process_spin_output( mask = self._to_xp(self.virtual_scale_mask, xp, out_tensor) else: mask = self._to_xp(self.spin_mask, xp, out_tensor) - atomic_mask = xp.reshape(mask[atype], (nframes, nloc, 1)) + atomic_mask = xp.reshape( + self._lookup_type_values(mask, atype, out_tensor), + (nframes, nloc, 1), + ) out_real, out_mag = out_tensor[:, :nloc], out_tensor[:, nloc:] if add_mag: out_real = out_real + out_mag @@ -249,7 +275,10 @@ def process_spin_output_lower( mask = self._to_xp(self.virtual_scale_mask, xp, extended_out_tensor) else: mask = self._to_xp(self.spin_mask, xp, extended_out_tensor) - atomic_mask = xp.reshape(mask[extended_atype], (nframes, nall, 1)) + atomic_mask = xp.reshape( + self._lookup_type_values(mask, extended_atype, extended_out_tensor), + (nframes, nall, 1), + ) extended_out_real = xp.concat( [ extended_out_tensor[:, :nloc], @@ -698,8 +727,10 @@ def call_common( if "mask_mag" not in model_ret: xp = array_api_compat.array_namespace(atype) nframes_m, nloc_m = atype.shape[:2] - vsm = self._to_xp(self.virtual_scale_mask, xp, atype) - atomic_mask = xp.reshape(vsm[atype], (nframes_m, nloc_m, 1)) + atomic_mask = xp.reshape( + self._lookup_type_values(self.virtual_scale_mask, atype, atype), + (nframes_m, nloc_m, 1), + ) model_ret["mask_mag"] = atomic_mask > 0.0 return model_ret @@ -881,8 +912,12 @@ def call_common_lower( if "mask_mag" not in model_ret: xp = array_api_compat.array_namespace(extended_atype) nall = extended_atype.shape[1] - vsm = self._to_xp(self.virtual_scale_mask, xp, extended_atype) - atomic_mask = xp.reshape(vsm[extended_atype], (nframes, nall, 1)) + atomic_mask = xp.reshape( + self._lookup_type_values( + self.virtual_scale_mask, extended_atype, extended_atype + ), + (nframes, nall, 1), + ) model_ret["mask_mag"] = atomic_mask > 0.0 return model_ret diff --git a/source/tests/common/dpmodel/test_spin_model_virtual_types.py b/source/tests/common/dpmodel/test_spin_model_virtual_types.py new file mode 100644 index 0000000000..35246f4e03 --- /dev/null +++ b/source/tests/common/dpmodel/test_spin_model_virtual_types.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Regression tests for virtual placeholders in dpmodel SpinModel inputs.""" + +import numpy as np +import pytest + +from deepmd.dpmodel.common import ( + to_numpy_array, +) +from deepmd.dpmodel.model.model import ( + get_model, +) + +MODEL_CONFIG = { + "type_map": ["A", "B", "C"], + "descriptor": { + "type": "se_e2_a", + "sel": [4, 4, 4], + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [3, 6], + "axis_neuron": 2, + "precision": "float64", + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "type": "ener", + "neuron": [5, 5], + "precision": "float64", + "seed": 1, + }, + # Keep the final real type magnetic so an accidental ``mask[-1]`` lookup + # is observable instead of being hidden by a zero scale. + "spin": {"use_spin": [False, False, True], "virtual_scale": [0.5]}, +} + + +@pytest.fixture +def model(): + """Build a small NumPy spin model with a magnetic final real type.""" + return get_model(MODEL_CONFIG) + + +def test_dense_spin_expansion_preserves_virtual_types(model) -> None: + """A dense placeholder and its spin partner must both stay virtual.""" + coord = np.arange(9, dtype=np.float64).reshape(1, 3, 3) + atype = np.array([[0, -1, 2]], dtype=np.int32) + spin = np.ones_like(coord) + + coord_updated, atype_updated, coord_corr = model.process_spin_input( + coord, atype, spin + ) + + np.testing.assert_array_equal(atype_updated, [[0, -1, 2, 3, -1, 5]]) + # The placeholder's virtual partner has neither a displacement nor a + # virial correction, even though the final real type is magnetic. + np.testing.assert_array_equal(coord_updated[:, 4], coord[:, 1]) + np.testing.assert_array_equal(coord_corr[:, 4], 0.0) + + _, magnetic_output, magnetic_mask = model.process_spin_output( + atype, np.ones((1, 6, 3), dtype=np.float64) + ) + np.testing.assert_array_equal(magnetic_output[:, 1], 0.0) + np.testing.assert_array_equal(magnetic_mask[:, 1], False) + + prediction = model(coord, atype, spin, box=None) + np.testing.assert_array_equal(prediction["atom_energy"][:, 1], 0.0) + np.testing.assert_array_equal(prediction["mask_mag"][:, 1], False) + + changed_coord = coord.copy() + changed_spin = spin.copy() + changed_coord[:, 1] += 100.0 + changed_spin[:, 1] += 100.0 + changed_prediction = model(changed_coord, atype, changed_spin, box=None) + np.testing.assert_allclose(changed_prediction["energy"], prediction["energy"]) + np.testing.assert_allclose( + changed_prediction["atom_energy"][:, [0, 2]], + prediction["atom_energy"][:, [0, 2]], + ) + + +def test_lower_spin_expansion_preserves_virtual_types(model) -> None: + """Local and ghost placeholders remain negative in the switched layout.""" + extended_coord = np.arange(12, dtype=np.float64).reshape(1, 4, 3) + extended_atype = np.array([[0, -1, 2, -1]], dtype=np.int32) + extended_spin = np.ones_like(extended_coord) + nlist = np.array([[[2, -1], [0, -1]]], dtype=np.int32) + mapping = np.array([[0, 1, 0, 1]], dtype=np.int32) + + ( + coord_updated, + atype_updated, + _, + _, + coord_corr, + ) = model.process_spin_input_lower( + extended_coord, + extended_atype, + extended_spin, + nlist, + mapping=mapping, + ) + + np.testing.assert_array_equal(atype_updated, [[0, -1, 3, -1, 2, -1, 5, -1]]) + for real_index, virtual_index in ((1, 3), (3, 7)): + np.testing.assert_array_equal( + coord_updated[:, virtual_index], extended_coord[:, real_index] + ) + np.testing.assert_array_equal(coord_corr[:, virtual_index], 0.0) + + _, magnetic_output, magnetic_mask = model.process_spin_output_lower( + extended_atype, + np.ones((1, 8, 3), dtype=np.float64), + nloc=2, + ) + np.testing.assert_array_equal(magnetic_output[:, [1, 3]], 0.0) + np.testing.assert_array_equal(magnetic_mask[:, [1, 3]], False) + + +def test_virtual_type_lookup_supports_array_api_strict(model) -> None: + """The masked lookup must not rely on NumPy negative-index semantics.""" + xp = pytest.importorskip("array_api_strict") + coord = xp.asarray(np.arange(9, dtype=np.float64).reshape(1, 3, 3)) + atype = xp.asarray(np.array([[0, -1, 2]], dtype=np.int64)) + spin = xp.ones_like(coord) + + coord_updated, atype_updated, coord_corr = model.process_spin_input( + coord, atype, spin + ) + np.testing.assert_array_equal(to_numpy_array(atype_updated), [[0, -1, 2, 3, -1, 5]]) + np.testing.assert_array_equal(to_numpy_array(coord_updated)[:, 4], [[3, 4, 5]]) + np.testing.assert_array_equal(to_numpy_array(coord_corr)[:, 4], 0.0)