fix: canonical-bridging review round 3

- dpmodel LinearEnergyModel gains a composite update_sel (twin of
  pt_expt): update learned children, skip inner_potential/pairtab,
  aggregate min_nbor_dist; BaseModel.update_sel dispatch no longer
  crashes with KeyError 'descriptor' on normalized linear configs.
- the shared linear builder rejects a bridged learned child whose
  type_map differs from the composition's (the graph route rejects the
  non-identity remap on every forward; fail at construction like pt).
- pt SeZM pair-exclusion reconciliation is factored into ONE helper
  used by all three builders (plain, native spin, virtual spin), so a
  pair_exclude_types vs descriptor.exclude_types mismatch fails fast on
  the spin routes too instead of being silently overwritten.
- canonical top-level learned-model options (data_stat_protect,
  preset_out_bias, use_compile, ...) are routed to the learned child by
  route_canonical_learned_options with conflict checks, in both the pt
  bridged builder and the shared factory; silently dropped before.
- pt_expt get_model rejects trainer-owned lora after bridge expansion
  (pt_expt has no LoRA support; it silently trained a plain model).
- module-level assert pins the routing tables disjoint (also resolves
  the CodeQL unused-variable finding on _LEARNED_CHILD_KEYS).

Author: Han Wang

deepmd/dpmodel/model/dp_linear_model.py

@@ -22,6 +22,9 @@
 from deepmd.dpmodel.model.make_model import (
     make_model,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 
 DPLinearModel_ = make_model(LinearEnergyAtomicModel, T_Bases=(NativeOP, BaseModel))
 
@@ -45,3 +48,48 @@ def __init__(
     ) -> None:
         DPModelCommon.__init__(self)
         DPLinearModel_.__init__(self, *args, **kwargs)
+
+    @classmethod
+    def update_sel(
+        cls,
+        train_data: DeepmdDataSystem,
+        type_map: list[str] | None,
+        local_jdata: dict,
+    ) -> tuple[dict, float | None]:
+        """Update the selection and perform neighbor statistics.
+
+        Updates each learned child in place, skipping analytical
+        (``inner_potential``) and pair-table children, and aggregates the
+        minimum neighbor distance (twin of the pt_expt implementation).
+
+        Parameters
+        ----------
+        train_data : DeepmdDataSystem
+            data used to do neighbor statistics
+        type_map : list[str], optional
+            The name of each type of atoms
+        local_jdata : dict
+            The local data refer to the current class
+
+        Returns
+        -------
+        dict
+            The updated local data
+        float
+            The minimum distance between two atoms
+        """
+        local_jdata_cpy = local_jdata.copy()
+        type_map = local_jdata_cpy["type_map"]
+        min_nbor_dist = None
+        for idx, sub_model in enumerate(local_jdata_cpy["models"]):
+            if sub_model.get("type") == "inner_potential":
+                # analytical child: no descriptor, no selection to update
+                continue
+            if "tab_file" not in sub_model:
+                sub_model, temp_min = DPModelCommon.update_sel(
+                    train_data, type_map, local_jdata_cpy["models"][idx]
+                )
+                local_jdata_cpy["models"][idx] = sub_model
+                if min_nbor_dist is None or temp_min <= min_nbor_dist:
+                    min_nbor_dist = temp_min
+        return local_jdata_cpy, min_nbor_dist

deepmd/dpmodel/model/model_factory.py

@@ -10,6 +10,9 @@
     Any,
 )
 
+from deepmd.utils.bridging import (
+    route_canonical_learned_options,
+)
 from deepmd.utils.spin import (
     Spin,
 )
@@ -233,13 +236,22 @@ def get_linear_atomic_model(
         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))
+        route_canonical_learned_options(data, children[learned_indices[0]])
 
     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)
+        elif inner_indices and i == learned_indices[0] and sub["type_map"] != type_map:
+            # The analytical child always uses the composition's type_map,
+            # and the graph route rejects a non-identity remap at forward
+            # time; fail at construction like the pt builder does.
+            raise ValueError(
+                "A bridged linear_ener composition requires the learned "
+                "child's type_map to match the composition type_map."
+            )
         if "descriptor" in sub:
             child = None
             if descriptor_child_builder is not None:

deepmd/pt/model/model/__init__.py

