feat(pt_expt): implementing energy hessian model (#5287)

### Summary

Add hessian (second derivative of energy w.r.t. coordinates) support to
the pt_expt backend, mirroring the JAX backend's approach.

- Compute hessian on extended coordinates in forward_common_atomic (via
torch.autograd.functional.hessian), then let dpmodel's
communicate_extended_output naturally map from nall×nall to nloc×nloc
- make_hessian_model only overrides atomic_output_def() to set
r_hessian=True — no forward_common override or r_hessian toggle hack
needed
- Hessian is enabled at runtime via EnergyModel.enable_hessian() and
returned through the user-facing forward() interface

### Changed files

- deepmd/pt_expt/model/make_model.py — add _cal_hessian_ext and
_WrapperForwardEnergy for hessian computation in forward_common_atomic
- deepmd/pt_expt/model/make_hessian_model.py — minimal wrapper:
__init__, requires_hessian, atomic_output_def
- deepmd/pt_expt/model/ener_model.py — enable_hessian(), hessian output
in forward() and translated_output_def()
  - deepmd/pt_expt/model/__init__.py — export make_hessian_model
- source/tests/pt_expt/model/test_ener_hessian_model.py — unit test:
autograd hessian vs finite-difference, parametrized over nv={1,2}
- source/tests/consistent/model/test_ener_hessian.py — cross-backend
consistency test covering PT, pt_expt, and JAX
  

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

* **New Features**
  * Adds Hessian (second-order derivative) support for energy models.
* Provides a mechanism to enable Hessian calculations per model
instance.
* When enabled, Hessian tensors are included in model prediction
outputs.

* **Tests**
* Added comprehensive multi-backend validation of Hessian functionality.
* Added unit tests comparing analytical Hessians to finite-difference
references.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

Author: wanghan-iapcm

deepmd/dpmodel/model/make_hessian_model.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import copy
+from typing import (
+    Any,
+)
+
+from deepmd.dpmodel.output_def import (
+    FittingOutputDef,
+)
+
+
+def make_hessian_model(T_Model: type) -> type:
+    """Make a model that can compute Hessian.
+
+    With the JAX-mirrored approach, hessian is computed in
+    ``forward_common_atomic`` (in make_model.py) on extended coordinates.
+    This wrapper only needs to override ``atomic_output_def()`` to set
+    ``r_hessian=True``, and ``communicate_extended_output`` in dpmodel
+    naturally maps it from nall to nloc.
+
+    Parameters
+    ----------
+    T_Model
+        The model. Should provide the ``atomic_output_def`` method.
+
+    Returns
+    -------
+    The model that computes hessian.
+
+    """
+
+    class CM(T_Model):
+        def __init__(
+            self,
+            *args: Any,
+            **kwargs: Any,
+        ) -> None:
+            super().__init__(
+                *args,
+                **kwargs,
+            )
+            self.hess_fitting_def = copy.deepcopy(super().atomic_output_def())
+
+        def requires_hessian(
+            self,
+            keys: str | list[str],
+        ) -> None:
+            """Set which output variable(s) requires hessian."""
+            if isinstance(keys, str):
+                keys = [keys]
+            for kk in self.hess_fitting_def.keys():
+                if kk in keys:
+                    self.hess_fitting_def[kk].r_hessian = True
+
+        def atomic_output_def(self) -> FittingOutputDef:
+            """Get the fitting output def."""
+            return self.hess_fitting_def
+
+    return CM

deepmd/pt_expt/model/__init__.py

@@ -1,4 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.dpmodel.model.make_hessian_model import (
+    make_hessian_model,
+)
+
 from .dipole_model import (
     DipoleModel,
 )
@@ -33,4 +37,5 @@
     "PolarModel",
     "PropertyModel",
     "get_model",
+    "make_hessian_model",
 ]

deepmd/pt_expt/model/ener_model.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import copy
 from typing import (
     Any,
 )
@@ -14,6 +15,9 @@
 from deepmd.dpmodel.model.dp_model import (
     DPModelCommon,
 )
