From fcf717806d115615c6f178177b0e96c501917418 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 11 Aug 2026 01:06:37 +0800 Subject: [PATCH 01/11] refactor(model): express analytical bridging as an explicit linear_ener composition Close #5948. A bridged model IS a linear composition, but it was spelled as a bridging_method flag on a non-composite model type, so the type requested was not the type returned, and every builder accepting the flag re-implemented the composition (and drifted, see #5947). - Register inner_potential as a config-level model type (argcheck), so it can be named as a linear_ener child: {type: inner_potential, mode: zbl, r_inner, r_outer}. - The linear builder derives the learned sibling descriptor's InnerClamp/BridgingSwitch radii from the inner_potential child at build time: the radii are written once, one source of truth. - Keep bridging_method as sugar expanded by ONE shared normalizer (deepmd.utils.bridging.expand_bridging_method) at every backend's get_model entry; the non-composite builders (get_standard_model, get_sezm_model) now fail fast on the flag instead of composing or silently dropping it (pt's standard route used to drop it). - dpmodel gains a real linear_ener config builder (it had none); the linear child-parsing core is shared between dpmodel and pt_expt. - pt realizes the canonical form via its existing SeZMModel implementation, so pt checkpoints and physics are unchanged; both spellings serialize to the identical wire dict (pinned by tests in all three backends). - Migrate examples/water/dpa4/input-zbl.json and doc/model/dpa4.md to the canonical spelling. --- deepmd/dpmodel/model/model.py | 212 ++++++++++++---- deepmd/pt/model/model/__init__.py | 86 ++++++- deepmd/pt_expt/model/get_model.py | 195 ++++++--------- deepmd/utils/argcheck.py | 40 ++- deepmd/utils/bridging.py | 129 ++++++++++ doc/model/dpa4.md | 36 ++- examples/water/dpa4/input-zbl.json | 227 +++++++++--------- .../tests/common/dpmodel/test_zbl_bridging.py | 104 ++++++++ source/tests/common/test_bridging.py | 162 +++++++++++++ .../tests/pt/model/test_get_model_bridging.py | 176 ++++++++++++++ .../pt_expt/model/test_get_model_bridging.py | 115 +++++++-- 11 files changed, 1182 insertions(+), 300 deletions(-) create mode 100644 deepmd/utils/bridging.py create mode 100644 source/tests/common/test_bridging.py create mode 100644 source/tests/pt/model/test_get_model_bridging.py diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index a66f40e7f3..418e910b00 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -1,5 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import copy +from typing import ( + Any, +) from deepmd.dpmodel.atomic_model.dp_atomic_model import ( DPAtomicModel, @@ -31,6 +34,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,51 +64,173 @@ 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. + 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, + ) + + 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 = _build_linear_atomic_model( + data, + model_components_factory=_model_factory.get_model_components, + dp_atomic_model=DPAtomicModel, + pairtab_atomic_model=PairTabAtomicModel, + ) + 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) + + +def _build_linear_atomic_model( + data: dict, + *, + model_components_factory: Any, + dp_atomic_model: type, + pairtab_atomic_model: type, +) -> Any: + """Build the ``LinearEnergyAtomicModel`` composition from a config. + + Shared between the dpmodel and pt_expt linear builders: the caller + supplies its backend's component factory and atomic-model classes. + ``data`` is mutated (callers pass a private deep copy). + + Parameters + ---------- + data : dict + The ``linear_ener`` model configuration. + model_components_factory : callable + Backend factory building (descriptor, fitting, type_map) from a + standard sub-model config. + dp_atomic_model : type + Backend learned atomic-model class. + pairtab_atomic_model : type + Backend pair-tabulation atomic-model class. + """ from deepmd.dpmodel.atomic_model.inner_potential import ( InnerPotentialAtomicModel, ) from deepmd.dpmodel.atomic_model.linear_atomic_model import ( LinearEnergyAtomicModel, ) - 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", + type_map = data["type_map"] + ntypes = len(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] + 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: + raise ValueError( + "An `inner_potential` sub-model bridges exactly one learned " + f"sibling, but got {len(learned_indices)} sub-models with a " + "descriptor." + ) + # 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: + sub["descriptor"]["ntypes"] = ntypes + descriptor, fitting, _ = model_components_factory(sub) + built[i] = dp_atomic_model(descriptor, fitting, type_map=sub["type_map"]) + 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_atomic_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=atom_exclude_types, - pair_exclude_types=pair_exclude_types, + atom_exclude_types=data.get("atom_exclude_types", []), + pair_exclude_types=data.get("pair_exclude_types", []), ) - return LinearEnergyModel(atomic_model_=composed) def get_spin_model(data: dict) -> SpinModel: @@ -135,14 +263,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 +320,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, + }, ) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 8671a1e94e..4095ccb4ee 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -36,6 +36,9 @@ from deepmd.pt.utils.multi_task import ( preprocess_shared_params, ) +from deepmd.utils.bridging import ( + expand_bridging_method, +) from deepmd.utils.spin import ( Spin, ) @@ -193,7 +196,11 @@ def get_spin_model(model_params: dict) -> SpinModel: return SpinEnergyModel(backbone_model=backbone_model, spin=spin) -def get_linear_model(model_params: dict) -> LinearEnergyModel: +def get_linear_model(model_params: dict) -> BaseModel: + if any( + sub.get("type") == "inner_potential" for sub in model_params.get("models", []) + ): + return _get_bridged_linear_model(model_params) model_params = copy.deepcopy(model_params) weights = model_params.get("weights", "mean") shared_links = None @@ -265,6 +272,73 @@ def get_linear_model(model_params: dict) -> LinearEnergyModel: return model +def _get_bridged_linear_model(model_params: dict) -> BaseModel: + """Realize a bridged ``linear_ener`` composition in the pt backend. + + The pt backend implements analytical bridging inside ``SeZMModel`` + (the ``InnerPotential`` term is a sub-module of the learned model), + so a canonical ``linear_ener`` composition over + ``[learned, inner_potential]`` maps onto the ``SeZMModel`` + constructor arguments. The physics and the checkpoint format are + identical to the legacy ``bridging_method`` flag form. + + Parameters + ---------- + model_params : dict + A ``linear_ener`` config with exactly one ``inner_potential`` + sub-model and exactly one learned (descriptor-bearing) sub-model. + + Raises + ------ + ValueError + If the composition shape is not the bridging one (child counts, + ``weights``, ``shared_dict``). + NotImplementedError + If the learned sibling is not of the DPA4/SeZM family: the pt + backend has no bridging implementation for other descriptors. + """ + model_params = copy.deepcopy(model_params) + children = model_params.get("models", []) + inner_cfgs = [sub for sub in children if sub.get("type") == "inner_potential"] + learned_cfgs = [sub for sub in children if "descriptor" in sub] + if len(inner_cfgs) > 1: + raise ValueError( + "A linear_ener composition supports at most one " + "`inner_potential` sub-model." + ) + if len(learned_cfgs) != 1 or len(children) != 2: + raise ValueError( + "An `inner_potential` sub-model bridges exactly one learned " + "sibling: expected a linear_ener composition over " + "[learned, inner_potential]." + ) + if str(model_params.get("weights", "mean")) != "sum": + raise ValueError('A bridged linear_ener composition requires `weights: "sum"`.') + if model_params.get("shared_dict"): + raise NotImplementedError( + "`shared_dict` is not supported with an `inner_potential` sub-model." + ) + learned = copy.deepcopy(learned_cfgs[0]) + descriptor_type = str(learned.get("descriptor", {}).get("type", "dpa4")) + if descriptor_type not in ("dpa4", "DPA4", "sezm", "SeZM"): + raise NotImplementedError( + "The pt backend implements `inner_potential` bridging only for " + f"the DPA4/SeZM descriptor family, but got {descriptor_type!r}." + ) + inner_cfg = inner_cfgs[0] + learned["type"] = "dpa4" + learned["type_map"] = copy.deepcopy(model_params["type_map"]) + learned["atom_exclude_types"] = model_params.get("atom_exclude_types", []) + learned["pair_exclude_types"] = model_params.get("pair_exclude_types", []) + learned["bridging_method"] = inner_cfg.get("mode", "zbl") + learned["bridging_r_inner"] = float(inner_cfg.get("r_inner", 0.5)) + learned["bridging_r_outer"] = float(inner_cfg.get("r_outer", 0.8)) + if "spin" in model_params: + learned["spin"] = model_params["spin"] + return get_sezm_spin_model(learned) + return get_sezm_model(learned) + + def get_zbl_model(model_params: dict) -> DPZBLModel: model_params = copy.deepcopy(model_params) ntypes = len(model_params["type_map"]) @@ -331,6 +405,15 @@ def _convert_preset_out_bias_to_array( def get_standard_model(model_params: dict) -> BaseModel: + bridging_method = str(model_params.get("bridging_method", "none")) + 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.' + ) model_params_old = model_params model_params = copy.deepcopy(model_params) ntypes = len(model_params["type_map"]) @@ -657,6 +740,7 @@ def _get_sezm_virtual_spin_model(model_params: dict) -> BaseModel: def get_model(model_params: dict) -> Any: + model_params = expand_bridging_method(model_params) model_type = model_params.get("type", "standard") if model_type == "standard": if "spin" in model_params: diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index 50f60ecf49..1928036ded 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -8,9 +8,6 @@ import copy import logging -from typing import ( - TYPE_CHECKING, -) from deepmd.dpmodel.atomic_model.dp_atomic_model import ( DPAtomicModel, @@ -45,16 +42,14 @@ from deepmd.pt_expt.model.spin_ener_model import ( SpinEnergyModel, ) +from deepmd.utils.bridging import ( + expand_bridging_method, +) from deepmd.utils.spin import ( Spin, normalize_spin_use_spin, ) -if TYPE_CHECKING: - from deepmd.pt_expt.model.dp_linear_model import ( - LinearEnergyModel, - ) - log = logging.getLogger(__name__) # Warn at most once per process for backend-ignored switches (keyed by name). @@ -80,9 +75,11 @@ def get_sezm_model(data: dict) -> BaseModel: training configs are interchangeable between the pt and pt_expt backends. In addition to the ``SeZM``/``sezm``/``dpa4`` aliases accepted by pt, pt_expt also accepts ``DPA4``. - Supported SeZM extensions: analytical bridging (e.g. ZBL), composed by - :func:`_compose_bridging`, and native-scheme spin, routed to - :func:`get_native_spin_model`; the two combine. + Supported SeZM extension: native-scheme spin, routed to + :func:`get_native_spin_model`. 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`); + this builder rejects the flag. Still unsupported here, each raising ``NotImplementedError``: the virtual-atom (``deepspin``) spin scheme, ``lora``, ``use_compile``, and @@ -109,11 +106,15 @@ def get_sezm_model(data: dict) -> BaseModel: "scheme 'native' instead." ) return get_native_spin_model(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 at construction below. bridging_method = str(data.get("bridging_method", "none")) - bridging_enabled = bridging_method.lower() not in ("none", "") + if bridging_method.lower() not in ("none", ""): + raise ValueError( + "`bridging_method` is not supported by the DPA4/SeZM builder: " + "analytical bridging builds a linear composition. 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." + ) if data.get("lora") is not None: raise NotImplementedError( "`lora` is not supported for DPA4/SeZM in the pt_expt backend." @@ -129,9 +130,6 @@ def get_sezm_model(data: dict) -> BaseModel: data.pop("type", None) data.setdefault("descriptor", {}) data.setdefault("fitting_net", {}) - 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) data["descriptor"].setdefault("type", "dpa4") data["fitting_net"].setdefault("type", "dpa4_ener") # the DPA4/SeZM model type is a fixed descriptor/fitting contract; reject @@ -164,74 +162,13 @@ def get_sezm_model(data: dict) -> BaseModel: data["descriptor"]["exclude_types"] = copy.deepcopy(pair_exclude_types) descriptor, fitting, _ = _model_factory.get_model_components(data) - model = DPA4EnergyModel( + return DPA4EnergyModel( descriptor=descriptor, fitting=fitting, type_map=data["type_map"], atom_exclude_types=data.get("atom_exclude_types", []), pair_exclude_types=pair_exclude_types, ) - if bridging_enabled: - return _compose_bridging(model, data, bridging_method) - return model - - -def _compose_bridging( - model: BaseModel, data: dict, bridging_method: str -) -> "LinearEnergyModel": - """Compose the learned model with its analytical bridging term. - - 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. The ONE owner of the - composition build for this backend: :func:`get_sezm_model` - (``type: "dpa4"``/``"sezm"``) is its only caller, because bridging - yields a composition and so is not expressible on a non-composite - model type -- :func:`get_standard_model` rejects it. Issue #5948 - tracks spelling the composition explicitly as ``linear_ener``. - - Parameters - ---------- - model - The learned backbone model (its descriptor already carries the - bridging radii injected by the caller). - data - The model config (``type_map`` and the exclusion lists are read). - bridging_method - The analytical bridging mode (e.g. ``"ZBL"``). - - Returns - ------- - LinearEnergyModel - A composition over ``[learned, InnerPotential]``. - """ - from deepmd.dpmodel.atomic_model.inner_potential import ( - InnerPotentialAtomicModel, - ) - from deepmd.dpmodel.atomic_model.linear_atomic_model import ( - LinearEnergyAtomicModel, - ) - from deepmd.pt_expt.model.dp_linear_model import ( - LinearEnergyModel, - ) - - descriptor = model.atomic_model.descriptor - 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=data.get("atom_exclude_types", []), - pair_exclude_types=data.get("pair_exclude_types", []), - ) - return LinearEnergyModel(atomic_model_=composed) def get_standard_model(data: dict) -> BaseModel: @@ -240,15 +177,15 @@ def get_standard_model(data: dict) -> BaseModel: ``bridging_method`` is rejected here rather than honored. Analytical bridging is a COMPOSITION -- it yields a ``LinearEnergyModel`` over ``[learned, InnerPotential]`` -- so a builder that accepted it would - return a model of a different kind than the one requested. pt_expt - keeps exactly one bridging owner, :func:`get_sezm_model` - (``type: "dpa4"``/``"sezm"``), so the composition and its - ``exclude_types`` reconciliation cannot drift between two builders. + return a model of a different kind than the one requested. The ONE + owner of the flag is the shared + :func:`deepmd.utils.bridging.expand_bridging_method` normalizer, + applied in :func:`get_model`; the composition itself is built by + :func:`get_linear_model`. Rejecting is deliberate over silently ignoring: dropping a bridging term without a word yields a physically different model than the config - asks for. Issue #5948 tracks replacing the flag with an explicit - ``linear_ener`` composition, at which point this restriction is moot. + asks for. Parameters ---------- @@ -271,8 +208,10 @@ def get_standard_model(data: dict) -> BaseModel: raise ValueError( "`bridging_method` is not supported for a standard model in the " "pt_expt backend: analytical bridging builds a linear " - 'composition, not a standard model. Use model `type: "dpa4"` ' - '(or `"sezm"`) with the same descriptor and fitting net.' + "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) @@ -342,54 +281,63 @@ def get_native_spin_model(data: dict) -> NativeSpinEnergyModel: def get_linear_model(model_params: dict) -> BaseModel: - """Get a linear energy model from a config dictionary. + """Get a linear energy model from a ``linear_ener`` config dictionary. + + Children with a ``descriptor`` build as learned atomic models; + ``pairtab`` children build as pair-tabulation atomic models; an + ``inner_potential`` child builds the analytical bridging term, with + the learned sibling descriptor's ``inner_clamp_r_inner``/``_outer`` + derived from the child's ``r_inner``/``r_outer`` (issue #5948). A + top-level ``spin`` section (scheme ``native``) wraps the composition + as a :class:`NativeSpinEnergyModel`. Parameters ---------- model_params : dict The model parameters. """ + from deepmd.dpmodel.model.model import ( + _build_linear_atomic_model, + ) + from .dp_linear_model import ( LinearEnergyModel, ) model_params = copy.deepcopy(model_params) - weights = model_params.get("weights", "mean") - list_of_models = [] - ntypes = len(model_params["type_map"]) - for sub_model_params in model_params["models"]: - if "type_map" not in sub_model_params: - sub_model_params["type_map"] = model_params["type_map"] - if "descriptor" in sub_model_params: - sub_model_params["descriptor"]["ntypes"] = ntypes - descriptor, fitting, _ = _model_factory.get_model_components( - sub_model_params - ) - list_of_models.append( - DPAtomicModel(descriptor, fitting, type_map=model_params["type_map"]) - ) - else: - assert ( - "type" in sub_model_params and sub_model_params["type"] == "pairtab" - ), "Sub-models in LinearEnergyModel must be a DPModel or a PairTable Model" - list_of_models.append( - PairTabAtomicModel( - sub_model_params["tab_file"], - sub_model_params["rcut"], - sub_model_params["sel"], - type_map=model_params["type_map"], - ) + spin = None + if "spin" in model_params: + spin_cfg = model_params.pop("spin") + if str(spin_cfg.get("scheme", "deepspin")) != "native": + raise NotImplementedError( + "Spin linear_ener models support only spin scheme 'native' " + "in the pt_expt backend." ) - - atom_exclude_types = model_params.get("atom_exclude_types", []) - pair_exclude_types = model_params.get("pair_exclude_types", []) - return LinearEnergyModel( - models=list_of_models, - type_map=model_params["type_map"], - weights=weights, - atom_exclude_types=atom_exclude_types, - pair_exclude_types=pair_exclude_types, + use_spin = normalize_spin_use_spin( + spin_cfg["use_spin"], model_params["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 model_params["models"]: + if "descriptor" in sub: + sub["descriptor"]["use_spin"] = use_spin + composed = _build_linear_atomic_model( + model_params, + model_components_factory=_model_factory.get_model_components, + dp_atomic_model=DPAtomicModel, + pairtab_atomic_model=PairTabAtomicModel, ) + 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) def get_spin_model(data: dict) -> SpinEnergyModel: @@ -414,6 +362,7 @@ 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, diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index eae723bc3d..bce9f57080 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -3488,7 +3488,10 @@ def sezm_model_args() -> Argument: ) doc_bridging_method = ( "Short-range bridging method. Currently supports 'ZBL'. " - "The value is case-insensitive; set it to 'None' to disable bridging." + "The value is case-insensitive; set it to 'None' to disable bridging. " + "This flag is sugar: it expands to the canonical `linear_ener` " + "composition over the learned model and an `inner_potential` " + "sub-model." ) doc_bridging_r_inner = ( "Inner clamping radius in Å. ML descriptor distances below this radius are frozen. " @@ -3704,6 +3707,41 @@ def pairtab_model_args() -> Argument: return ca +@model_args_plugin.register("inner_potential") +def inner_potential_model_args() -> Argument: + doc_mode = ( + "The analytical pair-potential formula. Currently supports 'zbl' " + "(case-insensitive)." + ) + doc_r_inner = ( + "Inner clamping radius in Å, applied to the learned sibling's " + "descriptor: ML descriptor distances below this radius are frozen. " + "For ZBL bridging, set `training.training_data.min_pair_dist` to the " + "same value so frames with atom pairs closer than `r_inner` are " + "skipped during training." + ) + doc_r_outer = ( + "Outer clamping radius in Å, applied to the learned sibling's " + "descriptor. The transition zone `[r_inner, r_outer]` uses a " + "C^3-continuous septic Hermite polynomial." + ) + ca = Argument( + "inner_potential", + dict, + [ + Argument("mode", str, optional=True, default="zbl", doc=doc_mode), + Argument("r_inner", float, optional=True, default=0.5, doc=doc_r_inner), + Argument("r_outer", float, optional=True, default=0.8, doc=doc_r_outer), + ], + doc=supported_backends("pt", "pt_expt") + + "Analytical short-range bridging pair potential (e.g. ZBL), usable " + "only as a sub-model of a `linear_ener` composition; the clamping " + "radii are derived onto the learned sibling's descriptor at build " + "time.", + ) + return ca + + @hybrid_model_args_plugin.register("linear_ener") def linear_ener_model_args() -> Argument: doc_weights = ( diff --git a/deepmd/utils/bridging.py b/deepmd/utils/bridging.py new file mode 100644 index 0000000000..6959cddcaf --- /dev/null +++ b/deepmd/utils/bridging.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Expansion of the ``bridging_method`` sugar into its canonical config form. + +A bridged model IS a linear composition: the learned model plus an +analytical inner potential, summed by ``linear_ener``. The canonical +config spelling is therefore:: + + "model": { + "type": "linear_ener", "weights": "sum", "type_map": [...], + "models": [ + {"type": "dpa4", "descriptor": {...}, "fitting_net": {...}}, + {"type": "inner_potential", "mode": "zbl", + "r_inner": 0.5, "r_outer": 0.8} + ] + } + +The legacy spelling -- a ``bridging_method`` flag on a non-composite model +type -- is kept as sugar. This module is the ONE owner of that expansion +(issue #5948): every backend's ``get_model`` entry point calls +:func:`expand_bridging_method` before dispatch, so no per-builder flag +handling can drift. +""" + +import copy + +__all__ = [ + "expand_bridging_method", +] + +# Top-level keys that belong to the composition, not to the learned child. +_COMPOSITION_KEYS = ( + "type", + "type_map", + "spin", + "atom_exclude_types", + "pair_exclude_types", + "bridging_method", + "bridging_r_inner", + "bridging_r_outer", +) + + +def expand_bridging_method(data: dict) -> dict: + """Expand the ``bridging_method`` sugar into a ``linear_ener`` config. + + A config without an active ``bridging_method`` is returned unchanged + (the same object, not a copy). A config with an active method is + deep-copied and rewritten to the canonical composition form: a + ``linear_ener`` model with ``weights: "sum"`` over the learned + sub-model and an ``inner_potential`` sub-model. The exclusion lists + move to the composition level; a top-level ``spin`` section stays at + the top level; every other key stays on the learned child. + + For backward compatibility with the legacy pt ``type: "dpa4"`` + builder, ``descriptor.exclude_types`` is promoted to the composition's + ``pair_exclude_types`` (and the two must match when both are given). + Hand-written canonical configs get no such promotion. + + Parameters + ---------- + data : dict + The model section of a training config. + + Returns + ------- + dict + The canonical config; ``data`` itself when no expansion applies. + + Raises + ------ + ValueError + If ``bridging_method`` is set on a model type that does not + support it, or if ``pair_exclude_types`` and + ``descriptor.exclude_types`` are both given and differ. + """ + method = str(data.get("bridging_method", "none")) + if method.lower() in ("none", ""): + return data + model_type = str(data.get("type", "standard")) + if model_type.lower() not in ("standard", "dpa4", "sezm"): + raise ValueError( + "`bridging_method` is only supported on the 'standard' and " + f"'dpa4'/'sezm' model types, but got type {model_type!r}. " + 'Spell the composition explicitly with `type: "linear_ener"` ' + "and an `inner_potential` sub-model instead." + ) + data = copy.deepcopy(data) + r_inner = float(data.get("bridging_r_inner", 0.5)) + r_outer = float(data.get("bridging_r_outer", 0.8)) + + # Legacy promotion (pt `type: "dpa4"` semantics): a descriptor-scoped + # exclusion also governs the analytical term of a bridged model. + descriptor_exclude_types = [ + list(pair) for pair in (data.get("descriptor", {}).get("exclude_types") or []) + ] + if "pair_exclude_types" in data: + pair_exclude_types = [list(pair) for pair in (data["pair_exclude_types"] or [])] + if descriptor_exclude_types and descriptor_exclude_types != pair_exclude_types: + raise ValueError( + "SeZM `pair_exclude_types` and `descriptor.exclude_types` must match " + "when both are provided." + ) + else: + pair_exclude_types = descriptor_exclude_types + + learned = { + key: value for key, value in data.items() if key not in _COMPOSITION_KEYS + } + learned["type"] = model_type + learned["type_map"] = copy.deepcopy(data["type_map"]) + canonical = { + "type": "linear_ener", + "type_map": data["type_map"], + "weights": "sum", + "models": [ + learned, + { + "type": "inner_potential", + "mode": method, + "r_inner": r_inner, + "r_outer": r_outer, + }, + ], + "atom_exclude_types": data.get("atom_exclude_types", []), + "pair_exclude_types": pair_exclude_types, + } + if "spin" in data: + canonical["spin"] = data["spin"] + return canonical diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index 8cce7b9faf..178e58980c 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -322,23 +322,45 @@ learned energy in a protected region: E_i = E_i^{\mathrm{DPA4/SeZM}} + E_i^{\mathrm{ZBL}}. ``` -Below `bridging_r_inner` the distance seen by the descriptor is clamped, with a -smooth transition back to the true distance up to `bridging_r_outer`; a source +Below `r_inner` the distance seen by the descriptor is clamped, with a +smooth transition back to the true distance up to `r_outer`; a source gate additionally blocks the learned model from leaking information about the -frozen short-range pairs. Enable it with: +frozen short-range pairs. + +A bridged model is a linear composition: the learned model plus the +analytical `inner_potential` term, summed by `linear_ener`. Spell it as: ```json { "model": { - "bridging_method": "zbl", - "bridging_r_inner": 0.5, - "bridging_r_outer": 0.8 + "type": "linear_ener", + "weights": "sum", + "type_map": ["O", "H"], + "models": [ + { + "type": "dpa4", + "descriptor": { "...": "..." }, + "fitting_net": { "...": "..." } + }, + { + "type": "inner_potential", + "mode": "zbl", + "r_inner": 0.5, + "r_outer": 0.8 + } + ] } } ``` +The composition derives the learned descriptor's clamping window from the +`inner_potential` child, so the radii are written once. The legacy +spelling -- a `bridging_method` / `bridging_r_inner` / `bridging_r_outer` +flag set on the `dpa4` (or `standard`) model type -- is kept as sugar and +expands to exactly the composition above. + When ZBL bridging is enabled, set `training.training_data.min_pair_dist` to the -same value as `bridging_r_inner` so frames with shorter atom pairs are excluded +same value as `r_inner` so frames with shorter atom pairs are excluded from training. See `examples/water/dpa4/input-zbl.json` for a complete example. ## Performance and precision diff --git a/examples/water/dpa4/input-zbl.json b/examples/water/dpa4/input-zbl.json index ddd010233c..5c0dc93c4c 100644 --- a/examples/water/dpa4/input-zbl.json +++ b/examples/water/dpa4/input-zbl.json @@ -1,115 +1,124 @@ { - "_comment": "DPA4-Mini energy-training example with ZBL zone bridging.", - "model": { - "type": "dpa4", - "type_map": [ - "O", - "H" - ], - "descriptor": { - "rcut": 6.0, - "channels": 32, - "n_radial": 16, - "edge_norm": false, - "use_env_seed": true, - "lmax": 2, - "mmax": 1, - "n_blocks": 2, - "mixing_layers": 3, - "radial_so2_mode": "degree_channel", - "radial_so2_rank": 1, - "n_focus": 1, - "focus_dim": 0, - "n_atten_head": 1, - "message_node_so3": true, - "ffn_neurons": 0, - "ffn_so3_grid": true, - "grid_mlp": false, - "grid_branch": [ - 0, - 0, - 1 - ], - "ffn_blocks": 1, - "so3_readout": "mlp", - "use_amp": true, - "precision": "float32", - "seed": 42 + "_comment": "DPA4-Mini energy-training example with ZBL zone bridging.", + "model": { + "type": "linear_ener", + "weights": "sum", + "type_map": [ + "O", + "H" + ], + "models": [ + { + "type": "dpa4", + "descriptor": { + "rcut": 6.0, + "channels": 32, + "n_radial": 16, + "edge_norm": false, + "use_env_seed": true, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "mixing_layers": 3, + "radial_so2_mode": "degree_channel", + "radial_so2_rank": 1, + "n_focus": 1, + "focus_dim": 0, + "n_atten_head": 1, + "message_node_so3": true, + "ffn_neurons": 0, + "ffn_so3_grid": true, + "grid_mlp": false, + "grid_branch": [ + 0, + 0, + 1 + ], + "ffn_blocks": 1, + "so3_readout": "mlp", + "use_amp": true, + "precision": "float32", + "seed": 42 + }, + "fitting_net": { + "neuron": [ + 0 + ], + "precision": "float32", + "seed": 42 + }, + "use_compile": false, + "enable_tf32": true + }, + { + "type": "inner_potential", + "mode": "zbl", + "r_inner": 0.5, + "r_outer": 0.8 + } + ] }, - "fitting_net": { - "neuron": [ - 0 - ], - "precision": "float32", - "seed": 42 + "learning_rate": { + "type": "wsd", + "start_lr": 0.00045, + "stop_lr": 1e-06, + "warmup_ratio": 0.003, + "warmup_start_factor": 0.2, + "decay_phase_ratio": 0.65, + "decay_type": "cosine" }, - "use_compile": false, - "enable_tf32": true, - "bridging_method": "zbl", - "bridging_r_inner": 0.5, - "bridging_r_outer": 0.8 - }, - "learning_rate": { - "type": "wsd", - "start_lr": 4.5e-4, - "stop_lr": 1e-6, - "warmup_ratio": 0.003, - "warmup_start_factor": 0.2, - "decay_phase_ratio": 0.65, - "decay_type": "cosine" - }, - "loss": { - "type": "ener", - "loss_func": "mae", - "f_use_norm": true, - "start_pref_e": 20, - "limit_pref_e": 20, - "start_pref_f": 20, - "limit_pref_f": 20, - "start_pref_v": 5, - "limit_pref_v": 5 - }, - "optimizer": { - "type": "HybridMuon", - "weight_decay": 0.001 - }, - "training": { - "stat_file": "./dpa4.hdf5", - "training_data": { - "systems": [ - "../data/data_0", - "../data/data_1", - "../data/data_2" - ], - "batch_size": 1, - "min_pair_dist": 0.5 + "loss": { + "type": "ener", + "loss_func": "mae", + "f_use_norm": true, + "start_pref_e": 20, + "limit_pref_e": 20, + "start_pref_f": 20, + "limit_pref_f": 20, + "start_pref_v": 5, + "limit_pref_v": 5 }, - "validation_data": { - "systems": [ - "../data/data_3" - ], - "batch_size": 1, - "numb_batch": 1 + "optimizer": { + "type": "HybridMuon", + "weight_decay": 0.001 }, - "numb_steps": 2000000, - "gradient_max_norm": 5.0, - "save_freq": 2000, - "max_ckpt_keep": 3, - "enable_ema": true, - "ema_decay": 0.999, - "ema_ckpt_keep": 3, - "disp_file": "lcurve.out", - "disp_freq": 1000, - "disp_avg": true, - "disp_training": true, - "time_training": true, - "tensorboard": false, - "enable_profiler": false, - "tensorboard_freq": 1000, - "tensorboard_log_dir": "tb_log", - "profiling": false, - "profiling_file": "timeline.json", - "zero_stage": 1, - "seed": 42 - } + "training": { + "stat_file": "./dpa4.hdf5", + "training_data": { + "systems": [ + "../data/data_0", + "../data/data_1", + "../data/data_2" + ], + "batch_size": 1, + "min_pair_dist": 0.5 + }, + "validation_data": { + "systems": [ + "../data/data_3" + ], + "batch_size": 1, + "numb_batch": 1 + }, + "numb_steps": 2000000, + "gradient_max_norm": 5.0, + "save_freq": 2000, + "max_ckpt_keep": 3, + "enable_ema": true, + "ema_decay": 0.999, + "ema_ckpt_keep": 3, + "disp_file": "lcurve.out", + "disp_freq": 1000, + "disp_avg": true, + "disp_training": true, + "time_training": true, + "tensorboard": false, + "enable_profiler": false, + "tensorboard_freq": 1000, + "tensorboard_log_dir": "tb_log", + "profiling": false, + "profiling_file": "timeline.json", + "zero_stage": 1, + "seed": 42 + } } diff --git a/source/tests/common/dpmodel/test_zbl_bridging.py b/source/tests/common/dpmodel/test_zbl_bridging.py index 0db497dd84..126b74337c 100644 --- a/source/tests/common/dpmodel/test_zbl_bridging.py +++ b/source/tests/common/dpmodel/test_zbl_bridging.py @@ -629,3 +629,107 @@ def test_forwarded_from_children(self) -> None: assert bridged.atomic_model.get_compute_stats_distinguish_types() == any( c.get_compute_stats_distinguish_types() for c in children ) + + +def _canonical_config() -> dict: + """``ZBL_CONFIG`` spelled canonically (issue #5948): an explicit + ``linear_ener`` composition with an ``inner_potential`` sub-model. + """ + cfg = copy.deepcopy(ZBL_CONFIG) + cfg["fitting_net"]["seed"] = 7 + return { + "type": "linear_ener", + "weights": "sum", + "type_map": cfg["type_map"], + "models": [ + { + "type": "standard", + "descriptor": cfg["descriptor"], + "fitting_net": cfg["fitting_net"], + }, + { + "type": "inner_potential", + "mode": "ZBL", + "r_inner": 0.8, + "r_outer": 1.2, + }, + ], + } + + +class TestCanonicalComposition: + """The canonical ``linear_ener`` + ``inner_potential`` spelling.""" + + def test_canonical_config_composes(self) -> None: + model = get_model(_canonical_config()) + assert type(model) is LinearEnergyModel + am = model.atomic_model + assert isinstance(am, LinearEnergyAtomicModel) + assert am.weights == "sum" + assert isinstance(am.models[1], InnerPotentialAtomicModel) + # the composition derives the learned sibling's clamp window from + # the inner_potential child: one source of truth for the radii + dp_child = am.models[0] + assert dp_child.descriptor.inner_clamp is not None + assert float(dp_child.descriptor.inner_clamp.r_inner) == 0.8 + assert dp_child.descriptor.bridging_switch is not None + + def test_canonical_matches_sugar_energy(self) -> None: + """Same seeds, both spellings: bit-identical construction, so the + energies must be exactly equal. + """ + sugar = copy.deepcopy(ZBL_CONFIG) + sugar["fitting_net"]["seed"] = 7 + m_sugar = get_model(sugar) + m_canon = get_model(_canonical_config()) + coord, atype, box = _close_pair_inputs() + e_sugar = m_sugar.call_common( + coord, atype, box=box, neighbor_graph_method="dense" + )["energy_redu"] + e_canon = m_canon.call_common( + coord, atype, box=box, neighbor_graph_method="dense" + )["energy_redu"] + np.testing.assert_array_equal(e_canon, e_sugar) + + def test_canonical_serialize_matches_sugar(self) -> None: + """Both spellings serialize to the same wire dict: the flag is + sugar, not a different model. + """ + sugar = copy.deepcopy(ZBL_CONFIG) + sugar["fitting_net"]["seed"] = 7 + d_sugar = get_model(sugar).serialize() + d_canon = get_model(_canonical_config()).serialize() + + def _strip_arrays(obj): + if isinstance(obj, dict): + return {k: _strip_arrays(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_strip_arrays(v) for v in obj] + if isinstance(obj, np.ndarray): + return ("ndarray", obj.shape) + return obj + + assert _strip_arrays(d_canon) == _strip_arrays(d_sugar) + + def test_two_inner_children_raise(self) -> None: + cfg = _canonical_config() + cfg["models"].append(dict(cfg["models"][1])) + with pytest.raises(ValueError, match="at most one"): + get_model(cfg) + + def test_inner_without_learned_sibling_raises(self) -> None: + cfg = _canonical_config() + cfg["models"] = [cfg["models"][1]] + with pytest.raises(ValueError, match="exactly one learned"): + get_model(cfg) + + def test_standard_builder_rejects_the_flag(self) -> None: + """Direct standard construction with the flag fails fast instead of + silently dropping the analytical term. + """ + from deepmd.dpmodel.model.model import ( + get_standard_model, + ) + + with pytest.raises(ValueError, match="bridging_method"): + get_standard_model(copy.deepcopy(ZBL_CONFIG)) diff --git a/source/tests/common/test_bridging.py b/source/tests/common/test_bridging.py new file mode 100644 index 0000000000..88dcf37e53 --- /dev/null +++ b/source/tests/common/test_bridging.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Unit tests for the ``bridging_method`` sugar expansion (issue #5948). + +``expand_bridging_method`` is the ONE owner of the sugar: it rewrites a +flag-form config into the canonical ``linear_ener`` composition over the +learned model and an ``inner_potential`` sub-model. These tests pin the +key routing, the legacy exclusion promotion, and the rejections. +""" + +import copy + +import pytest + +from deepmd.utils.bridging import ( + expand_bridging_method, +) + + +def _flag_config() -> dict: + return { + "type": "dpa4", + "type_map": ["Ni", "O"], + "descriptor": {"type": "dpa4", "rcut": 4.0, "sel": 8}, + "fitting_net": {"type": "dpa4_ener", "neuron": [8, 8]}, + "bridging_method": "ZBL", + "bridging_r_inner": 0.8, + "bridging_r_outer": 1.2, + } + + +@pytest.mark.parametrize( + "method", + [ + None, # key absent + "none", # lower-case disable spelling + "None", # argcheck default spelling + "", # empty string disables too + ], +) +def test_inactive_flag_returns_config_unchanged(method) -> None: + data = _flag_config() + if method is None: + del data["bridging_method"] + else: + data["bridging_method"] = method + assert expand_bridging_method(data) is data + + +def test_expansion_shape() -> None: + out = expand_bridging_method(_flag_config()) + assert out["type"] == "linear_ener" + assert out["weights"] == "sum" + assert out["type_map"] == ["Ni", "O"] + learned, inner = out["models"] + assert learned["type"] == "dpa4" + assert learned["descriptor"]["type"] == "dpa4" + assert learned["fitting_net"]["type"] == "dpa4_ener" + assert inner == { + "type": "inner_potential", + "mode": "ZBL", + "r_inner": 0.8, + "r_outer": 1.2, + } + # the flag keys must not leak into the canonical config + for key in ("bridging_method", "bridging_r_inner", "bridging_r_outer"): + assert key not in out + assert key not in learned + + +def test_default_radii() -> None: + data = _flag_config() + del data["bridging_r_inner"] + del data["bridging_r_outer"] + inner = expand_bridging_method(data)["models"][1] + assert inner["r_inner"] == 0.5 + assert inner["r_outer"] == 0.8 + + +def test_input_is_not_mutated() -> None: + data = _flag_config() + ref = copy.deepcopy(data) + expand_bridging_method(data) + assert data == ref + + +def test_spin_stays_top_level() -> None: + data = _flag_config() + data["spin"] = {"scheme": "native", "use_spin": [True, False]} + out = expand_bridging_method(data) + assert out["spin"] == {"scheme": "native", "use_spin": [True, False]} + assert "spin" not in out["models"][0] + + +def test_other_model_keys_stay_on_learned_child() -> None: + data = _flag_config() + data["data_stat_protect"] = 1e-3 + data["preset_out_bias"] = {"energy": [None, 1.0]} + out = expand_bridging_method(data) + learned = out["models"][0] + assert learned["data_stat_protect"] == 1e-3 + assert learned["preset_out_bias"] == {"energy": [None, 1.0]} + assert "data_stat_protect" not in out + assert "preset_out_bias" not in out + + +def test_exclusions_move_to_composition_level() -> None: + data = _flag_config() + data["pair_exclude_types"] = [[0, 1]] + data["atom_exclude_types"] = [1] + out = expand_bridging_method(data) + assert out["pair_exclude_types"] == [[0, 1]] + assert out["atom_exclude_types"] == [1] + learned = out["models"][0] + assert "pair_exclude_types" not in learned + assert "atom_exclude_types" not in learned + + +def test_descriptor_exclude_types_promotion() -> None: + """Legacy pt semantics: a descriptor-scoped exclusion on a bridged + model also governs the analytical term. + """ + data = _flag_config() + data["descriptor"]["exclude_types"] = [[0, 1]] + out = expand_bridging_method(data) + assert out["pair_exclude_types"] == [[0, 1]] + + +def test_descriptor_exclude_types_mismatch_raises() -> None: + data = _flag_config() + data["descriptor"]["exclude_types"] = [[0, 1]] + data["pair_exclude_types"] = [[0, 0]] + with pytest.raises(ValueError, match="must match"): + expand_bridging_method(data) + + +def test_matching_exclusions_pass() -> None: + data = _flag_config() + data["descriptor"]["exclude_types"] = [[0, 1]] + data["pair_exclude_types"] = [[0, 1]] + out = expand_bridging_method(data) + assert out["pair_exclude_types"] == [[0, 1]] + + +@pytest.mark.parametrize( + "model_type", + [ + "linear_ener", # composition types must spell inner_potential directly + "frozen", # unrelated model type + ], +) +def test_unsupported_model_type_raises(model_type: str) -> None: + data = _flag_config() + data["type"] = model_type + with pytest.raises(ValueError, match="linear_ener"): + expand_bridging_method(data) + + +def test_standard_type_is_supported() -> None: + data = _flag_config() + data["type"] = "standard" + out = expand_bridging_method(data) + assert out["models"][0]["type"] == "standard" diff --git a/source/tests/pt/model/test_get_model_bridging.py b/source/tests/pt/model/test_get_model_bridging.py new file mode 100644 index 0000000000..d7abe5a55c --- /dev/null +++ b/source/tests/pt/model/test_get_model_bridging.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""pt realization of the canonical bridging composition (issue #5948). + +The canonical config spelling is ``type: "linear_ener"`` with an +``inner_potential`` sub-model. The pt backend implements bridging inside +``SeZMModel``, so its linear builder maps the canonical form onto the +``SeZMModel`` constructor arguments; the legacy ``bridging_method`` flag +is sugar expanded by the shared normalizer at the ``get_model`` entry. +Both spellings must therefore build the same model. +""" + +import copy + +import pytest + +from deepmd.pt.model.model import ( + SeZMModel, + get_model, + get_standard_model, +) + + +def _descriptor() -> dict: + return { + "type": "dpa4", + "rcut": 4.0, + "rcut_smth": 0.5, + "sel": 20, + "n_dim": 8, + "e_dim": 8, + "precision": "float64", + "seed": 7, + } + + +def _fitting() -> dict: + return { + "type": "dpa4_ener", + "neuron": [4, 4], + "precision": "float64", + "seed": 7, + } + + +def _sugar_config() -> dict: + return { + "type": "dpa4", + "type_map": ["Ni", "O"], + "descriptor": _descriptor(), + "fitting_net": _fitting(), + "bridging_method": "ZBL", + "bridging_r_inner": 0.8, + "bridging_r_outer": 1.2, + } + + +def _canonical_config() -> dict: + return { + "type": "linear_ener", + "weights": "sum", + "type_map": ["Ni", "O"], + "models": [ + { + "type": "dpa4", + "descriptor": _descriptor(), + "fitting_net": _fitting(), + }, + { + "type": "inner_potential", + "mode": "ZBL", + "r_inner": 0.8, + "r_outer": 1.2, + }, + ], + } + + +def test_canonical_builds_sezm_model() -> None: + model = get_model(_canonical_config()) + assert isinstance(model, SeZMModel) + assert model.bridging_method == "ZBL" + assert model.bridging_r_inner == 0.8 + assert model.bridging_r_outer == 1.2 + + +def test_canonical_matches_sugar_serialize() -> None: + """Same seeds, both spellings: the serialized models must agree.""" + import numpy as np + + d_sugar = get_model(_sugar_config()).serialize() + d_canon = get_model(_canonical_config()).serialize() + + def _strip_arrays(obj): + if isinstance(obj, dict): + return {k: _strip_arrays(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_strip_arrays(v) for v in obj] + if isinstance(obj, np.ndarray): + return ("ndarray", obj.shape) + if hasattr(obj, "detach"): # torch tensor + return ("tensor", tuple(obj.shape)) + return obj + + assert _strip_arrays(d_canon) == _strip_arrays(d_sugar) + + +def test_standard_builder_rejects_the_flag() -> None: + """Fail fast: the pt standard builder used to silently DROP the + bridging term. + """ + cfg = _sugar_config() + cfg["type"] = "standard" + with pytest.raises(ValueError, match="bridging_method"): + get_standard_model(cfg) + + +def test_get_model_expands_the_flag_on_standard_type() -> None: + """Through the dispatcher the flag is sugar on any supported type.""" + cfg = _sugar_config() + cfg["type"] = "standard" + model = get_model(cfg) + assert isinstance(model, SeZMModel) + assert model.bridging_method == "ZBL" + + +def test_canonical_requires_sum_weights() -> None: + cfg = _canonical_config() + cfg["weights"] = "mean" + with pytest.raises(ValueError, match="sum"): + get_model(cfg) + + +def test_canonical_rejects_non_dpa4_learned_sibling() -> None: + cfg = _canonical_config() + cfg["models"][0]["descriptor"] = { + "type": "se_e2_a", + "rcut": 4.0, + "rcut_smth": 0.5, + "sel": [20, 20], + "neuron": [4, 8], + } + with pytest.raises(NotImplementedError, match="DPA4/SeZM"): + get_model(cfg) + + +def test_canonical_rejects_two_inner_children() -> None: + cfg = _canonical_config() + cfg["models"].append(copy.deepcopy(cfg["models"][1])) + with pytest.raises(ValueError, match="at most one"): + get_model(cfg) + + +def test_plain_linear_ener_is_unaffected() -> None: + """A linear_ener composition without an inner_potential child keeps + the pre-existing builder path. + """ + sub = { + "descriptor": { + "type": "se_atten", + "rcut": 4.0, + "rcut_smth": 0.5, + "sel": 20, + "neuron": [4, 8], + "attn_layer": 0, + "seed": 1, + }, + "fitting_net": {"neuron": [5, 5], "seed": 1}, + } + cfg = { + "type": "linear_ener", + "weights": "mean", + "type_map": ["Ni", "O"], + "models": [copy.deepcopy(sub), copy.deepcopy(sub)], + } + model = get_model(cfg) + assert not isinstance(model, SeZMModel) diff --git a/source/tests/pt_expt/model/test_get_model_bridging.py b/source/tests/pt_expt/model/test_get_model_bridging.py index d9e02492dd..1cd5c5bb64 100644 --- a/source/tests/pt_expt/model/test_get_model_bridging.py +++ b/source/tests/pt_expt/model/test_get_model_bridging.py @@ -1,19 +1,14 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Analytical bridging has exactly ONE owner per backend. +"""The ``bridging_method`` sugar has exactly ONE owner. Bridging builds a COMPOSITION (``LinearEnergyModel`` over -``[learned, InnerPotential]``), so it is not expressible on a non-composite -model type: ``type: "standard"`` would have to return a model of a -different kind than the one requested. pt_expt therefore owns bridging on -the DPA4/SeZM route only and REJECTS it in the standard builder -- loudly, -because silently dropping the term yields a physically different model. - -Two builders accepting the flag is exactly how the routes drifted: -``get_sezm_model`` promotes ``descriptor.exclude_types`` to model-level -``pair_exclude_types`` and the standard route never did, which changes a -0.9 A Ni-O dimer by ~80 eV (issue #5947). Issue #5948 replaces the flag -with an explicit ``linear_ener`` composition, after which this restriction -becomes moot. +``[learned, InnerPotential]``). Since issue #5948 the canonical spelling is +``type: "linear_ener"`` with an ``inner_potential`` sub-model, and the +``bridging_method`` flag is sugar expanded by the shared +``deepmd.utils.bridging.expand_bridging_method`` normalizer at the +``get_model`` entry. The non-composite builders (``get_standard_model``, +``get_sezm_model``) REJECT the flag -- loudly, because silently dropping +the term yields a physically different model. """ import copy @@ -68,12 +63,29 @@ def test_standard_builder_rejects_bridging() -> None: get_standard_model(_bridged(_dpa4_standard_config())) -def test_get_model_rejects_bridging_without_dpa4_model_type() -> None: - """Same contract through the dispatcher: an omitted model type defaults - to the standard route, so it must reject rather than compose. +def test_get_model_expands_bridging_without_dpa4_model_type() -> None: + """Through the dispatcher the flag is sugar: an omitted model type + defaults to 'standard', and the normalizer expands the flag into the + canonical composition instead of rejecting it. """ + model = get_model(_bridged(_dpa4_standard_config())) + assert isinstance(model.atomic_model, LinearEnergyAtomicModel) + assert len(model.atomic_model.models) == 2 + assert model.atomic_model.models[0].descriptor.bridging_switch is not None + + +def test_sezm_builder_rejects_bridging() -> None: + """The DPA4/SeZM builder must not hand back a composition either: the + flag's one owner is the shared normalizer at the get_model entry. + """ + from deepmd.pt_expt.model.get_model import ( + get_sezm_model, + ) + + data = _bridged(_dpa4_standard_config()) + data["type"] = "dpa4" with pytest.raises(ValueError, match="bridging_method"): - get_model(_bridged(_dpa4_standard_config())) + get_sezm_model(data) def test_standard_builder_without_bridging_is_unaffected() -> None: @@ -157,3 +169,72 @@ def test_compile_attention_probe_tolerates_composition() -> None: assert isinstance(model.atomic_model, LinearEnergyAtomicModel) # must not raise _warn_compiled_attention(model, "Default") + + +def _canonical_config() -> dict: + """The bridged config spelled canonically (issue #5948).""" + base = _dpa4_standard_config() + return { + "type": "linear_ener", + "weights": "sum", + "type_map": base["type_map"], + "models": [ + { + "type": "dpa4", + "descriptor": base["descriptor"], + "fitting_net": base["fitting_net"], + }, + { + "type": "inner_potential", + "mode": "ZBL", + "r_inner": 0.8, + "r_outer": 1.2, + }, + ], + } + + +def test_canonical_composition_builds() -> None: + """The canonical spelling composes [learned, InnerPotential] with the + clamp radii derived onto the learned child's descriptor. + """ + model = get_model(_canonical_config()) + assert isinstance(model.atomic_model, LinearEnergyAtomicModel) + assert len(model.atomic_model.models) == 2 + learned = model.atomic_model.models[0] + assert learned.descriptor.bridging_switch is not None + assert float(learned.descriptor.inner_clamp.r_inner) == 0.8 + + +def test_canonical_matches_sugar_serialize() -> None: + """Both spellings serialize to the same wire dict.""" + import numpy as np + + data = _bridged(_dpa4_standard_config()) + data["type"] = "dpa4" + d_sugar = get_model(data).serialize() + d_canon = get_model(_canonical_config()).serialize() + + def _strip_arrays(obj): + if isinstance(obj, dict): + return {k: _strip_arrays(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_strip_arrays(v) for v in obj] + if isinstance(obj, np.ndarray): + return ("ndarray", obj.shape) + return obj + + assert _strip_arrays(d_canon) == _strip_arrays(d_sugar) + + +def test_canonical_native_spin_composition() -> None: + """A top-level native-spin section wraps the canonical composition.""" + from deepmd.pt_expt.model.native_spin_model import ( + NativeSpinEnergyModel, + ) + + cfg = _canonical_config() + cfg["spin"] = {"scheme": "native", "use_spin": [True, False]} + model = get_model(cfg) + assert isinstance(model, NativeSpinEnergyModel) + assert isinstance(model.atomic_model, LinearEnergyAtomicModel) From fb52e0fa776a2cddd3c98940f3db0a3c27f23788 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 10 Aug 2026 17:08:46 +0000 Subject: [PATCH 02/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- doc/model/dpa4.md | 13 +- examples/water/dpa4/input-zbl.json | 240 ++++++++++++++--------------- 2 files changed, 130 insertions(+), 123 deletions(-) diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index 178e58980c..3bdcb66c9a 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -335,12 +335,19 @@ analytical `inner_potential` term, summed by `linear_ener`. Spell it as: "model": { "type": "linear_ener", "weights": "sum", - "type_map": ["O", "H"], + "type_map": [ + "O", + "H" + ], "models": [ { "type": "dpa4", - "descriptor": { "...": "..." }, - "fitting_net": { "...": "..." } + "descriptor": { + "...": "..." + }, + "fitting_net": { + "...": "..." + } }, { "type": "inner_potential", diff --git a/examples/water/dpa4/input-zbl.json b/examples/water/dpa4/input-zbl.json index 5c0dc93c4c..6d4c7f2e00 100644 --- a/examples/water/dpa4/input-zbl.json +++ b/examples/water/dpa4/input-zbl.json @@ -1,124 +1,124 @@ { - "_comment": "DPA4-Mini energy-training example with ZBL zone bridging.", - "model": { - "type": "linear_ener", - "weights": "sum", - "type_map": [ - "O", - "H" - ], - "models": [ - { - "type": "dpa4", - "descriptor": { - "rcut": 6.0, - "channels": 32, - "n_radial": 16, - "edge_norm": false, - "use_env_seed": true, - "lmax": 2, - "mmax": 1, - "n_blocks": 2, - "mixing_layers": 3, - "radial_so2_mode": "degree_channel", - "radial_so2_rank": 1, - "n_focus": 1, - "focus_dim": 0, - "n_atten_head": 1, - "message_node_so3": true, - "ffn_neurons": 0, - "ffn_so3_grid": true, - "grid_mlp": false, - "grid_branch": [ - 0, - 0, - 1 - ], - "ffn_blocks": 1, - "so3_readout": "mlp", - "use_amp": true, - "precision": "float32", - "seed": 42 - }, - "fitting_net": { - "neuron": [ - 0 - ], - "precision": "float32", - "seed": 42 - }, - "use_compile": false, - "enable_tf32": true - }, - { - "type": "inner_potential", - "mode": "zbl", - "r_inner": 0.5, - "r_outer": 0.8 - } - ] - }, - "learning_rate": { - "type": "wsd", - "start_lr": 0.00045, - "stop_lr": 1e-06, - "warmup_ratio": 0.003, - "warmup_start_factor": 0.2, - "decay_phase_ratio": 0.65, - "decay_type": "cosine" - }, - "loss": { - "type": "ener", - "loss_func": "mae", - "f_use_norm": true, - "start_pref_e": 20, - "limit_pref_e": 20, - "start_pref_f": 20, - "limit_pref_f": 20, - "start_pref_v": 5, - "limit_pref_v": 5 - }, - "optimizer": { - "type": "HybridMuon", - "weight_decay": 0.001 - }, - "training": { - "stat_file": "./dpa4.hdf5", - "training_data": { - "systems": [ - "../data/data_0", - "../data/data_1", - "../data/data_2" - ], - "batch_size": 1, - "min_pair_dist": 0.5 + "_comment": "DPA4-Mini energy-training example with ZBL zone bridging.", + "model": { + "type": "linear_ener", + "weights": "sum", + "type_map": [ + "O", + "H" + ], + "models": [ + { + "type": "dpa4", + "descriptor": { + "rcut": 6.0, + "channels": 32, + "n_radial": 16, + "edge_norm": false, + "use_env_seed": true, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "mixing_layers": 3, + "radial_so2_mode": "degree_channel", + "radial_so2_rank": 1, + "n_focus": 1, + "focus_dim": 0, + "n_atten_head": 1, + "message_node_so3": true, + "ffn_neurons": 0, + "ffn_so3_grid": true, + "grid_mlp": false, + "grid_branch": [ + 0, + 0, + 1 + ], + "ffn_blocks": 1, + "so3_readout": "mlp", + "use_amp": true, + "precision": "float32", + "seed": 42 }, - "validation_data": { - "systems": [ - "../data/data_3" - ], - "batch_size": 1, - "numb_batch": 1 + "fitting_net": { + "neuron": [ + 0 + ], + "precision": "float32", + "seed": 42 }, - "numb_steps": 2000000, - "gradient_max_norm": 5.0, - "save_freq": 2000, - "max_ckpt_keep": 3, - "enable_ema": true, - "ema_decay": 0.999, - "ema_ckpt_keep": 3, - "disp_file": "lcurve.out", - "disp_freq": 1000, - "disp_avg": true, - "disp_training": true, - "time_training": true, - "tensorboard": false, - "enable_profiler": false, - "tensorboard_freq": 1000, - "tensorboard_log_dir": "tb_log", - "profiling": false, - "profiling_file": "timeline.json", - "zero_stage": 1, - "seed": 42 - } + "use_compile": false, + "enable_tf32": true + }, + { + "type": "inner_potential", + "mode": "zbl", + "r_inner": 0.5, + "r_outer": 0.8 + } + ] + }, + "learning_rate": { + "type": "wsd", + "start_lr": 0.00045, + "stop_lr": 1e-06, + "warmup_ratio": 0.003, + "warmup_start_factor": 0.2, + "decay_phase_ratio": 0.65, + "decay_type": "cosine" + }, + "loss": { + "type": "ener", + "loss_func": "mae", + "f_use_norm": true, + "start_pref_e": 20, + "limit_pref_e": 20, + "start_pref_f": 20, + "limit_pref_f": 20, + "start_pref_v": 5, + "limit_pref_v": 5 + }, + "optimizer": { + "type": "HybridMuon", + "weight_decay": 0.001 + }, + "training": { + "stat_file": "./dpa4.hdf5", + "training_data": { + "systems": [ + "../data/data_0", + "../data/data_1", + "../data/data_2" + ], + "batch_size": 1, + "min_pair_dist": 0.5 + }, + "validation_data": { + "systems": [ + "../data/data_3" + ], + "batch_size": 1, + "numb_batch": 1 + }, + "numb_steps": 2000000, + "gradient_max_norm": 5.0, + "save_freq": 2000, + "max_ckpt_keep": 3, + "enable_ema": true, + "ema_decay": 0.999, + "ema_ckpt_keep": 3, + "disp_file": "lcurve.out", + "disp_freq": 1000, + "disp_avg": true, + "disp_training": true, + "time_training": true, + "tensorboard": false, + "enable_profiler": false, + "tensorboard_freq": 1000, + "tensorboard_log_dir": "tb_log", + "profiling": false, + "profiling_file": "timeline.json", + "zero_stage": 1, + "seed": 42 + } } From ba3f746439bc9e48a9f6ea9157711644559172f0 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 11 Aug 2026 09:00:41 +0800 Subject: [PATCH 03/11] docs(dpa4): recommend the concise bridging_method form; linear_ener is the explicit equivalent Both spellings stay supported: the concise 'type: dpa4' + bridging_method form is the recommended user interface, and the explicit linear_ener + inner_potential composition is the canonical internal form it expands to (one shared normalizer defines the equivalence). Revert the example to the concise form; document both. --- deepmd/utils/argcheck.py | 6 +- deepmd/utils/bridging.py | 7 +- doc/model/dpa4.md | 41 ++++-- examples/water/dpa4/input-zbl.json | 227 ++++++++++++++--------------- 4 files changed, 143 insertions(+), 138 deletions(-) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index bce9f57080..cf52d914fa 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -3489,9 +3489,9 @@ def sezm_model_args() -> Argument: doc_bridging_method = ( "Short-range bridging method. Currently supports 'ZBL'. " "The value is case-insensitive; set it to 'None' to disable bridging. " - "This flag is sugar: it expands to the canonical `linear_ener` " - "composition over the learned model and an `inner_potential` " - "sub-model." + "This concise form is the recommended interface; it expands to the " + "equivalent explicit `linear_ener` composition over the learned " + "model and an `inner_potential` sub-model." ) doc_bridging_r_inner = ( "Inner clamping radius in Å. ML descriptor distances below this radius are frozen. " diff --git a/deepmd/utils/bridging.py b/deepmd/utils/bridging.py index 6959cddcaf..90576188bf 100644 --- a/deepmd/utils/bridging.py +++ b/deepmd/utils/bridging.py @@ -14,9 +14,10 @@ ] } -The legacy spelling -- a ``bridging_method`` flag on a non-composite model -type -- is kept as sugar. This module is the ONE owner of that expansion -(issue #5948): every backend's ``get_model`` entry point calls +The concise spelling -- a ``bridging_method`` flag on the ``dpa4`` (or +``standard``) model type -- is the recommended user interface; it is pure +sugar over the canonical form. This module is the ONE owner of that +expansion (issue #5948): every backend's ``get_model`` entry point calls :func:`expand_bridging_method` before dispatch, so no per-builder flag handling can drift. """ diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index 178e58980c..76bd596af8 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -322,13 +322,30 @@ learned energy in a protected region: E_i = E_i^{\mathrm{DPA4/SeZM}} + E_i^{\mathrm{ZBL}}. ``` -Below `r_inner` the distance seen by the descriptor is clamped, with a -smooth transition back to the true distance up to `r_outer`; a source +Below `bridging_r_inner` the distance seen by the descriptor is clamped, with a +smooth transition back to the true distance up to `bridging_r_outer`; a source gate additionally blocks the learned model from leaking information about the -frozen short-range pairs. +frozen short-range pairs. The recommended way to enable it is the concise +form, set directly on the `dpa4` model: -A bridged model is a linear composition: the learned model plus the -analytical `inner_potential` term, summed by `linear_ener`. Spell it as: +```json +{ + "model": { + "type": "dpa4", + "bridging_method": "zbl", + "bridging_r_inner": 0.5, + "bridging_r_outer": 0.8 + } +} +``` + +When ZBL bridging is enabled, set `training.training_data.min_pair_dist` to the +same value as `bridging_r_inner` so frames with shorter atom pairs are excluded +from training. See `examples/water/dpa4/input-zbl.json` for a complete example. + +Internally, a bridged model is a linear composition: the learned model plus +the analytical `inner_potential` term, summed by `linear_ener`. The concise +form above expands to exactly this equivalent explicit form: ```json { @@ -353,15 +370,11 @@ analytical `inner_potential` term, summed by `linear_ener`. Spell it as: } ``` -The composition derives the learned descriptor's clamping window from the -`inner_potential` child, so the radii are written once. The legacy -spelling -- a `bridging_method` / `bridging_r_inner` / `bridging_r_outer` -flag set on the `dpa4` (or `standard`) model type -- is kept as sugar and -expands to exactly the composition above. - -When ZBL bridging is enabled, set `training.training_data.min_pair_dist` to the -same value as `r_inner` so frames with shorter atom pairs are excluded -from training. See `examples/water/dpa4/input-zbl.json` for a complete example. +Both spellings build the same model (one shared normalizer defines the +equivalence). The explicit form exposes the composition machinery directly: +use it when you combine models beyond the standard bridged pair. In either +form, the composition derives the learned descriptor's clamping window from +the analytical term, so the radii are written once. ## Performance and precision diff --git a/examples/water/dpa4/input-zbl.json b/examples/water/dpa4/input-zbl.json index 5c0dc93c4c..ddd010233c 100644 --- a/examples/water/dpa4/input-zbl.json +++ b/examples/water/dpa4/input-zbl.json @@ -1,124 +1,115 @@ { - "_comment": "DPA4-Mini energy-training example with ZBL zone bridging.", - "model": { - "type": "linear_ener", - "weights": "sum", - "type_map": [ - "O", - "H" - ], - "models": [ - { - "type": "dpa4", - "descriptor": { - "rcut": 6.0, - "channels": 32, - "n_radial": 16, - "edge_norm": false, - "use_env_seed": true, - "lmax": 2, - "mmax": 1, - "n_blocks": 2, - "mixing_layers": 3, - "radial_so2_mode": "degree_channel", - "radial_so2_rank": 1, - "n_focus": 1, - "focus_dim": 0, - "n_atten_head": 1, - "message_node_so3": true, - "ffn_neurons": 0, - "ffn_so3_grid": true, - "grid_mlp": false, - "grid_branch": [ - 0, - 0, - 1 - ], - "ffn_blocks": 1, - "so3_readout": "mlp", - "use_amp": true, - "precision": "float32", - "seed": 42 - }, - "fitting_net": { - "neuron": [ - 0 - ], - "precision": "float32", - "seed": 42 - }, - "use_compile": false, - "enable_tf32": true - }, - { - "type": "inner_potential", - "mode": "zbl", - "r_inner": 0.5, - "r_outer": 0.8 - } - ] + "_comment": "DPA4-Mini energy-training example with ZBL zone bridging.", + "model": { + "type": "dpa4", + "type_map": [ + "O", + "H" + ], + "descriptor": { + "rcut": 6.0, + "channels": 32, + "n_radial": 16, + "edge_norm": false, + "use_env_seed": true, + "lmax": 2, + "mmax": 1, + "n_blocks": 2, + "mixing_layers": 3, + "radial_so2_mode": "degree_channel", + "radial_so2_rank": 1, + "n_focus": 1, + "focus_dim": 0, + "n_atten_head": 1, + "message_node_so3": true, + "ffn_neurons": 0, + "ffn_so3_grid": true, + "grid_mlp": false, + "grid_branch": [ + 0, + 0, + 1 + ], + "ffn_blocks": 1, + "so3_readout": "mlp", + "use_amp": true, + "precision": "float32", + "seed": 42 }, - "learning_rate": { - "type": "wsd", - "start_lr": 0.00045, - "stop_lr": 1e-06, - "warmup_ratio": 0.003, - "warmup_start_factor": 0.2, - "decay_phase_ratio": 0.65, - "decay_type": "cosine" + "fitting_net": { + "neuron": [ + 0 + ], + "precision": "float32", + "seed": 42 }, - "loss": { - "type": "ener", - "loss_func": "mae", - "f_use_norm": true, - "start_pref_e": 20, - "limit_pref_e": 20, - "start_pref_f": 20, - "limit_pref_f": 20, - "start_pref_v": 5, - "limit_pref_v": 5 + "use_compile": false, + "enable_tf32": true, + "bridging_method": "zbl", + "bridging_r_inner": 0.5, + "bridging_r_outer": 0.8 + }, + "learning_rate": { + "type": "wsd", + "start_lr": 4.5e-4, + "stop_lr": 1e-6, + "warmup_ratio": 0.003, + "warmup_start_factor": 0.2, + "decay_phase_ratio": 0.65, + "decay_type": "cosine" + }, + "loss": { + "type": "ener", + "loss_func": "mae", + "f_use_norm": true, + "start_pref_e": 20, + "limit_pref_e": 20, + "start_pref_f": 20, + "limit_pref_f": 20, + "start_pref_v": 5, + "limit_pref_v": 5 + }, + "optimizer": { + "type": "HybridMuon", + "weight_decay": 0.001 + }, + "training": { + "stat_file": "./dpa4.hdf5", + "training_data": { + "systems": [ + "../data/data_0", + "../data/data_1", + "../data/data_2" + ], + "batch_size": 1, + "min_pair_dist": 0.5 }, - "optimizer": { - "type": "HybridMuon", - "weight_decay": 0.001 + "validation_data": { + "systems": [ + "../data/data_3" + ], + "batch_size": 1, + "numb_batch": 1 }, - "training": { - "stat_file": "./dpa4.hdf5", - "training_data": { - "systems": [ - "../data/data_0", - "../data/data_1", - "../data/data_2" - ], - "batch_size": 1, - "min_pair_dist": 0.5 - }, - "validation_data": { - "systems": [ - "../data/data_3" - ], - "batch_size": 1, - "numb_batch": 1 - }, - "numb_steps": 2000000, - "gradient_max_norm": 5.0, - "save_freq": 2000, - "max_ckpt_keep": 3, - "enable_ema": true, - "ema_decay": 0.999, - "ema_ckpt_keep": 3, - "disp_file": "lcurve.out", - "disp_freq": 1000, - "disp_avg": true, - "disp_training": true, - "time_training": true, - "tensorboard": false, - "enable_profiler": false, - "tensorboard_freq": 1000, - "tensorboard_log_dir": "tb_log", - "profiling": false, - "profiling_file": "timeline.json", - "zero_stage": 1, - "seed": 42 - } + "numb_steps": 2000000, + "gradient_max_norm": 5.0, + "save_freq": 2000, + "max_ckpt_keep": 3, + "enable_ema": true, + "ema_decay": 0.999, + "ema_ckpt_keep": 3, + "disp_file": "lcurve.out", + "disp_freq": 1000, + "disp_avg": true, + "disp_training": true, + "time_training": true, + "tensorboard": false, + "enable_profiler": false, + "tensorboard_freq": 1000, + "tensorboard_log_dir": "tb_log", + "profiling": false, + "profiling_file": "timeline.json", + "zero_stage": 1, + "seed": 42 + } } From 2fde06a0948b10e7438f0154e406f91764ea0982 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 11 Aug 2026 09:09:30 +0800 Subject: [PATCH 04/11] refactor(dpmodel): move the linear atomic-model builder into BackendModelFactory The linear child-parsing core was a private function in deepmd/dpmodel/model/model.py that pt_expt imported cross-module. The factory is the established home for registry-parameterized composition builders (get_zbl_model already builds the srtab two-child composition there), so get_linear_atomic_model joins it: dpmodel and pt_expt now call _model_factory.get_linear_atomic_model(data), with the backend classes bound once at factory construction. No behavior change. --- deepmd/dpmodel/model/model.py | 110 +-------------------- deepmd/dpmodel/model/model_factory.py | 134 ++++++++++++++++++++++++++ deepmd/pt_expt/model/get_model.py | 11 +-- 3 files changed, 136 insertions(+), 119 deletions(-) diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index 418e910b00..245e5987ee 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import copy -from typing import ( - Any, -) from deepmd.dpmodel.atomic_model.dp_atomic_model import ( DPAtomicModel, @@ -117,12 +114,7 @@ def get_linear_model(data: dict) -> BaseModel: for sub in data["models"]: if "descriptor" in sub: sub["descriptor"]["use_spin"] = use_spin - composed = _build_linear_atomic_model( - data, - model_components_factory=_model_factory.get_model_components, - dp_atomic_model=DPAtomicModel, - pairtab_atomic_model=PairTabAtomicModel, - ) + composed = _model_factory.get_linear_atomic_model(data) if spin is not None: if not composed.supports_native_spin(): raise NotImplementedError( @@ -133,106 +125,6 @@ def get_linear_model(data: dict) -> BaseModel: return LinearEnergyModel(atomic_model_=composed) -def _build_linear_atomic_model( - data: dict, - *, - model_components_factory: Any, - dp_atomic_model: type, - pairtab_atomic_model: type, -) -> Any: - """Build the ``LinearEnergyAtomicModel`` composition from a config. - - Shared between the dpmodel and pt_expt linear builders: the caller - supplies its backend's component factory and atomic-model classes. - ``data`` is mutated (callers pass a private deep copy). - - Parameters - ---------- - data : dict - The ``linear_ener`` model configuration. - model_components_factory : callable - Backend factory building (descriptor, fitting, type_map) from a - standard sub-model config. - dp_atomic_model : type - Backend learned atomic-model class. - pairtab_atomic_model : type - Backend pair-tabulation atomic-model class. - """ - from deepmd.dpmodel.atomic_model.inner_potential import ( - InnerPotentialAtomicModel, - ) - from deepmd.dpmodel.atomic_model.linear_atomic_model import ( - LinearEnergyAtomicModel, - ) - - type_map = data["type_map"] - ntypes = len(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] - 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: - raise ValueError( - "An `inner_potential` sub-model bridges exactly one learned " - f"sibling, but got {len(learned_indices)} sub-models with a " - "descriptor." - ) - # 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: - sub["descriptor"]["ntypes"] = ntypes - descriptor, fitting, _ = model_components_factory(sub) - built[i] = dp_atomic_model(descriptor, fitting, type_map=sub["type_map"]) - 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_atomic_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) -> SpinModel: """Get a spin model from a dictionary. diff --git a/deepmd/dpmodel/model/model_factory.py b/deepmd/dpmodel/model/model_factory.py index b40ef7d937..9d6b7832eb 100644 --- a/deepmd/dpmodel/model/model_factory.py +++ b/deepmd/dpmodel/model/model_factory.py @@ -124,6 +124,127 @@ 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, +) -> 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. + + Raises + ------ + ValueError + If more than one ``inner_potential`` child is given, if an + ``inner_potential`` child has no unique learned sibling, 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] + 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: + raise ValueError( + "An `inner_potential` sub-model bridges exactly one learned " + f"sibling, but got {len(learned_indices)} sub-models with a " + "descriptor." + ) + # 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: + descriptor, fitting, _ = get_model_components( + sub, + descriptor_base=descriptor_base, + fitting_base=fitting_base, + backend_name=backend_name, + ) + built[i] = atomic_model(descriptor, fitting, type_map=sub["type_map"]) + 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 +378,19 @@ def get_standard_model(self, data: dict) -> Any: backend_name=self.backend_name, ) + def get_linear_atomic_model(self, data: dict) -> 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, + ) + def get_zbl_model(self, data: dict) -> Any: """Construct a ZBL model for this backend.""" if ( diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index 1928036ded..989c59c505 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -296,10 +296,6 @@ def get_linear_model(model_params: dict) -> BaseModel: model_params : dict The model parameters. """ - from deepmd.dpmodel.model.model import ( - _build_linear_atomic_model, - ) - from .dp_linear_model import ( LinearEnergyModel, ) @@ -324,12 +320,7 @@ def get_linear_model(model_params: dict) -> BaseModel: for sub in model_params["models"]: if "descriptor" in sub: sub["descriptor"]["use_spin"] = use_spin - composed = _build_linear_atomic_model( - model_params, - model_components_factory=_model_factory.get_model_components, - dp_atomic_model=DPAtomicModel, - pairtab_atomic_model=PairTabAtomicModel, - ) + composed = _model_factory.get_linear_atomic_model(model_params) if spin is not None: if not composed.supports_native_spin(): raise NotImplementedError( From b25bd482a5a8dbe8c9b98f09d8e9540477b0bc2c Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 11 Aug 2026 10:13:33 +0800 Subject: [PATCH 05/11] fix(bridging): address review findings on the linear composition builders Guards (shared builder, dpmodel + pt_expt via BackendModelFactory): - a bridged composition requires weights: "sum" ("mean" silently halved both energy terms; pt already rejected it) - an inner_potential child must not carry a descriptor (was a KeyError) - a bridging_method flag on a linear child is rejected instead of silently dropped (also guarded in the pt linear builder) pt canonical realization: - reject lora on the learned child (the trainer reads top-level lora only, so it would silently train without adapters) - reject a child type_map that differs from the composition's instead of silently overwriting it pt checkpoint consumers: _is_sezm_model_params, is_sezm_checkpoint and freeze_sezm_to_pt2 now recognize the canonical bridged spelling via the shared is_bridged_sezm_config predicate, so DeepPot no-jit routing and the .pt2 freeze route work for canonical checkpoints. update_sel (pt + pt_expt linear models): skip inner_potential children instead of crashing with KeyError: 'descriptor' on the default CLI path. pt_expt: DPA4/SeZM-family linear children now route through get_sezm_model via a descriptor_child_builder hook on the shared builder, restoring the family defaults and the loud rejections of lora / use_compile / preset_out_bias. argcheck: bridging args shared with the standard variant (the documented standard sugar previously failed strict normalization); inner_potential is now child-only (top-level model.type: inner_potential is rejected at normalization, not at build). --- deepmd/dpmodel/model/model_factory.py | 65 +++++++++-- deepmd/pt/entrypoints/freeze_pt2.py | 18 ++- deepmd/pt/infer/deep_eval.py | 7 ++ deepmd/pt/model/model/__init__.py | 21 ++++ deepmd/pt/model/model/dp_linear_model.py | 3 + deepmd/pt_expt/model/dp_linear_model.py | 3 + deepmd/pt_expt/model/get_model.py | 27 ++++- deepmd/utils/argcheck.py | 105 +++++++++++------- deepmd/utils/bridging.py | 38 +++++++ .../tests/common/dpmodel/test_zbl_bridging.py | 28 +++++ source/tests/common/test_bridging.py | 48 ++++++++ .../tests/pt/model/test_get_model_bridging.py | 102 +++++++++++++++++ .../pt_expt/model/test_get_model_bridging.py | 65 +++++++++++ 13 files changed, 473 insertions(+), 57 deletions(-) diff --git a/deepmd/dpmodel/model/model_factory.py b/deepmd/dpmodel/model/model_factory.py index 9d6b7832eb..561f3950be 100644 --- a/deepmd/dpmodel/model/model_factory.py +++ b/deepmd/dpmodel/model/model_factory.py @@ -132,6 +132,7 @@ def get_linear_atomic_model( 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. @@ -158,13 +159,22 @@ def get_linear_atomic_model( 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, or if a - child is of an unsupported kind. + ``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, @@ -179,7 +189,25 @@ def get_linear_atomic_model( 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] + 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( @@ -192,6 +220,10 @@ def get_linear_atomic_model( f"sibling, but got {len(learned_indices)} sub-models with a " "descriptor." ) + 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]] @@ -206,13 +238,18 @@ def get_linear_atomic_model( if "type_map" not in sub: sub["type_map"] = copy.deepcopy(type_map) if "descriptor" in sub: - descriptor, fitting, _ = get_model_components( - sub, - descriptor_base=descriptor_base, - fitting_base=fitting_base, - backend_name=backend_name, - ) - built[i] = atomic_model(descriptor, fitting, type_map=sub["type_map"]) + 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( @@ -378,7 +415,12 @@ def get_standard_model(self, data: dict) -> Any: backend_name=self.backend_name, ) - def get_linear_atomic_model(self, data: dict) -> Any: + 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.") @@ -389,6 +431,7 @@ def get_linear_atomic_model(self, data: dict) -> Any: 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: diff --git a/deepmd/pt/entrypoints/freeze_pt2.py b/deepmd/pt/entrypoints/freeze_pt2.py index 8c88e64105..e22d640fbb 100644 --- a/deepmd/pt/entrypoints/freeze_pt2.py +++ b/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}." ) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index dffdbee5e1..1083d5e498 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -87,6 +87,9 @@ from deepmd.utils.econf_embd import ( sort_element_type, ) +from deepmd.utils.bridging import ( + is_bridged_sezm_config, +) from deepmd.utils.model_branch_dict import ( get_model_dict, ) @@ -112,6 +115,10 @@ def _is_sezm_model_params(model_params: dict[str, Any]) -> bool: model_type = str(model_params.get("type", "")).lower() if model_type in {"sezm", "dpa4", "sezm_spin"}: return True + # canonical bridged spelling: linear_ener over [dpa4, inner_potential], + # realized by the pt backend as a SeZMModel + if is_bridged_sezm_config(model_params): + return True descriptor = model_params.get("descriptor") if isinstance(descriptor, dict): descriptor_type = str(descriptor.get("type", "")).lower() diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 4095ccb4ee..4c1ba74fe1 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -197,6 +197,13 @@ def get_spin_model(model_params: dict) -> SpinModel: def get_linear_model(model_params: dict) -> BaseModel: + for sub in model_params.get("models", []): + 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 any( sub.get("type") == "inner_potential" for sub in model_params.get("models", []) ): @@ -325,6 +332,20 @@ def _get_bridged_linear_model(model_params: dict) -> BaseModel: "The pt backend implements `inner_potential` bridging only for " f"the DPA4/SeZM descriptor family, but got {descriptor_type!r}." ) + if learned.get("lora") is not None: + raise NotImplementedError( + "`lora` on the learned child of a bridged linear_ener " + "composition is not supported: the pt trainer reads `lora` " + "from the top-level model section only. Use the concise " + '`type: "dpa4"` form with top-level `lora` and ' + "`bridging_method` instead." + ) + learned_type_map = learned.get("type_map", model_params["type_map"]) + if learned_type_map != model_params["type_map"]: + raise NotImplementedError( + "The pt backend requires the learned child's `type_map` to " + "match the bridged linear_ener composition's `type_map`." + ) inner_cfg = inner_cfgs[0] learned["type"] = "dpa4" learned["type_map"] = copy.deepcopy(model_params["type_map"]) diff --git a/deepmd/pt/model/model/dp_linear_model.py b/deepmd/pt/model/model/dp_linear_model.py index a0a3c6f2c7..7a99b401b8 100644 --- a/deepmd/pt/model/model/dp_linear_model.py +++ b/deepmd/pt/model/model/dp_linear_model.py @@ -349,6 +349,9 @@ def update_sel( 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_type_map = sub_model.get("type_map", type_map) local_jdata_cpy["models"][idx], temp_min = DPModelCommon.update_sel( diff --git a/deepmd/pt_expt/model/dp_linear_model.py b/deepmd/pt_expt/model/dp_linear_model.py index 1ae8255f84..85cc927f4e 100644 --- a/deepmd/pt_expt/model/dp_linear_model.py +++ b/deepmd/pt_expt/model/dp_linear_model.py @@ -224,6 +224,9 @@ def update_sel( 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] diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index 989c59c505..0d87f25af3 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -280,6 +280,28 @@ def get_native_spin_model(data: dict) -> NativeSpinEnergyModel: return NativeSpinEnergyModel(atomic_model_=backbone_model.atomic_model, spin=spin) +def _dpa4_family_child_builder(sub: dict) -> "BaseModel | None": + """Route DPA4/SeZM linear children through the family builder. + + A ``linear_ener`` child of the DPA4/SeZM model type must get exactly + the semantics of a standalone ``type: "dpa4"`` model -- the + descriptor/fitting type defaults, the exclusion consistency check, and + the loud rejections of unsupported options (``lora``, ``use_compile``, + ``preset_out_bias``) -- instead of the generic component build that + would silently ignore them. Returns ``None`` for non-DPA4-family + children so the shared builder uses its generic path. + + Parameters + ---------- + sub : dict + The sub-model config (``type_map`` and any derived clamp radii + already injected by the shared linear builder). + """ + if str(sub.get("type", "standard")) not in ("dpa4", "DPA4", "sezm", "SeZM"): + return None + return get_sezm_model(sub).atomic_model + + def get_linear_model(model_params: dict) -> BaseModel: """Get a linear energy model from a ``linear_ener`` config dictionary. @@ -320,7 +342,10 @@ def get_linear_model(model_params: dict) -> BaseModel: for sub in model_params["models"]: if "descriptor" in sub: sub["descriptor"]["use_spin"] = use_spin - composed = _model_factory.get_linear_atomic_model(model_params) + composed = _model_factory.get_linear_atomic_model( + model_params, + descriptor_child_builder=_dpa4_family_child_builder, + ) if spin is not None: if not composed.supports_native_spin(): raise NotImplementedError( diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index cf52d914fa..1ce58ab310 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -3245,7 +3245,10 @@ def model_compression_type_args() -> Variant: hybrid_model_args_plugin = ArgsPlugin() -def model_args(exclude_hybrid: bool = False) -> list[Argument]: +def model_args( + exclude_hybrid: bool = False, + extra_model_types: "list[Argument] | None" = None, +) -> list[Argument]: doc_type_map = "A list of strings. Give the name to each type of atoms. It is noted that the number of atom type of training system must be less than 128 in a GPU environment. If not given, type.raw in each system should use the same type indexes, and type_map.raw will take no effect." doc_data_stat_nbatch = "The model determines the normalization from the statistics of the data. This key specifies the number of `frames` in each `system` used for statistics." doc_data_stat_protect = "Protect parameter for atomic energy regression." @@ -3387,6 +3390,7 @@ def model_args(exclude_hybrid: bool = False) -> list[Argument]: [ *model_args_plugin.get_all_argument(), *hybrid_models, + *(extra_model_types or []), ], optional=True, default_tag="standard", @@ -3443,6 +3447,7 @@ def standard_model_args() -> Argument: default={}, doc=supported_backends("pt", "jax", "pd", "pt_expt", "tf2") + doc_info, ), + *_bridging_method_args(), ], doc=supported_backends("tf", "pt", "jax", "pd", "pt_expt", "tf2") + "Standard model, which contains a descriptor and a fitting.", @@ -3450,6 +3455,53 @@ def standard_model_args() -> Argument: return ca +def _bridging_method_args() -> list[Argument]: + """The concise analytical-bridging arguments, shared by the model types + that accept the ``bridging_method`` sugar (``dpa4`` and ``standard``). + """ + doc_bridging_method = ( + "Short-range bridging method. Currently supports 'ZBL'. " + "The value is case-insensitive; set it to 'None' to disable bridging. " + "This concise form is the recommended interface; it expands to the " + "equivalent explicit `linear_ener` composition over the learned " + "model and an `inner_potential` sub-model." + ) + doc_bridging_r_inner = ( + "Inner clamping radius in Å. ML descriptor distances below this radius are frozen. " + "Only used when `bridging_method` is enabled. " + "For ZBL bridging, set `training.training_data.min_pair_dist` to the same value " + "so frames with atom pairs closer than `bridging_r_inner` are skipped during training." + ) + doc_bridging_r_outer = ( + "Outer clamping radius in Å. The transition zone " + "`[bridging_r_inner, bridging_r_outer]` uses a C^3-continuous " + "septic Hermite polynomial. Only used when `bridging_method` is enabled." + ) + return [ + Argument( + "bridging_method", + str, + optional=True, + default="None", + doc=supported_backends("pt", "pt_expt") + doc_bridging_method, + ), + Argument( + "bridging_r_inner", + float, + optional=True, + default=0.5, + doc=supported_backends("pt", "pt_expt") + doc_bridging_r_inner, + ), + Argument( + "bridging_r_outer", + float, + optional=True, + default=0.8, + doc=supported_backends("pt", "pt_expt") + doc_bridging_r_outer, + ), + ] + + @model_args_plugin.register( "dpa4", alias=["DPA4", "SeZM", "sezm"], @@ -3486,24 +3538,6 @@ def sezm_model_args() -> Argument: "TF32 is controlled separately by `validating.tf32_infer` or " "`DP_TF32_INFER`." ) - doc_bridging_method = ( - "Short-range bridging method. Currently supports 'ZBL'. " - "The value is case-insensitive; set it to 'None' to disable bridging. " - "This concise form is the recommended interface; it expands to the " - "equivalent explicit `linear_ener` composition over the learned " - "model and an `inner_potential` sub-model." - ) - doc_bridging_r_inner = ( - "Inner clamping radius in Å. ML descriptor distances below this radius are frozen. " - "Only used when `bridging_method` is enabled. " - "For ZBL bridging, set `training.training_data.min_pair_dist` to the same value " - "so frames with atom pairs closer than `bridging_r_inner` are skipped during training." - ) - doc_bridging_r_outer = ( - "Outer clamping radius in Å. The transition zone " - "`[bridging_r_inner, bridging_r_outer]` uses a C^3-continuous " - "septic Hermite polynomial. Only used when `bridging_method` is enabled." - ) doc_lora_rank = "LoRA rank; adapters are injected on every SO3Linear and SO2Linear." doc_lora_alpha = ( "LoRA scaling numerator; effective scaling is alpha / rank. " @@ -3603,27 +3637,7 @@ def sezm_model_args() -> Argument: default={}, doc=supported_backends("pt", "pt_expt") + doc_info, ), - Argument( - "bridging_method", - str, - optional=True, - default="None", - doc=supported_backends("pt", "pt_expt") + doc_bridging_method, - ), - Argument( - "bridging_r_inner", - float, - optional=True, - default=0.5, - doc=supported_backends("pt", "pt_expt") + doc_bridging_r_inner, - ), - Argument( - "bridging_r_outer", - float, - optional=True, - default=0.8, - doc=supported_backends("pt", "pt_expt") + doc_bridging_r_outer, - ), + *_bridging_method_args(), Argument( "lora", dict, @@ -3707,8 +3721,11 @@ def pairtab_model_args() -> Argument: return ca -@model_args_plugin.register("inner_potential") def inner_potential_model_args() -> Argument: + """Child-only model type: NOT registered in ``model_args_plugin``, so + ``model.type: "inner_potential"`` is rejected at the top level; it is + injected only into the ``linear_ener`` ``models`` variant. + """ doc_mode = ( "The analytical pair-potential formula. Currently supports 'zbl' " "(case-insensitive)." @@ -3750,7 +3767,11 @@ def linear_ener_model_args() -> Argument: 'If "sum", the weights are set to be 1.' ) doc_shared_dict = "The definition of the shared parameters used in the `models` within linear model." - models_args = model_args(exclude_hybrid=True) + models_args = model_args( + exclude_hybrid=True, + # child-only model type: valid inside `models`, rejected at top level + extra_model_types=[inner_potential_model_args()], + ) models_args.name = "models" models_args.fold_subdoc = True models_args.set_dtype(list) diff --git a/deepmd/utils/bridging.py b/deepmd/utils/bridging.py index 90576188bf..7f2bda605c 100644 --- a/deepmd/utils/bridging.py +++ b/deepmd/utils/bridging.py @@ -26,8 +26,46 @@ __all__ = [ "expand_bridging_method", + "is_bridged_sezm_config", ] +_DPA4_FAMILY_TYPES = ("dpa4", "sezm") + + +def is_bridged_sezm_config(data: dict) -> bool: + """Return whether a config is a bridged DPA4/SeZM linear composition. + + True for a ``linear_ener`` config whose children contain an + ``inner_potential`` sub-model and a DPA4/SeZM-family learned + sub-model -- the canonical form the ``bridging_method`` sugar expands + to. Checkpoint consumers that route DPA4/SeZM models specially (e.g. + the ``.pt2`` freeze path) must recognize this shape too, because the + persisted model params keep the canonical spelling while the pt + backend realizes it as a ``SeZMModel``. + + Parameters + ---------- + data : dict + The model section of a training config. + """ + if str(data.get("type", "standard")).lower() != "linear_ener": + return False + children = [sub for sub in (data.get("models") or []) if isinstance(sub, dict)] + if not any(sub.get("type") == "inner_potential" for sub in children): + return False + + def _is_dpa4_family(sub: dict) -> bool: + if str(sub.get("type", "standard")).lower() in _DPA4_FAMILY_TYPES: + return True + descriptor = sub.get("descriptor") + return ( + isinstance(descriptor, dict) + and str(descriptor.get("type", "")).lower() in _DPA4_FAMILY_TYPES + ) + + return any(_is_dpa4_family(sub) for sub in children) + + # Top-level keys that belong to the composition, not to the learned child. _COMPOSITION_KEYS = ( "type", diff --git a/source/tests/common/dpmodel/test_zbl_bridging.py b/source/tests/common/dpmodel/test_zbl_bridging.py index 126b74337c..38e2ae592c 100644 --- a/source/tests/common/dpmodel/test_zbl_bridging.py +++ b/source/tests/common/dpmodel/test_zbl_bridging.py @@ -733,3 +733,31 @@ def test_standard_builder_rejects_the_flag(self) -> None: with pytest.raises(ValueError, match="bridging_method"): get_standard_model(copy.deepcopy(ZBL_CONFIG)) + + +class TestCanonicalCompositionGuards: + """Fail-fast guards of the shared linear builder.""" + + def test_mean_weights_with_inner_child_raise(self) -> None: + """`weights: "mean"` would silently halve both energy terms.""" + cfg = _canonical_config() + cfg["weights"] = "mean" + with pytest.raises(ValueError, match="sum"): + get_model(cfg) + + def test_nested_bridging_flag_on_child_raises(self) -> None: + """A `bridging_method` flag on a linear child must not be dropped.""" + cfg = _canonical_config() + cfg["models"] = [cfg["models"][0]] + cfg["models"][0]["bridging_method"] = "ZBL" + with pytest.raises(ValueError, match="sub-model"): + get_model(cfg) + + def test_inner_child_with_descriptor_raises_cleanly(self) -> None: + """A child carrying both `type: inner_potential` and a descriptor is + a configuration error, not a KeyError. + """ + cfg = _canonical_config() + cfg["models"][1]["descriptor"] = {"type": "dpa4"} + with pytest.raises(ValueError, match="must not carry"): + get_model(cfg) diff --git a/source/tests/common/test_bridging.py b/source/tests/common/test_bridging.py index 88dcf37e53..2476eb4448 100644 --- a/source/tests/common/test_bridging.py +++ b/source/tests/common/test_bridging.py @@ -160,3 +160,51 @@ def test_standard_type_is_supported() -> None: data["type"] = "standard" out = expand_bridging_method(data) assert out["models"][0]["type"] == "standard" + + +class TestIsBridgedSezmConfig: + """The canonical-shape predicate used by pt checkpoint consumers.""" + + def _canonical(self) -> dict: + return { + "type": "linear_ener", + "weights": "sum", + "type_map": ["Ni", "O"], + "models": [ + {"type": "dpa4", "descriptor": {"type": "dpa4"}}, + {"type": "inner_potential", "mode": "zbl"}, + ], + } + + def test_canonical_shape_is_recognized(self) -> None: + from deepmd.utils.bridging import is_bridged_sezm_config + + assert is_bridged_sezm_config(self._canonical()) + + def test_descriptor_type_alone_is_recognized(self) -> None: + from deepmd.utils.bridging import is_bridged_sezm_config + + cfg = self._canonical() + cfg["models"][0] = {"type": "standard", "descriptor": {"type": "dpa4"}} + assert is_bridged_sezm_config(cfg) + + def test_non_linear_type_is_not(self) -> None: + from deepmd.utils.bridging import is_bridged_sezm_config + + cfg = self._canonical() + cfg["type"] = "dpa4" + assert not is_bridged_sezm_config(cfg) + + def test_linear_without_inner_child_is_not(self) -> None: + from deepmd.utils.bridging import is_bridged_sezm_config + + cfg = self._canonical() + cfg["models"] = [cfg["models"][0]] + assert not is_bridged_sezm_config(cfg) + + def test_non_dpa4_learned_child_is_not(self) -> None: + from deepmd.utils.bridging import is_bridged_sezm_config + + cfg = self._canonical() + cfg["models"][0] = {"type": "standard", "descriptor": {"type": "se_e2_a"}} + assert not is_bridged_sezm_config(cfg) diff --git a/source/tests/pt/model/test_get_model_bridging.py b/source/tests/pt/model/test_get_model_bridging.py index d7abe5a55c..a3dd0974fa 100644 --- a/source/tests/pt/model/test_get_model_bridging.py +++ b/source/tests/pt/model/test_get_model_bridging.py @@ -174,3 +174,105 @@ def test_plain_linear_ener_is_unaffected() -> None: } model = get_model(cfg) assert not isinstance(model, SeZMModel) + + +def test_canonical_rejects_lora_on_child() -> None: + """The pt trainer reads `lora` from the top level only; a child-level + `lora` must fail fast instead of silently training without adapters. + """ + cfg = _canonical_config() + cfg["models"][0]["lora"] = {"rank": 2} + with pytest.raises(NotImplementedError, match="lora"): + get_model(cfg) + + +def test_canonical_rejects_mismatched_child_type_map() -> None: + """An explicit child type_map that differs from the composition's must + not be silently overwritten. + """ + cfg = _canonical_config() + cfg["models"][0]["type_map"] = ["O", "Ni"] + with pytest.raises(NotImplementedError, match="type_map"): + get_model(cfg) + + +def test_nested_bridging_flag_on_child_raises() -> None: + """A `bridging_method` flag on a linear child must not be dropped.""" + cfg = _canonical_config() + cfg["models"] = [cfg["models"][0]] + cfg["models"][0]["bridging_method"] = "ZBL" + with pytest.raises(ValueError, match="sub-model"): + get_model(cfg) + + +def test_deep_eval_recognizes_canonical_params() -> None: + """`_is_sezm_model_params` must route the canonical bridged spelling + like the flag spelling (both realize a SeZMModel). + """ + from deepmd.pt.infer.deep_eval import ( + _is_sezm_model_params, + ) + + assert _is_sezm_model_params(_canonical_config()) + assert not _is_sezm_model_params( + { + "type": "linear_ener", + "models": [ + {"descriptor": {"type": "se_atten"}}, + {"descriptor": {"type": "se_atten"}}, + ], + } + ) + + +def test_is_sezm_checkpoint_recognizes_canonical_params(tmp_path) -> None: + """The `.pt2` freeze router must recognize a checkpoint whose persisted + model params keep the canonical bridged spelling. + """ + import torch + + from deepmd.pt.entrypoints.freeze_pt2 import ( + is_sezm_checkpoint, + ) + + ckpt = str(tmp_path / "canonical.pt") + torch.save({"model": {"_extra_state": {"model_params": _canonical_config()}}}, ckpt) + assert is_sezm_checkpoint(ckpt) + ckpt2 = str(tmp_path / "multitask.pt") + torch.save( + { + "model": { + "_extra_state": { + "model_params": {"model_dict": {"branch": _canonical_config()}} + } + } + }, + ckpt2, + ) + assert is_sezm_checkpoint(ckpt2) + + +def test_update_sel_skips_inner_potential_child(monkeypatch) -> None: + """Neighbor-stat selection must skip the analytical child instead of + crashing on its missing descriptor. + """ + from deepmd.pt.model.model import ( + LinearEnergyModel, + ) + from deepmd.pt.model.model.dp_model import ( + DPModelCommon, + ) + + 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 = _canonical_config() + updated, min_dist = LinearEnergyModel.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" diff --git a/source/tests/pt_expt/model/test_get_model_bridging.py b/source/tests/pt_expt/model/test_get_model_bridging.py index 1cd5c5bb64..545989fbc7 100644 --- a/source/tests/pt_expt/model/test_get_model_bridging.py +++ b/source/tests/pt_expt/model/test_get_model_bridging.py @@ -238,3 +238,68 @@ def test_canonical_native_spin_composition() -> None: model = get_model(cfg) assert isinstance(model, NativeSpinEnergyModel) assert isinstance(model.atomic_model, LinearEnergyAtomicModel) + + +def test_canonical_requires_sum_weights() -> None: + """`weights: "mean"` would silently halve both energy terms.""" + cfg = _canonical_config() + cfg["weights"] = "mean" + with pytest.raises(ValueError, match="sum"): + get_model(cfg) + + +def test_canonical_rejects_lora_on_child() -> None: + """The DPA4-family child routes through get_sezm_model, so unsupported + options are rejected loudly instead of silently ignored. + """ + cfg = _canonical_config() + cfg["models"][0]["lora"] = {"rank": 2} + with pytest.raises(NotImplementedError, match="lora"): + get_model(cfg) + + +def test_canonical_child_gets_dpa4_defaults() -> None: + """A dpa4-family child may omit the descriptor/fitting `type` keys; the + family builder fills them like a standalone `type: "dpa4"` model. + """ + cfg = _canonical_config() + del cfg["models"][0]["descriptor"]["type"] + del cfg["models"][0]["fitting_net"]["type"] + model = get_model(cfg) + assert isinstance(model.atomic_model, LinearEnergyAtomicModel) + assert model.atomic_model.models[0].descriptor.bridging_switch is not None + + +def test_nested_bridging_flag_on_child_raises() -> None: + """A `bridging_method` flag on a linear child must not be dropped.""" + cfg = _canonical_config() + cfg["models"] = [cfg["models"][0]] + cfg["models"][0]["bridging_method"] = "ZBL" + with pytest.raises(ValueError, match="sub-model"): + get_model(cfg) + + +def test_update_sel_skips_inner_potential_child(monkeypatch) -> None: + """Neighbor-stat selection must skip the analytical child instead of + crashing on its missing descriptor. + """ + from deepmd.dpmodel.model.dp_model import ( + DPModelCommon, + ) + from deepmd.pt_expt.model.dp_linear_model import ( + LinearEnergyModel, + ) + + 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 = _canonical_config() + updated, min_dist = LinearEnergyModel.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" From 508ac203a6b416622138d7b01f07cf0d7984eee0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 11 Aug 2026 02:14:37 +0000 Subject: [PATCH 06/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/infer/deep_eval.py | 6 +++--- source/tests/common/test_bridging.py | 20 +++++++++++++++----- 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 1083d5e498..f468379565 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -84,12 +84,12 @@ from deepmd.utils.batch_size import ( RetrySignal, ) -from deepmd.utils.econf_embd import ( - sort_element_type, -) from deepmd.utils.bridging import ( is_bridged_sezm_config, ) +from deepmd.utils.econf_embd import ( + sort_element_type, +) from deepmd.utils.model_branch_dict import ( get_model_dict, ) diff --git a/source/tests/common/test_bridging.py b/source/tests/common/test_bridging.py index 2476eb4448..52ec8ba900 100644 --- a/source/tests/common/test_bridging.py +++ b/source/tests/common/test_bridging.py @@ -177,33 +177,43 @@ def _canonical(self) -> dict: } def test_canonical_shape_is_recognized(self) -> None: - from deepmd.utils.bridging import is_bridged_sezm_config + from deepmd.utils.bridging import ( + is_bridged_sezm_config, + ) assert is_bridged_sezm_config(self._canonical()) def test_descriptor_type_alone_is_recognized(self) -> None: - from deepmd.utils.bridging import is_bridged_sezm_config + from deepmd.utils.bridging import ( + is_bridged_sezm_config, + ) cfg = self._canonical() cfg["models"][0] = {"type": "standard", "descriptor": {"type": "dpa4"}} assert is_bridged_sezm_config(cfg) def test_non_linear_type_is_not(self) -> None: - from deepmd.utils.bridging import is_bridged_sezm_config + from deepmd.utils.bridging import ( + is_bridged_sezm_config, + ) cfg = self._canonical() cfg["type"] = "dpa4" assert not is_bridged_sezm_config(cfg) def test_linear_without_inner_child_is_not(self) -> None: - from deepmd.utils.bridging import is_bridged_sezm_config + from deepmd.utils.bridging import ( + is_bridged_sezm_config, + ) cfg = self._canonical() cfg["models"] = [cfg["models"][0]] assert not is_bridged_sezm_config(cfg) def test_non_dpa4_learned_child_is_not(self) -> None: - from deepmd.utils.bridging import is_bridged_sezm_config + from deepmd.utils.bridging import ( + is_bridged_sezm_config, + ) cfg = self._canonical() cfg["models"][0] = {"type": "standard", "descriptor": {"type": "se_e2_a"}} From 1abe20dfb80946aeaf2c2699f8e20e18c1f40a4f Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 11 Aug 2026 12:11:33 +0800 Subject: [PATCH 07/11] fix(pt_expt): route descriptor-typed DPA4 linear children through the family builder A child spelled 'type: standard' with a DPA4/SeZM descriptor (the shape the sugar on 'type: standard' expands to) bypassed get_sezm_model and so lost the family defaults and the loud lora/use_compile/ preset_out_bias rejections. The child-builder hook now keys on the descriptor type too, matching the pt backend's routing. --- deepmd/pt_expt/model/get_model.py | 8 +++++++- .../tests/pt_expt/model/test_get_model_bridging.py | 12 ++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index 0d87f25af3..5f5d2f817a 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -297,7 +297,13 @@ def _dpa4_family_child_builder(sub: dict) -> "BaseModel | None": The sub-model config (``type_map`` and any derived clamp radii already injected by the shared linear builder). """ - if str(sub.get("type", "standard")) not in ("dpa4", "DPA4", "sezm", "SeZM"): + family_types = ("dpa4", "sezm") + model_type = str(sub.get("type", "standard")).lower() + descriptor = sub.get("descriptor") + descriptor_type = ( + str(descriptor.get("type", "")).lower() if isinstance(descriptor, dict) else "" + ) + if model_type not in family_types and descriptor_type not in family_types: return None return get_sezm_model(sub).atomic_model diff --git a/source/tests/pt_expt/model/test_get_model_bridging.py b/source/tests/pt_expt/model/test_get_model_bridging.py index 545989fbc7..c52e11ba37 100644 --- a/source/tests/pt_expt/model/test_get_model_bridging.py +++ b/source/tests/pt_expt/model/test_get_model_bridging.py @@ -303,3 +303,15 @@ 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_descriptor_typed_child_routes_through_family_builder() -> None: + """A child of `type: "standard"` with a DPA4 descriptor is a DPA4-family + child (the sugar on `type: "standard"` expands to this shape), so it must + get the family builder's rejections, not the silent generic path. + """ + cfg = _canonical_config() + cfg["models"][0]["type"] = "standard" + cfg["models"][0]["lora"] = {"rank": 2} + with pytest.raises(NotImplementedError, match="lora"): + get_model(cfg) From d725272644b38f46a5e516cbf8548ea17f8e1e7a Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 12 Aug 2026 17:22:16 +0800 Subject: [PATCH 08/11] fix: three canonical-bridging P1s from review round 2 1. Sugar expansion keeps trainer-owned top-level `lora` at the composition level instead of forwarding it to the learned child (which the pt bridge builder rejects), restoring the concise dpa4+bridging_method+lora form. The key routing is now a four-way table (composition / consumed / trainer / learned-child) whose coverage of the standard+dpa4 argcheck schemas is pinned by test_routing_covers_the_argcheck_schema: a new argcheck model key fails the test until it gets an explicit routing decision, instead of silently landing on the child. The guard immediately surfaced use_compile and enable_tf32 as previously-undecided keys (both routed to the learned child, preserving behavior). 2. pt LinearEnergyModel.update_sel: the shared-config reconstruction loop now skips the `inner_potential` child too. Normalization always inserts `shared_dict: {}`, so the default CLI path entered the reconstruction loop and dereferenced the analytical child's missing descriptor (KeyError). New test runs on a NORMALIZED config. 3. The shared dpmodel/pt_expt linear builder rejects bridged compositions with a third child (e.g. pairtab): no common execution route exists (pairtab is dense-only, the bridged pair graph-only), matching the pt builder's exact-two-child constraint and message. --- deepmd/dpmodel/model/model_factory.py | 9 ++-- deepmd/pt/model/model/dp_linear_model.py | 3 ++ deepmd/utils/bridging.py | 54 ++++++++++++++++--- .../tests/common/dpmodel/test_zbl_bridging.py | 23 ++++++++ source/tests/common/test_bridging.py | 54 +++++++++++++++++++ .../tests/pt/model/test_get_model_bridging.py | 51 ++++++++++++++++++ .../pt_expt/model/test_get_model_bridging.py | 13 +++++ 7 files changed, 198 insertions(+), 9 deletions(-) diff --git a/deepmd/dpmodel/model/model_factory.py b/deepmd/dpmodel/model/model_factory.py index 561f3950be..5aa0475012 100644 --- a/deepmd/dpmodel/model/model_factory.py +++ b/deepmd/dpmodel/model/model_factory.py @@ -214,11 +214,14 @@ def get_linear_atomic_model( "A linear_ener composition supports at most one " "`inner_potential` sub-model." ) - if len(learned_indices) != 1: + 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 " - f"sibling, but got {len(learned_indices)} sub-models with a " - "descriptor." + "sibling: expected a linear_ener composition over " + "[learned, inner_potential]." ) if str(data.get("weights", "mean")) != "sum": raise ValueError( diff --git a/deepmd/pt/model/model/dp_linear_model.py b/deepmd/pt/model/model/dp_linear_model.py index 7a99b401b8..94403166da 100644 --- a/deepmd/pt/model/model/dp_linear_model.py +++ b/deepmd/pt/model/model/dp_linear_model.py @@ -370,6 +370,9 @@ def get_shared_key(shared_ref: str) -> str: if "type_map" not in ret_jdata: ret_jdata["type_map"] = deepcopy(type_map) for idx, original_sub_model in enumerate(original_models): + if original_sub_model.get("type") == "inner_potential": + # analytical child: no descriptor to write back + continue if "tab_file" in original_sub_model: continue updated_sub_model = local_jdata_cpy["models"][idx] diff --git a/deepmd/utils/bridging.py b/deepmd/utils/bridging.py index 7f2bda605c..73901be19b 100644 --- a/deepmd/utils/bridging.py +++ b/deepmd/utils/bridging.py @@ -66,17 +66,57 @@ def _is_dpa4_family(sub: dict) -> bool: return any(_is_dpa4_family(sub) for sub in children) -# Top-level keys that belong to the composition, not to the learned child. +# Routing of the concise-form top-level keys during sugar expansion. Every +# key the `standard`/`dpa4` argcheck schemas declare must appear in exactly +# one tuple below: a schema-coverage test derives the key universe from +# `deepmd.utils.argcheck` and fails when a new key is left unrouted, so +# adding a model key forces an explicit routing decision here. + +# Keys that belong to the composition, not to the learned child. _COMPOSITION_KEYS = ( "type", "type_map", "spin", "atom_exclude_types", "pair_exclude_types", +) +# Keys consumed by the expansion itself; they appear in neither the +# composition nor the learned child. +_CONSUMED_KEYS = ( "bridging_method", "bridging_r_inner", "bridging_r_outer", ) +# Training-owned keys: the trainer reads them from the top level of the +# model section, so they stay at the composition level and must never be +# forwarded to a sub-model. +_TRAINER_KEYS = ("lora",) +# Keys that configure the learned model and are forwarded to the learned +# child. This tuple is not consulted at expansion time (the child receives +# every key not routed above); it exists so the schema-coverage test can +# assert that every argcheck key has an explicit routing decision. +_LEARNED_CHILD_KEYS = ( + "descriptor", + "fitting_net", + "model_branch_alias", + "info", + "use_compile", + "enable_tf32", + "data_stat_nbatch", + "data_stat_protect", + "data_bias_nsample", + "use_srtab", + "smin_alpha", + "sw_rmin", + "sw_rmax", + "preset_out_bias", + "srtab_add_bias", + "type_embedding", + "modifier", + "compress", + "finetune_head", +) +_NON_CHILD_KEYS = _COMPOSITION_KEYS + _CONSUMED_KEYS + _TRAINER_KEYS def expand_bridging_method(data: dict) -> dict: @@ -87,8 +127,9 @@ def expand_bridging_method(data: dict) -> dict: deep-copied and rewritten to the canonical composition form: a ``linear_ener`` model with ``weights: "sum"`` over the learned sub-model and an ``inner_potential`` sub-model. The exclusion lists - move to the composition level; a top-level ``spin`` section stays at - the top level; every other key stays on the learned child. + move to the composition level; a top-level ``spin`` section and the + training-owned keys (``lora``) stay at the top level; every other key + stays on the learned child. For backward compatibility with the legacy pt ``type: "dpa4"`` builder, ``descriptor.exclude_types`` is promoted to the composition's @@ -142,9 +183,7 @@ def expand_bridging_method(data: dict) -> dict: else: pair_exclude_types = descriptor_exclude_types - learned = { - key: value for key, value in data.items() if key not in _COMPOSITION_KEYS - } + learned = {key: value for key, value in data.items() if key not in _NON_CHILD_KEYS} learned["type"] = model_type learned["type_map"] = copy.deepcopy(data["type_map"]) canonical = { @@ -165,4 +204,7 @@ def expand_bridging_method(data: dict) -> dict: } if "spin" in data: canonical["spin"] = data["spin"] + for key in _TRAINER_KEYS: + if key in data: + canonical[key] = data[key] return canonical diff --git a/source/tests/common/dpmodel/test_zbl_bridging.py b/source/tests/common/dpmodel/test_zbl_bridging.py index 38e2ae592c..9f78a0bc79 100644 --- a/source/tests/common/dpmodel/test_zbl_bridging.py +++ b/source/tests/common/dpmodel/test_zbl_bridging.py @@ -76,6 +76,29 @@ def test_builder_composes_linear_model(): assert float(dp_child.descriptor.inner_clamp.r_inner) == 0.8 +def test_third_child_without_common_route_raises(): + """[learned, inner_potential, pairtab] has no common execution route + (pairtab is dense-only, the bridged pair is graph-only): the builder + must reject it at construction like the pt backend does. + """ + cfg = { + "type": "linear_ener", + "weights": "sum", + "type_map": ["Ni", "O"], + "models": [ + { + "type": "dpa4", + "descriptor": copy.deepcopy(ZBL_CONFIG["descriptor"]), + "fitting_net": copy.deepcopy(ZBL_CONFIG["fitting_net"]), + }, + {"type": "inner_potential", "mode": "ZBL"}, + {"type": "pairtab", "tab_file": "unused.txt", "rcut": 4.0, "sel": 8}, + ], + } + with pytest.raises(ValueError, match="exactly one learned"): + get_model(cfg) + + def test_zbl_child_equals_composition_minus_learned(): """Composition energy == learned child + analytical child (exact sum).""" model = get_model(copy.deepcopy(ZBL_CONFIG)) diff --git a/source/tests/common/test_bridging.py b/source/tests/common/test_bridging.py index 52ec8ba900..c887b8f94f 100644 --- a/source/tests/common/test_bridging.py +++ b/source/tests/common/test_bridging.py @@ -103,6 +103,60 @@ def test_other_model_keys_stay_on_learned_child() -> None: assert "preset_out_bias" not in out +def test_lora_stays_top_level() -> None: + """`lora` is training-owned: the pt trainer reads it from the top + level of the model section, so the expansion must keep it there and + never forward it to the learned child (which the pt bridge builder + rejects). + """ + data = _flag_config() + data["lora"] = {"rank": 2, "alpha": None} + out = expand_bridging_method(data) + assert out["lora"] == {"rank": 2, "alpha": None} + assert "lora" not in out["models"][0] + + +def test_routing_covers_the_argcheck_schema() -> None: + """Every key the `standard`/`dpa4` argcheck schemas declare must have + an explicit routing decision in the expansion. Adding a model key to + argcheck without deciding its routing fails here instead of silently + landing on the learned child (how top-level `lora` once broke). + """ + from deepmd.utils.argcheck import ( + model_args, + sezm_model_args, + standard_model_args, + ) + from deepmd.utils.bridging import ( + _COMPOSITION_KEYS, + _CONSUMED_KEYS, + _LEARNED_CHILD_KEYS, + _TRAINER_KEYS, + ) + + schema_keys = {"type"} + schema_keys |= set(model_args(exclude_hybrid=True).sub_fields) + schema_keys |= set(standard_model_args().sub_fields) + schema_keys |= set(sezm_model_args().sub_fields) + + routing = [ + set(_COMPOSITION_KEYS), + set(_CONSUMED_KEYS), + set(_TRAINER_KEYS), + set(_LEARNED_CHILD_KEYS), + ] + routed = set().union(*routing) + assert sum(len(s) for s in routing) == len(routed), ( + "a key is routed to more than one destination" + ) + assert schema_keys == routed, ( + f"unrouted argcheck keys: {sorted(schema_keys - routed)}; " + f"routed keys absent from the schema: {sorted(routed - schema_keys)}. " + "Decide the routing in deepmd.utils.bridging and update the " + "corresponding tuple." + ) + + def test_exclusions_move_to_composition_level() -> None: data = _flag_config() data["pair_exclude_types"] = [[0, 1]] diff --git a/source/tests/pt/model/test_get_model_bridging.py b/source/tests/pt/model/test_get_model_bridging.py index a3dd0974fa..a6eb6d7a33 100644 --- a/source/tests/pt/model/test_get_model_bridging.py +++ b/source/tests/pt/model/test_get_model_bridging.py @@ -252,6 +252,57 @@ def test_is_sezm_checkpoint_recognizes_canonical_params(tmp_path) -> None: assert is_sezm_checkpoint(ckpt2) +def test_sugar_with_top_level_lora_builds() -> None: + """The concise dpa4+bridging form with trainer-owned top-level `lora` + must keep building: the expansion routes `lora` to the composition + level, so the bridge builder never sees it on the learned child. + """ + from deepmd.pt.train.training import ( + get_model_for_wrapper, + ) + + cfg = _sugar_config() + cfg["lora"] = {"rank": 2, "alpha": None} + model = get_model_for_wrapper(copy.deepcopy(cfg)) + assert isinstance(model, SeZMModel) + # The trainer injects the adapters later by reading the top level of + # its own (unexpanded) config; expansion must not have mutated it. + assert cfg["lora"] == {"rank": 2, "alpha": None} + + +def test_update_sel_normalized_config_skips_inner_potential_child( + monkeypatch, +) -> None: + """The default CLI path hands `update_sel` a NORMALIZED config, where + argcheck always inserts `shared_dict: {}`. Both the update loop and + the shared-config reconstruction loop must skip the analytical child. + """ + from deepmd.pt.model.model import ( + LinearEnergyModel, + ) + from deepmd.pt.model.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="_*") + assert cfg["shared_dict"] == {} # inserted by normalization + updated, min_dist = LinearEnergyModel.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" + + def test_update_sel_skips_inner_potential_child(monkeypatch) -> None: """Neighbor-stat selection must skip the analytical child instead of crashing on its missing descriptor. diff --git a/source/tests/pt_expt/model/test_get_model_bridging.py b/source/tests/pt_expt/model/test_get_model_bridging.py index c52e11ba37..6302aa1ddb 100644 --- a/source/tests/pt_expt/model/test_get_model_bridging.py +++ b/source/tests/pt_expt/model/test_get_model_bridging.py @@ -194,6 +194,19 @@ def _canonical_config() -> dict: } +def test_third_child_without_common_route_raises() -> None: + """[learned, inner_potential, pairtab] has no common execution route + (pairtab is dense-only, the bridged pair is graph-only): the shared + builder must reject it at construction like the pt backend does. + """ + cfg = _canonical_config() + cfg["models"].append( + {"type": "pairtab", "tab_file": "unused.txt", "rcut": 4.0, "sel": 8} + ) + with pytest.raises(ValueError, match="exactly one learned"): + get_model(cfg) + + def test_canonical_composition_builds() -> None: """The canonical spelling composes [learned, InnerPotential] with the clamp radii derived onto the learned child's descriptor. From 187ad365e3cd7d626239b5ca209508878b35bb01 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Fri, 14 Aug 2026 09:09:07 +0800 Subject: [PATCH 09/11] 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). --- deepmd/dpmodel/model/dp_linear_model.py | 48 +++++++++++++++++ deepmd/dpmodel/model/model_factory.py | 12 +++++ deepmd/pt/model/model/__init__.py | 54 +++++++++++++++---- deepmd/pt_expt/model/get_model.py | 8 +++ deepmd/utils/bridging.py | 42 +++++++++++++++ .../tests/common/dpmodel/test_zbl_bridging.py | 46 ++++++++++++++++ source/tests/common/test_bridging.py | 30 +++++++++++ .../tests/pt/model/test_get_model_bridging.py | 44 +++++++++++++++ .../pt_expt/model/test_get_model_bridging.py | 31 +++++++++++ 9 files changed, 304 insertions(+), 11 deletions(-) diff --git a/deepmd/dpmodel/model/dp_linear_model.py b/deepmd/dpmodel/model/dp_linear_model.py index 3aecf8e911..2255edbede 100644 --- a/deepmd/dpmodel/model/dp_linear_model.py +++ b/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 diff --git a/deepmd/dpmodel/model/model_factory.py b/deepmd/dpmodel/model/model_factory.py index 5aa0475012..08d03fdabb 100644 --- a/deepmd/dpmodel/model/model_factory.py +++ b/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,6 +236,7 @@ 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): @@ -240,6 +244,14 @@ def get_linear_atomic_model( 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: diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 4c1ba74fe1..de3dba3e63 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/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 diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index 5f5d2f817a..d68eb27b0e 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/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, diff --git a/deepmd/utils/bridging.py b/deepmd/utils/bridging.py index 73901be19b..43848be820 100644 --- a/deepmd/utils/bridging.py +++ b/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: diff --git a/source/tests/common/dpmodel/test_zbl_bridging.py b/source/tests/common/dpmodel/test_zbl_bridging.py index 9f78a0bc79..b5e79f68b5 100644 --- a/source/tests/common/dpmodel/test_zbl_bridging.py +++ b/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" diff --git a/source/tests/common/test_bridging.py b/source/tests/common/test_bridging.py index c887b8f94f..d416234cd1 100644 --- a/source/tests/common/test_bridging.py +++ b/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) diff --git a/source/tests/pt/model/test_get_model_bridging.py b/source/tests/pt/model/test_get_model_bridging.py index a6eb6d7a33..e0cbd56704 100644 --- a/source/tests/pt/model/test_get_model_bridging.py +++ b/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) diff --git a/source/tests/pt_expt/model/test_get_model_bridging.py b/source/tests/pt_expt/model/test_get_model_bridging.py index 6302aa1ddb..36b5e9c3bc 100644 --- a/source/tests/pt_expt/model/test_get_model_bridging.py +++ b/source/tests/pt_expt/model/test_get_model_bridging.py @@ -328,3 +328,34 @@ def test_descriptor_typed_child_routes_through_family_builder() -> None: cfg["models"][0]["lora"] = {"rank": 2} with pytest.raises(NotImplementedError, match="lora"): get_model(cfg) + + +def test_canonical_rejects_mismatched_learned_type_map() -> None: + """A remapped learned-child type_map builds a model the graph route + rejects on every forward; the shared builder fails at construction. + """ + 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_expanded_sugar_with_lora_raises() -> None: + """The expansion keeps trainer-owned `lora` at the composition top + level; pt_expt has no LoRA support and must reject it instead of + silently training a plain full model (covers the normalized path, + where absent `lora` normalizes to None and must NOT trigger). + """ + from deepmd.utils.argcheck import ( + model_args, + ) + + cfg = _dpa4_standard_config() + cfg["type"] = "dpa4" + cfg["bridging_method"] = "ZBL" + cfg = model_args().normalize_value(cfg, trim_pattern="_*") + assert cfg.get("lora") is None # normalization default stays buildable + get_model(copy.deepcopy(cfg)) + cfg["lora"] = {"rank": 2} + with pytest.raises(NotImplementedError, match="lora"): + get_model(cfg) From 0cb7b3fc834232065c81c8ae369451ae9c286e50 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Fri, 14 Aug 2026 23:23:36 +0800 Subject: [PATCH 10/11] fix: close the remaining consume-or-reject and guard-parity gaps Self-audit of the linear_ener key x backend x route matrix (the same probe pattern the review rounds used) found four residual holes; all are closed by explicit rejections instead of silent drops: - the shared linear builder (dpmodel/pt_expt) now rejects top-level `lora`, non-empty `shared_dict`, and child-level `lora` -- the pt backend's consumers (trainer lora, linear shared_dict) do not exist in these backends, so all three silently built a different model than the config asked for (empty shared_dict from strict normalization still passes). - the shared builder mirrors pt's DPA4/SeZM family restriction on the bridged learned sibling: other families died on an obscure unknown-kwarg TypeError from the clamp injection. - tf registers linear_ener but cannot build inner_potential children; its update_sel now rejects them explicitly before any neighbor statistics run, instead of an 'unknown model type' dispatch error. - a table-walk test runs EVERY _LEARNED_CHILD_KEYS entry through route_canonical_learned_options (copy-down and conflict branches), so a future per-key special case cannot land untested. All rejections are pinned by tests verified to fail without the guards. --- deepmd/dpmodel/model/model_factory.py | 32 ++++++++++ deepmd/tf/model/linear.py | 8 +++ .../tests/common/dpmodel/test_zbl_bridging.py | 62 +++++++++++++++++++ source/tests/common/test_bridging.py | 19 ++++++ .../pt_expt/model/test_get_model_bridging.py | 10 +++ source/tests/tf/test_linear_model.py | 31 ++++++++++ 6 files changed, 162 insertions(+) diff --git a/deepmd/dpmodel/model/model_factory.py b/deepmd/dpmodel/model/model_factory.py index 08d03fdabb..f0442828ae 100644 --- a/deepmd/dpmodel/model/model_factory.py +++ b/deepmd/dpmodel/model/model_factory.py @@ -204,6 +204,21 @@ def get_linear_atomic_model( "`descriptor`: the analytical term has no learned " "component." ) + # Consume-or-reject: this builder has no consumer for these keys, so + # accepting them silently would train/evaluate a different model than + # the config asks for. (The pt backend consumes top-level `lora` in its + # trainer and `shared_dict` in its own linear builder; this shared + # builder serves backends without either consumer.) + if data.get("lora") is not None: + raise NotImplementedError( + f"`lora` on a linear_ener composition is not supported in the " + f"{backend_name} backend." + ) + if data.get("shared_dict"): + raise NotImplementedError( + f"`shared_dict` is not supported for linear_ener in the " + f"{backend_name} backend." + ) for sub in children: if str(sub.get("bridging_method", "none")).lower() not in ("none", ""): raise ValueError( @@ -211,6 +226,11 @@ def get_linear_atomic_model( "sub-model: add an `inner_potential` sub-model to the " "composition instead." ) + if sub.get("lora") is not None: + raise NotImplementedError( + "`lora` on a linear_ener sub-model is not supported in the " + f"{backend_name} backend." + ) if inner_indices: if len(inner_indices) > 1: raise ValueError( @@ -230,6 +250,18 @@ def get_linear_atomic_model( raise ValueError( 'A bridged linear_ener composition requires `weights: "sum"`.' ) + learned_descriptor_type = str( + children[learned_indices[0]]["descriptor"].get("type", "dpa4") + ) + if learned_descriptor_type not in ("dpa4", "DPA4", "sezm", "SeZM"): + # same family restriction as the pt builder: the clamp window + # below only exists on DPA4/SeZM descriptors, so any other + # family would die on an obscure unknown-kwarg TypeError + raise NotImplementedError( + f"The {backend_name} backend implements `inner_potential` " + "bridging only for the DPA4/SeZM descriptor family, but got " + f"{learned_descriptor_type!r}." + ) # 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]] diff --git a/deepmd/tf/model/linear.py b/deepmd/tf/model/linear.py index 4d147d5169..9534371f7f 100644 --- a/deepmd/tf/model/linear.py +++ b/deepmd/tf/model/linear.py @@ -173,6 +173,14 @@ def update_sel( float The minimum distance between two atoms """ + if any(sub.get("type") == "inner_potential" for sub in local_jdata["models"]): + # reject explicitly (and before any neighbor statistics run): + # the generic dispatch below would only report an obscure + # "unknown model type" for this child + raise NotImplementedError( + "`inner_potential` sub-models (analytical bridging) are " + "not supported in the TensorFlow backend." + ) local_jdata_cpy = local_jdata.copy() new_list = [] min_nbor_dist = None diff --git a/source/tests/common/dpmodel/test_zbl_bridging.py b/source/tests/common/dpmodel/test_zbl_bridging.py index ec8185a6ac..e4379c1b12 100644 --- a/source/tests/common/dpmodel/test_zbl_bridging.py +++ b/source/tests/common/dpmodel/test_zbl_bridging.py @@ -876,3 +876,65 @@ 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" + + +class TestConsumeOrRejectGuards: + """Every key this route accepts is either consumed or loudly rejected; + the pt backend's consumers (trainer `lora`, linear `shared_dict`) do + not exist here, so silence would build a different model than asked. + """ + + def test_top_level_lora_raises(self) -> None: + cfg = _canonical_config() + cfg["lora"] = {"rank": 2} + with pytest.raises(NotImplementedError, match="lora"): + get_model(cfg) + + def test_expanded_sugar_with_lora_raises(self) -> None: + """The sugar expansion keeps trainer-owned `lora` at the top level; + dpmodel has no trainer to consume it. + """ + from deepmd.utils.bridging import ( + expand_bridging_method, + ) + + cfg = copy.deepcopy(ZBL_CONFIG) + cfg["type"] = "dpa4" + cfg["bridging_method"] = "ZBL" + cfg["lora"] = {"rank": 2} + with pytest.raises(NotImplementedError, match="lora"): + get_model(expand_bridging_method(cfg)) + + def test_nonempty_shared_dict_raises(self) -> None: + cfg = _canonical_config() + cfg["shared_dict"] = {"my_descriptor": "descriptor"} + with pytest.raises(NotImplementedError, match="shared_dict"): + get_model(cfg) + + def test_empty_shared_dict_is_fine(self) -> None: + """Strict normalization always inserts `shared_dict: {}`; the + default CLI path must keep building. + """ + cfg = _canonical_config() + cfg["shared_dict"] = {} + get_model(cfg) + + def test_child_level_lora_raises(self) -> None: + cfg = _canonical_config() + cfg["models"][0]["lora"] = {"rank": 2} + with pytest.raises(NotImplementedError, match="lora"): + get_model(cfg) + + def test_non_dpa4_learned_sibling_raises_cleanly(self) -> None: + """Same family restriction as the pt builder: without it the clamp + injection dies on an obscure unknown-kwarg TypeError. + """ + cfg = _canonical_config() + cfg["models"][0]["descriptor"] = { + "type": "se_e2_a", + "rcut": 4.0, + "rcut_smth": 3.5, + "sel": [8, 8], + } + with pytest.raises(NotImplementedError, match="DPA4/SeZM"): + get_model(cfg) diff --git a/source/tests/common/test_bridging.py b/source/tests/common/test_bridging.py index d416234cd1..1893aa4678 100644 --- a/source/tests/common/test_bridging.py +++ b/source/tests/common/test_bridging.py @@ -302,3 +302,22 @@ def test_route_canonical_learned_options_copies_and_conflicts() -> None: learned["data_stat_protect"] = 0.456 with pytest.raises(ValueError, match="data_stat_protect"): route_canonical_learned_options(composition, learned) + + +def test_routing_helper_handles_every_learned_key_uniformly() -> None: + """Walk the WHOLE learned-key table through the canonical-route helper: + each key copies down when the child lacks it and conflicts when the two + levels differ. Pins uniformity, so a future per-key special case in the + helper cannot land untested. + """ + from deepmd.utils.bridging import ( + _LEARNED_CHILD_KEYS, + route_canonical_learned_options, + ) + + for key in _LEARNED_CHILD_KEYS: + learned = {} + route_canonical_learned_options({key: "sentinel-a"}, learned) + assert learned[key] == "sentinel-a", key + with pytest.raises(ValueError, match=key): + route_canonical_learned_options({key: "sentinel-a"}, {key: "sentinel-b"}) diff --git a/source/tests/pt_expt/model/test_get_model_bridging.py b/source/tests/pt_expt/model/test_get_model_bridging.py index 36b5e9c3bc..6a174753b5 100644 --- a/source/tests/pt_expt/model/test_get_model_bridging.py +++ b/source/tests/pt_expt/model/test_get_model_bridging.py @@ -359,3 +359,13 @@ def test_expanded_sugar_with_lora_raises() -> None: cfg["lora"] = {"rank": 2} with pytest.raises(NotImplementedError, match="lora"): get_model(cfg) + + +def test_nonempty_shared_dict_raises() -> None: + """pt_expt has no `shared_dict` consumer for linear compositions: reject + loudly instead of silently building without parameter sharing. + """ + cfg = _canonical_config() + cfg["shared_dict"] = {"my_descriptor": "descriptor"} + with pytest.raises(NotImplementedError, match="shared_dict"): + get_model(cfg) diff --git a/source/tests/tf/test_linear_model.py b/source/tests/tf/test_linear_model.py index 1392a97820..81efd0a21b 100644 --- a/source/tests/tf/test_linear_model.py +++ b/source/tests/tf/test_linear_model.py @@ -156,3 +156,34 @@ def get_loss(self, loss, lr): self.assertEqual(result, "ener-loss") # the original config must be preserved for other consumers self.assertEqual(loss_config, {"type": "ener", "start_pref_e": 1.0}) + + +class TestLinearUpdateSelRejectsInnerPotential(unittest.TestCase): + def test_inner_potential_child_raises_cleanly(self) -> None: + """The TF backend does not implement analytical bridging: the + neighbor-stat phase (the first CLI touchpoint) must say so + explicitly instead of dying on a generic unknown-type dispatch. + """ + from deepmd.tf.model.linear import ( + LinearEnergyModel, + ) + + cfg = { + "type": "linear_ener", + "type_map": ["O", "H"], + "models": [ + { + "type": "standard", + "descriptor": {"type": "se_e2_a", "sel": [10, 10], "rcut": 4.0}, + "fitting_net": {"neuron": [4]}, + }, + { + "type": "inner_potential", + "mode": "ZBL", + "r_inner": 0.5, + "r_outer": 0.8, + }, + ], + } + with self.assertRaises(NotImplementedError): + LinearEnergyModel.update_sel(None, cfg["type_map"], cfg) From edf02fb7c6771eae50eb1a998ab5459844bb4c07 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 15 Aug 2026 12:53:28 +0800 Subject: [PATCH 11/11] fix: resolve canonical-option conflicts against argcheck defaults Strict normalization injects schema defaults on BOTH the composition top level and the learned child before any builder runs, so the previous strict conflict check rejected every explicitly set top-level learned-owned option on the normal CLI path (top-level 0.123 vs the injected child default 0.01) -- the advertised composition-level routing was unusable. route_canonical_learned_options now recovers the lost provenance by comparing each level against the key's argcheck default (collected once from the model schema, with a consistency check): a level holding exactly the default is treated as not explicitly configured and the other level wins; only two explicit non-default values still raise. Known residual ambiguity: explicitly setting a level to the default value is indistinguishable from not setting it. Regression per review: normalize a canonical config with a top-level data_stat_protect=0.123 (asserting the child default injection), then get_model() must build with 0.123; verified to fail before the fix. --- deepmd/utils/bridging.py | 98 ++++++++++++++++--- source/tests/common/test_bridging.py | 24 +++++ .../tests/pt/model/test_get_model_bridging.py | 17 ++++ 3 files changed, 128 insertions(+), 11 deletions(-) diff --git a/deepmd/utils/bridging.py b/deepmd/utils/bridging.py index 43848be820..88467d66ab 100644 --- a/deepmd/utils/bridging.py +++ b/deepmd/utils/bridging.py @@ -121,6 +121,63 @@ def _is_dpa4_family(sub: dict) -> bool: assert not set(_LEARNED_CHILD_KEYS) & set(_NON_CHILD_KEYS) +_NO_DEFAULT = object() +_SCHEMA_DEFAULTS: dict | None = None + + +def _learned_key_schema_defaults() -> dict: + """Collect the argcheck defaults of the learned-owned keys (cached). + + Strict normalization injects these defaults on BOTH the composition + top level and the learned child, erasing the "did the user set this?" + provenance. The conflict resolution in + :func:`route_canonical_learned_options` recovers it by comparing a + value against its schema default: a level holding exactly the default + is treated as not explicitly configured. + + Returns + ------- + dict + Mapping from key name to its argcheck default, for every + ``_LEARNED_CHILD_KEYS`` entry that declares one. + + Raises + ------ + RuntimeError + If a key is declared with two different defaults anywhere in the + model schema: the recovery above then has no single reference + value and must not guess. + """ + global _SCHEMA_DEFAULTS + if _SCHEMA_DEFAULTS is None: + from deepmd.utils.argcheck import ( # deferred: heavy import + model_args, + ) + + defaults: dict = {} + + def _walk(arg: object) -> None: + for field in getattr(arg, "sub_fields", {}).values(): + if field.name in _LEARNED_CHILD_KEYS and field.optional: + if field.name in defaults and defaults[field.name] != ( + field.default + ): + raise RuntimeError( + f"`{field.name}` is declared with inconsistent " + "argcheck defaults; the canonical-route conflict " + "resolution relies on a single one." + ) + defaults[field.name] = field.default + _walk(field) + for variant in getattr(arg, "sub_variants", {}).values(): + for choice in variant.choice_dict.values(): + _walk(choice) + + _walk(model_args()) + _SCHEMA_DEFAULTS = defaults + return _SCHEMA_DEFAULTS + + def route_canonical_learned_options(composition: dict, learned: dict) -> None: """Route learned-model options from a canonical composition to its child. @@ -129,9 +186,18 @@ def route_canonical_learned_options(composition: dict, learned: dict) -> None: 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. + place) when the child does not set it. + + When the two levels disagree, the argcheck default decides: strict + normalization injects defaults on both levels, so a level holding + exactly the schema default is treated as not explicitly configured + and the other level wins (in particular, a user-set top-level value + survives the child default injected on the normal CLI path). Only two + explicitly configured (non-default) values raise — a silent drop or + a silent override there would unpin the ownership contract. The one + unrecoverable ambiguity: explicitly setting a level to exactly the + default value is indistinguishable from not setting it, and loses to + an explicit non-default on the other level. Parameters ---------- @@ -143,20 +209,30 @@ def route_canonical_learned_options(composition: dict, learned: dict) -> None: Raises ------ ValueError - If a learned-owned key is set at both levels with different - values. + If a learned-owned key is set to two different non-default values + at the two levels. """ 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." - ) + default = _learned_key_schema_defaults().get(key, _NO_DEFAULT) + if learned[key] == default: + # argcheck-injected child default: the explicit + # top-level value wins + learned[key] = copy.deepcopy(composition[key]) + elif composition[key] == default: + # top-level default: the explicit child value wins + pass + else: + 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]) diff --git a/source/tests/common/test_bridging.py b/source/tests/common/test_bridging.py index 1893aa4678..24480ffddf 100644 --- a/source/tests/common/test_bridging.py +++ b/source/tests/common/test_bridging.py @@ -321,3 +321,27 @@ def test_routing_helper_handles_every_learned_key_uniformly() -> None: assert learned[key] == "sentinel-a", key with pytest.raises(ValueError, match=key): route_canonical_learned_options({key: "sentinel-a"}, {key: "sentinel-b"}) + + +def test_routing_resolves_argcheck_default_conflicts() -> None: + """Strict normalization injects schema defaults on BOTH levels, erasing + the set-by-user provenance; the helper recovers it by comparing against + the argcheck default, so only two explicit non-default values conflict. + """ + from deepmd.utils.bridging import ( + route_canonical_learned_options, + ) + + # child holds the injected default -> the explicit top-level value wins + learned = {"data_stat_protect": 0.01} + route_canonical_learned_options({"data_stat_protect": 0.123}, learned) + assert learned["data_stat_protect"] == 0.123 + # the top level holds the injected default -> the explicit child wins + learned = {"data_stat_protect": 0.123} + route_canonical_learned_options({"data_stat_protect": 0.01}, learned) + assert learned["data_stat_protect"] == 0.123 + # two explicit non-default values are a REAL conflict + with pytest.raises(ValueError, match="data_stat_protect"): + route_canonical_learned_options( + {"data_stat_protect": 0.2}, {"data_stat_protect": 0.3} + ) diff --git a/source/tests/pt/model/test_get_model_bridging.py b/source/tests/pt/model/test_get_model_bridging.py index e0cbd56704..5c851f111d 100644 --- a/source/tests/pt/model/test_get_model_bridging.py +++ b/source/tests/pt/model/test_get_model_bridging.py @@ -371,3 +371,20 @@ def test_canonical_spin_rejects_mismatched_pair_exclusions(scheme: str) -> None: } with pytest.raises(ValueError, match="must match"): get_model(cfg) + + +def test_normalized_canonical_top_level_option_survives_child_defaults() -> None: + """The normal CLI path normalizes BEFORE building, which injects the + schema default on the learned child; an explicit top-level value must + survive that injection instead of being rejected as a conflict. + """ + from deepmd.utils.argcheck import ( + model_args, + ) + + cfg = _canonical_config() + cfg["data_stat_protect"] = 0.123 + cfg = model_args().normalize_value(cfg, trim_pattern="_*") + assert cfg["models"][0]["data_stat_protect"] == 0.01 # injected default + model = get_model(cfg) + assert model.atomic_model.data_stat_protect == 0.123