refactor: the atomic model answers supports_native_spin, not a descriptor hunt

The native-spin builders reached into the atomic model for a descriptor
(``learned_descriptor``) to decide spin eligibility -- duck-typing that
assumes a descriptor+fitting architecture, guesses which child of a
composition is the learned one, and cannot answer at all for an atomic model
with no descriptor.  Same mistake the ``supports_graph_route`` detour made:
the capability belongs ON the atomic model.

``supports_native_spin()`` is now the twin of ``uses_graph_lower()``:
concrete ``False`` on ``BaseAtomicModel``; ``DPAtomicModel`` delegates to its
descriptor; ``LinearEnergyAtomicModel`` answers from its children.  Both
builders just ask the atomic model, and ``learned_descriptor`` is gone.

``DPAtomicModel`` already cached the answer, but as a bool ATTRIBUTE -- which
is exactly why reaching for the descriptor looked necessary, since an
attribute cannot be overridden polymorphically by the base or by a
composition.  It becomes a method backed by the private
``_supports_native_spin``, so the per-forward hot path still avoids a
descriptor call.

The composition uses ANY, deliberately unlike ``uses_graph_lower``'s ALL:
every child must run on the shared graph, but spin only has to reach ONE
consumer -- analytical terms accept and ignore it.  With no consumer the
magnetic force is identically zero, which is not a spin model; pinned by a
test that a two-ZBL composition reports False.

Author: Han Wang

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -179,6 +179,17 @@ def uses_graph_lower(self) -> bool:
         """
         return False
 
+    def supports_native_spin(self) -> bool:
+        """Returns whether this atomic model consumes a per-atom spin input.
+
+        Generic capability (concrete default ``False``), the twin of
+        :meth:`uses_graph_lower`: the model layer asks the atomic model
+        directly instead of reaching into it for a descriptor, so the answer
+        stays correct for architectures with no descriptor at all (analytical
+        terms) and for compositions.
+        """
+        return False
+
     def get_default_fparam(self) -> list[float] | None:
         """Get the default frame parameters."""
         return None

deepmd/dpmodel/atomic_model/dp_atomic_model.py

@@ -133,7 +133,7 @@ def __init__(
         # (that would be a ``TypeError``, not a no-op). Queried via the
         # ``supports_native_spin`` capability method declared on
         # ``BaseDescriptor`` (concrete default ``False``; DPA4 overrides).
-        self.supports_native_spin: bool = self.descriptor.supports_native_spin()
+        self._supports_native_spin: bool = self.descriptor.supports_native_spin()
         # Same capability method as ``supports_native_spin`` above, for the
         # frame-level ``charge_spin`` FiLM kwarg: only DPA4's ``call_graph``
         # declares it; other descriptors' ``call_graph`` would ``TypeError``
@@ -167,6 +167,10 @@ def uses_graph_lower(self) -> bool:
         """Delegates to this model's own descriptor."""
         return bool(self.descriptor.uses_graph_lower())
 
+    def supports_native_spin(self) -> bool:
+        """Delegates to this model's own descriptor (cached at construction)."""
+        return self._supports_native_spin
+
     def fitting_output_def(self) -> FittingOutputDef:
         """Get the output def of the fitting net."""
         return self.fitting_net.output_def()
