feat(pt/dpa4): Support property fitting in DPA4/SeZM models

deepmd/pt/model/atomic_model/sezm_atomic_model.py

@@ -149,6 +149,22 @@ def get_active_mode(self) -> str:
         """Return the current SeZM execution mode."""
         return str(getattr(self, "_active_mode", "ener"))
 
+    def get_compute_stats_distinguish_types(self) -> bool:
+        """Return whether output statistics are type-resolved."""
+        active_fitting = self.get_active_fitting_net()
+        if active_fitting is not None and hasattr(
+            active_fitting, "get_distinguish_types"
+        ):
+            return bool(active_fitting.get_distinguish_types())
+        return super().get_compute_stats_distinguish_types()
+
+    def get_intensive(self) -> bool:
+        """Return whether the active reducible output is intensive."""
+        active_fitting = self.get_active_fitting_net()
+        if active_fitting is not None and hasattr(active_fitting, "get_intensive"):
+            return bool(active_fitting.get_intensive())
+        return super().get_intensive()
+
     def _compute_or_load_dens_force_stat(
         self,
         sampled_func: Any,
@@ -595,9 +611,16 @@ def apply_out_stat(
         dict[str, torch.Tensor]
             Outputs after SeZM output-stat post-processing.
         """
-        if "energy" in ret:
-            out_bias, _ = self._fetch_out_stat(["energy"])
-            ret["energy"] = ret["energy"] + out_bias["energy"][atype]
+        out_bias, out_std = self._fetch_out_stat(self.bias_keys)
+        for key in self.bias_keys:
+            if key not in ret:
+                continue
+            if key == "energy":
+                ret[key] = ret[key] + out_bias[key][atype]
+            elif self.get_compute_stats_distinguish_types():
+                ret[key] = ret[key] * out_std[key][atype] + out_bias[key][atype]
+            else:
+                ret[key] = ret[key] * out_std[key][0] + out_bias[key][0]
         return ret
 
     def get_dim_fparam(self) -> int:

deepmd/pt/model/model/__init__.py

@@ -31,6 +31,7 @@
 )
 from deepmd.pt.model.task.sezm_ener import (
     SeZMEnergyFittingNet,
+    _resolve_auto_neuron,
 )
 from deepmd.utils.spin import (
     Spin,
@@ -75,6 +76,9 @@
 from .sezm_model import (
     SeZMModel,
 )
+from .sezm_property_model import (
+    SeZMPropertyModel,
+)
 from .sezm_spin_model import (
     SeZMSpinModel,
 )
@@ -341,12 +345,37 @@ def get_sezm_model(model_params: dict) -> BaseModel:
     descriptor = BaseDescriptor(**model_params["descriptor"])
 
     fitting_net = copy.deepcopy(model_params["fitting_net"])
-    fitting_net.pop("type", None)
+    fitting_net.setdefault("type", "dpa4_ener")
     fitting_net["ntypes"] = descriptor.get_ntypes()
     fitting_net["type_map"] = copy.deepcopy(model_params["type_map"])
     fitting_net["mixed_types"] = descriptor.mixed_types()
     fitting_net["dim_descrpt"] = descriptor.get_dim_out()
-    fitting = SeZMEnergyFittingNet(**fitting_net)
+    if fitting_net["type"] in ("dpa4_ener", "sezm_ener"):
+        fitting = BaseFitting(**fitting_net)
+        modelcls = SeZMModel
+    elif fitting_net["type"] == "property":
+        if bridging_method != "NONE":
+            raise ValueError(
+                "DPA4/SeZM property fitting does not support analytical bridging "
+                "potentials; set `bridging_method` to `none`."
+            )
+        # Share the SeZM auto-width convention
+        fitting_net["neuron"] = _resolve_auto_neuron(
+            fitting_net.get("neuron"),
+            dim_descrpt=fitting_net["dim_descrpt"],
+            numb_fparam=fitting_net.get("numb_fparam", 0),
+            numb_aparam=fitting_net.get("numb_aparam", 0),
+            dim_case_embd=fitting_net.get("dim_case_embd", 0),
+            case_film_embd=fitting_net.get("case_film_embd", False),
+            use_aparam_as_mask=fitting_net.get("use_aparam_as_mask", False),
+        )
+        fitting = BaseFitting(**fitting_net)
+        modelcls = SeZMPropertyModel
+    else:
+        raise ValueError(
+            "DPA4/SeZM model supports `dpa4_ener`, `sezm_ener`, or `property` "
+            f"fitting, but got `{fitting_net['type']}`."
+        )
     atom_exclude_types = model_params.get("atom_exclude_types", [])
     preset_out_bias = model_params.get("preset_out_bias")
     preset_out_bias = _convert_preset_out_bias_to_array(
@@ -356,7 +385,7 @@ def get_sezm_model(model_params: dict) -> BaseModel:
     use_compile = bool(model_params.get("use_compile", False))
     enable_tf32 = bool(model_params.get("enable_tf32", True))
 
-    model = SeZMModel(
+    model = modelcls(
         descriptor=descriptor,
         fitting=fitting,
         type_map=model_params["type_map"],
@@ -416,6 +445,12 @@ def get_sezm_spin_model(model_params: dict) -> BaseModel:
     descriptor = BaseDescriptor(**model_params["descriptor"])
 
     fitting_net = copy.deepcopy(model_params["fitting_net"])
+    fitting_net_type = fitting_net.get("type", "dpa4_ener")
+    if fitting_net_type not in ("dpa4_ener", "sezm_ener"):
+        raise ValueError(
+            "Spin DPA4/SeZM currently supports only `dpa4_ener` or `sezm_ener` "
+            f"fitting, but got `{fitting_net_type}`."
+        )
     fitting_net.pop("type", None)
     fitting_net["ntypes"] = descriptor.get_ntypes()
     fitting_net["type_map"] = copy.deepcopy(model_params["type_map"])
@@ -461,7 +496,7 @@ def get_model(model_params: dict) -> Any:
             return get_standard_model(model_params)
     elif model_type == "linear_ener":
         return get_linear_model(model_params)
-    elif model_type in ("SeZM", "sezm", "dpa4"):
+    elif model_type in ("SeZM", "sezm", "DPA4", "dpa4"):
         if "spin" in model_params:
             return get_sezm_spin_model(model_params)
         return get_sezm_model(model_params)
@@ -480,6 +515,7 @@ def get_model(model_params: dict) -> Any:
     "LinearEnergyModel",
     "PolarModel",
     "SeZMModel",
+    "SeZMPropertyModel",
     "SeZMSpinModel",
     "SpinEnergyModel",
     "SpinModel",

deepmd/pt/model/model/sezm_model.py

@@ -487,6 +487,7 @@
 )
 from deepmd.pt.model.model.transform_output import (
     edge_energy_deriv,
+    fit_output_to_model_output,
 )
 from deepmd.pt.utils import (
     env,
@@ -1363,6 +1364,7 @@ def core_compute(
         extended_atype: torch.Tensor | None = None,
         extended_coord_corr: torch.Tensor | None = None,
         embedding_only: bool = False,
+        conservative: bool = True,
     ) -> dict[str, torch.Tensor]:
         """
         Compute SeZM lower outputs from the unified edge-vector schema.
@@ -1406,6 +1408,11 @@ def core_compute(
         embedding_only
             When ``True``, return only the embedding outputs and skip the
             force/virial autograd entirely.
+        conservative
+            Whether to run the conservative energy derivative path. Energy
+            fitting keeps this enabled. Non-conservative property fitting
+            disables it, so fitting outputs are reduced by their output
+            definition without constructing edge-force gradients.
 
         Returns
         -------
@@ -1428,8 +1435,9 @@ def core_compute(
         # This keeps coordinate gathering and shift application outside the
         # differentiated region while preserving conservative forces through the
         # scatter indices below.  The embedding path produces no force, so it
-        # keeps ``edge_vec`` detached and never allocates an autograd leaf.
-        if not embedding_only:
+        # keeps ``edge_vec`` detached and never allocates an autograd leaf. The
+        # same forward-only treatment is used by non-conservative property heads.
+        if conservative and not embedding_only:
             edge_vec = edge_vec.detach().requires_grad_(True)
 
         # === Step 2. Descriptor forward ===
@@ -1502,10 +1510,20 @@ def core_compute(
             ).view(out_shape)
         fit_ret["mask"] = atom_mask
 
+        if not conservative:
+            return fit_output_to_model_output(
+                fit_ret,
+                self.atomic_output_def(),
+                coord,
+                create_graph=False,
+                mask=fit_ret["mask"],
+                extended_coord_corr=extended_coord_corr,
+            )
+
         # === Step 5. Inject analytical pair potential (edge form) ===
         # ZBL is evaluated from ``edge_vec`` (the autograd leaf) so its force
         # and virial flow through the same edge backward as the learned energy.
-        if self.inter_potential is not None:
+        if self.inter_potential is not None and "energy" in fit_ret:
             fit_ret["energy"] = fit_ret["energy"] + self.inter_potential(
                 edge_vec=edge_vec,
                 edge_index=edge_index,
@@ -2091,8 +2109,9 @@ def compute_fn(  # type: ignore[misc]
             traced = rebuild_graph_module(traced)
 
         # The conservative Inductor option set that keeps the dynamic edge
-        # graph lowerable is centralised in ``deepmd.pt.utils.compile_compat``.
-        compile_options = build_inductor_compile_options()
+        # graph lowerable is centralised in ``deepmd.pt.utils.compile_compat``;
+        # subclasses may augment it via ``_inductor_compile_options``.
+        compile_options = self._inductor_compile_options()
 
         # NOTE: Store the compiled callable inside the plain-``dict``
         # cache ``compiled_core_compute_cache``.  The dict itself was installed
@@ -2275,6 +2294,14 @@ def should_use_compile(self) -> bool:
             return self.use_compile
         return bool(self._env_use_compile_infer)
 
+    def _inductor_compile_options(self) -> dict[str, Any]:
+        """Return the Inductor lowering options for this model's compiled core.
+
+        Subclasses may override this to augment the shared option set from
+        :func:`build_inductor_compile_options` with model-specific entries.
+        """
+        return build_inductor_compile_options()
+
     # =========================================================================
     # Export Utilities
     # =========================================================================