@@ -38,6 +38,7 @@
 )
 from deepmd.utils.bridging import (
     expand_bridging_method,
+    route_canonical_learned_options,
 )
 from deepmd.utils.spin import (
     Spin,
@@ -346,6 +347,7 @@ def _get_bridged_linear_model(model_params: dict) -> BaseModel:
             "The pt backend requires the learned child's `type_map` to "
             "match the bridged linear_ener composition's `type_map`."
         )
+    route_canonical_learned_options(model_params, learned)
     inner_cfg = inner_cfgs[0]
     learned["type"] = "dpa4"
     learned["type_map"] = copy.deepcopy(model_params["type_map"])
@@ -479,16 +481,31 @@ def get_standard_model(model_params: dict) -> BaseModel:
     return model
 
 
-def get_sezm_model(model_params: dict) -> BaseModel:
-    model_params_old = model_params
-    model_params = copy.deepcopy(model_params)
-    model_params.setdefault("descriptor", {})
-    model_params.setdefault("fitting_net", {})
-    model_params["descriptor"].setdefault("type", "dpa4")
+def _reconcile_sezm_pair_exclude_types(model_params: dict) -> list[list[int]]:
+    """Reconcile ``pair_exclude_types`` with ``descriptor.exclude_types``.
 
-    ntypes = len(model_params["type_map"])
-    model_params["descriptor"]["ntypes"] = ntypes
-    model_params["descriptor"]["type_map"] = copy.deepcopy(model_params["type_map"])
+    A DPA4/SeZM config may spell pair exclusions at the model level
+    (``pair_exclude_types``) or on the descriptor (``exclude_types``).
+    Every SeZM builder (plain, native spin, virtual spin) resolves the two
+    through this ONE helper so that a mismatch always fails fast instead
+    of one spelling silently overwriting the other.
+
+    Parameters
+    ----------
+    model_params : dict
+        The DPA4/SeZM model config; ``model_params["descriptor"]`` must
+        exist.
+
+    Returns
+    -------
+    list[list[int]]
+        The reconciled real-type pair exclusion list.
+
+    Raises
+    ------
+    ValueError
+        If both spellings are given and differ.
+    """
     descriptor_exclude_types = [
         list(pair) for pair in (model_params["descriptor"].get("exclude_types") or [])
     ]
@@ -503,6 +520,20 @@ def get_sezm_model(model_params: dict) -> BaseModel:
             )
     else:
         pair_exclude_types = descriptor_exclude_types
+    return pair_exclude_types
+
+
+def get_sezm_model(model_params: dict) -> BaseModel:
+    model_params_old = model_params
+    model_params = copy.deepcopy(model_params)
+    model_params.setdefault("descriptor", {})
+    model_params.setdefault("fitting_net", {})
+    model_params["descriptor"].setdefault("type", "dpa4")
+
+    ntypes = len(model_params["type_map"])
+    model_params["descriptor"]["ntypes"] = ntypes
+    model_params["descriptor"]["type_map"] = copy.deepcopy(model_params["type_map"])
+    pair_exclude_types = _reconcile_sezm_pair_exclude_types(model_params)
     model_params["pair_exclude_types"] = pair_exclude_types
     model_params["descriptor"]["exclude_types"] = copy.deepcopy(pair_exclude_types)
 
@@ -621,7 +652,7 @@ def _get_sezm_native_spin_model(model_params: dict) -> BaseModel:
     model_params["descriptor"]["type_map"] = copy.deepcopy(model_params["type_map"])
     model_params["descriptor"]["use_spin"] = use_spin
 
-    pair_exclude_types = model_params.get("pair_exclude_types", [])
+    pair_exclude_types = _reconcile_sezm_pair_exclude_types(model_params)
     model_params["pair_exclude_types"] = pair_exclude_types
     if pair_exclude_types:
         model_params["descriptor"]["exclude_types"] = copy.deepcopy(pair_exclude_types)
@@ -694,9 +725,10 @@ def _get_sezm_virtual_spin_model(model_params: dict) -> BaseModel:
         virtual_scale=model_params["spin"]["virtual_scale"],
         allow_missing_label=model_params["spin"].get("allow_missing_label", False),
     )
+    real_pair_exclude_types = _reconcile_sezm_pair_exclude_types(model_params)
     model_params["type_map"] += [item + "_spin" for item in model_params["type_map"]]
     pair_exclude_types = spin.get_pair_exclude_types(
-        exclude_types=model_params.get("pair_exclude_types", None)
+        exclude_types=real_pair_exclude_types or None
     )
     model_params["pair_exclude_types"] = pair_exclude_types
     model_params["descriptor"]["exclude_types"] = pair_exclude_types