+from deepmd.dpmodel.model.make_hessian_model import (
+    make_hessian_model,
+)
 
 from .make_model import (
     make_model,
@@ -34,6 +38,17 @@ def __init__(
     ) -> None:
         DPModelCommon.__init__(self)
         DPEnergyModel_.__init__(self, *args, **kwargs)
+        self._hessian_enabled = False
+
+    def enable_hessian(self) -> None:
+        if self._hessian_enabled:
+            return
+        self.__class__ = make_hessian_model(type(self))
+        self.hess_fitting_def = copy.deepcopy(
+            super(type(self), self).atomic_output_def()
+        )
+        self.requires_hessian("energy")
+        self._hessian_enabled = True
 
     def forward(
         self,
@@ -63,6 +78,8 @@ def forward(
                 model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-2)
         if "mask" in model_ret:
             model_predict["mask"] = model_ret["mask"]
+        if self.atomic_output_def()["energy"].r_hessian:
+            model_predict["hessian"] = model_ret["energy_derv_r_derv_r"].squeeze(-3)
         return model_predict
 
     def forward_lower(
@@ -115,6 +132,8 @@ def translated_output_def(self) -> dict[str, Any]:
             output_def["atom_virial"].squeeze(-2)
         if "mask" in out_def_data:
             output_def["mask"] = out_def_data["mask"]
+        if self.atomic_output_def()["energy"].r_hessian:
+            output_def["hessian"] = out_def_data["energy_derv_r_derv_r"]
         return output_def
 
     def forward_lower_exportable(

deepmd/pt_expt/model/make_model.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+import math
 from typing import (
     Any,
 )
@@ -8,10 +9,16 @@
     make_fx,
 )
 
+from deepmd.dpmodel import (
+    get_hessian_name,
+)
 from deepmd.dpmodel.atomic_model.base_atomic_model import (
     BaseAtomicModel,
 )
 from deepmd.dpmodel.model.make_model import make_model as make_model_dp
+from deepmd.dpmodel.output_def import (
+    OutputVariableDef,
+)
 from deepmd.pt_expt.common import (
     torch_module,
 )
@@ -21,6 +28,136 @@
 )
 
 
+def _cal_hessian_ext(
+    model: Any,
+    kk: str,
+    vdef: OutputVariableDef,
+    extended_coord: torch.Tensor,
+    extended_atype: torch.Tensor,
+    nlist: torch.Tensor,
+    mapping: torch.Tensor | None,
+    fparam: torch.Tensor | None,
+    aparam: torch.Tensor | None,
+    create_graph: bool = False,
+) -> torch.Tensor:
+    """Compute hessian of reduced output w.r.t. extended coordinates.
+
+    Mirrors the JAX approach: compute hessian on extended coordinates,
+    then let communicate_extended_output map nall->nloc.
+
+    Parameters
+    ----------
+    model
+        The model (CM instance). Must have ``atomic_model.forward_common_atomic``.
+    kk
+        The output key (e.g. "energy").
+    vdef
+        The output variable definition.
+    extended_coord
+        Extended coordinates. Shape: [nf, nall, 3].
+    extended_atype
+        Extended atom types. Shape: [nf, nall].
+    nlist
+        Neighbor list. Shape: [nf, nloc, nsel].
+    mapping
+        Mapping from extended to local. Shape: [nf, nall] or None.
+    fparam
+        Frame parameters. Shape: [nf, nfp] or None.
+    aparam
+        Atomic parameters. Shape: [nf, nloc, nap] or None.
+    create_graph
+        Whether to create graph for higher-order derivatives.
+
+    Returns
+    -------
+    torch.Tensor
+        Hessian on extended coordinates. Shape: [nf, *def, nall, 3, nall, 3].
+    """
+    nf, nall, _ = extended_coord.shape
+    vsize = math.prod(vdef.shape)
+    coord_flat = extended_coord.reshape(nf, nall * 3)
+    hessians = []
+    for ii in range(nf):
+        for ci in range(vsize):
+            wrapper = _WrapperForwardEnergy(
+                model,
+                kk,
+                ci,
+                nall,
+                extended_atype[ii],
+                nlist[ii],
+                mapping[ii] if mapping is not None else None,
+                fparam[ii] if fparam is not None else None,
+                aparam[ii] if aparam is not None else None,
+            )
+            hess = torch.autograd.functional.hessian(
+                wrapper,
+                coord_flat[ii],
+                create_graph=create_graph,
+            )
+            hessians.append(hess)
+    # [nf * vsize, nall*3, nall*3] -> [nf, *vshape, nall, 3, nall, 3]
+    result = torch.stack(hessians).reshape(nf, *vdef.shape, nall, 3, nall, 3)
+    return result
+
+
+class _WrapperForwardEnergy:
+    """Callable wrapper for torch.autograd.functional.hessian.
+
+    Given flattened extended coordinates, recomputes the reduced energy
+    for one frame and one output component.
+    """
+
+    def __init__(
+        self,
+        model: Any,
+        kk: str,
+        ci: int,
+        nall: int,
+        atype: torch.Tensor,
+        nlist: torch.Tensor,
+        mapping: torch.Tensor | None,
+        fparam: torch.Tensor | None,
+        aparam: torch.Tensor | None,
+    ) -> None:
+        self.model = model
+        self.kk = kk
+        self.ci = ci
+        self.nall = nall
+        self.atype = atype
+        self.nlist = nlist
+        self.mapping = mapping
+        self.fparam = fparam
+        self.aparam = aparam
+
+    def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor:
+        """Compute scalar reduced energy for one frame, one component.
+
+        Parameters
+        ----------
+        coord_flat
+            Flattened extended coordinates for one frame. Shape: [nall * 3].
+
+        Returns
+        -------
+        torch.Tensor
+            Scalar energy component.
+        """
+        cc_3d = coord_flat.reshape(1, self.nall, 3)
+        atomic_ret = self.model.atomic_model.forward_common_atomic(
+            cc_3d,
+            self.atype.unsqueeze(0),
+            self.nlist.unsqueeze(0),
+            mapping=self.mapping.unsqueeze(0) if self.mapping is not None else None,
+            fparam=self.fparam.unsqueeze(0) if self.fparam is not None else None,
+            aparam=self.aparam.unsqueeze(0) if self.aparam is not None else None,
+        )
+        # atomic_ret[kk]: [1, nloc, *def]
+        atom_energy = atomic_ret[self.kk][0]  # [nloc, *def]
+        energy_redu = atom_energy.sum(dim=0).reshape(-1)[self.ci]
+        return energy_redu
+
+
 def make_model(
     T_AtomicModel: type[BaseAtomicModel],
     T_Bases: tuple[type, ...] = (),
@@ -84,14 +221,35 @@ def forward_common_atomic(
                 fparam=fparam,
                 aparam=aparam,
             )
-            return fit_output_to_model_output(
+            model_ret = fit_output_to_model_output(
                 atomic_ret,
                 self.atomic_output_def(),
                 extended_coord,
                 do_atomic_virial=do_atomic_virial,
                 create_graph=self.training,
                 mask=atomic_ret.get("mask"),
             )
+            # Hessian computation (mirrors JAX's forward_common_atomic).
+            # Produces hessian on extended coords [nf, *def, nall, 3, nall, 3],
+            # then communicate_extended_output maps it to nloc x nloc.
+            aod = self.atomic_output_def()
+            for kk in aod.keys():
+                vdef = aod[kk]
+                if vdef.reducible and vdef.r_hessian:
+                    kk_hess = get_hessian_name(kk)
+                    model_ret[kk_hess] = _cal_hessian_ext(
+                        self,
+                        kk,
+                        vdef,
+                        extended_coord,
+                        extended_atype,
+                        nlist,
+                        mapping,
+                        fparam,
+                        aparam,
+                        create_graph=self.training,
+                    )
+            return model_ret
 
         def forward_common_lower_exportable(
             self,