refactor(model): express analytical bridging as an explicit linear_ener composition

deepmd/dpmodel/model/model.py

@@ -31,6 +31,9 @@
 from deepmd.dpmodel.model.spin_model import (
     SpinModel,
 )
+from deepmd.utils.bridging import (
+    expand_bridging_method,
+)
 from deepmd.utils.spin import (
     Spin,
     normalize_spin_use_spin,
@@ -58,50 +61,67 @@ def get_standard_model(data: dict) -> BaseModel:
     data : dict
         The data to construct the model.
     """
-    data = copy.deepcopy(data)
-    # Analytical bridging (e.g. ZBL): the radii feed the DESCRIPTOR's
-    # InnerClamp/BridgingSwitch (mirrors pt's builder); the method builds the
-    # atomic model's InnerPotential below.
     bridging_method = str(data.get("bridging_method", "none"))
-    bridging_enabled = bridging_method.lower() not in ("none", "")
-    if bridging_enabled:
-        data["descriptor"]["inner_clamp_r_inner"] = data.get("bridging_r_inner", 0.5)
-        data["descriptor"]["inner_clamp_r_outer"] = data.get("bridging_r_outer", 0.8)
-    model = _model_factory.get_standard_model(data)
-    if not bridging_enabled:
-        return model
-
-    descriptor = model.atomic_model.descriptor
-    atom_exclude_types = data.get("atom_exclude_types", [])
-    pair_exclude_types = data.get("pair_exclude_types", [])
-    # Composition, not a flag (first-principles design): the analytical
-    # bridging term is its own atomic model, summed with the learned one by the
-    # existing linear composition machinery.
-    from deepmd.dpmodel.atomic_model.inner_potential import (
-        InnerPotentialAtomicModel,
-    )
-    from deepmd.dpmodel.atomic_model.linear_atomic_model import (
-        LinearEnergyAtomicModel,
-    )
+    if bridging_method.lower() not in ("none", ""):
+        raise ValueError(
+            "`bridging_method` is not supported for a standard model: "
+            "analytical bridging builds a linear composition, not a "
+            "standard model. Route the config through `get_model` (which "
+            "expands the flag), or spell the composition explicitly with "
+            '`type: "linear_ener"` and an `inner_potential` sub-model.'
+        )
+    return _model_factory.get_standard_model(data)
+
+
+def get_linear_model(data: dict) -> BaseModel:
+    """Build a linear energy model from a ``linear_ener`` config.
+
+    Children with a ``descriptor`` build as standard learned atomic
+    models; ``pairtab`` children build as pair-tabulation atomic models;
+    an ``inner_potential`` child builds the analytical bridging term. The
+    composition is the ONE owner of the bridging coupling: it derives the
+    learned sibling descriptor's ``inner_clamp_r_inner``/``_outer`` from
+    the ``inner_potential`` child's ``r_inner``/``r_outer``, so the radii
+    are written once in the config (issue #5948, task 2).
+
+    A top-level ``spin`` section (scheme ``native``) wraps the composed
+    atomic model as a :class:`NativeSpinEnergyModel`, with ``use_spin``
+    injected into every learned child's descriptor.
+
+    Parameters
+    ----------
+    data : dict
+        The model configuration.
+    """
     from deepmd.dpmodel.model.dp_linear_model import (
         LinearEnergyModel,
     )
 
-    zbl_atomic = InnerPotentialAtomicModel(
-        type_map=data["type_map"],
-        mode=bridging_method,
-        rcut=descriptor.get_rcut(),
-        sel=descriptor.get_sel(),
-    )
-    composed = LinearEnergyAtomicModel(
-        models=[model.atomic_model, zbl_atomic],
-        type_map=data["type_map"],
-        weights="sum",
-        # Both exclusions belong to the composition: its children share one
-        # graph, so "excluded" must cover the analytical term too.
-        atom_exclude_types=atom_exclude_types,
-        pair_exclude_types=pair_exclude_types,
-    )
+    data = copy.deepcopy(data)
+    spin = None
+    if "spin" in data:
+        spin_cfg = data.pop("spin")
+        if str(spin_cfg.get("scheme", "deepspin")) != "native":
+            raise NotImplementedError(
+                "Spin linear_ener models support only spin scheme 'native'."
+            )
+        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),
+        )
+        for sub in data["models"]:
+            if "descriptor" in sub:
+                sub["descriptor"]["use_spin"] = use_spin
+    composed = _model_factory.get_linear_atomic_model(data)
+    if spin is not None:
+        if not composed.supports_native_spin():
+            raise NotImplementedError(
+                "spin scheme 'native' requires an atomic model declaring "
+                "supports_native_spin()"
+            )
+        return NativeSpinEnergyModel(atomic_model_=composed, spin=spin)
     return LinearEnergyModel(atomic_model_=composed)
 
 
@@ -135,14 +155,10 @@ def get_native_spin_model(data: dict) -> NativeSpinEnergyModel:
     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
-    everything about assembling the atomic model -- descriptor/fitting,
-    exclusions and the analytical-bridging composition -- so ``spin`` and
-    ``bridging_method`` combine for free: the wrapper re-classes whatever
-    atomic model came back, be it a single learned model or a
-    ``LinearEnergyAtomicModel`` over ``[learned, InnerPotential]`` (the
-    analytical child accepts and ignores ``spin``; the learned child consumes
-    it).
+    The non-spin backbone is built by :func:`get_standard_model`. A spin
+    model with analytical bridging is a ``linear_ener`` composition and
+    routes through :func:`get_linear_model` instead (the ``bridging_method``
+    sugar expands to that form in :func:`get_model`).
 
     Parameters
     ----------
@@ -196,9 +212,13 @@ def get_model(data: dict) -> BaseModel:
     data : dict
         The data to construct the model.
     """
+    data = expand_bridging_method(data)
     return _model_factory.get_model(
         data,
         standard_model_factory=get_standard_model,
         spin_model_factory=get_spin_model,
         native_spin_model_factory=get_native_spin_model,
+        model_factories={
+            "linear_ener": get_linear_model,
+        },
     )

deepmd/dpmodel/model/model_factory.py

@@ -124,6 +124,167 @@ def get_zbl_model(
     )
 
 
+def get_linear_atomic_model(
+    data: dict,
+    *,
+    descriptor_base: type,
+    fitting_base: type,
+    backend_name: str,
+    atomic_model: type,
+    pairtab_model: type,
+    descriptor_child_builder: "Callable[[dict], Any | None] | None" = None,
+) -> Any:
+    """Build the ``LinearEnergyAtomicModel`` composition from a config.
+
+    Children with a ``descriptor`` build as learned atomic models through
+    the backend registries; ``pairtab`` children build as pair-tabulation
+    atomic models; an ``inner_potential`` child builds the analytical
+    bridging term. The composition is the ONE owner of the bridging
+    coupling: it derives the learned sibling descriptor's
+    ``inner_clamp_r_inner``/``_outer`` from the ``inner_potential``
+    child's ``r_inner``/``r_outer``, so the radii are written once in the
+    config (issue #5948).
+
+    Parameters
+    ----------
+    data : dict
+        The ``linear_ener`` model configuration.
+    descriptor_base : type
+        Backend descriptor registry base class.
+    fitting_base : type
+        Backend fitting registry base class.
+    backend_name : str
+        Backend name used in error messages.
+    atomic_model : type
+        Backend learned atomic-model class.
+    pairtab_model : type
+        Backend pair-tabulation atomic-model class.
+    descriptor_child_builder : callable, optional
+        Backend hook for descriptor-bearing children: called with the
+        child config (``type_map`` and derived clamp radii already
+        injected) and returns the child atomic model, or ``None`` to fall
+        back to the generic registry build. Backends use it to route
+        family-specific model types (e.g. DPA4/SeZM) through their
+        validated builders.
+
+    Raises
+    ------
+    ValueError
+        If more than one ``inner_potential`` child is given, if an
+        ``inner_potential`` child has no unique learned sibling, if a
+        bridged composition does not use ``weights: "sum"``, if a child
+        carries a ``bridging_method`` flag, or if a child is of an
+        unsupported kind.
+    """
+    from deepmd.dpmodel.atomic_model.inner_potential import (
+        InnerPotentialAtomicModel,
+    )
+    from deepmd.dpmodel.atomic_model.linear_atomic_model import (
+        LinearEnergyAtomicModel,
+    )
+
+    data = copy.deepcopy(data)
+    type_map = data["type_map"]
+    children = data["models"]
+    inner_indices = [
+        i for i, sub in enumerate(children) if sub.get("type") == "inner_potential"
+    ]
+    learned_indices = [
+        i
+        for i, sub in enumerate(children)
+        if "descriptor" in sub and i not in inner_indices
+    ]
+    for i in inner_indices:
+        if "descriptor" in children[i]:
+            raise ValueError(
+                "An `inner_potential` sub-model must not carry a "
+                "`descriptor`: the analytical term has no learned "
+                "component."
+            )
+    for sub in children:
+        if str(sub.get("bridging_method", "none")).lower() not in ("none", ""):
+            raise ValueError(
+                "`bridging_method` is not supported on a linear_ener "
+                "sub-model: add an `inner_potential` sub-model to the "
+                "composition instead."
+            )
+    if inner_indices:
+        if len(inner_indices) > 1:
+            raise ValueError(
+                "A linear_ener composition supports at most one "
+                "`inner_potential` sub-model."
+            )
+        if len(learned_indices) != 1 or len(children) != 2:
+            # A third child (e.g. pairtab) has no common execution route
+            # with the graph-only bridged pair; reject at construction
+            # like the pt builder does.
+            raise ValueError(
+                "An `inner_potential` sub-model bridges exactly one learned "
+                "sibling: expected a linear_ener composition over "
+                "[learned, inner_potential]."
+            )
+        if str(data.get("weights", "mean")) != "sum":
+            raise ValueError(
+                'A bridged linear_ener composition requires `weights: "sum"`.'
+            )
+        # The composition derives the sibling descriptor's clamp window from
+        # the inner_potential child: one source of truth for the radii.
+        inner_cfg = children[inner_indices[0]]
+        learned_descriptor = children[learned_indices[0]]["descriptor"]
+        learned_descriptor["inner_clamp_r_inner"] = float(inner_cfg.get("r_inner", 0.5))
+        learned_descriptor["inner_clamp_r_outer"] = float(inner_cfg.get("r_outer", 0.8))
+
+    built: dict[int, Any] = {}
+    for i, sub in enumerate(children):
+        if i in inner_indices:
+            continue
+        if "type_map" not in sub:
+            sub["type_map"] = copy.deepcopy(type_map)
+        if "descriptor" in sub:
+            child = None
+            if descriptor_child_builder is not None:
+                child = descriptor_child_builder(sub)
+            if child is None:
+                descriptor, fitting, _ = get_model_components(
+                    sub,
+                    descriptor_base=descriptor_base,
+                    fitting_base=fitting_base,
+                    backend_name=backend_name,
+                )
+                child = atomic_model(descriptor, fitting, type_map=sub["type_map"])
+            built[i] = child
+        else:
+            if sub.get("type") != "pairtab":
+                raise ValueError(
+                    "Sub-models in LinearEnergyModel must be a standard model, "
+                    "a pairtab model, or an inner_potential model, but got "
+                    f"type {sub.get('type')!r}."
+                )
+            built[i] = pairtab_model(
+                sub["tab_file"],
+                sub["rcut"],
+                sub["sel"],
+                type_map=copy.deepcopy(type_map),
+            )
+    for i in inner_indices:
+        learned_descriptor_obj = built[learned_indices[0]].descriptor
+        built[i] = InnerPotentialAtomicModel(
+            type_map=copy.deepcopy(type_map),
+            mode=children[i].get("mode", "zbl"),
+            rcut=learned_descriptor_obj.get_rcut(),
+            sel=learned_descriptor_obj.get_sel(),
+        )
+    return LinearEnergyAtomicModel(
+        models=[built[i] for i in range(len(children))],
+        type_map=type_map,
+        weights=data.get("weights", "mean"),
+        # Both exclusions belong to the composition: its children share one
+        # graph, so "excluded" must cover the analytical term too.
+        atom_exclude_types=data.get("atom_exclude_types", []),
+        pair_exclude_types=data.get("pair_exclude_types", []),
+    )
+
+
 def get_spin_model(
     data: dict,
     *,
@@ -257,6 +418,25 @@ def get_standard_model(self, data: dict) -> Any:
             backend_name=self.backend_name,
         )
 
+    def get_linear_atomic_model(
+        self,
+        data: dict,
+        *,
+        descriptor_child_builder: "Callable[[dict], Any | None] | None" = None,
+    ) -> Any:
+        """Construct the linear atomic-model composition for this backend."""
+        if self.atomic_model is None or self.pairtab_model is None:
+            raise NotImplementedError("Linear model is not implemented yet.")
+        return get_linear_atomic_model(
+            data,
+            descriptor_base=self.descriptor_base,
+            fitting_base=self.fitting_base,
+            backend_name=self.backend_name,
+            atomic_model=self.atomic_model,
+            pairtab_model=self.pairtab_model,
+            descriptor_child_builder=descriptor_child_builder,
+        )
+
     def get_zbl_model(self, data: dict) -> Any:
         """Construct a ZBL model for this backend."""
         if (

deepmd/pt/entrypoints/freeze_pt2.py

@@ -68,6 +68,9 @@
 from deepmd.pt_expt.utils.edge_schema import (
     edge_schema_from_extended,
 )
+from deepmd.utils.bridging import (
+    is_bridged_sezm_config,
+)
 from deepmd.utils.model_branch_dict import (
     get_model_dict,
 )
@@ -159,12 +162,21 @@ def is_sezm_checkpoint(ckpt_path: str) -> bool:
         _, params = _extract_state_and_params(raw)
     except ValueError:
         return False
+
+    def _is_sezm_params(branch_params: dict[str, Any]) -> bool:
+        # the flag spelling and the canonical bridged linear spelling both
+        # realize a SeZMModel in the pt backend
+        return str(branch_params.get("type", "")).lower() in (
+            "sezm",
+            "dpa4",
+        ) or is_bridged_sezm_config(branch_params)
+
     if "model_dict" in params:
         return any(
-            str(branch_params.get("type", "")).lower() in ("sezm", "dpa4")
+            _is_sezm_params(branch_params)
             for branch_params in params["model_dict"].values()
         )
-    return str(params.get("type", "")).lower() in ("sezm", "dpa4")
+    return _is_sezm_params(params)
 
 
 def _select_model_head(
@@ -870,7 +882,7 @@ def freeze_sezm_to_pt2(
     state_dict, params = _select_model_head(state_dict, params, head)
 
     model_type = str(params.get("type", "")).lower()
-    if model_type not in ("sezm", "dpa4"):
+    if model_type not in ("sezm", "dpa4") and not is_bridged_sezm_config(params):
         raise ValueError(
             f"freeze_sezm_to_pt2 expects a SeZM/DPA4 checkpoint, got type={params.get('type')!r}."
         )