deepmd/pt_expt/model/get_model.py

@@ -385,6 +385,14 @@ def get_model(data: dict) -> BaseModel:
         The data to construct the model.
     """
     data = expand_bridging_method(data)
+    if data.get("lora") is not None:
+        # The expansion keeps trainer-owned `lora` at the composition top
+        # level (the pt trainer reads it there); pt_expt has no LoRA
+        # support, so reject it here instead of silently training a plain
+        # full model.
+        raise NotImplementedError(
+            "`lora` is not supported for DPA4/SeZM in the pt_expt backend."
+        )
     return _model_factory.get_model(
         data,
         standard_model_factory=get_standard_model,

deepmd/utils/bridging.py

@@ -117,6 +117,48 @@ def _is_dpa4_family(sub: dict) -> bool:
     "finetune_head",
 )
 _NON_CHILD_KEYS = _COMPOSITION_KEYS + _CONSUMED_KEYS + _TRAINER_KEYS
+# The routing tables are pairwise disjoint: a key has exactly one owner.
+assert not set(_LEARNED_CHILD_KEYS) & set(_NON_CHILD_KEYS)
+
+
+def route_canonical_learned_options(composition: dict, learned: dict) -> None:
+    """Route learned-model options from a canonical composition to its child.
+
+    A canonical ``linear_ener`` config accepts generic model options (e.g.
+    ``data_stat_protect``, ``preset_out_bias``) at the composition top
+    level, but the learned child is their one owner: a bridged builder
+    reads them from the child config only. This helper copies each
+    learned-owned key present at the top level onto ``learned`` (in
+    place) when the child does not set it, and raises when both levels
+    set different values — a silent drop or a silent override would both
+    unpin the ownership contract.
+
+    Parameters
+    ----------
+    composition : dict
+        The canonical ``linear_ener`` model config.
+    learned : dict
+        The learned child's config; modified in place.
+
+    Raises
+    ------
+    ValueError
+        If a learned-owned key is set at both levels with different
+        values.
+    """
+    for key in _LEARNED_CHILD_KEYS:
+        if key not in composition:
+            continue
+        if key in learned:
+            if learned[key] != composition[key]:
+                raise ValueError(
+                    f"`{key}` is set both on the linear_ener composition "
+                    f"({composition[key]!r}) and on its learned child "
+                    f"({learned[key]!r}) with different values. The learned "
+                    "child owns this option: set it on the child only."
+                )
+        else:
+            learned[key] = copy.deepcopy(composition[key])
 
 
 def expand_bridging_method(data: dict) -> dict:

source/tests/common/dpmodel/test_zbl_bridging.py

@@ -784,3 +784,49 @@ def test_inner_child_with_descriptor_raises_cleanly(self) -> None:
         cfg["models"][1]["descriptor"] = {"type": "dpa4"}
         with pytest.raises(ValueError, match="must not carry"):
             get_model(cfg)
+
+    def test_canonical_rejects_mismatched_learned_type_map(self) -> None:
+        """A remapped learned-child type_map builds a model the graph
+        route rejects on every forward; fail at construction instead.
+        """
+        cfg = _canonical_config()
+        cfg["models"][0]["type_map"] = list(reversed(cfg["type_map"]))
+        with pytest.raises(ValueError, match="type_map"):
+            get_model(cfg)
+
+    def test_canonical_conflicting_top_level_option_raises(self) -> None:
+        """A learned-owned option set differently at both levels must
+        fail loudly instead of one value silently winning.
+        """
+        cfg = _canonical_config()
+        cfg["data_stat_protect"] = 0.123
+        cfg["models"][0]["data_stat_protect"] = 0.456
+        with pytest.raises(ValueError, match="data_stat_protect"):
+            get_model(cfg)
+
+    def test_update_sel_dispatches_and_skips_inner_child(self, monkeypatch) -> None:
+        """``BaseModel.update_sel`` dispatches ``linear_ener`` to a
+        composite implementation that updates the learned child and
+        skips the analytical one (the default neighbor-stat phase would
+        otherwise crash with ``KeyError: 'descriptor'``).
+        """
+        from deepmd.dpmodel.model.dp_model import (
+            DPModelCommon,
+        )
+        from deepmd.utils.argcheck import (
+            model_args,
+        )
+
+        seen = []
+
+        def _fake_update_sel(train_data, type_map, sub):
+            seen.append(copy.deepcopy(sub))
+            return sub, 0.9
+
+        monkeypatch.setattr(DPModelCommon, "update_sel", staticmethod(_fake_update_sel))
+        cfg = model_args().normalize_value(_canonical_config(), trim_pattern="_*")
+        updated, min_dist = BaseModel.update_sel(None, cfg["type_map"], cfg)
+        assert min_dist == 0.9
+        assert len(seen) == 1  # only the learned child
+        assert "descriptor" in seen[0]
+        assert updated["models"][1]["type"] == "inner_potential"

source/tests/common/test_bridging.py

@@ -272,3 +272,33 @@ def test_non_dpa4_learned_child_is_not(self) -> None:
         cfg = self._canonical()
         cfg["models"][0] = {"type": "standard", "descriptor": {"type": "se_e2_a"}}
         assert not is_bridged_sezm_config(cfg)
+
+
+def test_route_canonical_learned_options_copies_and_conflicts() -> None:
+    """The learned child owns the generic model options: a top-level value
+    is copied onto the child when absent, accepted when equal, and
+    rejected when the two levels differ.
+    """
+    from deepmd.utils.bridging import (
+        route_canonical_learned_options,
+    )
+
+    composition = {
+        "type": "linear_ener",  # composition-owned: never routed
+        "type_map": ["Ni", "O"],  # composition-owned: never routed
+        "data_stat_protect": 0.123,  # learned-owned: routed
+        "preset_out_bias": {"energy": [1.0, 2.0]},  # learned-owned: routed
+    }
+    learned = {"descriptor": {"type": "dpa4"}}
+    route_canonical_learned_options(composition, learned)
+    assert learned["data_stat_protect"] == 0.123
+    assert learned["preset_out_bias"] == {"energy": [1.0, 2.0]}
+    assert learned["preset_out_bias"] is not composition["preset_out_bias"]
+    assert "type_map" not in learned
+
+    # equal values at both levels pass
+    route_canonical_learned_options(composition, learned)
+
+    learned["data_stat_protect"] = 0.456
+    with pytest.raises(ValueError, match="data_stat_protect"):
+        route_canonical_learned_options(composition, learned)

source/tests/pt/model/test_get_model_bridging.py

@@ -327,3 +327,47 @@ def _fake_update_sel(train_data, type_map, sub):
     assert len(seen) == 1  # only the learned child
     assert "descriptor" in seen[0]
     assert updated["models"][1]["type"] == "inner_potential"
+
+
+def test_canonical_top_level_option_routes_to_learned_child() -> None:
+    """Generic learned-model options at the canonical top level reach the
+    learned child (the child is their one owner).
+    """
+    cfg = _canonical_config()
+    cfg["data_stat_protect"] = 0.123
+    model = get_model(cfg)
+    assert model.atomic_model.data_stat_protect == 0.123
+
+
+def test_canonical_conflicting_top_level_option_raises() -> None:
+    """A learned-owned option set differently at both levels must fail
+    loudly instead of one value silently winning.
+    """
+    cfg = _canonical_config()
+    cfg["data_stat_protect"] = 0.123
+    cfg["models"][0]["data_stat_protect"] = 0.456
+    with pytest.raises(ValueError, match="data_stat_protect"):
+        get_model(cfg)
+
+
+@pytest.mark.parametrize(
+    "scheme",
+    [
+        "native",  # spin as an equivariant descriptor feature
+        "deepspin",  # classical virtual-atom representation
+    ],
+)
+def test_canonical_spin_rejects_mismatched_pair_exclusions(scheme: str) -> None:
+    """Both spin routes must fail fast on a pair-exclusion mismatch like
+    the no-spin route, not silently overwrite the descriptor's exclusions.
+    """
+    cfg = _canonical_config()
+    cfg["pair_exclude_types"] = [[0, 0]]
+    cfg["models"][0]["descriptor"]["exclude_types"] = [[0, 1]]
+    cfg["spin"] = {
+        "scheme": scheme,
+        "use_spin": [True, False],
+        "virtual_scale": 0.3,
+    }
+    with pytest.raises(ValueError, match="must match"):
+        get_model(cfg)