fix(dpmodel,pt_expt): normalize use_spin forms and forward allow_missing_label in the native-spin builders

Review: PR deepmodeling#5884 (OutisLi, P1 x2). The DPA4 native-spin builders did a
direct boolean conversion of spin.use_spin, breaking the public schema's
index and symbol forms (['Ni'] became [True] and failed the length check;
[0, 1] became [False, True]), and dropped spin.allow_missing_label (so
datasets without spin.npy were rejected despite the documented zero
default).

- New pure helper deepmd.utils.spin.normalize_spin_use_spin(use_spin,
  type_map) -> list[bool]: single owner of the three-form contract pt
  implements in its frozen tree (bool passthrough / index scatter /
  symbol lookup with ValueError on unknown symbols); direct unit tests
  cover all forms incl. the empty list and input purity.
- Both builder twins (dpmodel model.py get_dpa4_native_spin_model and
  pt_expt get_model._get_dpa4_native_spin_model) normalize via the
  helper and forward allow_missing_label into Spin(...).
- get_additional_data_requirement drops its getattr has_spin/spin
  dances for direct calls (has_spin is base-declared; spin models carry
  .spin).
- Tests: config-form parametrization (symbols/indices/booleans/unknown
  symbol) on both builders; the pt_expt data-requirement regression pins
  must=False + default=0.0 with allow_missing_label=True and must=True
  without.

Author: Han Wang

deepmd/dpmodel/model/model.py

@@ -48,6 +48,7 @@
 )
 from deepmd.utils.spin import (
     Spin,
+    normalize_spin_use_spin,
 )
 
 _DPA4_SEZM_DESCRIPTOR_TYPES = ("dpa4", "DPA4", "sezm", "SeZM")
@@ -223,10 +224,13 @@ def get_dpa4_native_spin_model(data: dict) -> DPA4NativeSpinModel:
             "is a follow-up; use one or the other"
         )
     spin_cfg = data.pop("spin")
-    use_spin = [bool(flag) for flag in spin_cfg["use_spin"]]
+    # Expand index/symbol forms of ``use_spin`` against ``type_map`` into the
+    # per-type boolean list (pure; validates symbols).
+    use_spin = normalize_spin_use_spin(spin_cfg["use_spin"], data["type_map"])
     spin = Spin(
         use_spin=use_spin,
         virtual_scale=spin_cfg.get("virtual_scale", 1.0),
+        allow_missing_label=spin_cfg.get("allow_missing_label", False),
     )
     data["descriptor"]["use_spin"] = use_spin
     backbone_model = get_standard_model(data)

deepmd/pt_expt/model/get_model.py

@@ -46,6 +46,7 @@
 )
 from deepmd.utils.spin import (
     Spin,
+    normalize_spin_use_spin,
 )
 
 log = logging.getLogger(__name__)
@@ -242,10 +243,13 @@ def _get_dpa4_native_spin_model(data: dict) -> DPA4NativeSpinModel:
             "charge-spin FiLM combined with native spin on the graph route "
             "is a follow-up; use one or the other"
         )
-    use_spin = [bool(flag) for flag in spin_cfg["use_spin"]]
+    # Expand index/symbol forms of ``use_spin`` against ``type_map`` into the
+    # per-type boolean list (pure; validates symbols).
+    use_spin = normalize_spin_use_spin(spin_cfg["use_spin"], data["type_map"])
     spin = Spin(
         use_spin=use_spin,
         virtual_scale=spin_cfg.get("virtual_scale", 1.0),
+        allow_missing_label=spin_cfg.get("allow_missing_label", False),
     )
     data["descriptor"]["use_spin"] = use_spin
     backbone_model = get_sezm_model(data)

deepmd/pt_expt/train/training.py

@@ -261,17 +261,13 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]:
                 "aparam", _model.get_dim_aparam(), atomic=True, must=True
             )
         )
-    has_spin = getattr(_model, "has_spin", False)
-    if callable(has_spin):
-        has_spin = has_spin()
-    if has_spin:
+    if _model.has_spin():
         # ``model.spin.allow_missing_label`` relaxes the spin label from
         # mandatory to optional with a zero default, so a system without a
         # ``spin`` file is filled with zeros rather than rejected. Mirrors
         # ``deepmd.pt.train.training.get_additional_data_requirement``.