@@ -375,7 +379,7 @@ def forward_atomic_graph(
         # See ``self.supports_native_spin``/``self.supports_charge_spin`` in
         # ``__init__``: only forward the ``spin``/``charge_spin`` keyword to
         # descriptors whose ``call_graph`` declares it.
-        spin_kwargs = {"spin": spin} if self.supports_native_spin else {}
+        spin_kwargs = {"spin": spin} if self._supports_native_spin else {}
         charge_spin_kwargs = (
             {"charge_spin": charge_spin} if self.supports_charge_spin else {}
         )

deepmd/dpmodel/atomic_model/linear_atomic_model.py

@@ -261,6 +261,17 @@ def uses_graph_lower(self) -> bool:
         """
         return all(m.uses_graph_lower() for m in self.models)
 
+    def supports_native_spin(self) -> bool:
+        """Spin-capable when ANY child consumes the spin input.
+
+        Unlike :meth:`uses_graph_lower` (every child must run on the shared
+        graph), spin only has to reach ONE consumer: analytical terms accept
+        and ignore it, so a composition of a spin-aware learned model with a
+        ZBL term is a valid native-spin model.  With no consumer at all the
+        magnetic force would be identically zero, which is not a spin model.
+        """
+        return any(m.supports_native_spin() for m in self.models)
+
     def forward_atomic_graph(
         self,
         graph: Any,

deepmd/dpmodel/model/model.py

@@ -240,44 +240,14 @@ def get_spin_model(data: dict) -> SpinModel:
     return SpinModel(backbone_model=backbone_model, spin=spin)
 
 
-def learned_descriptor(atomic_model: Any) -> Any:
-    """The descriptor of the LEARNED term of an atomic model.
-
-    A plain atomic model owns its descriptor directly; a composition (e.g.
-    analytical bridging, ``LinearEnergyAtomicModel`` over ``[learned,
-    InterPotential]``) has exactly one descriptor-bearing child, since the
-    analytical terms are descriptor-free.  Shared by every backend's builder
-    so descriptor capability gates read the same object regardless of whether
-    the model happens to be composed.
-
-    Parameters
-    ----------
-    atomic_model
-        The atomic model to inspect.
-
-    Returns
-    -------
-    descriptor or None
-        The learned descriptor, or ``None`` if no child has one.
-    """
-    descriptor = getattr(atomic_model, "descriptor", None)
-    if descriptor is not None:
-        return descriptor
-    for child in getattr(atomic_model, "models", []):
-        descriptor = getattr(child, "descriptor", None)
-        if descriptor is not None:
-            return descriptor
-    return None
-
-
 def get_native_spin_model(data: dict) -> NativeSpinEnergyModel:
     """Get a native (virtual-atom-free) spin model from a dictionary.
 
     Unlike :func:`get_spin_model`, no virtual atoms or doubled type map are
     introduced: ``spin`` is injected into the descriptor config as
     ``use_spin`` and consumed by the descriptor's equivariant spin
-    embedding. Any descriptor declaring ``supports_native_spin()`` is
-    eligible; the gate is the capability method, not a descriptor-type
+    embedding. Any atomic model declaring ``supports_native_spin()`` is
+    eligible; the gate is that capability method, not a descriptor-type
     list.
 
     The non-spin backbone is built by :func:`get_standard_model`, which OWNS
@@ -322,10 +292,12 @@ def get_native_spin_model(data: dict) -> NativeSpinEnergyModel:
             "support (supports_native_spin()); descriptor type "
             f"{data['descriptor'].get('type')!r} does not accept `use_spin`"
         ) from err
-    descriptor = learned_descriptor(backbone_model.atomic_model)
-    if descriptor is None or not descriptor.supports_native_spin():
+    # The ATOMIC MODEL answers the capability -- it knows its own structure,
+    # so this holds for a plain descriptor+fitting model and for a bridging
+    # composition alike, with no assumption here about either.
+    if not backbone_model.atomic_model.supports_native_spin():
         raise NotImplementedError(
-            "spin scheme 'native' requires a descriptor declaring "
+            "spin scheme 'native' requires an atomic model declaring "
             "supports_native_spin()"
         )
     return NativeSpinEnergyModel(atomic_model_=backbone_model.atomic_model, spin=spin)

deepmd/pt_expt/model/get_model.py

@@ -258,9 +258,9 @@ def get_native_spin_model(data: dict) -> NativeSpinEnergyModel:
     for the DPA4/SeZM family (keeping its bridging/lora/compile/
     preset_out_bias rejections and ``exclude_types`` consistency check),
     else :func:`get_standard_model` -- then re-classed through the
-    registered :class:`NativeSpinEnergyModel`. Eligibility is the
-    ``descriptor.supports_native_spin()`` capability, not a descriptor-type
-    list.
+    registered :class:`NativeSpinEnergyModel`. Eligibility is the atomic
+    model's own ``supports_native_spin()`` capability, not a descriptor-type
+    list -- so a bridging composition answers for itself.
 
     Parameters
     ----------
