fix(jax): add hessian energy loss

deepmd/dpmodel/loss/ener.py

@@ -13,6 +13,7 @@
 )
 from deepmd.dpmodel.loss.reduction import (
     masked_atom_mean,
+    masked_pair_mean,
     per_frame_component_mean,
 )
 from deepmd.utils.data import (
@@ -107,6 +108,10 @@ class EnergyLoss(Loss):
         The prefactor of generalized force loss at the end of the training.
     numb_generalized_coord : int
         The dimension of generalized coordinates.
+    start_pref_h : float
+        The prefactor of Hessian loss at the start of the training.
+    limit_pref_h : float
+        The prefactor of Hessian loss at the end of the training.
     use_default_pf : bool
         If true, use default atom_pref of 1.0 for all atoms when atom_pref data is not provided.
         This allows using the prefactor force loss (pf) without requiring atom_pref.npy files.
@@ -159,6 +164,8 @@ def __init__(
         start_pref_gf: float = 0.0,
         limit_pref_gf: float = 0.0,
         numb_generalized_coord: int = 0,
+        start_pref_h: float = 0.0,
+        limit_pref_h: float = 0.0,
         use_huber: bool = False,
         huber_delta: float | list[float] = 0.01,
         loss_func: str = "mse",
@@ -191,12 +198,15 @@ def __init__(
         self.start_pref_gf = start_pref_gf
         self.limit_pref_gf = limit_pref_gf
         self.numb_generalized_coord = numb_generalized_coord
+        self.start_pref_h = start_pref_h
+        self.limit_pref_h = limit_pref_h
         self.has_e = self.start_pref_e != 0.0 or self.limit_pref_e != 0.0
         self.has_f = self.start_pref_f != 0.0 or self.limit_pref_f != 0.0
         self.has_v = self.start_pref_v != 0.0 or self.limit_pref_v != 0.0
         self.has_ae = self.start_pref_ae != 0.0 or self.limit_pref_ae != 0.0
         self.has_pf = self.start_pref_pf != 0.0 or self.limit_pref_pf != 0.0
         self.has_gf = self.start_pref_gf != 0.0 or self.limit_pref_gf != 0.0
+        self.has_h = self.start_pref_h != 0.0 or self.limit_pref_h != 0.0
         if self.has_gf and self.numb_generalized_coord < 1:
             raise RuntimeError(
                 "When generalized force loss is used, the dimension of generalized coordinates should be larger than 0"
@@ -219,8 +229,10 @@ def __init__(
             self.has_pf or self.has_gf or self.relative_f is not None
         ):
             raise RuntimeError(
-                "Huber loss is not implemented for force with atom_pref, generalized force and relative force. "
+                "Huber loss is not implemented for force with atom_pref, generalized force and relative force."
             )
+        if self.use_huber and self.has_h:
+            raise RuntimeError("Huber loss is not implemented for hessian.")
 
     def call(
         self,
@@ -328,6 +340,7 @@ def call(
         pref_pf = find_atom_pref * (
             self.limit_pref_pf + (self.start_pref_pf - self.limit_pref_pf) * lr_ratio
         )
+        pref_h = self.limit_pref_h + (self.start_pref_h - self.limit_pref_h) * lr_ratio
 
         loss = 0
         more_loss = {}
@@ -741,6 +754,42 @@ def call(
             more_loss["rmse_gf"] = self.display_if_exist(
                 xp.sqrt(l2_gen_force_loss), find_drdq
             )
+        hessian = model_dict.get("hessian", model_dict.get("energy_derv_r_derv_r"))
+        if self.has_h and hessian is not None and "hessian" in label_dict:
+            find_hessian = label_dict.get("find_hessian", 0.0)
+            if maskf is not None:
+                hessian_shape = (_nf, _nloc * 3, _nloc * 3)
+                diff_h = xp.reshape(label_dict["hessian"], hessian_shape) - xp.reshape(
+                    hessian, hessian_shape
+                )
+                # A Hessian element couples two Cartesian atom components, so
+                # it is valid only when both corresponding atoms are real.
+                l2_hessian_loss = masked_pair_mean(xp.square(diff_h), maskf, ncomp=3)
+            else:
+                diff_h = xp.reshape(label_dict["hessian"], (-1,)) - xp.reshape(
+                    hessian,
+                    (-1,),
+                )
+                l2_hessian_loss = xp.mean(xp.square(diff_h))
+            mae_h = None
+            if self.loss_func == "mae" or mae:
+                if maskf is not None:
+                    mae_h = masked_pair_mean(xp.abs(diff_h), maskf, ncomp=3)
+                else:
+                    mae_h = xp.mean(xp.abs(diff_h))
+            if self.loss_func == "mse":
+                loss += pref_h * find_hessian * l2_hessian_loss
+            elif self.loss_func == "mae":
+                loss += pref_h * find_hessian * mae_h
+            else:
+                raise NotImplementedError(
+                    f"Loss type {self.loss_func} is not implemented for hessian loss."
+                )
+            more_loss["rmse_h"] = self.display_if_exist(
+                xp.sqrt(l2_hessian_loss), find_hessian
+            )
+            if mae:
+                more_loss["mae_h"] = self.display_if_exist(mae_h, find_hessian)
 
         self.l2_l = loss
         more_loss["rmse"] = xp.sqrt(loss)
@@ -819,6 +868,17 @@ def label_requirement(self) -> list[DataRequirementItem]:
                     default=1.0,
                 )
             )
+        if self.has_h:
+            label_requirement.append(
+                DataRequirementItem(
+                    "hessian",
+                    ndof=1,
+                    atomic=False,
+                    must=False,
+                    high_prec=False,
+                    special_shape="hessian",
+                )
+            )
         return label_requirement
 
     def serialize(self) -> dict:
@@ -829,9 +889,12 @@ def serialize(self) -> dict:
         dict
             The serialized loss module
         """
-        return {
+        data = {
             "@class": "EnergyLoss",
-            "@version": 4,
+            # Version 5 identifies the opt-in Hessian fields. Keep ordinary
+            # energy losses at version 4 so readers that already support the
+            # standard schema remain interoperable across backends.
+            "@version": 5 if self.has_h else 4,
             "starter_learning_rate": self.starter_learning_rate,
             "start_pref_e": self.start_pref_e,
             "limit_pref_e": self.limit_pref_e,
@@ -855,6 +918,12 @@ def serialize(self) -> dict:
             "use_default_pf": self.use_default_pf,
             "intensive_ener_virial": self.intensive_ener_virial,
         }
+        if self.has_h:
+            # Keep the established cross-backend serialization unchanged for
+            # ordinary energy losses; Hessian-only fields are an opt-in schema.
+            data["start_pref_h"] = self.start_pref_h
+            data["limit_pref_h"] = self.limit_pref_h
+        return data
 
     @classmethod
     def deserialize(cls, data: dict) -> "Loss":
@@ -872,9 +941,15 @@ def deserialize(cls, data: dict) -> "Loss":
         """
         data = data.copy()
         version = data.pop("@version")
-        check_version_compatibility(version, 4, 1)
+        check_version_compatibility(version, 5, 1)
         data.pop("@class")
         # Backward compatibility: version 1-2 used legacy normalization
         if version < 3:
             data.setdefault("intensive_ener_virial", False)
+        # Version 5 introduced explicit Hessian prefactors. Older payloads
+        # represent an ordinary energy loss unless these development fields
+        # were already present.
+        if version < 5:
+            data.setdefault("start_pref_h", 0.0)
+            data.setdefault("limit_pref_h", 0.0)
         return cls(**data)

deepmd/dpmodel/loss/reduction.py

@@ -67,6 +67,45 @@ def masked_atom_mean(elem: Array, maskf: Array, ncomp: int) -> Array:
     return xp.mean(per_frame)
 
 
+def masked_pair_mean(elem: Array, maskf: Array, ncomp: int) -> Array:
+    """Return a per-frame mean over valid atom-pair components.
+
+    Parameters
+    ----------
+    elem : Array
+        Non-negative pair contribution of shape
+        ``[nf, nloc * ncomp, nloc * ncomp]``. The contribution has already
+        been squared or converted to an absolute value, but is not masked.
+    maskf : Array
+        Per-atom real/placeholder mask of shape ``[nf, nloc]``.
+    ncomp : int
+        Number of components per atom on each pair axis. A Cartesian Hessian
+        uses three components on both axes.
+
+    Returns
+    -------
+    Array
+        ``mean_over_frames(sum(valid_pair_elem) / (real_natoms*ncomp)**2)``.
+        A pair is valid only when both atom indices are real. An all-padding
+        frame contributes a neutral zero.
+    """
+    xp = array_api_compat.array_namespace(elem, maskf)
+    nf, nloc = maskf.shape
+    component_mask = xp.reshape(
+        xp.broadcast_to(maskf[:, :, None], (nf, nloc, ncomp)),
+        (nf, nloc * ncomp),
+    )
+    masked = elem * component_mask[:, :, None] * component_mask[:, None, :]
+    per_frame_sum = xp.sum(xp.reshape(masked, (nf, -1)), axis=-1)
+    per_frame_dof = xp.square(xp.sum(component_mask, axis=-1))
+    has_dof = per_frame_dof > 0
+    safe_dof = xp.where(has_dof, per_frame_dof, xp.ones_like(per_frame_dof))
+    per_frame = xp.where(
+        has_dof, per_frame_sum / safe_dof, xp.zeros_like(per_frame_sum)
+    )
+    return xp.mean(per_frame)
+
+
 def per_frame_component_mean(err: Array) -> Array:
     """Idiom 2 primitive: per-frame mean over the flattened component axis.
 

deepmd/dpmodel/model/ener_model.py

@@ -73,6 +73,9 @@ def __init__(
         self.hess_fitting_def = None
 
     def enable_hessian(self) -> None:
+        """Enable Hessian outputs without changing an already-enabled model."""
+        if self._enable_hessian:
+            return
         self.hess_fitting_def = deepcopy(self.atomic_output_def())
         self.hess_fitting_def["energy"].r_hessian = True
         self._enable_hessian = True

deepmd/jax/train/trainer.py

@@ -95,6 +95,7 @@
 )
 from deepmd.utils.data import (
     DataRequirementItem,
+    has_data_requirement,
 )
 from deepmd.utils.data_system import (
     DeepmdDataSystem,
@@ -217,11 +218,24 @@ def __init__(
         learning_rate_param = jdata["learning_rate"]
         self.lr = self._get_lr_and_coef(learning_rate_param)
         self.losses = self._build_losses(jdata, learning_rate_param)
-        self.loss = self.losses if self.multi_task else self.losses[DEFAULT_TASK_KEY]
         self.data_requirements_by_task = {
             model_key: list(self.losses[model_key].label_requirement)
             for model_key in self.model_keys
         }
+        for model_key, requirements in self.data_requirements_by_task.items():
+            # Hessians are expensive and are only exposed by models whose task
+            # loss requests the corresponding label. Enable each multi-task
+            # branch independently so unrelated branches keep normal outputs.
+            if has_data_requirement(requirements, "hessian"):
+                enable_hessian = getattr(self.models[model_key], "enable_hessian", None)
+                if not callable(enable_hessian):
+                    raise RuntimeError(
+                        f"Model {type(self.models[model_key]).__name__} does not "
+                        "support Hessian supervision."
+                    )
+                enable_hessian()
+                self.model_params_by_task[model_key]["hessian_mode"] = True
+        self.loss = self.losses if self.multi_task else self.losses[DEFAULT_TASK_KEY]
 
         self.valid_numb_batch_by_task = self._valid_numb_batch_by_task()
         self.valid_numb_batch = (
@@ -1083,6 +1097,8 @@ def _evaluate_model_dict(
     model_dict["energy"] = model_dict["energy_redu"]
     model_dict["force"] = model_dict["energy_derv_r"].squeeze(-2)
     model_dict["virial"] = model_dict["energy_derv_c_redu"].squeeze(-2)
+    if model_dict.get("energy_derv_r_derv_r") is not None:
+        model_dict["hessian"] = model_dict["energy_derv_r_derv_r"].squeeze(-3)
     return model_dict
 
 

deepmd/pd/loss/ener.py

@@ -37,6 +37,33 @@ def custom_huber_loss(
     return paddle.mean(loss)
 
 
+def masked_pair_mean(
+    elem: paddle.Tensor, maskf: paddle.Tensor, ncomp: int
+) -> paddle.Tensor:
+    """Average pair contributions over real atoms on both pair axes.
+
+    Hessian padding is square: a component is valid only when both its row atom
+    and column atom are real. The reduction is normalized per frame so mixed
+    systems with different atom counts receive the same weighting.
+    """
+    nf, nloc = maskf.shape
+    component_mask = paddle.broadcast_to(
+        maskf.astype(elem.dtype).reshape([nf, nloc, 1]),
+        [nf, nloc, ncomp],
+    ).reshape([nf, nloc * ncomp])
+    masked = elem * component_mask[:, :, None] * component_mask[:, None, :]
+    per_frame_sum = paddle.sum(masked.reshape([nf, -1]), axis=-1)
+    per_frame_dof = paddle.square(paddle.sum(component_mask, axis=-1))
+    has_dof = per_frame_dof > 0
+    safe_dof = paddle.where(has_dof, per_frame_dof, paddle.ones_like(per_frame_dof))
+    per_frame = paddle.where(
+        has_dof,
+        per_frame_sum / safe_dof,
+        paddle.zeros_like(per_frame_sum),
+    )
+    return paddle.mean(per_frame)
+
+
 class EnergyStdLoss(TaskLoss):
     def __init__(
         self,
@@ -649,7 +676,11 @@ def __init__(
             Other keyword arguments.
         """
         super().__init__(**kwargs)
-        self.has_h = (start_pref_h != 0.0 and limit_pref_h != 0.0) or self.inference
+        # A scheduled term is active when either endpoint is nonzero. Requiring
+        # both endpoints silently disabled valid ramp-up and ramp-down inputs.
+        self.has_h = (start_pref_h != 0.0 or limit_pref_h != 0.0) or self.inference
+        if self.use_huber and self.has_h:
+            raise RuntimeError("Huber loss is not implemented for hessian.")
 
         self.start_pref_h = start_pref_h
         self.limit_pref_h = limit_pref_h
@@ -672,12 +703,21 @@ def forward(
         if self.has_h and "hessian" in model_pred and "hessian" in label:
             find_hessian = label.get("find_hessian", 0.0)
             pref_h = pref_h * find_hessian
-            diff_h = label["hessian"].reshape(
-                [-1],
-            ) - model_pred["hessian"].reshape(
-                [-1],
-            )
-            l2_hessian_loss = paddle.mean(paddle.square(diff_h))
+            maskf = model_pred.get("mask")
+            if maskf is not None:
+                nf, nloc = maskf.shape
+                hessian_shape = [nf, nloc * 3, nloc * 3]
+                diff_h = label["hessian"].reshape(hessian_shape) - model_pred[
+                    "hessian"
+                ].reshape(hessian_shape)
+                l2_hessian_loss = masked_pair_mean(
+                    paddle.square(diff_h), maskf, ncomp=3
+                )
+            else:
+                diff_h = label["hessian"].reshape([-1]) - model_pred["hessian"].reshape(
+                    [-1]
+                )
+                l2_hessian_loss = paddle.mean(paddle.square(diff_h))
             if not self.inference:
                 more_loss["l2_hessian_loss"] = self.display_if_exist(
                     l2_hessian_loss.detach(), find_hessian
@@ -686,7 +726,10 @@ def forward(
             rmse_h = l2_hessian_loss.sqrt()
             more_loss["rmse_h"] = self.display_if_exist(rmse_h.detach(), find_hessian)
             if mae:
-                mae_h = paddle.mean(paddle.abs(diff_h))
+                if maskf is not None:
+                    mae_h = masked_pair_mean(paddle.abs(diff_h), maskf, ncomp=3)
+                else:
+                    mae_h = paddle.mean(paddle.abs(diff_h))
                 more_loss["mae_h"] = self.display_if_exist(mae_h.detach(), find_hessian)
 
         if not self.inference:
@@ -701,10 +744,11 @@ def label_requirement(self) -> list[DataRequirementItem]:
             label_requirement.append(
                 DataRequirementItem(
                     "hessian",
-                    ndof=1,  # 9=3*3 --> 3N*3N=ndof*natoms*natoms
-                    atomic=True,
+                    ndof=1,
+                    atomic=False,
                     must=False,
                     high_prec=False,
+                    special_shape="hessian",
                 )
             )
         return label_requirement

deepmd/pd/train/training.py

@@ -1353,8 +1353,12 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]:
 
 
 def whether_hessian(loss_params: dict[str, Any]) -> bool:
+    """Return whether either Hessian schedule endpoint enables supervision."""
     loss_type = loss_params.get("type", "ener")
-    return loss_type == "ener" and loss_params.get("start_pref_h", 0.0) > 0.0
+    return loss_type in {"ener", "ener_hess"} and (
+        loss_params.get("start_pref_h", 0.0) != 0.0
+        or loss_params.get("limit_pref_h", 0.0) != 0.0
+    )
 
 
 def get_loss(