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

Author: njzjz-bot

deepmd/dpmodel/model/model.py

@@ -1,8 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import copy
-from typing import (
-    Any,
-)
 
 from deepmd.dpmodel.atomic_model.dp_atomic_model import (
     DPAtomicModel,
@@ -16,29 +13,24 @@
 from deepmd.dpmodel.fitting.base_fitting import (
     BaseFitting,
 )
-from deepmd.dpmodel.fitting.ener_fitting import (
-    EnergyFittingNet,
-)
+from deepmd.dpmodel.model import dipole_model as _dipole_model  # noqa: F401
+from deepmd.dpmodel.model import dos_model as _dos_model  # noqa: F401
+from deepmd.dpmodel.model import polar_model as _polar_model  # noqa: F401
 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 get_model as get_model_from_factory
+from deepmd.dpmodel.model.model_factory import (
+    get_standard_model as get_standard_model_from_factory,
 )
-from deepmd.dpmodel.model.property_model import (
-    PropertyModel,
+from deepmd.dpmodel.model.model_factory import (
+    get_zbl_model as get_zbl_model_from_factory,
 )
 from deepmd.dpmodel.model.spin_model import (
     SpinModel,
@@ -48,31 +40,6 @@
 )
 
 
-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.
 
@@ -81,78 +48,24 @@ def get_standard_model(data: dict) -> EnergyModel:
     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 get_standard_model_from_factory(
+        data,
+        descriptor_base=BaseDescriptor,
+        fitting_base=BaseFitting,
+        model_base=BaseModel,
+        backend_name="DP",
     )
-    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,
+    return get_zbl_model_from_factory(
+        data,
+        descriptor_base=BaseDescriptor,
+        fitting_base=BaseFitting,
+        atomic_model=DPAtomicModel,
+        pairtab_model=PairTabAtomicModel,
+        zbl_model=DPZBLModel,
+        backend_name="DP",
     )
 
 
@@ -198,13 +111,10 @@ 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 get_model_from_factory(
+        data,
+        base_model=BaseModel,
+        standard_model_factory=get_standard_model,
+        spin_model_factory=get_spin_model,
+        zbl_model_factory=get_zbl_model,
+    )

deepmd/dpmodel/model/model_factory.py

@@ -0,0 +1,174 @@
+# 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,
+)
+
+ModelFactory = 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"],
+        atom_exclude_types=data.get("atom_exclude_types", []),
+        pair_exclude_types=data.get("pair_exclude_types", []),
+    )
+
+
+def get_model(
+    data: dict,
+    *,
+    base_model: type,
+    standard_model_factory: ModelFactory,
+    spin_model_factory: ModelFactory | None = None,
+    zbl_model_factory: ModelFactory | None = None,
+    model_factories: Mapping[str, ModelFactory] | 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)