Skip to content

Commit 1f0ca6c

Browse files
authored
feat(pt/dpa4): Support property fitting in DPA4/SeZM models (deepmodeling#5587)
<!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Added support for invariant property fitting with the DPA4/SeZM model family. - Property outputs can now be trained and read back with dedicated labels, including intensive or non-intensive aggregation. - Expanded model selection to recognize the DPA4 path and property-fitting configuration. - **Bug Fixes** - Improved output-stat handling so non-energy outputs are normalized more consistently. - Added validation to prevent unsupported configuration combinations. - **Documentation** - Updated training guides and examples for property-based DPA4/SeZM workflows. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
1 parent dbdc9a3 commit 1f0ca6c

9 files changed

Lines changed: 630 additions & 22 deletions

File tree

deepmd/pt/model/atomic_model/sezm_atomic_model.py

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,22 @@ def get_active_mode(self) -> str:
149149
"""Return the current SeZM execution mode."""
150150
return str(getattr(self, "_active_mode", "ener"))
151151

152+
def get_compute_stats_distinguish_types(self) -> bool:
153+
"""Return whether output statistics are type-resolved."""
154+
active_fitting = self.get_active_fitting_net()
155+
if active_fitting is not None and hasattr(
156+
active_fitting, "get_distinguish_types"
157+
):
158+
return bool(active_fitting.get_distinguish_types())
159+
return super().get_compute_stats_distinguish_types()
160+
161+
def get_intensive(self) -> bool:
162+
"""Return whether the active reducible output is intensive."""
163+
active_fitting = self.get_active_fitting_net()
164+
if active_fitting is not None and hasattr(active_fitting, "get_intensive"):
165+
return bool(active_fitting.get_intensive())
166+
return super().get_intensive()
167+
152168
def _compute_or_load_dens_force_stat(
153169
self,
154170
sampled_func: Any,
@@ -595,9 +611,16 @@ def apply_out_stat(
595611
dict[str, torch.Tensor]
596612
Outputs after SeZM output-stat post-processing.
597613
"""
598-
if "energy" in ret:
599-
out_bias, _ = self._fetch_out_stat(["energy"])
600-
ret["energy"] = ret["energy"] + out_bias["energy"][atype]
614+
out_bias, out_std = self._fetch_out_stat(self.bias_keys)
615+
for key in self.bias_keys:
616+
if key not in ret:
617+
continue
618+
if key == "energy":
619+
ret[key] = ret[key] + out_bias[key][atype]
620+
elif self.get_compute_stats_distinguish_types():
621+
ret[key] = ret[key] * out_std[key][atype] + out_bias[key][atype]
622+
else:
623+
ret[key] = ret[key] * out_std[key][0] + out_bias[key][0]
601624
return ret
602625

603626
def get_dim_fparam(self) -> int:

deepmd/pt/model/model/__init__.py

Lines changed: 40 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
)
3232
from deepmd.pt.model.task.sezm_ener import (
3333
SeZMEnergyFittingNet,
34+
_resolve_auto_neuron,
3435
)
3536
from deepmd.utils.spin import (
3637
Spin,
@@ -78,6 +79,9 @@
7879
from .sezm_model import (
7980
SeZMModel,
8081
)
82+
from .sezm_property_model import (
83+
SeZMPropertyModel,
84+
)
8185
from .sezm_spin_model import (
8286
SeZMSpinModel,
8387
)
@@ -346,12 +350,37 @@ def get_sezm_model(model_params: dict) -> BaseModel:
346350
descriptor = BaseDescriptor(**model_params["descriptor"])
347351

348352
fitting_net = copy.deepcopy(model_params["fitting_net"])
349-
fitting_net.pop("type", None)
353+
fitting_net.setdefault("type", "dpa4_ener")
350354
fitting_net["ntypes"] = descriptor.get_ntypes()
351355
fitting_net["type_map"] = copy.deepcopy(model_params["type_map"])
352356
fitting_net["mixed_types"] = descriptor.mixed_types()
353357
fitting_net["dim_descrpt"] = descriptor.get_dim_out()
354-
fitting = SeZMEnergyFittingNet(**fitting_net)
358+
if fitting_net["type"] in ("dpa4_ener", "sezm_ener"):
359+
fitting = BaseFitting(**fitting_net)
360+
modelcls = SeZMModel
361+
elif fitting_net["type"] == "property":
362+
if bridging_method != "NONE":
363+
raise ValueError(
364+
"DPA4/SeZM property fitting does not support analytical bridging "
365+
"potentials; set `bridging_method` to `none`."
366+
)
367+
# Share the SeZM auto-width convention
368+
fitting_net["neuron"] = _resolve_auto_neuron(
369+
fitting_net.get("neuron"),
370+
dim_descrpt=fitting_net["dim_descrpt"],
371+
numb_fparam=fitting_net.get("numb_fparam", 0),
372+
numb_aparam=fitting_net.get("numb_aparam", 0),
373+
dim_case_embd=fitting_net.get("dim_case_embd", 0),
374+
case_film_embd=fitting_net.get("case_film_embd", False),
375+
use_aparam_as_mask=fitting_net.get("use_aparam_as_mask", False),
376+
)
377+
fitting = BaseFitting(**fitting_net)
378+
modelcls = SeZMPropertyModel
379+
else:
380+
raise ValueError(
381+
"DPA4/SeZM model supports `dpa4_ener`, `sezm_ener`, or `property` "
382+
f"fitting, but got `{fitting_net['type']}`."
383+
)
355384
atom_exclude_types = model_params.get("atom_exclude_types", [])
356385
preset_out_bias = model_params.get("preset_out_bias")
357386
preset_out_bias = _convert_preset_out_bias_to_array(
@@ -361,7 +390,7 @@ def get_sezm_model(model_params: dict) -> BaseModel:
361390
use_compile = bool(model_params.get("use_compile", False))
362391
enable_tf32 = bool(model_params.get("enable_tf32", True))
363392

364-
model = SeZMModel(
393+
model = modelcls(
365394
descriptor=descriptor,
366395
fitting=fitting,
367396
type_map=model_params["type_map"],
@@ -421,6 +450,12 @@ def get_sezm_spin_model(model_params: dict) -> BaseModel:
421450
descriptor = BaseDescriptor(**model_params["descriptor"])
422451

423452
fitting_net = copy.deepcopy(model_params["fitting_net"])
453+
fitting_net_type = fitting_net.get("type", "dpa4_ener")
454+
if fitting_net_type not in ("dpa4_ener", "sezm_ener"):
455+
raise ValueError(
456+
"Spin DPA4/SeZM currently supports only `dpa4_ener` or `sezm_ener` "
457+
f"fitting, but got `{fitting_net_type}`."
458+
)
424459
fitting_net.pop("type", None)
425460
fitting_net["ntypes"] = descriptor.get_ntypes()
426461
fitting_net["type_map"] = copy.deepcopy(model_params["type_map"])
@@ -466,7 +501,7 @@ def get_model(model_params: dict) -> Any:
466501
return get_standard_model(model_params)
467502
elif model_type == "linear_ener":
468503
return get_linear_model(model_params)
469-
elif model_type in ("SeZM", "sezm", "dpa4"):
504+
elif model_type in ("SeZM", "sezm", "DPA4", "dpa4"):
470505
if "spin" in model_params:
471506
return get_sezm_spin_model(model_params)
472507
return get_sezm_model(model_params)
@@ -486,6 +521,7 @@ def get_model(model_params: dict) -> Any:
486521
"PolarModel",
487522
"PopulationModel",
488523
"SeZMModel",
524+
"SeZMPropertyModel",
489525
"SeZMSpinModel",
490526
"SpinEnergyModel",
491527
"SpinModel",

deepmd/pt/model/model/sezm_model.py

Lines changed: 32 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -487,6 +487,7 @@
487487
)
488488
from deepmd.pt.model.model.transform_output import (
489489
edge_energy_deriv,
490+
fit_output_to_model_output,
490491
)
491492
from deepmd.pt.utils import (
492493
env,
@@ -1363,6 +1364,7 @@ def core_compute(
13631364
extended_atype: torch.Tensor | None = None,
13641365
extended_coord_corr: torch.Tensor | None = None,
13651366
embedding_only: bool = False,
1367+
conservative: bool = True,
13661368
) -> dict[str, torch.Tensor]:
13671369
"""
13681370
Compute SeZM lower outputs from the unified edge-vector schema.
@@ -1406,6 +1408,11 @@ def core_compute(
14061408
embedding_only
14071409
When ``True``, return only the embedding outputs and skip the
14081410
force/virial autograd entirely.
1411+
conservative
1412+
Whether to run the conservative energy derivative path. Energy
1413+
fitting keeps this enabled. Non-conservative property fitting
1414+
disables it, so fitting outputs are reduced by their output
1415+
definition without constructing edge-force gradients.
14091416
14101417
Returns
14111418
-------
@@ -1428,8 +1435,9 @@ def core_compute(
14281435
# This keeps coordinate gathering and shift application outside the
14291436
# differentiated region while preserving conservative forces through the
14301437
# scatter indices below. The embedding path produces no force, so it
1431-
# keeps ``edge_vec`` detached and never allocates an autograd leaf.
1432-
if not embedding_only:
1438+
# keeps ``edge_vec`` detached and never allocates an autograd leaf. The
1439+
# same forward-only treatment is used by non-conservative property heads.
1440+
if conservative and not embedding_only:
14331441
edge_vec = edge_vec.detach().requires_grad_(True)
14341442

14351443
# === Step 2. Descriptor forward ===
@@ -1502,10 +1510,20 @@ def core_compute(
15021510
).view(out_shape)
15031511
fit_ret["mask"] = atom_mask
15041512

1513+
if not conservative:
1514+
return fit_output_to_model_output(
1515+
fit_ret,
1516+
self.atomic_output_def(),
1517+
coord,
1518+
create_graph=False,
1519+
mask=fit_ret["mask"],
1520+
extended_coord_corr=extended_coord_corr,
1521+
)
1522+
15051523
# === Step 5. Inject analytical pair potential (edge form) ===
15061524
# ZBL is evaluated from ``edge_vec`` (the autograd leaf) so its force
15071525
# and virial flow through the same edge backward as the learned energy.
1508-
if self.inter_potential is not None:
1526+
if self.inter_potential is not None and "energy" in fit_ret:
15091527
fit_ret["energy"] = fit_ret["energy"] + self.inter_potential(
15101528
edge_vec=edge_vec,
15111529
edge_index=edge_index,
@@ -2091,8 +2109,9 @@ def compute_fn( # type: ignore[misc]
20912109
traced = rebuild_graph_module(traced)
20922110

20932111
# The conservative Inductor option set that keeps the dynamic edge
2094-
# graph lowerable is centralised in ``deepmd.pt.utils.compile_compat``.
2095-
compile_options = build_inductor_compile_options()
2112+
# graph lowerable is centralised in ``deepmd.pt.utils.compile_compat``;
2113+
# subclasses may augment it via ``_inductor_compile_options``.
2114+
compile_options = self._inductor_compile_options()
20962115

20972116
# NOTE: Store the compiled callable inside the plain-``dict``
20982117
# cache ``compiled_core_compute_cache``. The dict itself was installed
@@ -2275,6 +2294,14 @@ def should_use_compile(self) -> bool:
22752294
return self.use_compile
22762295
return bool(self._env_use_compile_infer)
22772296

2297+
def _inductor_compile_options(self) -> dict[str, Any]:
2298+
"""Return the Inductor lowering options for this model's compiled core.
2299+
2300+
Subclasses may override this to augment the shared option set from
2301+
:func:`build_inductor_compile_options` with model-specific entries.
2302+
"""
2303+
return build_inductor_compile_options()
2304+
22782305
# =========================================================================
22792306
# Export Utilities
22802307
# =========================================================================

0 commit comments

Comments
 (0)