@@ -300,16 +300,12 @@ def get_native_spin_model(data: dict) -> NativeSpinEnergyModel:
             "support (supports_native_spin()); descriptor type "
             f"{data['descriptor'].get('type')!r} does not accept `use_spin`"
         ) from err
-    # A bridged backbone is a LinearEnergyAtomicModel; the capability gate
-    # reads the LEARNED child's descriptor either way (shared helper).
-    from deepmd.dpmodel.model.model import (
-        learned_descriptor,
-    )
-
-    descriptor = learned_descriptor(backbone_model.atomic_model)
-    if descriptor is None or not descriptor.supports_native_spin():
+    # The ATOMIC MODEL answers the capability -- it knows its own structure,
+    # so this holds for a plain descriptor+fitting model and for a bridging
+    # composition alike, with no assumption here about either.
+    if not backbone_model.atomic_model.supports_native_spin():
         raise NotImplementedError(
-            "spin scheme 'native' requires a descriptor declaring "
+            "spin scheme 'native' requires an atomic model declaring "
             "supports_native_spin()"
         )
     return NativeSpinEnergyModel(atomic_model_=backbone_model.atomic_model, spin=spin)

source/tests/common/dpmodel/test_dp_atomic_model.py

@@ -58,7 +58,7 @@ def test_methods(self) -> None:
         # Base-default (False) branch of the descriptor spin capabilities,
         # cached at construction via direct method calls (the True branch is
         # pinned in test_dpa4_call_graph.py).
-        self.assertFalse(md0.supports_native_spin)
+        self.assertFalse(md0.supports_native_spin())
         self.assertFalse(md0.supports_charge_spin)
 
     def test_self_consistency(

source/tests/common/dpmodel/test_dpa4_call_graph.py

@@ -424,7 +424,7 @@ def test_supports_native_spin_capability_gate() -> None:
         seed=5,
     )
     dpa4_model = DPAtomicModel(dd, ft, type_map=["A", "B", "C"])
-    assert dpa4_model.supports_native_spin is True
+    assert dpa4_model.supports_native_spin() is True
     assert dpa4_model.supports_charge_spin is True
 
 

source/tests/common/dpmodel/test_zbl_bridging.py

@@ -312,3 +312,41 @@ def test_serialize_roundtrip_after_change_type_map(self) -> None:
         np.testing.assert_allclose(
             _pair_energy(restored), _pair_energy(model), rtol=1e-12
         )
+
+
+class TestNativeSpinCapabilityOnAtomicModel:
+    """``supports_native_spin`` is answered by the ATOMIC MODEL.
+
+    The model layer must not reach into an atomic model for a descriptor to
+    decide spin eligibility: an analytical term has no descriptor at all, and
+    a composition has several children.  Each atomic model answers from its
+    own structure, exactly like ``uses_graph_lower``.
+    """
+
+    def test_analytical_term_is_not_spin_capable(self) -> None:
+        zbl = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8])
+        # inherits the concrete base default -- no descriptor, no spin input
+        assert zbl.supports_native_spin() is False
+
+    def test_composition_is_capable_when_any_child_is(self) -> None:
+        """ANY, not ALL: analytical children accept and ignore ``spin``."""
+        learned = get_model(
+            {
+                **copy.deepcopy(ZBL_CONFIG),
+                "spin": {"use_spin": [True, False], "scheme": "native"},
+            }
+        ).atomic_model
+        assert learned.supports_native_spin() is True
+        kinds = [type(c).__name__ for c in learned.models]
+        assert kinds[1] == "InterPotentialAtomicModel", kinds
+        # ... and the spin-free analytical child alone is not capable
+        assert learned.models[1].supports_native_spin() is False
+
+    def test_composition_without_a_spin_consumer_is_not_capable(self) -> None:
+        """No consumer => the magnetic force would be identically zero."""
+        zbl_a = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8])
+        zbl_b = InterPotentialAtomicModel(type_map=["Ni", "O"], rcut=4.0, sel=[8])
+        composed = LinearEnergyAtomicModel(
+            [zbl_a, zbl_b], type_map=["Ni", "O"], weights="sum"
+        )
+        assert composed.supports_native_spin() is False