refactor(model): unify dpmodel backend factories (deepmodeling#5786)

## Summary

- introduce BackendModelFactory to bind backend descriptor, fitting,
model, atomic-model, pair-table, and ZBL registries once
- centralize descriptor/fitting parameter injection, standard-model
selection, ZBL assembly, legacy model routing, plugin fallback, and spin
preprocessing
- make JAX and TF2 export the bound shared factory methods directly;
keep only genuine backend special cases in dpmodel and pt_expt
- reuse the shared component path for pt_expt DPA4/SeZM and linear
submodels
- register/export all dpmodel model types through the package, removing
unused side-effect imports
- preserve configured ZBL smin_alpha values and route pt_expt use_srtab
configurations to its backend-native DPZBLModel
- use descriptor-derived cutoff and selection values for JAX/TF2 ZBL
pair tables, carrying the normalized deepmodeling#4339 behavior beyond dpmodel

## Validation

- ruff check .
- ruff format .
- all commit hooks, including isort, Ruff, Velin, and pylint
- shared routing, construction, input-immutability, and
unsupported-variant tests
- dpmodel standard/spin construction tests
- JAX standard and ZBL factory tests
- pt_expt standard, spin, ZBL, DPA4/SeZM, and linear-model tests
- isolated TF2 standard-model construction smoke test


Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added unified, backend-agnostic model construction with routing for
standard, spin, ZBL, linear-energy, and registered variants.
* Expanded the public model API to include linear-energy, DPA4 energy,
native-spin, ZBL, and additional property models (DOS, Dipole, Polar).
* **Bug Fixes**
* Improved configuration handling so caller-provided settings are
preserved.
* Ensured tabulated-potential (`use_srtab`) configurations correctly
produce ZBL models.
* **Tests**
* Added coverage for model routing, unsupported variants, shared
construction behavior, and configuration preservation.
  * Added a dedicated pt-expt ZBL selection test using `use_srtab`.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: njzjz-bot <njzjz.bot@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: njzjz-bot

deepmd/dpmodel/model/__init__.py

@@ -12,15 +12,36 @@
 Models generated by `make_model` have already done it.
 """
 
+from .dipole_model import (
+    DipoleModel,
+)
+from .dos_model import (
+    DOSModel,
+)
+from .dp_linear_model import (
+    LinearEnergyModel,
+)
 from .dp_model import (
     DPModelCommon,
 )
+from .dp_zbl_model import (
+    DPZBLModel,
+)
+from .dpa4_model import (
+    DPA4EnergyModel,
+)
 from .ener_model import (
     EnergyModel,
 )
 from .make_model import (
     make_model,
 )
+from .native_spin_model import (
+    NativeSpinEnergyModel,
+)
+from .polar_model import (
+    PolarModel,
+)
 from .property_model import (
     PropertyModel,
 )
@@ -29,8 +50,15 @@
 )
 
 __all__ = [
+    "DOSModel",
+    "DPA4EnergyModel",
     "DPModelCommon",
+    "DPZBLModel",
+    "DipoleModel",
     "EnergyModel",
+    "LinearEnergyModel",
+    "NativeSpinEnergyModel",
+    "PolarModel",
     "PropertyModel",
     "SpinModel",
     "make_model",

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,36 +13,21 @@
 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.dpa4_model import (
-    DPA4EnergyModel,
+from deepmd.dpmodel.model.model_factory import (
+    BackendModelFactory,
 )
-from deepmd.dpmodel.model.ener_model import (
-    EnergyModel,
+from deepmd.dpmodel.model.model_factory import (
+    get_spin_model as get_spin_model_from_factory,
 )
 from deepmd.dpmodel.model.native_spin_model import (
     NativeSpinEnergyModel,
 )
-from deepmd.dpmodel.model.polar_model import (
-    PolarModel,
-)
-from deepmd.dpmodel.model.property_model import (
-    PropertyModel,
-)
 from deepmd.dpmodel.model.spin_model import (
     SpinModel,
 )
@@ -54,48 +36,29 @@
     normalize_spin_use_spin,
 )
 
-_DPA4_SEZM_DESCRIPTOR_TYPES = ("dpa4", "DPA4", "sezm", "SeZM")
-
+_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_zbl_model = _model_factory.get_zbl_model
 
-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"]
+_DPA4_SEZM_DESCRIPTOR_TYPES = ("dpa4", "DPA4", "sezm", "SeZM")
 
 
-def get_standard_model(data: dict) -> EnergyModel:
-    """Get a EnergyModel from a dictionary.
+def get_standard_model(data: dict) -> BaseModel:
+    """Get a standard model 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"])
     # Analytical bridging (e.g. ZBL): the radii feed the DESCRIPTOR's
     # InnerClamp/BridgingSwitch (mirrors pt's builder); the method builds the
     # atomic model's InterPotential below.
@@ -104,105 +67,42 @@ def get_standard_model(data: dict) -> EnergyModel:
     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)
-    descriptor, fitting, fitting_net_type = _get_standard_model_components(data, ntypes)
+    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", [])
-
-    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 in ["dpa4_ener", "sezm_ener"]:
-        modelcls = DPA4EnergyModel
-    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,
+    # 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.inter_potential import (
+        InterPotentialAtomicModel,
+    )
+    from deepmd.dpmodel.atomic_model.linear_atomic_model import (
+        LinearEnergyAtomicModel,
+    )
+    from deepmd.dpmodel.model.dp_linear_model import (
+        LinearEnergyModel,
     )
-    if bridging_enabled:
-        # 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.inter_potential import (
-            InterPotentialAtomicModel,
-        )
-        from deepmd.dpmodel.atomic_model.linear_atomic_model import (
-            LinearEnergyAtomicModel,
-        )
-        from deepmd.dpmodel.model.dp_linear_model import (
-            LinearEnergyModel,
-        )
-
-        zbl_atomic = InterPotentialAtomicModel(
-            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,
-        )
-        return LinearEnergyModel(atomic_model_=composed)
-    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(),
+    zbl_atomic = InterPotentialAtomicModel(
         type_map=data["type_map"],
+        mode=bridging_method,
+        rcut=descriptor.get_rcut(),
+        sel=descriptor.get_sel(),
     )
-
-    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,
+    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,
     )
+    return LinearEnergyModel(atomic_model_=composed)
 
 
 def get_spin_model(data: dict) -> SpinModel:
@@ -218,30 +118,11 @@ def get_spin_model(data: dict) -> SpinModel:
             "the virtual-atom (deepspin) scheme is not supported for "
             "DPA4/SeZM; use spin scheme 'native'"
         )
-    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"],
-    )
-    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=SpinModel,
     )
-    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_native_spin_model(data: dict) -> NativeSpinEnergyModel:
@@ -315,15 +196,9 @@ 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:
-            if data["spin"].get("scheme", "deepspin") == "native":
-                return get_native_spin_model(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,
+        standard_model_factory=get_standard_model,
+        spin_model_factory=get_spin_model,
+        native_spin_model_factory=get_native_spin_model,
+    )