-        allow_missing_spin = getattr(
-            getattr(_model, "spin", None), "allow_missing_label", False
-        )
+        # Every spin model wrapper carries a ``spin`` attribute.
+        allow_missing_spin = _model.spin.allow_missing_label
         additional_data_requirement.append(
             DataRequirementItem(
                 "spin",

deepmd/utils/spin.py

@@ -8,6 +8,51 @@
 )
 
 
+def normalize_spin_use_spin(use_spin: list, type_map: list[str]) -> list[bool]:
+    """Normalize ``use_spin`` to a per-type boolean list.
+
+    Three equivalent forms are accepted: a per-type boolean list, a list of
+    magnetic type indices, or a list of magnetic element symbols. The index
+    and symbol forms are expanded against ``type_map``, so a large type map
+    only needs its magnetic species named. Pure: the inputs are not
+    modified.
+
+    Parameters
+    ----------
+    use_spin : list
+        The ``spin.use_spin`` configuration value, in any of the three
+        accepted forms.
+    type_map : list[str]
+        The model's type map, defining the per-type order and (for the
+        symbol form) the element names.
+
+    Returns
+    -------
+    list[bool]
+        The per-type boolean form, of length ``len(type_map)``.
+
+    Raises
+    ------
+    ValueError
+        If a symbol in ``use_spin`` is absent from ``type_map``.
+    """
+    if use_spin and isinstance(use_spin[0], str):
+        type_index = {name: idx for idx, name in enumerate(type_map)}
+        unknown = [name for name in use_spin if name not in type_index]
+        if unknown:
+            raise ValueError(
+                f"spin.use_spin references element(s) {unknown} absent from type_map."
+            )
+        use_spin = [type_index[name] for name in use_spin]
+    # ``bool`` is a subclass of ``int``; an already-boolean list is passed
+    # through while an index list is scattered into a per-type mask.
+    if not use_spin or not isinstance(use_spin[0], bool):
+        mask = np.full(len(type_map), False, dtype=bool)
+        mask[use_spin] = True
+        return mask.tolist()
+    return [bool(flag) for flag in use_spin]
+
+
 class Spin:
     """Class for spin, mainly processes the spin type-related information.
     Atom types can be split into three kinds:

source/tests/common/dpmodel/test_dpa4_native_spin_model.py

@@ -126,3 +126,47 @@ def test_translated_output_def_has_spin_keys(self):
         assert "force_mag" in out_def
         assert "energy" in out_def
         assert "force" in out_def
+
+
+class TestNativeSpinConfigForms:
+    """``spin.use_spin`` index/symbol forms and ``allow_missing_label``.
+
+    The public schema accepts a per-type boolean list, a list of magnetic
+    type indices, or a list of element symbols (expanded against
+    ``type_map`` by ``normalize_spin_use_spin``); ``allow_missing_label``
+    must be forwarded into the constructed :class:`Spin`.
+    """
+
+    @pytest.mark.parametrize(
+        ("use_spin_form", "expected"),
+        [
+            (["Ni"], [True, False]),  # element-symbol form
+            ([0], [True, False]),  # type-index form
+            ([0, 1], [True, True]),  # multiple type indices
+            ([True, False], [True, False]),  # canonical boolean passthrough
+        ],
+    )
+    def test_use_spin_forms(self, use_spin_form, expected):
+        config = copy.deepcopy(NATIVE_SPIN_CONFIG)
+        config["spin"] = {"use_spin": use_spin_form, "scheme": "native"}
+        model = get_model(config)
+        assert model.spin.use_spin.tolist() == expected
+        # The descriptor consumes the SAME normalized boolean list.
+        descriptor = model.backbone_model.atomic_model.descriptor
+        assert [bool(flag) for flag in descriptor.use_spin] == expected
+
+    def test_use_spin_unknown_symbol_raises(self):
+        config = copy.deepcopy(NATIVE_SPIN_CONFIG)
+        config["spin"] = {"use_spin": ["Fe"], "scheme": "native"}
+        with pytest.raises(ValueError, match="absent from type_map"):
+            get_model(config)
+
+    def test_allow_missing_label_forwarded(self):
+        config = copy.deepcopy(NATIVE_SPIN_CONFIG)
+        config["spin"]["allow_missing_label"] = True
+        model = get_model(config)
+        assert model.spin.allow_missing_label is True
+
+    def test_allow_missing_label_default_false(self):
+        model = get_model(copy.deepcopy(NATIVE_SPIN_CONFIG))
+        assert model.spin.allow_missing_label is False

source/tests/common/test_spin.py

@@ -6,11 +6,52 @@
 
 from deepmd.utils.spin import (
     Spin,
+    normalize_spin_use_spin,
 )
 
 CUR_DIR = os.path.dirname(__file__)
 
 
+class NormalizeUseSpinTest(unittest.TestCase):
+    """Unit tests for ``normalize_spin_use_spin`` (pure; all three forms)."""
+
+    def setUp(self) -> None:
+        self.type_map = ["Ni", "O", "H"]
+
+    def test_boolean_passthrough(self) -> None:
+        self.assertEqual(
+            normalize_spin_use_spin([True, False, True], self.type_map),
+            [True, False, True],
+        )
+
+    def test_index_form(self) -> None:
+        self.assertEqual(
+            normalize_spin_use_spin([0, 2], self.type_map),
+            [True, False, True],
+        )
+
+    def test_symbol_form(self) -> None:
+        self.assertEqual(
+            normalize_spin_use_spin(["Ni", "H"], self.type_map),
+            [True, False, True],
+        )
+
+    def test_empty_list_all_false(self) -> None:
+        self.assertEqual(
+            normalize_spin_use_spin([], self.type_map),
+            [False, False, False],
+        )
+
+    def test_unknown_symbol_raises(self) -> None:
+        with self.assertRaisesRegex(ValueError, "absent from type_map"):
+            normalize_spin_use_spin(["Fe"], self.type_map)
+
+    def test_pure_no_input_mutation(self) -> None:
+        use_spin = ["Ni"]
+        normalize_spin_use_spin(use_spin, self.type_map)
+        self.assertEqual(use_spin, ["Ni"])
+
+
 class SpinTest(unittest.TestCase):
     def setUp(self) -> None:
         type_map_1 = ["H", "O"]

source/tests/pt_expt/model/test_dpa4_native_spin.py

@@ -902,3 +902,60 @@ def test_training_smoke(self, tmp_path) -> None:
             )
         finally:
             os.chdir(old_cwd)
+
+
+class TestNativeSpinConfigFormsPtExpt:
+    """pt_expt twin of ``test_dpa4_native_spin_model.py::TestNativeSpinConfigForms``.
+
+    ``spin.use_spin`` index/symbol forms are expanded against ``type_map``
+    (``normalize_spin_use_spin``), and ``allow_missing_label`` is forwarded
+    into the constructed :class:`Spin` -- observable through the trainer's
+    spin data requirement.
+    """
+
+    @pytest.mark.parametrize(
+        ("use_spin_form", "expected"),
+        [
+            (["Ni"], [True, False]),  # element-symbol form
+            ([0], [True, False]),  # type-index form
+            ([0, 1], [True, True]),  # multiple type indices
+            ([True, False], [True, False]),  # canonical boolean passthrough
+        ],
+    )
+    def test_use_spin_forms(self, use_spin_form, expected) -> None:
+        config = copy.deepcopy(NATIVE_SPIN_CONFIG)
+        config["spin"] = {"use_spin": use_spin_form, "scheme": "native"}
+        model = get_model(config)
+        assert model.spin.use_spin.tolist() == expected
+        descriptor = model.backbone_model.atomic_model.descriptor
+        assert [bool(flag) for flag in descriptor.use_spin] == expected
+
+    def test_use_spin_unknown_symbol_raises(self) -> None:
+        config = copy.deepcopy(NATIVE_SPIN_CONFIG)
+        config["spin"] = {"use_spin": ["Fe"], "scheme": "native"}
+        with pytest.raises(ValueError, match="absent from type_map"):
+            get_model(config)
+
+    @pytest.mark.parametrize(
+        ("allow_missing", "expected_must"),
+        [
+            (True, False),  # relaxed: spin file optional, zero default
+            (False, True),  # default: spin file mandatory
+        ],
+    )
+    def test_allow_missing_label_data_requirement(
+        self, allow_missing, expected_must
+    ) -> None:
+        from deepmd.pt_expt.train.training import (
+            get_additional_data_requirement,
+        )
+
+        config = copy.deepcopy(NATIVE_SPIN_CONFIG)
+        if allow_missing:
+            config["spin"]["allow_missing_label"] = True
+        model = get_model(config)
+        assert model.spin.allow_missing_label is allow_missing
+        reqs = get_additional_data_requirement(model)
+        spin_req = next(rr for rr in reqs if rr.key == "spin")
+        assert spin_req.must is expected_must
+        assert spin_req.default == 0.0