From 6e5a4a91427b30f77fb01f00fea1f493a656996a Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Tue, 14 Jul 2026 07:12:18 +0800 Subject: [PATCH 1/3] refactor(model): unify dpmodel backend factories Centralize descriptor and fitting parameter injection, standard-model selection, ZBL assembly, and model-type routing across dpmodel, pt_expt, JAX, and TF2. Keep backend modules limited to their native class registries and supported special cases. Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/dpmodel/model/__init__.py | 16 + deepmd/dpmodel/model/model.py | 187 ++---------- deepmd/dpmodel/model/model_factory.py | 285 ++++++++++++++++++ deepmd/jax/model/model.py | 122 +------- deepmd/pt_expt/model/get_model.py | 167 +++------- deepmd/tf2/model/model.py | 120 +------- .../common/dpmodel/test_model_factory.py | 197 ++++++++++++ .../tests/pt_expt/model/test_get_model_zbl.py | 67 ++++ 8 files changed, 659 insertions(+), 502 deletions(-) create mode 100644 deepmd/dpmodel/model/model_factory.py create mode 100644 source/tests/common/dpmodel/test_model_factory.py create mode 100644 source/tests/pt_expt/model/test_get_model_zbl.py diff --git a/deepmd/dpmodel/model/__init__.py b/deepmd/dpmodel/model/__init__.py index 37ef57b38b..2816afc9e0 100644 --- a/deepmd/dpmodel/model/__init__.py +++ b/deepmd/dpmodel/model/__init__.py @@ -12,15 +12,27 @@ Models generated by `make_model` have already done it. """ +from .dipole_model import ( + DipoleModel, +) +from .dos_model import ( + DOSModel, +) from .dp_model import ( DPModelCommon, ) +from .dp_zbl_model import ( + DPZBLModel, +) from .ener_model import ( EnergyModel, ) from .make_model import ( make_model, ) +from .polar_model import ( + PolarModel, +) from .property_model import ( PropertyModel, ) @@ -29,8 +41,12 @@ ) __all__ = [ + "DOSModel", "DPModelCommon", + "DPZBLModel", + "DipoleModel", "EnergyModel", + "PolarModel", "PropertyModel", "SpinModel", "make_model", diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index 8f96e965b0..ed2de188fb 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -1,9 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import copy -from typing import ( - Any, -) - from deepmd.dpmodel.atomic_model.dp_atomic_model import ( DPAtomicModel, ) @@ -16,144 +11,33 @@ from deepmd.dpmodel.fitting.base_fitting import ( BaseFitting, ) -from deepmd.dpmodel.fitting.ener_fitting import ( - EnergyFittingNet, -) from deepmd.dpmodel.model.base_model import ( BaseModel, ) -from deepmd.dpmodel.model.dipole_model import ( - DipoleModel, -) -from deepmd.dpmodel.model.dos_model import ( - DOSModel, -) from deepmd.dpmodel.model.dp_zbl_model import ( DPZBLModel, ) -from deepmd.dpmodel.model.ener_model import ( - EnergyModel, -) -from deepmd.dpmodel.model.polar_model import ( - PolarModel, +from deepmd.dpmodel.model.model_factory import ( + BackendModelFactory, ) -from deepmd.dpmodel.model.property_model import ( - PropertyModel, +from deepmd.dpmodel.model.model_factory import ( + get_spin_model as get_spin_model_from_factory, ) from deepmd.dpmodel.model.spin_model import ( SpinModel, ) -from deepmd.utils.spin import ( - Spin, -) - -def _get_standard_model_components( - data: dict[str, Any], ntypes: int -) -> tuple[BaseDescriptor, BaseFitting, str]: - # descriptor - data["descriptor"]["ntypes"] = ntypes - data["descriptor"]["type_map"] = copy.deepcopy(data["type_map"]) - descriptor = BaseDescriptor(**data["descriptor"]) - # fitting - fitting_net = data.get("fitting_net", {}) - fitting_net["type"] = fitting_net.get("type", "ener") - fitting_net["ntypes"] = descriptor.get_ntypes() - fitting_net["type_map"] = copy.deepcopy(data["type_map"]) - fitting_net["mixed_types"] = descriptor.mixed_types() - if fitting_net["type"] in ["dipole", "polar"]: - fitting_net["embedding_width"] = descriptor.get_dim_emb() - fitting_net["dim_descrpt"] = descriptor.get_dim_out() - grad_force = "direct" not in fitting_net["type"] - if not grad_force: - fitting_net["out_dim"] = descriptor.get_dim_emb() - if "ener" in fitting_net["type"]: - fitting_net["return_energy"] = True - fitting = BaseFitting(**fitting_net) - return descriptor, fitting, fitting_net["type"] - - -def get_standard_model(data: dict) -> EnergyModel: - """Get a EnergyModel from a dictionary. - - Parameters - ---------- - data : dict - The data to construct the model. - """ - if "type_embedding" in data: - raise ValueError( - "In the DP backend, type_embedding is not at the model level, but within the descriptor. See type embedding documentation for details." - ) - data = copy.deepcopy(data) - ntypes = len(data["type_map"]) - descriptor, fitting, fitting_net_type = _get_standard_model_components(data, ntypes) - atom_exclude_types = data.get("atom_exclude_types", []) - pair_exclude_types = data.get("pair_exclude_types", []) - - if fitting_net_type == "dipole": - modelcls = DipoleModel - elif fitting_net_type == "polar": - modelcls = PolarModel - elif fitting_net_type == "dos": - modelcls = DOSModel - elif fitting_net_type in ["ener", "direct_force_ener"]: - modelcls = EnergyModel - elif fitting_net_type == "property": - modelcls = PropertyModel - else: - raise RuntimeError(f"Unknown fitting type: {fitting_net_type}") - - model = modelcls( - descriptor=descriptor, - fitting=fitting, - type_map=data["type_map"], - atom_exclude_types=atom_exclude_types, - pair_exclude_types=pair_exclude_types, - ) - return model - - -def get_zbl_model(data: dict) -> DPZBLModel: - data = copy.deepcopy(data) - data["descriptor"]["ntypes"] = len(data["type_map"]) - data["descriptor"]["type_map"] = data["type_map"] - descriptor = BaseDescriptor(**data["descriptor"]) - fitting_type = data["fitting_net"].pop("type") - data["fitting_net"]["type_map"] = data["type_map"] - if fitting_type == "ener": - fitting = EnergyFittingNet( - ntypes=descriptor.get_ntypes(), - dim_descrpt=descriptor.get_dim_out(), - mixed_types=descriptor.mixed_types(), - **data["fitting_net"], - ) - else: - raise ValueError(f"Unknown fitting type {fitting_type}") - - dp_model = DPAtomicModel(descriptor, fitting, type_map=data["type_map"]) - # pairtab - filepath = data["use_srtab"] - pt_model = PairTabAtomicModel( - filepath, - descriptor.get_rcut(), - descriptor.get_sel(), - type_map=data["type_map"], - ) - - rmin = data["sw_rmin"] - rmax = data["sw_rmax"] - atom_exclude_types = data.get("atom_exclude_types", []) - pair_exclude_types = data.get("pair_exclude_types", []) - return DPZBLModel( - dp_model, - pt_model, - rmin, - rmax, - type_map=data["type_map"], - atom_exclude_types=atom_exclude_types, - pair_exclude_types=pair_exclude_types, - ) +_model_factory = BackendModelFactory( + descriptor_base=BaseDescriptor, + fitting_base=BaseFitting, + model_base=BaseModel, + backend_name="DP", + atomic_model=DPAtomicModel, + pairtab_model=PairTabAtomicModel, + zbl_model=DPZBLModel, +) +get_standard_model = _model_factory.get_standard_model +get_zbl_model = _model_factory.get_zbl_model def get_spin_model(data: dict) -> SpinModel: @@ -164,30 +48,11 @@ def get_spin_model(data: dict) -> SpinModel: data : dict The data to construct the model. """ - data = copy.deepcopy(data) - # include virtual spin and placeholder types - data["type_map"] += [item + "_spin" for item in data["type_map"]] - spin = Spin( - use_spin=data["spin"]["use_spin"], - virtual_scale=data["spin"]["virtual_scale"], + return get_spin_model_from_factory( + data, + standard_model_factory=get_standard_model, + spin_model=SpinModel, ) - pair_exclude_types = spin.get_pair_exclude_types( - exclude_types=data.get("pair_exclude_types", None) - ) - data["pair_exclude_types"] = pair_exclude_types - # for descriptor data stat - data["descriptor"]["exclude_types"] = pair_exclude_types - atom_exclude_types = spin.get_atom_exclude_types( - exclude_types=data.get("atom_exclude_types", None) - ) - data["atom_exclude_types"] = atom_exclude_types - if "env_protection" not in data["descriptor"]: - data["descriptor"]["env_protection"] = 1e-6 - if data["descriptor"]["type"] in ["se_e2_a"]: - # only expand sel for se_e2_a - data["descriptor"]["sel"] += data["descriptor"]["sel"] - backbone_model = get_standard_model(data) - return SpinModel(backbone_model=backbone_model, spin=spin) def get_model(data: dict) -> BaseModel: @@ -198,13 +63,7 @@ def get_model(data: dict) -> BaseModel: data : dict The data to construct the model. """ - model_type = data.get("type", "standard") - if model_type == "standard": - if "spin" in data: - return get_spin_model(data) - elif "use_srtab" in data: - return get_zbl_model(data) - else: - return get_standard_model(data) - else: - return BaseModel.get_class_by_type(model_type).get_model(data) + return _model_factory.get_model( + data, + spin_model_factory=get_spin_model, + ) diff --git a/deepmd/dpmodel/model/model_factory.py b/deepmd/dpmodel/model/model_factory.py new file mode 100644 index 0000000000..3da519b4cf --- /dev/null +++ b/deepmd/dpmodel/model/model_factory.py @@ -0,0 +1,285 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Shared model-factory dispatch for dpmodel-driven backends.""" + +import copy +from collections.abc import ( + Callable, + Mapping, +) +from typing import ( + Any, +) + +from deepmd.utils.spin import ( + Spin, +) + +ModelBuilder = Callable[[dict], Any] + + +def get_model_components( + data: dict, + *, + descriptor_base: type, + fitting_base: type, + backend_name: str, +) -> tuple[Any, Any, str]: + """Construct a backend descriptor and fitting net from model config. + + The backend registries expose the same descriptor/fitting constructor + contract. Keeping the parameter injection here prevents subtle differences + in ``type_map``, ``ntypes``, embedding width, and direct-force handling. + """ + data = copy.deepcopy(data) + if "type_embedding" in data: + raise ValueError( + f"In the {backend_name} backend, type_embedding is not at the model " + "level, but within the descriptor. See type embedding documentation " + "for details." + ) + type_map = copy.deepcopy(data["type_map"]) + descriptor_data = data["descriptor"] + descriptor_type = descriptor_data.pop("type") + descriptor_data["ntypes"] = len(type_map) + descriptor_data["type_map"] = copy.deepcopy(type_map) + descriptor = descriptor_base.get_class_by_type(descriptor_type)(**descriptor_data) + + fitting_data = data.get("fitting_net", {}) + fitting_type = fitting_data.pop("type", "ener") + fitting_data["ntypes"] = descriptor.get_ntypes() + fitting_data["type_map"] = copy.deepcopy(type_map) + fitting_data["mixed_types"] = descriptor.mixed_types() + if fitting_type in {"dipole", "polar"}: + fitting_data["embedding_width"] = descriptor.get_dim_emb() + fitting_data["dim_descrpt"] = descriptor.get_dim_out() + if "direct" in fitting_type: + fitting_data["out_dim"] = descriptor.get_dim_emb() + if "ener" in fitting_type: + fitting_data["return_energy"] = True + fitting = fitting_base.get_class_by_type(fitting_type)(**fitting_data) + return descriptor, fitting, fitting_type + + +def get_standard_model( + data: dict, + *, + descriptor_base: type, + fitting_base: type, + model_base: type, + backend_name: str, +) -> Any: + """Construct a standard model through backend registries.""" + descriptor, fitting, fitting_type = get_model_components( + data, + descriptor_base=descriptor_base, + fitting_base=fitting_base, + backend_name=backend_name, + ) + model_type = "ener" if fitting_type == "direct_force_ener" else fitting_type + model_cls = model_base.get_class_by_type(model_type) + return model_cls( + descriptor=descriptor, + fitting=fitting, + type_map=data["type_map"], + atom_exclude_types=data.get("atom_exclude_types", []), + pair_exclude_types=data.get("pair_exclude_types", []), + ) + + +def get_zbl_model( + data: dict, + *, + descriptor_base: type, + fitting_base: type, + atomic_model: type, + pairtab_model: type, + zbl_model: type, + backend_name: str, +) -> Any: + """Construct a ZBL model from backend-native atomic model classes.""" + data = copy.deepcopy(data) + descriptor, fitting, fitting_type = get_model_components( + data, + descriptor_base=descriptor_base, + fitting_base=fitting_base, + backend_name=backend_name, + ) + if fitting_type != "ener": + raise ValueError(f"Unknown fitting type {fitting_type}") + dp_model = atomic_model(descriptor, fitting, type_map=data["type_map"]) + pairtab = pairtab_model( + data["use_srtab"], + descriptor.get_rcut(), + descriptor.get_sel(), + type_map=data["type_map"], + ) + return zbl_model( + dp_model, + pairtab, + data["sw_rmin"], + data["sw_rmax"], + type_map=data["type_map"], + smin_alpha=data.get("smin_alpha", 0.1), + atom_exclude_types=data.get("atom_exclude_types", []), + pair_exclude_types=data.get("pair_exclude_types", []), + ) + + +def get_spin_model( + data: dict, + *, + standard_model_factory: ModelBuilder, + spin_model: type, +) -> Any: + """Construct a legacy spin model using a backend standard-model factory.""" + data = copy.deepcopy(data) + data["type_map"] += [item + "_spin" for item in data["type_map"]] + spin = Spin( + use_spin=data["spin"]["use_spin"], + virtual_scale=data["spin"]["virtual_scale"], + ) + pair_exclude_types = spin.get_pair_exclude_types( + exclude_types=data.get("pair_exclude_types") + ) + data["pair_exclude_types"] = pair_exclude_types + data["descriptor"]["exclude_types"] = pair_exclude_types + data["atom_exclude_types"] = spin.get_atom_exclude_types( + exclude_types=data.get("atom_exclude_types") + ) + data["descriptor"].setdefault("env_protection", 1e-6) + if data["descriptor"]["type"] == "se_e2_a": + data["descriptor"]["sel"] += data["descriptor"]["sel"] + backbone_model = standard_model_factory(data) + return spin_model(backbone_model=backbone_model, spin=spin) + + +def get_model( + data: dict, + *, + base_model: type, + standard_model_factory: ModelBuilder, + spin_model_factory: ModelBuilder | None = None, + zbl_model_factory: ModelBuilder | None = None, + model_factories: Mapping[str, ModelBuilder] | None = None, +) -> Any: + """Construct a backend model using the shared model-type routing rules. + + Backend modules supply the concrete constructors while this function owns + the routing precedence. In particular, legacy ``standard`` configurations + select spin before ZBL, matching the established dpmodel and PyTorch input + contract. Explicit model types may be handled by backend-specific factories + before falling back to the backend model plugin registry. + + Parameters + ---------- + data : dict + Model configuration. + base_model : type + Backend model base class providing ``get_class_by_type``. + standard_model_factory : callable + Constructor for an ordinary standard model. + spin_model_factory : callable, optional + Constructor for a legacy standard model containing ``spin``. + zbl_model_factory : callable, optional + Constructor for a legacy standard model containing ``use_srtab``. + model_factories : mapping, optional + Backend-specific constructors keyed by explicit model type. + + Returns + ------- + Any + The backend-native model instance. + """ + model_type = data.get("type", "standard") + if model_type == "standard": + if "spin" in data: + if spin_model_factory is None: + raise NotImplementedError("Spin model is not implemented yet.") + return spin_model_factory(data) + if "use_srtab" in data: + if zbl_model_factory is None: + raise NotImplementedError("ZBL model is not implemented yet.") + return zbl_model_factory(data) + return standard_model_factory(data) + + if model_factories is not None and model_type in model_factories: + return model_factories[model_type](data) + return base_model.get_class_by_type(model_type).get_model(data) + + +class BackendModelFactory: + """Bind backend registries once and expose the shared factory operations.""" + + def __init__( + self, + *, + descriptor_base: type, + fitting_base: type, + model_base: type, + backend_name: str, + atomic_model: type | None = None, + pairtab_model: type | None = None, + zbl_model: type | None = None, + ) -> None: + """Store backend-native classes used by all model construction paths.""" + self.descriptor_base = descriptor_base + self.fitting_base = fitting_base + self.model_base = model_base + self.backend_name = backend_name + self.atomic_model = atomic_model + self.pairtab_model = pairtab_model + self.zbl_model = zbl_model + + def get_model_components(self, data: dict) -> tuple[Any, Any, str]: + """Construct descriptor and fitting objects for this backend.""" + return get_model_components( + data, + descriptor_base=self.descriptor_base, + fitting_base=self.fitting_base, + backend_name=self.backend_name, + ) + + def get_standard_model(self, data: dict) -> Any: + """Construct a standard model for this backend.""" + return get_standard_model( + data, + descriptor_base=self.descriptor_base, + fitting_base=self.fitting_base, + model_base=self.model_base, + backend_name=self.backend_name, + ) + + def get_zbl_model(self, data: dict) -> Any: + """Construct a ZBL model for this backend.""" + if ( + self.atomic_model is None + or self.pairtab_model is None + or self.zbl_model is None + ): + raise NotImplementedError("ZBL model is not implemented yet.") + return get_zbl_model( + data, + descriptor_base=self.descriptor_base, + fitting_base=self.fitting_base, + atomic_model=self.atomic_model, + pairtab_model=self.pairtab_model, + zbl_model=self.zbl_model, + backend_name=self.backend_name, + ) + + def get_model( + self, + data: dict, + *, + spin_model_factory: ModelBuilder | None = None, + model_factories: Mapping[str, ModelBuilder] | None = None, + ) -> Any: + """Construct a model using this backend and the shared routing rules.""" + return get_model( + data, + base_model=self.model_base, + standard_model_factory=self.get_standard_model, + spin_model_factory=spin_model_factory, + zbl_model_factory=self.get_zbl_model, + model_factories=model_factories, + ) diff --git a/deepmd/jax/model/model.py b/deepmd/jax/model/model.py index a3d067c636..3e9055deda 100644 --- a/deepmd/jax/model/model.py +++ b/deepmd/jax/model/model.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from copy import ( - deepcopy, +from deepmd.dpmodel.model.model_factory import ( + BackendModelFactory, ) - from deepmd.jax.atomic_model.dp_atomic_model import ( DPAtomicModel, ) @@ -15,9 +14,6 @@ from deepmd.jax.fitting.base_fitting import ( BaseFitting, ) -from deepmd.jax.fitting.fitting import ( - EnergyFittingNet, -) from deepmd.jax.model.base_model import ( BaseModel, ) @@ -25,105 +21,15 @@ DPZBLModel, ) - -def get_standard_model(data: dict) -> BaseModel: - """Get a Model from a dictionary. - - Parameters - ---------- - data : dict - The data to construct the model. - """ - data = deepcopy(data) - if "type_embedding" in data: - raise ValueError( - "In the JAX backend, type_embedding is not at the model level, but within the descriptor. See type embedding documentation for details." - ) - descriptor_type = data["descriptor"].pop("type") - data["descriptor"]["type_map"] = data["type_map"] - data["descriptor"]["ntypes"] = len(data["type_map"]) - # Default the fitting type to energy and tolerate a missing fitting_net - # block, matching the dpmodel and TF2 standard-model factories. - data["fitting_net"] = data.get("fitting_net", {}) - fitting_type = data["fitting_net"].pop("type", "ener") - data["fitting_net"]["type_map"] = data["type_map"] - descriptor = BaseDescriptor.get_class_by_type(descriptor_type)( - **data["descriptor"], - ) - if fitting_type in {"dipole", "polar"}: - data["fitting_net"]["embedding_width"] = descriptor.get_dim_emb() - fitting = BaseFitting.get_class_by_type(fitting_type)( - ntypes=descriptor.get_ntypes(), - dim_descrpt=descriptor.get_dim_out(), - mixed_types=descriptor.mixed_types(), - **data["fitting_net"], - ) - return BaseModel.get_class_by_type(fitting_type)( - descriptor=descriptor, - fitting=fitting, - type_map=data["type_map"], - atom_exclude_types=data.get("atom_exclude_types", []), - pair_exclude_types=data.get("pair_exclude_types", []), - ) - - -def get_zbl_model(data: dict) -> DPZBLModel: - data = deepcopy(data) - data["descriptor"]["ntypes"] = len(data["type_map"]) - data["descriptor"]["type_map"] = data["type_map"] - descriptor_type = data["descriptor"].pop("type") - descriptor = BaseDescriptor.get_class_by_type(descriptor_type)(**data["descriptor"]) - fitting_type = data["fitting_net"].pop("type") - data["fitting_net"]["type_map"] = data["type_map"] - if fitting_type == "ener": - fitting = EnergyFittingNet( - ntypes=descriptor.get_ntypes(), - dim_descrpt=descriptor.get_dim_out(), - mixed_types=descriptor.mixed_types(), - **data["fitting_net"], - ) - else: - raise ValueError(f"Unknown fitting type {fitting_type}") - - dp_model = DPAtomicModel(descriptor, fitting, type_map=data["type_map"]) - # pairtab - filepath = data["use_srtab"] - pt_model = PairTabAtomicModel( - filepath, - data["descriptor"]["rcut"], - data["descriptor"]["sel"], - type_map=data["type_map"], - ) - rmin = data["sw_rmin"] - rmax = data["sw_rmax"] - atom_exclude_types = data.get("atom_exclude_types", []) - pair_exclude_types = data.get("pair_exclude_types", []) - return DPZBLModel( - dp_model, - pt_model, - rmin, - rmax, - type_map=data["type_map"], - atom_exclude_types=atom_exclude_types, - pair_exclude_types=pair_exclude_types, - ) - - -def get_model(data: dict) -> BaseModel: - """Get a model from a dictionary. - - Parameters - ---------- - data : dict - The data to construct the model. - """ - model_type = data.get("type", "standard") - if model_type == "standard": - if "spin" in data: - raise NotImplementedError("Spin model is not implemented yet.") - elif "use_srtab" in data: - return get_zbl_model(data) - else: - return get_standard_model(data) - else: - return BaseModel.get_class_by_type(model_type).get_model(data) +_model_factory = BackendModelFactory( + descriptor_base=BaseDescriptor, + fitting_base=BaseFitting, + model_base=BaseModel, + backend_name="JAX", + atomic_model=DPAtomicModel, + pairtab_model=PairTabAtomicModel, + zbl_model=DPZBLModel, +) +get_standard_model = _model_factory.get_standard_model +get_zbl_model = _model_factory.get_zbl_model +get_model = _model_factory.get_model diff --git a/deepmd/pt_expt/model/get_model.py b/deepmd/pt_expt/model/get_model.py index 7efa904f23..9ec2b08259 100644 --- a/deepmd/pt_expt/model/get_model.py +++ b/deepmd/pt_expt/model/get_model.py @@ -8,23 +8,27 @@ import copy import logging -from typing import ( - Any, -) +from deepmd.dpmodel.atomic_model.dp_atomic_model import ( + DPAtomicModel, +) +from deepmd.dpmodel.atomic_model.pairtab_atomic_model import ( + PairTabAtomicModel, +) +from deepmd.dpmodel.model.model_factory import ( + BackendModelFactory, +) +from deepmd.dpmodel.model.model_factory import ( + get_spin_model as get_spin_model_from_factory, +) from deepmd.pt_expt.descriptor import ( BaseDescriptor, ) from deepmd.pt_expt.fitting import ( BaseFitting, ) - -# Import from submodules directly to avoid circular import via __init__.py -from deepmd.pt_expt.model.dipole_model import ( - DipoleModel, -) -from deepmd.pt_expt.model.dos_model import ( - DOSModel, +from deepmd.pt_expt.model.dp_zbl_model import ( + DPZBLModel, ) from deepmd.pt_expt.model.ener_model import ( EnergyModel, @@ -32,18 +36,9 @@ from deepmd.pt_expt.model.model import ( BaseModel, ) -from deepmd.pt_expt.model.polar_model import ( - PolarModel, -) -from deepmd.pt_expt.model.property_model import ( - PropertyModel, -) from deepmd.pt_expt.model.spin_ener_model import ( SpinEnergyModel, ) -from deepmd.utils.spin import ( - Spin, -) log = logging.getLogger(__name__) @@ -51,69 +46,17 @@ _WARNED_ONCE: set[str] = set() -def _get_standard_model_components( - data: dict[str, Any], - ntypes: int, -) -> tuple: - """Build descriptor and fitting from config dict.""" - # descriptor - data["descriptor"]["ntypes"] = ntypes - data["descriptor"]["type_map"] = copy.deepcopy(data["type_map"]) - descriptor = BaseDescriptor(**data["descriptor"]) - - # fitting - fitting_net = data.get("fitting_net", {}) - fitting_net["type"] = fitting_net.get("type", "ener") - fitting_net["ntypes"] = descriptor.get_ntypes() - fitting_net["type_map"] = copy.deepcopy(data["type_map"]) - fitting_net["mixed_types"] = descriptor.mixed_types() - if fitting_net["type"] in ["dipole", "polar"]: - fitting_net["embedding_width"] = descriptor.get_dim_emb() - fitting_net["dim_descrpt"] = descriptor.get_dim_out() - grad_force = "direct" not in fitting_net["type"] - if not grad_force: - fitting_net["out_dim"] = descriptor.get_dim_emb() - if "ener" in fitting_net["type"]: - fitting_net["return_energy"] = True - fitting = BaseFitting(**fitting_net) - return descriptor, fitting, fitting_net["type"] - - -def get_standard_model(data: dict) -> EnergyModel: - """Get a standard model from a config dictionary. - - Parameters - ---------- - data : dict - The data to construct the model. - """ - data = copy.deepcopy(data) - ntypes = len(data["type_map"]) - descriptor, fitting, fitting_net_type = _get_standard_model_components(data, ntypes) - atom_exclude_types = data.get("atom_exclude_types", []) - pair_exclude_types = data.get("pair_exclude_types", []) - - if fitting_net_type == "dipole": - modelcls = DipoleModel - elif fitting_net_type == "polar": - modelcls = PolarModel - elif fitting_net_type == "dos": - modelcls = DOSModel - elif fitting_net_type in ["ener", "direct_force_ener"]: - modelcls = EnergyModel - elif fitting_net_type == "property": - modelcls = PropertyModel - else: - raise RuntimeError(f"Unknown fitting type: {fitting_net_type}") - - model = modelcls( - descriptor=descriptor, - fitting=fitting, - type_map=data["type_map"], - atom_exclude_types=atom_exclude_types, - pair_exclude_types=pair_exclude_types, - ) - return model +_model_factory = BackendModelFactory( + descriptor_base=BaseDescriptor, + fitting_base=BaseFitting, + model_base=BaseModel, + backend_name="pt_expt", + atomic_model=DPAtomicModel, + pairtab_model=PairTabAtomicModel, + zbl_model=DPZBLModel, +) +get_standard_model = _model_factory.get_standard_model +get_zbl_model = _model_factory.get_zbl_model def get_sezm_model(data: dict) -> EnergyModel: @@ -194,8 +137,7 @@ def get_sezm_model(data: dict) -> EnergyModel: data["pair_exclude_types"] = pair_exclude_types data["descriptor"]["exclude_types"] = copy.deepcopy(pair_exclude_types) - ntypes = len(data["type_map"]) - descriptor, fitting, _ = _get_standard_model_components(data, ntypes) + descriptor, fitting, _ = _model_factory.get_model_components(data) return EnergyModel( descriptor=descriptor, fitting=fitting, @@ -213,13 +155,6 @@ def get_linear_model(model_params: dict) -> BaseModel: model_params : dict The model parameters. """ - from deepmd.dpmodel.atomic_model.dp_atomic_model import ( - DPAtomicModel, - ) - from deepmd.dpmodel.atomic_model.pairtab_atomic_model import ( - PairTabAtomicModel, - ) - from .dp_linear_model import ( LinearEnergyModel, ) @@ -233,8 +168,8 @@ def get_linear_model(model_params: dict) -> BaseModel: sub_model_params["type_map"] = model_params["type_map"] if "descriptor" in sub_model_params: sub_model_params["descriptor"]["ntypes"] = ntypes - descriptor, fitting, _ = _get_standard_model_components( - sub_model_params, ntypes + descriptor, fitting, _ = _model_factory.get_model_components( + sub_model_params ) list_of_models.append( DPAtomicModel(descriptor, fitting, type_map=model_params["type_map"]) @@ -270,27 +205,11 @@ def get_spin_model(data: dict) -> SpinEnergyModel: type map and descriptor sel for virtual spin atoms, then wraps the backbone EnergyModel as a :class:`SpinEnergyModel`. """ - data = copy.deepcopy(data) - data["type_map"] += [item + "_spin" for item in data["type_map"]] - spin = Spin( - use_spin=data["spin"]["use_spin"], - virtual_scale=data["spin"]["virtual_scale"], - ) - pair_exclude_types = spin.get_pair_exclude_types( - exclude_types=data.get("pair_exclude_types", None) + return get_spin_model_from_factory( + data, + standard_model_factory=get_standard_model, + spin_model=SpinEnergyModel, ) - data["pair_exclude_types"] = pair_exclude_types - data["descriptor"]["exclude_types"] = pair_exclude_types - atom_exclude_types = spin.get_atom_exclude_types( - exclude_types=data.get("atom_exclude_types", None) - ) - data["atom_exclude_types"] = atom_exclude_types - if "env_protection" not in data["descriptor"]: - data["descriptor"]["env_protection"] = 1e-6 - if data["descriptor"]["type"] in ["se_e2_a"]: - data["descriptor"]["sel"] += data["descriptor"]["sel"] - backbone_model = get_standard_model(data) - return SpinEnergyModel(backbone_model=backbone_model, spin=spin) def get_model(data: dict) -> BaseModel: @@ -301,14 +220,14 @@ def get_model(data: dict) -> BaseModel: data : dict The data to construct the model. """ - model_type = data.get("type", "standard") - if model_type == "standard": - if "spin" in data: - return get_spin_model(data) - return get_standard_model(data) - elif model_type == "linear_ener": - return get_linear_model(data) - elif model_type in ("dpa4", "DPA4", "sezm", "SeZM"): - return get_sezm_model(data) - else: - return BaseModel.get_class_by_type(model_type).get_model(data) + return _model_factory.get_model( + data, + spin_model_factory=get_spin_model, + model_factories={ + "linear_ener": get_linear_model, + "dpa4": get_sezm_model, + "DPA4": get_sezm_model, + "sezm": get_sezm_model, + "SeZM": get_sezm_model, + }, + ) diff --git a/deepmd/tf2/model/model.py b/deepmd/tf2/model/model.py index 8977c3357b..b1adc48632 100644 --- a/deepmd/tf2/model/model.py +++ b/deepmd/tf2/model/model.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from copy import ( - deepcopy, +from deepmd.dpmodel.model.model_factory import ( + BackendModelFactory, ) - from deepmd.tf2.atomic_model.dp_atomic_model import ( DPAtomicModel, ) @@ -15,9 +14,6 @@ from deepmd.tf2.fitting.base_fitting import ( BaseFitting, ) -from deepmd.tf2.fitting.fitting import ( - EnergyFittingNet, -) from deepmd.tf2.model.base_model import ( BaseModel, ) @@ -25,103 +21,15 @@ DPZBLModel, ) - -def get_standard_model(data: dict) -> BaseModel: - """Get a Model from a dictionary. - - Parameters - ---------- - data : dict - The data to construct the model. - """ - data = deepcopy(data) - if "type_embedding" in data: - raise ValueError( - "In the tf2 backend, type_embedding is not at the model level, but within the descriptor. See type embedding documentation for details." - ) - descriptor_type = data["descriptor"].pop("type") - data["descriptor"]["type_map"] = data["type_map"] - data["descriptor"]["ntypes"] = len(data["type_map"]) - data["fitting_net"] = data.get("fitting_net", {}) - fitting_type = data["fitting_net"].pop("type", "ener") - data["fitting_net"]["type_map"] = data["type_map"] - descriptor = BaseDescriptor.get_class_by_type(descriptor_type)( - **data["descriptor"], - ) - if fitting_type in {"dipole", "polar"}: - data["fitting_net"]["embedding_width"] = descriptor.get_dim_emb() - fitting = BaseFitting.get_class_by_type(fitting_type)( - ntypes=descriptor.get_ntypes(), - dim_descrpt=descriptor.get_dim_out(), - mixed_types=descriptor.mixed_types(), - **data["fitting_net"], - ) - return BaseModel.get_class_by_type(fitting_type)( - descriptor=descriptor, - fitting=fitting, - type_map=data["type_map"], - atom_exclude_types=data.get("atom_exclude_types", []), - pair_exclude_types=data.get("pair_exclude_types", []), - ) - - -def get_zbl_model(data: dict) -> DPZBLModel: - data = deepcopy(data) - data["descriptor"]["ntypes"] = len(data["type_map"]) - data["descriptor"]["type_map"] = data["type_map"] - descriptor_type = data["descriptor"].pop("type") - descriptor = BaseDescriptor.get_class_by_type(descriptor_type)(**data["descriptor"]) - fitting_type = data["fitting_net"].pop("type") - data["fitting_net"]["type_map"] = data["type_map"] - if fitting_type == "ener": - fitting = EnergyFittingNet( - ntypes=descriptor.get_ntypes(), - dim_descrpt=descriptor.get_dim_out(), - mixed_types=descriptor.mixed_types(), - **data["fitting_net"], - ) - else: - raise ValueError(f"Unknown fitting type {fitting_type}") - - dp_model = DPAtomicModel(descriptor, fitting, type_map=data["type_map"]) - # pairtab - filepath = data["use_srtab"] - pt_model = PairTabAtomicModel( - filepath, - data["descriptor"]["rcut"], - data["descriptor"]["sel"], - type_map=data["type_map"], - ) - rmin = data["sw_rmin"] - rmax = data["sw_rmax"] - atom_exclude_types = data.get("atom_exclude_types", []) - pair_exclude_types = data.get("pair_exclude_types", []) - return DPZBLModel( - dp_model, - pt_model, - rmin, - rmax, - type_map=data["type_map"], - atom_exclude_types=atom_exclude_types, - pair_exclude_types=pair_exclude_types, - ) - - -def get_model(data: dict) -> BaseModel: - """Get a model from a dictionary. - - Parameters - ---------- - data : dict - The data to construct the model. - """ - model_type = data.get("type", "standard") - if model_type == "standard": - if "spin" in data: - raise NotImplementedError("Spin model is not implemented yet.") - elif "use_srtab" in data: - return get_zbl_model(data) - else: - return get_standard_model(data) - else: - return BaseModel.get_class_by_type(model_type).get_model(data) +_model_factory = BackendModelFactory( + descriptor_base=BaseDescriptor, + fitting_base=BaseFitting, + model_base=BaseModel, + backend_name="TF2", + atomic_model=DPAtomicModel, + pairtab_model=PairTabAtomicModel, + zbl_model=DPZBLModel, +) +get_standard_model = _model_factory.get_standard_model +get_zbl_model = _model_factory.get_zbl_model +get_model = _model_factory.get_model diff --git a/source/tests/common/dpmodel/test_model_factory.py b/source/tests/common/dpmodel/test_model_factory.py new file mode 100644 index 0000000000..0974c0c5cb --- /dev/null +++ b/source/tests/common/dpmodel/test_model_factory.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for the shared dpmodel-driven model factory dispatch.""" + +import unittest + +from deepmd.dpmodel.model.model_factory import ( + get_model, + get_standard_model, +) + + +class _RegisteredModel: + @classmethod + def get_model(cls, data: dict) -> tuple[str, dict]: + return "registered", data + + +class _BaseModel: + @classmethod + def get_class_by_type(cls, model_type: str) -> type[_RegisteredModel]: + if model_type != "registered": + raise KeyError(model_type) + return _RegisteredModel + + +class _Descriptor: + def __init__(self, **kwargs) -> None: + self.kwargs = kwargs + + def get_ntypes(self) -> int: + return self.kwargs["ntypes"] + + def mixed_types(self) -> bool: + return True + + def get_dim_emb(self) -> int: + return 7 + + def get_dim_out(self) -> int: + return 11 + + +class _DescriptorBase: + @classmethod + def get_class_by_type(cls, descriptor_type: str) -> type[_Descriptor]: + if descriptor_type != "descriptor": + raise KeyError(descriptor_type) + return _Descriptor + + +class _Fitting: + def __init__(self, **kwargs) -> None: + self.kwargs = kwargs + + +class _FittingBase: + @classmethod + def get_class_by_type(cls, fitting_type: str) -> type[_Fitting]: + if fitting_type != "dipole": + raise KeyError(fitting_type) + return _Fitting + + +class _StandardModel: + def __init__(self, **kwargs) -> None: + self.kwargs = kwargs + + +class _StandardModelBase: + @classmethod + def get_class_by_type(cls, model_type: str) -> type[_StandardModel]: + if model_type != "dipole": + raise KeyError(model_type) + return _StandardModel + + +def _factory(name: str): + def factory(data: dict) -> tuple[str, dict]: + return name, data + + return factory + + +class TestModelFactory(unittest.TestCase): + """Verify common routing and backend extension points.""" + + def test_standard_routes(self) -> None: + standard = _factory("standard") + spin = _factory("spin") + zbl = _factory("zbl") + + self.assertEqual( + get_model( + {}, + base_model=_BaseModel, + standard_model_factory=standard, + spin_model_factory=spin, + zbl_model_factory=zbl, + )[0], + "standard", + ) + self.assertEqual( + get_model( + {"spin": {}, "use_srtab": "table"}, + base_model=_BaseModel, + standard_model_factory=standard, + spin_model_factory=spin, + zbl_model_factory=zbl, + )[0], + "spin", + ) + self.assertEqual( + get_model( + {"use_srtab": "table"}, + base_model=_BaseModel, + standard_model_factory=standard, + spin_model_factory=spin, + zbl_model_factory=zbl, + )[0], + "zbl", + ) + + def test_unsupported_standard_variant(self) -> None: + with self.assertRaisesRegex( + NotImplementedError, "Spin model is not implemented yet" + ): + get_model( + {"spin": {}}, + base_model=_BaseModel, + standard_model_factory=_factory("standard"), + ) + + def test_unsupported_zbl_variant(self) -> None: + with self.assertRaisesRegex( + NotImplementedError, "ZBL model is not implemented yet" + ): + get_model( + {"use_srtab": "table"}, + base_model=_BaseModel, + standard_model_factory=_factory("standard"), + ) + + def test_explicit_factory_precedes_registry(self) -> None: + result = get_model( + {"type": "custom"}, + base_model=_BaseModel, + standard_model_factory=_factory("standard"), + model_factories={"custom": _factory("custom")}, + ) + self.assertEqual(result[0], "custom") + + def test_registry_fallback(self) -> None: + data = {"type": "registered"} + self.assertEqual( + get_model( + data, + base_model=_BaseModel, + standard_model_factory=_factory("standard"), + ), + ("registered", data), + ) + + def test_standard_construction_is_shared_and_non_mutating(self) -> None: + data = { + "type_map": ["O", "H"], + "descriptor": {"type": "descriptor", "custom": 3}, + "fitting_net": {"type": "dipole", "custom": 5}, + "atom_exclude_types": [1], + "pair_exclude_types": [[0, 1]], + } + expected = { + "type_map": ["O", "H"], + "descriptor": {"type": "descriptor", "custom": 3}, + "fitting_net": {"type": "dipole", "custom": 5}, + "atom_exclude_types": [1], + "pair_exclude_types": [[0, 1]], + } + model = get_standard_model( + data, + descriptor_base=_DescriptorBase, + fitting_base=_FittingBase, + model_base=_StandardModelBase, + backend_name="test", + ) + + self.assertEqual(data, expected) + self.assertEqual(model.kwargs["descriptor"].kwargs["ntypes"], 2) + self.assertEqual(model.kwargs["descriptor"].kwargs["type_map"], ["O", "H"]) + self.assertEqual(model.kwargs["fitting"].kwargs["ntypes"], 2) + self.assertEqual(model.kwargs["fitting"].kwargs["dim_descrpt"], 11) + self.assertEqual(model.kwargs["fitting"].kwargs["embedding_width"], 7) + self.assertEqual(model.kwargs["atom_exclude_types"], [1]) + self.assertEqual(model.kwargs["pair_exclude_types"], [[0, 1]]) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt_expt/model/test_get_model_zbl.py b/source/tests/pt_expt/model/test_get_model_zbl.py new file mode 100644 index 0000000000..a41cc70e82 --- /dev/null +++ b/source/tests/pt_expt/model/test_get_model_zbl.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Test legacy ZBL routing in the pt_expt model factory.""" + +import os +import unittest + +from deepmd.pt_expt.model import ( + DPZBLModel, + get_model, +) + +from ...seed import ( + GLOBAL_SEED, +) + +TESTS_DIR = os.path.dirname(os.path.dirname(os.path.dirname(__file__))) +TAB_FILE = os.path.join( + TESTS_DIR, + "pt", + "model", + "water", + "data", + "zbl_tab_potential", + "H2O_tab_potential.txt", +) + + +class TestGetModelZBL(unittest.TestCase): + """Ensure ``use_srtab`` selects the backend-native ZBL model.""" + + def test_get_model_routes_use_srtab_to_zbl(self) -> None: + model = get_model( + { + "type_map": ["O", "H", "B"], + "use_srtab": TAB_FILE, + "smin_alpha": 0.37, + "sw_rmin": 0.2, + "sw_rmax": 4.0, + "descriptor": { + "type": "dpa1", + "rcut_smth": 0.5, + "rcut": 4.0, + "sel": 20, + "neuron": [3, 6], + "axis_neuron": 2, + "attn": 4, + "attn_layer": 2, + "attn_dotr": True, + "attn_mask": False, + "activation_function": "tanh", + "set_davg_zero": True, + "type_one_side": True, + "seed": GLOBAL_SEED, + }, + "fitting_net": { + "type": "ener", + "neuron": [], + "seed": GLOBAL_SEED, + }, + } + ) + self.assertIsInstance(model, DPZBLModel) + self.assertEqual(model.atomic_model.smin_alpha, 0.37) + + +if __name__ == "__main__": + unittest.main() From b1392f5f34b9918c7f2c8097d43fbf1b01add504 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Thu, 16 Jul 2026 18:02:17 +0800 Subject: [PATCH 2/3] test(model): cover shared ZBL factory paths Cover shared model-factory error branches, preserve non-default ZBL softmin values across the dpmodel-driven backends, and avoid an unnecessary ZBL config copy. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/dpmodel/model/model_factory.py | 1 - .../common/dpmodel/test_model_factory.py | 93 ++++++++++++++++++- source/tests/consistent/test_tf2_zbl_model.py | 3 +- source/tests/jax/test_zbl_model.py | 3 +- 4 files changed, 96 insertions(+), 4 deletions(-) diff --git a/deepmd/dpmodel/model/model_factory.py b/deepmd/dpmodel/model/model_factory.py index 3da519b4cf..0ababa77ab 100644 --- a/deepmd/dpmodel/model/model_factory.py +++ b/deepmd/dpmodel/model/model_factory.py @@ -97,7 +97,6 @@ def get_zbl_model( backend_name: str, ) -> Any: """Construct a ZBL model from backend-native atomic model classes.""" - data = copy.deepcopy(data) descriptor, fitting, fitting_type = get_model_components( data, descriptor_base=descriptor_base, diff --git a/source/tests/common/dpmodel/test_model_factory.py b/source/tests/common/dpmodel/test_model_factory.py index 0974c0c5cb..d73798cde7 100644 --- a/source/tests/common/dpmodel/test_model_factory.py +++ b/source/tests/common/dpmodel/test_model_factory.py @@ -6,6 +6,7 @@ from deepmd.dpmodel.model.model_factory import ( get_model, get_standard_model, + get_zbl_model, ) @@ -39,6 +40,12 @@ def get_dim_emb(self) -> int: def get_dim_out(self) -> int: return 11 + def get_rcut(self) -> float: + return 5.0 + + def get_sel(self) -> list[int]: + return [4, 8] + class _DescriptorBase: @classmethod @@ -56,7 +63,7 @@ def __init__(self, **kwargs) -> None: class _FittingBase: @classmethod def get_class_by_type(cls, fitting_type: str) -> type[_Fitting]: - if fitting_type != "dipole": + if fitting_type not in {"dipole", "ener"}: raise KeyError(fitting_type) return _Fitting @@ -74,6 +81,30 @@ def get_class_by_type(cls, model_type: str) -> type[_StandardModel]: return _StandardModel +class _AtomicModel: + def __init__(self, descriptor, fitting, **kwargs) -> None: + self.descriptor = descriptor + self.fitting = fitting + self.kwargs = kwargs + + +class _PairTabModel: + def __init__(self, table, rcut, sel, **kwargs) -> None: + self.table = table + self.rcut = rcut + self.sel = sel + self.kwargs = kwargs + + +class _ZBLModel: + def __init__(self, dp_model, pairtab, sw_rmin, sw_rmax, **kwargs) -> None: + self.dp_model = dp_model + self.pairtab = pairtab + self.sw_rmin = sw_rmin + self.sw_rmax = sw_rmax + self.kwargs = kwargs + + def _factory(name: str): def factory(data: dict) -> tuple[str, dict]: return name, data @@ -192,6 +223,66 @@ def test_standard_construction_is_shared_and_non_mutating(self) -> None: self.assertEqual(model.kwargs["atom_exclude_types"], [1]) self.assertEqual(model.kwargs["pair_exclude_types"], [[0, 1]]) + def test_model_level_type_embedding_is_rejected(self) -> None: + """Cover the shared validation used by every dpmodel-driven backend.""" + with self.assertRaisesRegex(ValueError, "type_embedding is not at the model"): + get_standard_model( + { + "type_map": ["O", "H"], + "type_embedding": {}, + "descriptor": {"type": "descriptor"}, + "fitting_net": {"type": "dipole"}, + }, + descriptor_base=_DescriptorBase, + fitting_base=_FittingBase, + model_base=_StandardModelBase, + backend_name="test", + ) + + def test_zbl_rejects_non_energy_fitting(self) -> None: + """ZBL construction accepts only an energy fitting network.""" + with self.assertRaisesRegex(ValueError, "Unknown fitting type dipole"): + get_zbl_model( + { + "type_map": ["O", "H"], + "descriptor": {"type": "descriptor"}, + "fitting_net": {"type": "dipole"}, + "use_srtab": "table", + "sw_rmin": 0.2, + "sw_rmax": 4.0, + }, + descriptor_base=_DescriptorBase, + fitting_base=_FittingBase, + atomic_model=_AtomicModel, + pairtab_model=_PairTabModel, + zbl_model=_ZBLModel, + backend_name="test", + ) + + def test_zbl_forwards_nondefault_softmin_and_descriptor_cutoff(self) -> None: + """Preserve configured softmin values and normalized descriptor geometry.""" + model = get_zbl_model( + { + "type_map": ["O", "H"], + "descriptor": {"type": "descriptor"}, + "fitting_net": {"type": "ener"}, + "use_srtab": "table", + "sw_rmin": 0.2, + "sw_rmax": 4.0, + "smin_alpha": 0.37, + }, + descriptor_base=_DescriptorBase, + fitting_base=_FittingBase, + atomic_model=_AtomicModel, + pairtab_model=_PairTabModel, + zbl_model=_ZBLModel, + backend_name="test", + ) + + self.assertEqual(model.kwargs["smin_alpha"], 0.37) + self.assertEqual(model.pairtab.rcut, 5.0) + self.assertEqual(model.pairtab.sel, [4, 8]) + if __name__ == "__main__": unittest.main() diff --git a/source/tests/consistent/test_tf2_zbl_model.py b/source/tests/consistent/test_tf2_zbl_model.py index 636d74c5d4..3e295048c4 100644 --- a/source/tests/consistent/test_tf2_zbl_model.py +++ b/source/tests/consistent/test_tf2_zbl_model.py @@ -35,7 +35,7 @@ def _zbl_config() -> dict: "use_srtab": SRTAB, "sw_rmin": 0.2, "sw_rmax": 4.0, - "smin_alpha": 0.1, + "smin_alpha": 0.37, # ZBL wraps a linear atomic model, which requires a mixed-type descriptor "descriptor": { "type": "se_atten", @@ -74,6 +74,7 @@ def test_injects_type_map_into_subconfigs(self) -> None: dp_atomic = model.atomic_model.models[0] self.assertEqual(list(dp_atomic.descriptor.get_type_map()), data["type_map"]) self.assertEqual(list(dp_atomic.fitting_net.get_type_map()), data["type_map"]) + self.assertEqual(model.atomic_model.smin_alpha, data["smin_alpha"]) if __name__ == "__main__": diff --git a/source/tests/jax/test_zbl_model.py b/source/tests/jax/test_zbl_model.py index befa240fa7..df96ec4ce2 100644 --- a/source/tests/jax/test_zbl_model.py +++ b/source/tests/jax/test_zbl_model.py @@ -30,7 +30,7 @@ def _zbl_config() -> dict: "use_srtab": SRTAB, "sw_rmin": 0.2, "sw_rmax": 4.0, - "smin_alpha": 0.1, + "smin_alpha": 0.37, # ZBL wraps a linear atomic model, which requires a mixed-type descriptor "descriptor": { "type": "se_atten", @@ -68,6 +68,7 @@ def test_injects_type_map_into_subconfigs(self) -> None: dp_atomic = model.atomic_model.models[0] self.assertEqual(dp_atomic.descriptor.get_type_map(), data["type_map"]) self.assertEqual(dp_atomic.fitting_net.get_type_map(), data["type_map"]) + self.assertEqual(model.atomic_model.smin_alpha, data["smin_alpha"]) if __name__ == "__main__": From b03a169b379dc0545f3de8acc5561f9115790a11 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 28 Jul 2026 05:47:55 +0000 Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/model/__init__.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/deepmd/dpmodel/model/__init__.py b/deepmd/dpmodel/model/__init__.py index f9889e66b9..462ee802d2 100644 --- a/deepmd/dpmodel/model/__init__.py +++ b/deepmd/dpmodel/model/__init__.py @@ -18,15 +18,15 @@ from .dos_model import ( DOSModel, ) +from .dp_linear_model import ( + LinearEnergyModel, +) from .dp_model import ( DPModelCommon, ) from .dp_zbl_model import ( DPZBLModel, ) -from .dp_linear_model import ( - LinearEnergyModel, -) from .dpa4_model import ( DPA4EnergyModel, )