Skip to content

refactor(model): unify dpmodel backend factories #5786

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
njzjz merged 4 commits into from
Jul 28, 2026
Merged

refactor(model): unify dpmodel backend factories #5786

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
6e5a4a9
refactor(model): unify dpmodel backend factories
njzjz-bot Jul 13, 2026
b1392f5
test(model): cover shared ZBL factory paths
njzjz-bot Jul 16, 2026
757ccb3
chore(model): merge master into factory refactor
njzjz-bot Jul 28, 2026
b03a169
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jul 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 28 additions & 0 deletions deepmd/dpmodel/model/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand All @@ -29,8 +50,15 @@
)

__all__ = [
"DOSModel",
"DPA4EnergyModel",
"DPModelCommon",
"DPZBLModel",
"DipoleModel",
"EnergyModel",
"LinearEnergyModel",
"NativeSpinEnergyModel",
"PolarModel",
"PropertyModel",
"SpinModel",
"make_model",
Expand Down
231 changes: 53 additions & 178 deletions deepmd/dpmodel/model/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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,
)
Expand All @@ -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.
Expand All @@ -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:
Expand All @@ -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:
Expand Down Expand Up @@ -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,
)
Loading
Loading