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.py b/deepmd/dpmodel/model/model.py index 2d378525dc..90e30b4b1c 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -31,6 +31,9 @@ from deepmd.dpmodel.model.spin_model import ( SpinModel, ) +from deepmd.utils.bridging import ( + expand_bridging_method, +) from deepmd.utils.spin import ( Spin, normalize_spin_use_spin, @@ -58,50 +61,67 @@ def get_standard_model(data: dict) -> BaseModel: data : dict The data to construct the model. """ - data = copy.deepcopy(data) - # Analytical bridging (e.g. ZBL): the radii feed the DESCRIPTOR's - # InnerClamp/BridgingSwitch (mirrors pt's builder); the method builds the - # atomic model's InnerPotential below. bridging_method = str(data.get("bridging_method", "none")) - bridging_enabled = bridging_method.lower() not in ("none", "") - if bridging_enabled: - data["descriptor"]["inner_clamp_r_inner"] = data.get("bridging_r_inner", 0.5) - data["descriptor"]["inner_clamp_r_outer"] = data.get("bridging_r_outer", 0.8) - model = _model_factory.get_standard_model(data) - if not bridging_enabled: - return model - - descriptor = model.atomic_model.descriptor - atom_exclude_types = data.get("atom_exclude_types", []) - pair_exclude_types = data.get("pair_exclude_types", []) - # Composition, not a flag (first-principles design): the analytical - # bridging term is its own atomic model, summed with the learned one by the - # existing linear composition machinery. - from deepmd.dpmodel.atomic_model.inner_potential import ( - InnerPotentialAtomicModel, - ) - from deepmd.dpmodel.atomic_model.linear_atomic_model import ( - LinearEnergyAtomicModel, - ) + if bridging_method.lower() not in ("none", ""): + raise ValueError( + "`bridging_method` is not supported for a standard model: " + "analytical bridging builds a linear composition, not a " + "standard model. Route the config through `get_model` (which " + "expands the flag), or spell the composition explicitly with " + '`type: "linear_ener"` and an `inner_potential` sub-model.' + ) + return _model_factory.get_standard_model(data) + + +def get_linear_model(data: dict) -> BaseModel: + """Build a linear energy model from a ``linear_ener`` config. + + Children with a ``descriptor`` build as standard learned atomic + models; ``pairtab`` children build as pair-tabulation atomic models; + an ``inner_potential`` child builds the analytical bridging term. The + composition is the ONE owner of the bridging coupling: it derives the + learned sibling descriptor's ``inner_clamp_r_inner``/``_outer`` from + the ``inner_potential`` child's ``r_inner``/``r_outer``, so the radii + are written once in the config (issue #5948, task 2). + + A top-level ``spin`` section (scheme ``native``) wraps the composed + atomic model as a :class:`NativeSpinEnergyModel`, with ``use_spin`` + injected into every learned child's descriptor. + + Parameters + ---------- + data : dict + The model configuration. + """ from deepmd.dpmodel.model.dp_linear_model import ( LinearEnergyModel, ) - zbl_atomic = InnerPotentialAtomicModel( - type_map=data["type_map"], - mode=bridging_method, - rcut=descriptor.get_rcut(), - sel=descriptor.get_sel(), - ) - composed = LinearEnergyAtomicModel( - models=[model.atomic_model, zbl_atomic], - type_map=data["type_map"], - weights="sum", - # Both exclusions belong to the composition: its children share one - # graph, so "excluded" must cover the analytical term too. - atom_exclude_types=atom_exclude_types, - pair_exclude_types=pair_exclude_types, - ) + data = copy.deepcopy(data) + spin = None + if "spin" in data: + spin_cfg = data.pop("spin") + if str(spin_cfg.get("scheme", "deepspin")) != "native": + raise NotImplementedError( + "Spin linear_ener models support only spin scheme 'native'." + ) + use_spin = normalize_spin_use_spin(spin_cfg["use_spin"], data["type_map"]) + spin = Spin( + use_spin=use_spin, + virtual_scale=spin_cfg.get("virtual_scale", 1.0), + allow_missing_label=spin_cfg.get("allow_missing_label", False), + ) + for sub in data["models"]: + if "descriptor" in sub: + sub["descriptor"]["use_spin"] = use_spin + composed = _model_factory.get_linear_atomic_model(data) + if spin is not None: + if not composed.supports_native_spin(): + raise NotImplementedError( + "spin scheme 'native' requires an atomic model declaring " + "supports_native_spin()" + ) + return NativeSpinEnergyModel(atomic_model_=composed, spin=spin) return LinearEnergyModel(atomic_model_=composed) @@ -135,14 +155,10 @@ def get_native_spin_model(data: dict) -> NativeSpinEnergyModel: eligible; the gate is that capability method, not a descriptor-type list. - The non-spin backbone is built by :func:`get_standard_model`, which OWNS - everything about assembling the atomic model -- descriptor/fitting, - exclusions and the analytical-bridging composition -- so ``spin`` and - ``bridging_method`` combine for free: the wrapper re-classes whatever - atomic model came back, be it a single learned model or a - ``LinearEnergyAtomicModel`` over ``[learned, InnerPotential]`` (the - analytical child accepts and ignores ``spin``; the learned child consumes - it). + The non-spin backbone is built by :func:`get_standard_model`. A spin + model with analytical bridging is a ``linear_ener`` composition and + routes through :func:`get_linear_model` instead (the ``bridging_method`` + sugar expands to that form in :func:`get_model`). Parameters ---------- @@ -195,9 +211,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/dpmodel/model/model_factory.py b/deepmd/dpmodel/model/model_factory.py index b40ef7d937..f0442828ae 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, ) @@ -124,6 +127,208 @@ def get_zbl_model( ) +def get_linear_atomic_model( + data: dict, + *, + descriptor_base: type, + fitting_base: type, + backend_name: str, + atomic_model: type, + pairtab_model: type, + descriptor_child_builder: "Callable[[dict], Any | None] | None" = None, +) -> Any: + """Build the ``LinearEnergyAtomicModel`` composition from a config. + + Children with a ``descriptor`` build as learned atomic models through + the backend registries; ``pairtab`` children build as pair-tabulation + atomic models; an ``inner_potential`` child builds the analytical + bridging term. The composition is the ONE owner of the bridging + coupling: it derives the learned sibling descriptor's + ``inner_clamp_r_inner``/``_outer`` from the ``inner_potential`` + child's ``r_inner``/``r_outer``, so the radii are written once in the + config (issue #5948). + + Parameters + ---------- + data : dict + The ``linear_ener`` model configuration. + descriptor_base : type + Backend descriptor registry base class. + fitting_base : type + Backend fitting registry base class. + backend_name : str + Backend name used in error messages. + atomic_model : type + Backend learned atomic-model class. + pairtab_model : type + Backend pair-tabulation atomic-model class. + descriptor_child_builder : callable, optional + Backend hook for descriptor-bearing children: called with the + child config (``type_map`` and derived clamp radii already + injected) and returns the child atomic model, or ``None`` to fall + back to the generic registry build. Backends use it to route + family-specific model types (e.g. DPA4/SeZM) through their + validated builders. + + Raises + ------ + ValueError + If more than one ``inner_potential`` child is given, if an + ``inner_potential`` child has no unique learned sibling, if a + bridged composition does not use ``weights: "sum"``, if a child + carries a ``bridging_method`` flag, or if a child is of an + unsupported kind. + """ + from deepmd.dpmodel.atomic_model.inner_potential import ( + InnerPotentialAtomicModel, + ) + from deepmd.dpmodel.atomic_model.linear_atomic_model import ( + LinearEnergyAtomicModel, + ) + + data = copy.deepcopy(data) + type_map = data["type_map"] + children = data["models"] + inner_indices = [ + i for i, sub in enumerate(children) if sub.get("type") == "inner_potential" + ] + learned_indices = [ + i + for i, sub in enumerate(children) + if "descriptor" in sub and i not in inner_indices + ] + for i in inner_indices: + if "descriptor" in children[i]: + raise ValueError( + "An `inner_potential` sub-model must not carry a " + "`descriptor`: the analytical term has no learned " + "component." + ) + # 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( + "`bridging_method` is not supported on a linear_ener " + "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( + "A linear_ener composition supports at most one " + "`inner_potential` sub-model." + ) + if len(learned_indices) != 1 or len(children) != 2: + # A third child (e.g. pairtab) has no common execution route + # with the graph-only bridged pair; reject at construction + # like the pt builder does. + raise ValueError( + "An `inner_potential` sub-model bridges exactly one learned " + "sibling: expected a linear_ener composition over " + "[learned, inner_potential]." + ) + if str(data.get("weights", "mean")) != "sum": + raise ValueError( + 'A bridged linear_ener composition requires `weights: "sum"`.' + ) + 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]] + learned_descriptor = children[learned_indices[0]]["descriptor"] + learned_descriptor["inner_clamp_r_inner"] = float(inner_cfg.get("r_inner", 0.5)) + learned_descriptor["inner_clamp_r_outer"] = float(inner_cfg.get("r_outer", 0.8)) + route_canonical_learned_options(data, children[learned_indices[0]]) + + built: dict[int, Any] = {} + for i, sub in enumerate(children): + if i in inner_indices: + continue + if "type_map" not in sub: + sub["type_map"] = copy.deepcopy(type_map) + elif inner_indices and i == learned_indices[0] and sub["type_map"] != type_map: + # The analytical child always uses the composition's type_map, + # and the graph route rejects a non-identity remap at forward + # time; fail at construction like the pt builder does. + raise ValueError( + "A bridged linear_ener composition requires the learned " + "child's type_map to match the composition type_map." + ) + if "descriptor" in sub: + child = None + if descriptor_child_builder is not None: + child = descriptor_child_builder(sub) + if child is None: + descriptor, fitting, _ = get_model_components( + sub, + descriptor_base=descriptor_base, + fitting_base=fitting_base, + backend_name=backend_name, + ) + child = atomic_model(descriptor, fitting, type_map=sub["type_map"]) + built[i] = child + else: + if sub.get("type") != "pairtab": + raise ValueError( + "Sub-models in LinearEnergyModel must be a standard model, " + "a pairtab model, or an inner_potential model, but got " + f"type {sub.get('type')!r}." + ) + built[i] = pairtab_model( + sub["tab_file"], + sub["rcut"], + sub["sel"], + type_map=copy.deepcopy(type_map), + ) + for i in inner_indices: + learned_descriptor_obj = built[learned_indices[0]].descriptor + built[i] = InnerPotentialAtomicModel( + type_map=copy.deepcopy(type_map), + mode=children[i].get("mode", "zbl"), + rcut=learned_descriptor_obj.get_rcut(), + sel=learned_descriptor_obj.get_sel(), + ) + return LinearEnergyAtomicModel( + models=[built[i] for i in range(len(children))], + type_map=type_map, + weights=data.get("weights", "mean"), + # Both exclusions belong to the composition: its children share one + # graph, so "excluded" must cover the analytical term too. + atom_exclude_types=data.get("atom_exclude_types", []), + pair_exclude_types=data.get("pair_exclude_types", []), + ) + + def get_spin_model( data: dict, *, @@ -257,6 +462,25 @@ def get_standard_model(self, data: dict) -> Any: backend_name=self.backend_name, ) + def get_linear_atomic_model( + self, + data: dict, + *, + descriptor_child_builder: "Callable[[dict], Any | None] | None" = None, + ) -> Any: + """Construct the linear atomic-model composition for this backend.""" + if self.atomic_model is None or self.pairtab_model is None: + raise NotImplementedError("Linear model is not implemented yet.") + return get_linear_atomic_model( + data, + descriptor_base=self.descriptor_base, + fitting_base=self.fitting_base, + backend_name=self.backend_name, + atomic_model=self.atomic_model, + pairtab_model=self.pairtab_model, + descriptor_child_builder=descriptor_child_builder, + ) + def get_zbl_model(self, data: dict) -> Any: """Construct a ZBL model for this backend.""" if ( 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..f468379565 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -84,6 +84,9 @@ from deepmd.utils.batch_size import ( RetrySignal, ) +from deepmd.utils.bridging import ( + is_bridged_sezm_config, +) from deepmd.utils.econf_embd import ( sort_element_type, ) @@ -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 887ce1754c..18225dbdc8 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -36,6 +36,10 @@ from deepmd.pt.utils.multi_task import ( preprocess_shared_params, ) +from deepmd.utils.bridging import ( + expand_bridging_method, + route_canonical_learned_options, +) from deepmd.utils.spin import ( Spin, ) @@ -193,7 +197,18 @@ 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: + 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", []) + ): + return _get_bridged_linear_model(model_params) model_params = copy.deepcopy(model_params) weights = model_params.get("weights", "mean") shared_links = None @@ -265,6 +280,88 @@ 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}." + ) + 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`." + ) + route_canonical_learned_options(model_params, learned) + 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 +428,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"]) @@ -375,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 []) ] @@ -399,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) @@ -516,7 +651,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) @@ -589,9 +724,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 @@ -656,6 +792,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/model/model/dp_linear_model.py b/deepmd/pt/model/model/dp_linear_model.py index a0a3c6f2c7..94403166da 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( @@ -367,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/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 f0d87a9118..a994958df4 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-relocated 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``, and @@ -107,11 +104,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) @@ -340,8 +279,44 @@ 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). + """ + 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 + + 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 ---------- @@ -353,42 +328,37 @@ def get_linear_model(model_params: dict) -> BaseModel: ) 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 = _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( + "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: @@ -413,6 +383,15 @@ def get_model(data: dict) -> BaseModel: data : dict 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/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/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index bbbf15acb4..828cf639c2 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,21 +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." - ) - 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. " @@ -3600,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, @@ -3704,6 +3721,44 @@ def pairtab_model_args() -> Argument: return ca +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)." + ) + 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 = ( @@ -3712,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 new file mode 100644 index 0000000000..88467d66ab --- /dev/null +++ b/deepmd/utils/bridging.py @@ -0,0 +1,328 @@ +# 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 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. +""" + +import copy + +__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) + + +# 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 +# The routing tables are pairwise disjoint: a key has exactly one owner. +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. + + 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. + + 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 + ---------- + 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 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]: + 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]) + + +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 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 + ``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 _NON_CHILD_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"] + for key in _TRAINER_KEYS: + if key in data: + canonical[key] = data[key] + return canonical diff --git a/doc/model/dpa4.md b/doc/model/dpa4.md index 3acee5ec35..63f1745e62 100644 --- a/doc/model/dpa4.md +++ b/doc/model/dpa4.md @@ -325,11 +325,13 @@ 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 gate additionally blocks the learned model from leaking information about the -frozen short-range pairs. Enable it with: +frozen short-range pairs. The recommended way to enable it is the concise +form, set directly on the `dpa4` model: ```json { "model": { + "type": "dpa4", "bridging_method": "zbl", "bridging_r_inner": 0.5, "bridging_r_outer": 0.8 @@ -352,6 +354,46 @@ from training. See `examples/water/dpa4/input-zbl.json` for a complete example. > prediction. See [change-bias](change-bias.md) for the precise > definitions. +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 +{ + "model": { + "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 + } + ] + } +} +``` + +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 ### Training-time settings diff --git a/source/tests/common/dpmodel/test_zbl_bridging.py b/source/tests/common/dpmodel/test_zbl_bridging.py index 9d9ddba47d..e4379c1b12 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)) @@ -675,3 +698,243 @@ def test_set_by_statistic_fits_raw_labels_by_definition(): model.atomic_model.compute_or_load_out_stat(samples) repeated_bias = np.asarray(model.atomic_model.out_bias).reshape(-1)[:2] np.testing.assert_allclose(repeated_bias, raw_fit, atol=1.0e-8) + + +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)) + + +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) + + 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" + + +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 new file mode 100644 index 0000000000..24480ffddf --- /dev/null +++ b/source/tests/common/test_bridging.py @@ -0,0 +1,347 @@ +# 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_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]] + 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" + + +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) + + +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) + + +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"}) + + +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 new file mode 100644 index 0000000000..5c851f111d --- /dev/null +++ b/source/tests/pt/model/test_get_model_bridging.py @@ -0,0 +1,390 @@ +# 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) + + +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_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. + """ + 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" + + +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) + + +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 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..6a174753b5 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,203 @@ 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_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. + """ + 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) + + +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" + + +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) + + +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) + + +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)