refactor(model): unify dpmodel backend factories

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",

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,
+    )