diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py index 7c7271cdfe..34f7858b83 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/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) diff --git a/deepmd/dpmodel/loss/reduction.py b/deepmd/dpmodel/loss/reduction.py index 068c11e2bc..2568913538 100644 --- a/deepmd/dpmodel/loss/reduction.py +++ b/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. diff --git a/deepmd/dpmodel/model/ener_model.py b/deepmd/dpmodel/model/ener_model.py index a8280dbebf..7e89b73972 100644 --- a/deepmd/dpmodel/model/ener_model.py +++ b/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 diff --git a/deepmd/jax/train/trainer.py b/deepmd/jax/train/trainer.py index ce2fb8d3f2..d78dbc38ed 100644 --- a/deepmd/jax/train/trainer.py +++ b/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 diff --git a/deepmd/pd/loss/ener.py b/deepmd/pd/loss/ener.py index 4ab9d71ab9..1a9b2c38d0 100644 --- a/deepmd/pd/loss/ener.py +++ b/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 diff --git a/deepmd/pd/train/training.py b/deepmd/pd/train/training.py index b2a8e6f5e5..53f768fded 100644 --- a/deepmd/pd/train/training.py +++ b/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( diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py index bff23e2ba2..468b95f1eb 100644 --- a/deepmd/pt/loss/ener.py +++ b/deepmd/pt/loss/ener.py @@ -8,6 +8,7 @@ from deepmd.dpmodel.loss.reduction import ( masked_atom_mean, + masked_pair_mean, per_frame_component_mean, ) from deepmd.pt.loss.loss import ( @@ -60,6 +61,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, loss_func: str = "mse", inference: bool = False, use_huber: bool = False, @@ -107,6 +110,10 @@ def __init__( 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. loss_func : str Loss function type. Options: 'mse' (Mean Squared Error, L2 loss, default) or 'mae' (Mean Absolute Error, L1 loss). MAE loss is less sensitive to outliers compared to MSE loss. @@ -156,6 +163,10 @@ def __init__( self.has_ae = (start_pref_ae != 0.0 and limit_pref_ae != 0.0) or inference self.has_pf = (start_pref_pf != 0.0 and limit_pref_pf != 0.0) or inference self.has_gf = start_pref_gf != 0.0 and limit_pref_gf != 0.0 + # Hessian labels scale quadratically with the atom count. Unlike the + # linear-size labels above, do not request them merely because a mock + # inference loss is collecting requirements for ``dp change-bias``. + self.has_h = start_pref_h != 0.0 or limit_pref_h != 0.0 self.start_pref_e = start_pref_e self.limit_pref_e = limit_pref_e @@ -169,6 +180,8 @@ def __init__( self.limit_pref_pf = limit_pref_pf self.start_pref_gf = start_pref_gf self.limit_pref_gf = limit_pref_gf + self.start_pref_h = start_pref_h + self.limit_pref_h = limit_pref_h self.use_default_pf = use_default_pf self.relative_f = relative_f self.enable_atom_ener_coeff = enable_atom_ener_coeff @@ -195,8 +208,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 forward( self, @@ -237,6 +252,7 @@ def forward( pref_ae = self.limit_pref_ae + (self.start_pref_ae - self.limit_pref_ae) * coef pref_pf = self.limit_pref_pf + (self.start_pref_pf - self.limit_pref_pf) * coef pref_gf = self.limit_pref_gf + (self.start_pref_gf - self.limit_pref_gf) * coef + pref_h = self.limit_pref_h + (self.start_pref_h - self.limit_pref_h) * coef loss = torch.zeros(1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)[0] more_loss = {} @@ -775,6 +791,46 @@ def forward( f"Loss type {self.loss_func} is not implemented for atomic energy loss." ) + 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 + if maskf is not None: + hessian_shape = (_nf, _nloc * 3, _nloc * 3) + diff_h = label["hessian"].reshape(hessian_shape) - model_pred[ + "hessian" + ].reshape(hessian_shape) + # Both Cartesian axes must refer to real atoms for a Hessian + # element to contribute to the loss or display denominator. + l2_hessian_loss = masked_pair_mean(torch.square(diff_h), maskf, ncomp=3) + else: + diff_h = label["hessian"].reshape(-1) - model_pred["hessian"].reshape( + -1 + ) + l2_hessian_loss = torch.mean(torch.square(diff_h)) + if not self.inference: + more_loss["l2_hessian_loss"] = self.display_if_exist( + l2_hessian_loss.detach(), find_hessian + ) + mae_h = None + if self.loss_func == "mae" or mae: + if maskf is not None: + mae_h = masked_pair_mean(torch.abs(diff_h), maskf, ncomp=3) + else: + mae_h = torch.mean(torch.abs(diff_h)) + if self.loss_func == "mse": + loss += pref_h * l2_hessian_loss + elif self.loss_func == "mae": + loss += pref_h * 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( + l2_hessian_loss.sqrt().detach(), find_hessian + ) + if mae: + more_loss["mae_h"] = self.display_if_exist(mae_h.detach(), find_hessian) + if not self.inference: more_loss["rmse"] = torch.sqrt(loss.detach()) return model_pred, loss, more_loss @@ -856,6 +912,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: @@ -866,9 +933,11 @@ def serialize(self) -> dict: dict The serialized loss module """ - return { + data = { "@class": "EnergyLoss", - "@version": 4, + # Only Hessian-bearing payloads need the version-5 schema. Keeping + # ordinary energy losses at version 4 preserves existing readers. + "@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, @@ -892,6 +961,10 @@ def serialize(self) -> dict: "use_default_pf": self.use_default_pf, "intensive_ener_virial": self.intensive_ener_virial, } + if self.start_pref_h != 0.0 or self.limit_pref_h != 0.0: + data["start_pref_h"] = self.start_pref_h + data["limit_pref_h"] = self.limit_pref_h + return data @classmethod def deserialize(cls, data: dict) -> "TaskLoss": @@ -909,89 +982,18 @@ def deserialize(cls, data: dict) -> "TaskLoss": """ data = data.copy() version = data.pop("@version") - check_version_compatibility(version, 4, 1) + check_version_compatibility(version, 5, 1) data.pop("@class") # Handle backward compatibility for older versions without intensive_ener_virial if version < 3: data.setdefault("intensive_ener_virial", False) + # Version 5 introduced explicit Hessian prefactors. Version 1-4 + # payloads therefore default to the standard non-Hessian loss. + if version < 5: + data.setdefault("start_pref_h", 0.0) + data.setdefault("limit_pref_h", 0.0) return cls(**data) class EnergyHessianStdLoss(EnergyStdLoss): - def __init__( - self, - start_pref_h: float = 0.0, - limit_pref_h: float = 0.0, - **kwargs: Any, - ) -> None: - r"""Enable the layer to compute loss on hessian. - - Parameters - ---------- - 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. - **kwargs - Other keyword arguments. - """ - super().__init__(**kwargs) - self.has_h = (start_pref_h != 0.0 and limit_pref_h != 0.0) or self.inference - - self.start_pref_h = start_pref_h - self.limit_pref_h = limit_pref_h - - def forward( - self, - input_dict: dict[str, torch.Tensor], - model: torch.nn.Module, - label: dict[str, torch.Tensor], - natoms: int, - learning_rate: float, - mae: bool = False, - ) -> tuple[dict[str, torch.Tensor], torch.Tensor, dict[str, torch.Tensor]]: - model_pred, loss, more_loss = super().forward( - input_dict, model, label, natoms, learning_rate, mae=mae - ) - coef = learning_rate / self.starter_learning_rate - pref_h = self.limit_pref_h + (self.start_pref_h - self.limit_pref_h) * coef - - 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 = torch.mean(torch.square(diff_h)) - if not self.inference: - more_loss["l2_hessian_loss"] = self.display_if_exist( - l2_hessian_loss.detach(), find_hessian - ) - loss += pref_h * l2_hessian_loss - rmse_h = l2_hessian_loss.sqrt() - more_loss["rmse_h"] = self.display_if_exist(rmse_h.detach(), find_hessian) - if mae: - mae_h = torch.mean(torch.abs(diff_h)) - more_loss["mae_h"] = self.display_if_exist(mae_h.detach(), find_hessian) - - if not self.inference: - more_loss["rmse"] = torch.sqrt(loss.detach()) - return model_pred, loss, more_loss - - @property - def label_requirement(self) -> list[DataRequirementItem]: - """Add hessian label requirement needed for this loss calculation.""" - label_requirement = super().label_requirement - if self.has_h: - label_requirement.append( - DataRequirementItem( - "hessian", - ndof=1, # 9=3*3 --> 3N*3N=ndof*natoms*natoms - atomic=True, - must=False, - high_prec=False, - ) - ) - return label_requirement + """Backward-compatible name for the unified energy loss.""" diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py index 28387553fb..bc858b6419 100644 --- a/deepmd/pt/model/model/ener_model.py +++ b/deepmd/pt/model/model/ener_model.py @@ -39,6 +39,9 @@ def __init__( self._hessian_enabled = False def enable_hessian(self) -> None: + """Enable Hessian outputs without changing an already-enabled model.""" + if self._hessian_enabled: + return self.__class__ = make_hessian_model(type(self)) self.hess_fitting_def = super(type(self), self).atomic_output_def() self.requires_hessian("energy") diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 40fddc51f1..9c55bea6cc 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -44,7 +44,6 @@ DenoiseLoss, DeNSLoss, DOSLoss, - EnergyHessianStdLoss, EnergySpinLoss, EnergyStdLoss, PopulationLoss, @@ -120,6 +119,7 @@ ) from deepmd.utils.data import ( DataRequirementItem, + has_data_requirement, ) from deepmd.utils.finetune import ( warn_configuration_mismatch_during_finetune, @@ -458,12 +458,12 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: if self.zero_stage > 0 and self.opt_type == "LKF": raise ValueError("training.zero_stage does not support LKF optimizer.") - # loss_param_tmp for Hessian activation - loss_param_tmp = None + # Loss parameters are also used to select SeZM/DeNS execution modes. + loss_params_for_model = None if not self.multi_task: - loss_param_tmp = config["loss"] + loss_params_for_model = config["loss"] else: - loss_param_tmp = { + loss_params_for_model = { model_key: config["loss_dict"][model_key] for model_key in self.model_keys } @@ -475,10 +475,9 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: self.model = get_model_for_wrapper( model_params, resuming=resuming, - _loss_params=loss_param_tmp, ) # SeZM specific process for DeNS training - prepare_model_for_loss(self.model, loss_param_tmp) + prepare_model_for_loss(self.model, loss_params_for_model) # Loss if not self.multi_task: @@ -498,6 +497,22 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: loss_param, lr_param, ntypes, self.model[model_key] ) + # Losses own the interpretation of their prefactors. The trainer only + # consumes their data contract when selecting expensive model outputs. + loss_data_requirements = ( + { + model_key: self.loss[model_key].label_requirement + for model_key in self.model_keys + } + if self.multi_task + else self.loss.label_requirement + ) + prepare_model_for_data_requirements( + self.model, + loss_data_requirements, + model_params, + ) + # Data if not self.multi_task: # add data requirement for labels @@ -2499,11 +2514,6 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: return additional_data_requirement -def whether_hessian(loss_params: dict[str, Any]) -> bool: - loss_type = loss_params.get("type", "ener") - return loss_type == "ener" and loss_params.get("start_pref_h", 0.0) > 0.0 - - def prepare_model_for_loss( model: Any, loss_params: dict[str, Any] | None, @@ -2518,17 +2528,54 @@ def prepare_model_for_loss( prepare_model_for_loss(sub_model, sub_loss) return if hasattr(model, "set_active_mode_from_loss"): - model.set_active_mode_from_loss(loss_params.get("type", "ener")) + loss_type = loss_params.get("type", "ener") + model.set_active_mode_from_loss( + "ener" if loss_type == "ener_hess" else loss_type + ) + + +def prepare_model_for_data_requirements( + model: Any, + data_requirements: list[DataRequirementItem] | dict[str, list[DataRequirementItem]], + model_params: dict[str, Any], +) -> None: + """Enable model outputs requested by a loss's data requirements. + + Hessian prefactors are deliberately not inspected here. The presence of a + ``hessian`` requirement is the loss-to-trainer contract for enabling that + expensive output. The mode is persisted in the model definition so reload + and freeze paths reconstruct the same model interface. + """ + if isinstance(model, dict): + if not isinstance(data_requirements, dict): + raise TypeError("Multi-task models require per-task data requirements.") + for model_key, sub_model in model.items(): + prepare_model_for_data_requirements( + sub_model, + data_requirements[model_key], + model_params["model_dict"][model_key], + ) + return + if isinstance(data_requirements, dict): + raise TypeError("Single-task models require a list of data requirements.") + if not has_data_requirement(data_requirements, "hessian"): + return + enable_hessian = getattr(model, "enable_hessian", None) + if not callable(enable_hessian): + raise RuntimeError( + f"Model {type(model).__name__} does not support Hessian supervision." + ) + enable_hessian() + model_params["hessian_mode"] = True + if hasattr(model, "model_def_script"): + model.model_def_script = json.dumps(model_params) def get_loss( loss_params: dict[str, Any], start_lr: float, _ntypes: int, _model: Any ) -> TaskLoss: loss_type = loss_params.get("type", "ener") - if whether_hessian(loss_params): - loss_params["starter_learning_rate"] = start_lr - return EnergyHessianStdLoss(**loss_params) - elif loss_type == "ener": + if loss_type in {"ener", "ener_hess"}: loss_params["starter_learning_rate"] = start_lr return EnergyStdLoss(**loss_params) elif loss_type == "dens": @@ -2585,11 +2632,8 @@ def get_single_model( def get_model_for_wrapper( _model_params: dict[str, Any], resuming: bool = False, - _loss_params: dict[str, Any] | None = None, ) -> Any: if "model_dict" not in _model_params: - if _loss_params is not None and whether_hessian(_loss_params): - _model_params["hessian_mode"] = True _model = get_single_model( _model_params, ) @@ -2598,8 +2642,6 @@ def get_model_for_wrapper( model_keys = list(_model_params["model_dict"]) do_case_embd, case_embd_index = get_case_embd_config(_model_params) for _model_key in model_keys: - if _loss_params is not None and whether_hessian(_loss_params[_model_key]): - _model_params["model_dict"][_model_key]["hessian_mode"] = True _model[_model_key] = get_single_model( _model_params["model_dict"][_model_key], ) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 407b66ff82..a0c847125f 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -4423,7 +4423,9 @@ def limit_pref(item: str) -> str: @loss_args_plugin.register( - "ener", doc=supported_backends("tf", "pt", "jax", "pd", "pt_expt", "tf2") + "ener", + alias=["ener_hess"], + doc=supported_backends("tf", "pt", "jax", "pd", "pt_expt", "tf2"), ) def loss_ener() -> list[Argument]: doc_start_pref_e = start_pref("energy", abbr="e") @@ -4530,14 +4532,14 @@ def loss_ener() -> list[Argument]: [float, int], optional=True, default=0.00, - doc=supported_backends("pt", "pd") + doc_start_pref_h, + doc=supported_backends("pt", "jax", "pd") + doc_start_pref_h, ), Argument( "limit_pref_h", [float, int], optional=True, default=0.00, - doc=supported_backends("pt", "pd") + doc_limit_pref_h, + doc=supported_backends("pt", "jax", "pd") + doc_limit_pref_h, ), Argument( "start_pref_ae", @@ -5119,7 +5121,7 @@ def loss_tensor() -> list[Argument]: def loss_variant_type_args() -> Variant: - doc_loss = "The type of the loss. When the fitting type is `ener`, the loss type should be set to `ener`, `dens` (Only DPA4/SeZM supported), or left unset. When the fitting type is `property`, the loss type should be set to `property`. When the fitting type is `dipole` or `polar`, the loss type should be set to `tensor`." + doc_loss = "The type of the loss. When the fitting type is `ener`, the loss type should be set to `ener`, its legacy alias `ener_hess`, `dens` (Only DPA4/SeZM supported), or left unset. Hessian supervision is configured through `start_pref_h` and `limit_pref_h` on the `ener` loss. When the fitting type is `property`, the loss type should be set to `property`. When the fitting type is `dipole` or `polar`, the loss type should be set to `tensor`." return Variant( "type", @@ -5131,7 +5133,7 @@ def loss_variant_type_args() -> Variant: def loss_args() -> list[Argument]: - doc_loss = "The definition of loss function. The loss type should be set to `tensor`, `property`, `ener`, `dens` or left unset." + doc_loss = "The definition of loss function. The loss type should be set to `tensor`, `property`, `ener`, `dens` or left unset. The legacy `ener_hess` type is normalized to `ener`." ca = Argument( "loss", dict, [], [loss_variant_type_args()], optional=True, doc=doc_loss ) diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py index 0db472dca9..c7a682c921 100644 --- a/deepmd/utils/data.py +++ b/deepmd/utils/data.py @@ -5,6 +5,9 @@ import copy import functools import logging +from collections.abc import ( + Iterable, +) from concurrent.futures import ( ThreadPoolExecutor, as_completed, @@ -159,6 +162,7 @@ def add( default: float = 0.0, dtype: np.dtype | None = None, output_natoms_for_type_sel: bool = False, + special_shape: str | None = None, ) -> "DeepmdData": """Add a data item that to be loaded. @@ -187,11 +191,14 @@ def add( the dtype of data, overwrites `high_prec` if provided output_natoms_for_type_sel : bool, optional if True and type_sel is True, the atomic dimension will be natoms instead of nsel + special_shape : str, optional + Name of a loader-defined non-standard shape contract. ``"hessian"`` + stores one full-frame ``(3 * natoms) x (3 * natoms)`` matrix per frame. """ # normalize key: "atomic_" prefix -> "atom_", same convention as _load_set output if key.startswith("atomic_"): key = "atom_" + key[7:] - self.data_dict[key] = { + data_config = { "ndof": ndof, "atomic": atomic, "must": must, @@ -203,6 +210,11 @@ def add( "dtype": dtype, "output_natoms_for_type_sel": output_natoms_for_type_sel, } + if special_shape is not None: + # Preserve the established dictionary schema for ordinary labels; + # only non-standard tensors need this extra shape contract. + data_config["special_shape"] = special_shape + self.data_dict[key] = data_config return self def reduce(self, key_out: str, key_in: str) -> "DeepmdData": @@ -735,6 +747,7 @@ def _load_set(self, set_name: DPPath) -> dict[str, Any]: output_natoms_for_type_sel=self.data_dict[kk][ "output_natoms_for_type_sel" ], + special_shape=self.data_dict[kk].get("special_shape"), ) for kk in self.data_dict.keys(): if self.data_dict[kk]["reduce"] is not None: @@ -811,7 +824,9 @@ def _load_data( default: float = 0.0, dtype: np.dtype | None = None, output_natoms_for_type_sel: bool = False, + special_shape: str | None = None, ) -> np.ndarray: + is_hessian = special_shape == "hessian" or key == "hessian" if atomic: natoms = self.natoms idx_map = self.idx_map @@ -839,7 +854,19 @@ def _load_data( if path.is_file(): data = path.load_numpy().astype(dtype) try: # YWolfeee: deal with data shape error - if atomic: + if is_hessian: + natoms = self.natoms + idx_map = self.idx_map + data = data.reshape(nframes, 3 * natoms, 3 * natoms) + num_chunks, chunk_size = len(idx_map), 3 + idx_map_hess = np.arange(num_chunks * chunk_size) # pylint: disable=no-explicit-dtype + idx_map_hess = idx_map_hess.reshape(num_chunks, chunk_size) + idx_map_hess = idx_map_hess[idx_map].flatten() + data = data[:, idx_map_hess, :] + data = data[:, :, idx_map_hess] + data = data.reshape([nframes, -1]) + ndof = 9 * natoms * natoms + elif atomic: if type_sel is not None: # check the data shape is nsel or natoms if data.size == nframes * natoms_sel * ndof_: @@ -871,24 +898,9 @@ def _load_data( f"({nframes}, {natoms_sel}, {ndof_}) or" f"({nframes}, {natoms}, {ndof_})" ) - if key == "hessian": - data = data.reshape(nframes, 3 * natoms, 3 * natoms) - # get idx_map for hessian - num_chunks, chunk_size = len(idx_map), 3 - idx_map_hess = np.arange(num_chunks * chunk_size) # pylint: disable=no-explicit-dtype - idx_map_hess = idx_map_hess.reshape(num_chunks, chunk_size) - idx_map_hess = idx_map_hess[idx_map] - idx_map_hess = idx_map_hess.flatten() - data = data[:, idx_map_hess, :] - data = data[:, :, idx_map_hess] - data = data.reshape([nframes, -1]) - ndof = ( - 3 * ndof * 3 * ndof - ) # size of hessian is 3Natoms * 3Natoms - else: - data = data.reshape([nframes, natoms, -1]) - data = data[:, idx_map, :] - data = data.reshape([nframes, -1]) + data = data.reshape([nframes, natoms, -1]) + data = data[:, idx_map, :] + data = data.reshape([nframes, -1]) data = np.reshape(data, [nframes, ndof]) except ValueError as err_message: explanation = "This error may occur when your label mismatch its name, i.e. you might store global tensor in `atomic_tensor.npy` or atomic tensor in `tensor.npy`." @@ -901,7 +913,9 @@ def _load_data( elif must: raise RuntimeError(f"{path} not found!") else: - if atomic and type_sel is not None and not output_natoms_for_type_sel: + if is_hessian: + ndof = 9 * self.natoms * self.natoms + elif atomic and type_sel is not None and not output_natoms_for_type_sel: ndof = ndof_ * natoms_sel data = np.full([nframes, ndof], default, dtype=dtype) if repeat != 1: @@ -928,8 +942,12 @@ def _load_single_data( """ vv = self.data_dict[key] path = self._get_data_path(set_dir, key) + is_hessian = vv.get("special_shape") == "hessian" or key == "hessian" - if vv["atomic"]: + if is_hessian: + natoms = self.natoms + idx_map = self.idx_map + elif vv["atomic"]: natoms = self.natoms idx_map = self.idx_map # if type_sel, then revise natoms and idx_map @@ -962,7 +980,9 @@ def _load_single_data( raise RuntimeError(f"{path} not found!") # Create a default array based on requirements - if vv["atomic"]: + if is_hessian: + data = np.full([9 * natoms * natoms], vv["default"], dtype=dtype) + elif vv["atomic"]: if vv["type_sel"] is not None and not vv["output_natoms_for_type_sel"]: natoms = natoms_sel data = np.full([natoms, ndof], vv["default"], dtype=dtype) @@ -986,7 +1006,17 @@ def _load_single_data( data = mmap_obj[frame_idx].copy().astype(dtype, copy=False) try: - if vv["atomic"]: + if is_hessian: + data = data.reshape(3 * natoms, 3 * natoms) + num_chunks, chunk_size = len(idx_map), 3 + idx_map_hess = np.arange(num_chunks * chunk_size, dtype=int).reshape( + num_chunks, chunk_size + ) + idx_map_hess = idx_map_hess[idx_map].flatten() + data = data[idx_map_hess, :] + data = data[:, idx_map_hess] + data = data.reshape(-1) + elif vv["atomic"]: # Handle type_sel logic if vv["type_sel"] is not None: if mmap_obj.shape[1] == natoms_sel * ndof: @@ -1011,23 +1041,9 @@ def _load_single_data( f"The shape of the data {key} in {set_dir} has width {mmap_obj.shape[1]}, which doesn't match either ({natoms_sel * ndof}) or ({natoms * ndof})" ) - # Handle special case for Hessian - if key == "hessian": - data = data.reshape(3 * natoms, 3 * natoms) - num_chunks, chunk_size = len(idx_map), 3 - idx_map_hess = np.arange( - num_chunks * chunk_size, dtype=int - ).reshape(num_chunks, chunk_size) - idx_map_hess = idx_map_hess[idx_map].flatten() - data = data[idx_map_hess, :] - data = data[:, idx_map_hess] - data = data.reshape(-1) - # size of hessian is 3Natoms * 3Natoms - # ndof = 3 * ndof * 3 * ndof - else: - # data should be 2D here: [natoms, ndof] - data = data.reshape([natoms, -1]) - data = data[idx_map, :] + # data should be 2D here: [natoms, ndof] + data = data.reshape([natoms, -1]) + data = data[idx_map, :] else: data = data.reshape([ndof]) @@ -1140,6 +1156,9 @@ class DataRequirementItem: the dtype of data, overwrites `high_prec` if provided output_natoms_for_type_sel : bool, optional if True and type_sel is True, the atomic dimension will be natoms instead of nsel + special_shape : str, optional + Name of a loader-defined non-standard shape contract. ``"hessian"`` + stores one full-frame ``(3 * natoms) x (3 * natoms)`` matrix per frame. """ def __init__( @@ -1154,6 +1173,7 @@ def __init__( default: float = 0.0, dtype: np.dtype | None = None, output_natoms_for_type_sel: bool = False, + special_shape: str | None = None, ) -> None: self.key = key self.ndof = ndof @@ -1165,10 +1185,11 @@ def __init__( self.default = default self.dtype = dtype self.output_natoms_for_type_sel = output_natoms_for_type_sel + self.special_shape = special_shape self.dict = self.to_dict() def to_dict(self) -> dict: - return { + data = { "key": self.key, "ndof": self.ndof, "atomic": self.atomic, @@ -1180,6 +1201,9 @@ def to_dict(self) -> dict: "dtype": self.dtype, "output_natoms_for_type_sel": self.output_natoms_for_type_sel, } + if self.special_shape is not None: + data["special_shape"] = self.special_shape + return data def __getitem__(self, key: str) -> np.ndarray: if key not in self.dict: @@ -1193,3 +1217,12 @@ def __eq__(self, value: object, /) -> bool: def __repr__(self) -> str: return f"DataRequirementItem({self.dict})" + + +def has_data_requirement(requirements: Iterable[DataRequirementItem], key: str) -> bool: + """Return whether a collection requests data identified by ``key``. + + Consumers should use requirement items as the contract with losses instead + of reinterpreting loss-specific configuration such as prefactors. + """ + return any(item.key == key for item in requirements) diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 9713a43d26..dfc4076a9f 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -333,6 +333,7 @@ def add_dict(self, adict: dict[str, dict[str, Any]]) -> None: output_natoms_for_type_sel=adict[kk].get( "output_natoms_for_type_sel", False ), + special_shape=adict[kk].get("special_shape"), ) def add_data_requirements( @@ -353,6 +354,7 @@ def add( default: float = 0.0, dtype: np.dtype | None = None, output_natoms_for_type_sel: bool = False, + special_shape: str | None = None, ) -> None: """Add a data item that to be loaded. @@ -381,6 +383,8 @@ def add( The dtype of data, overwrites `high_prec` if provided output_natoms_for_type_sel : bool If True and type_sel is True, the atomic dimension will be natoms instead of nsel + special_shape : str, optional + Name of a loader-defined non-standard shape contract. """ for ii in self.data_systems: ii.add( @@ -394,6 +398,7 @@ def add( default=default, dtype=dtype, output_natoms_for_type_sel=output_natoms_for_type_sel, + special_shape=special_shape, ) def reduce(self, key_out: str, key_in: str) -> None: @@ -549,7 +554,25 @@ def _merge_batch_data(self, batch_data: list[dict]) -> dict: if kk not in batch_data[0]: continue b_data["find_" + kk] = batch_data[0]["find_" + kk] - if not vv["atomic"]: + if vv.get("special_shape") == "hessian" or kk == "hessian": + # A Hessian is a (3 * natoms, 3 * natoms) matrix, so neither + # branch below pads it correctly: concatenating raises on + # ragged systems and copying a flat prefix would scatter the + # rows. Embed each frame's block in the top-left corner of the + # padded square instead; the padded rows and columns stay zero + # and are dropped by the loss mask. + padded_dof = max_natoms * 3 + merged = np.zeros( + (len(batch_data), padded_dof, padded_dof), + dtype=batch_data[0][kk].dtype, + ) + for ii, bb in enumerate(batch_data): + frame_dof = bb["natoms_vec"][0] * 3 + merged[ii, :frame_dof, :frame_dof] = bb[kk][0].reshape( + frame_dof, frame_dof + ) + b_data[kk] = merged.reshape(len(batch_data), -1) + elif not vv["atomic"]: b_data[kk] = np.concatenate([bb[kk] for bb in batch_data], axis=0) else: b_data[kk] = np.zeros( diff --git a/doc/model/train-energy-hessian.md b/doc/model/train-energy-hessian.md index dc1ca38fe8..d84fc754fe 100644 --- a/doc/model/train-energy-hessian.md +++ b/doc/model/train-energy-hessian.md @@ -1,13 +1,13 @@ -# Fit energy Hessian {{ pytorch_icon }} +# Fit energy Hessian {{ pytorch_icon }} {{ jax_icon }} > [!NOTE] -> **Supported backends**: PyTorch {{ pytorch_icon }} +> **Supported backends**: PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }} To train a model that takes Hessian matrices, i.e., the second order derivatives of energies w.r.t coordinates as input, you only need to prepare full Hessian matrices and modify the `loss` section to define the Hessian-specific settings, keeping other sections the same as the normal energy model's input script. ## Energy Hessian Loss -If you want to train with Hessians, you are expected to add the start and limit prefactors of Hessians, i.e., {ref}`start_pref_h ` and {ref}`limit_pref_h ` to the {ref}`loss ` section in the `input.json`: +If you want to train with Hessians, add the start and limit prefactors of Hessians, i.e., {ref}`start_pref_h ` and {ref}`limit_pref_h ` to the standard energy {ref}`loss ` section in the `input.json`: ```json "loss": { @@ -23,13 +23,17 @@ If you want to train with Hessians, you are expected to add the start and limit }, ``` -The options {ref}`start_pref_e `, {ref}`limit_pref_e `, {ref}`start_pref_f `, {ref}`limit_pref_f `, {ref}`start_pref_v ` and {ref}`limit_pref_v ` determine the start and limit prefactors of energy, force, and virial, respectively. The calculation and definition of Hessian loss are the same as for the other terms. +The legacy loss type `"ener_hess"` remains accepted as an alias of `"ener"`, but new input files should use the canonical `"ener"` type. -If one does not want to train with virial, then he/she may set the virial prefactors {ref}`start_pref_v ` and {ref}`limit_pref_v ` to 0. +Setting either Hessian prefactor to a nonzero value enables Hessian supervision, so schedules may ramp the term up from zero or down to zero. Both prefactors must be zero to disable the term. Earlier releases could silently disable the term on some backends when exactly one endpoint was zero. -## Hessian Format in PyTorch +The options {ref}`start_pref_e `, {ref}`limit_pref_e `, {ref}`start_pref_f `, {ref}`limit_pref_f `, {ref}`start_pref_v ` and {ref}`limit_pref_v ` determine the start and limit prefactors of energy, force, and virial, respectively. The calculation and definition of Hessian loss are the same as for the other terms. -In the PyTorch backend, Hessian matrices are listed in `hessian.npy` files, and the data format may contain the following files: +If one does not want to train with virial, set the virial prefactors {ref}`start_pref_v ` and {ref}`limit_pref_v ` to 0. + +## Hessian Data Format + +In the PyTorch and JAX backends, Hessian matrices are listed in `hessian.npy` files, and the data format may contain the following files: ``` type.raw @@ -50,7 +54,7 @@ Note that the `hessian.npy` should contain the **full** Hessian matrices with sh ## Train the Model -There are two approaches to training a Hessian model. The first method involves training the model from scratch using the same command as in the `ener` mode within the PyTorch backend: +There are two approaches to training a Hessian model. The first method involves training the model from scratch using the same command as in the `ener` mode: ::::{tab-set} @@ -61,9 +65,16 @@ dp --pt train input.json ``` ::: +:::{tab-item} JAX {{ jax_icon }} + +```bash +dp --jax train input.json +``` +::: + :::: -The second approach is to train a Hessian model from a pretrained energy model, following the same command as the `finetune` strategy within the PyTorch backend: +The second approach is to train a Hessian model from a pretrained energy model, following the same command as the `finetune` strategy: ::::{tab-set} @@ -74,6 +85,13 @@ dp --pt train input.json --finetune pretrained_energy.pt ``` ::: +:::{tab-item} JAX {{ jax_icon }} + +```bash +dp --jax train input.json --finetune pretrained_energy.jax +``` +::: + :::: The detailed loss can be found in `lcurve.out`: @@ -91,9 +109,9 @@ The detailed loss can be found in `lcurve.out`: ## Test the Model > [!WARNING] -> A model trained with Hessian cannot be frozen. If freezing is enforced, the model will be treated as a standard energy model, and the frozen one will no longer be able to output Hessian predictions. +> A PyTorch model trained with Hessian cannot currently be frozen with its Hessian output. If freezing is enforced, the frozen model is treated as a standard energy model. -If one do freeze and test a Hessian model using the commands: +PyTorch can test such a frozen model as a standard energy model. JAX can preserve the Hessian output during freezing by passing `--hessian`: ::::{tab-set} @@ -107,6 +125,15 @@ dp --pt test -m frozen_model.pth -s test_system -d ${output_prefix} -a -n 1 ``` ::: +:::{tab-item} JAX {{ jax_icon }} + +```bash +dp --jax freeze -c . -o frozen_model.hlo --hessian + +dp --jax test -m frozen_model.hlo -s test_system -d ${output_prefix} -a -n 1 +``` +::: + :::: If `dp --pt test -d ${output_prefix} -a` is specified, the output files will be the same as those in the `ener` mode, i.e., @@ -116,7 +143,7 @@ ${output_prefix}.e.out ${output_prefix}.e_peratom.out ${output_prefix}.f.out ${output_prefix}.v.out ${output_prefix}.v_peratom.out ``` -If one intends to use the trained model for Hessian predictions, then he/she is supposed to test the model directly without performing a freezing operation: +The backend checkpoints can also be tested directly without freezing: ::::{tab-set} @@ -128,9 +155,16 @@ dp --pt test -m model.pt -s test_system -d ${output_prefix} -a -n 1 ``` ::: +:::{tab-item} JAX {{ jax_icon }} + +```bash +dp --jax test -m model.ckpt.jax -s test_system -d ${output_prefix} -a -n 1 +``` +::: + :::: -If `dp --pt test -d ${output_prefix} -a` is specified, the predicted Hessian for each frame are output in an additional file in the working directory: +When either backend tests a Hessian-capable checkpoint with `-d ${output_prefix} -a`, the predicted Hessian for each frame is written to an additional file in the working directory: ``` ${output_prefix}.h.out diff --git a/source/tests/common/dpmodel/test_loss_ener.py b/source/tests/common/dpmodel/test_loss_ener.py index ebf9ba0a64..761bc39768 100644 --- a/source/tests/common/dpmodel/test_loss_ener.py +++ b/source/tests/common/dpmodel/test_loss_ener.py @@ -1,11 +1,17 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import unittest +from pathlib import ( + Path, +) import numpy as np from deepmd.dpmodel.loss.ener import ( EnergyLoss, ) +from deepmd.utils.data import ( + DeepmdData, +) from ...seed import ( GLOBAL_SEED, @@ -15,7 +21,14 @@ class TestEnergyLossBase(unittest.TestCase): """Base class providing common setup for dpmodel EnergyLoss tests.""" - def _make_data(self, natoms=5, nframes=2, numb_generalized_coord=0): + def _make_data( + self, + natoms=5, + nframes=2, + numb_generalized_coord=0, + hessian=False, + hessian_key="hessian", + ): """Generate fake model predictions and labels.""" rng = np.random.default_rng(GLOBAL_SEED) model_dict = { @@ -43,6 +56,10 @@ def _make_data(self, natoms=5, nframes=2, numb_generalized_coord=0): label_dict["find_drdq"] = 1.0 if hasattr(self, "enable_atom_ener_coeff") and self.enable_atom_ener_coeff: label_dict["atom_ener_coeff"] = rng.random((nframes, natoms, 1)) + if hessian: + model_dict[hessian_key] = rng.random((nframes, 3 * natoms, 3 * natoms)) + label_dict["hessian"] = rng.random((nframes, 3 * natoms, 3 * natoms)) + label_dict["find_hessian"] = 1.0 return model_dict, label_dict, natoms @@ -145,6 +162,87 @@ def test_forward(self) -> None: self.assertIsNotNone(loss) +class TestEnergyLossHessian(TestEnergyLossBase): + """Test Hessian loss inside the dpmodel energy loss.""" + + def test_disabled_hessian_preserves_standard_serialization(self) -> None: + """Default losses must keep the pre-Hessian serialization contract.""" + loss_fn = EnergyLoss(starter_learning_rate=1.0) + + data = loss_fn.serialize() + + self.assertEqual(data["@version"], 4) + self.assertNotIn("start_pref_h", data) + self.assertNotIn("limit_pref_h", data) + self.assertNotIn( + "hessian", + {item.key for item in loss_fn.label_requirement}, + ) + self.assertTrue( + all( + "special_shape" not in item.to_dict() + for item in loss_fn.label_requirement + ) + ) + + def test_forward_hessian(self) -> None: + loss_fn = EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_h=2.0, + limit_pref_h=1.0, + ) + model_dict, label_dict, natoms = self._make_data( + hessian=True, + hessian_key="energy_derv_r_derv_r", + ) + loss, more_loss = loss_fn.call(1.0, natoms, model_dict, label_dict) + diff_h = label_dict["hessian"].reshape(-1) - model_dict[ + "energy_derv_r_derv_r" + ].reshape(-1) + l2_hessian_loss = np.mean(np.square(diff_h)) + np.testing.assert_allclose(loss, 2.0 * l2_hessian_loss) + np.testing.assert_allclose(more_loss["rmse_h"], np.sqrt(l2_hessian_loss)) + self.assertIn( + "hessian", + {item.key for item in loss_fn.label_requirement}, + ) + + def test_hessian_label_requirement_loads_full_matrix(self) -> None: + loss_fn = EnergyLoss( + starter_learning_rate=1.0, + start_pref_h=1.0, + limit_pref_h=1.0, + ) + hessian_req = next( + item for item in loss_fn.label_requirement if item.key == "hessian" + ) + self.assertFalse(hessian_req.atomic) + self.assertEqual(hessian_req.special_shape, "hessian") + + system = ( + Path(__file__).resolve().parents[2] / "pt" / "hessian" / "data" / "H8C4N2O" + ) + data = DeepmdData(str(system), type_map=["C", "H", "N", "O"]) + data.add( + hessian_req.key, + hessian_req.ndof, + atomic=hessian_req.atomic, + must=hessian_req.must, + high_prec=hessian_req.high_prec, + special_shape=hessian_req.special_shape, + ) + + batch = data.get_batch(2) + + self.assertEqual(batch["hessian"].shape, (2, (3 * data.natoms) ** 2)) + + class TestEnergyLossSerialize(TestEnergyLossBase): """Test serialize/deserialize round-trip.""" @@ -160,16 +258,34 @@ def test_serialize_deserialize(self) -> None: start_pref_gf=1.0, limit_pref_gf=0.5, numb_generalized_coord=2, + start_pref_h=2.0, + limit_pref_h=0.5, ) data = loss_fn.serialize() + self.assertEqual(data["@version"], 5) + self.assertEqual(data["start_pref_h"], 2.0) + self.assertEqual(data["limit_pref_h"], 0.5) loss_fn2 = EnergyLoss.deserialize(data) - model_dict, label_dict, natoms = self._make_data(numb_generalized_coord=2) + model_dict, label_dict, natoms = self._make_data( + numb_generalized_coord=2, + hessian=True, + ) loss1, more1 = loss_fn.call(1.0, natoms, model_dict, label_dict) loss2, more2 = loss_fn2.call(1.0, natoms, model_dict, label_dict) np.testing.assert_allclose(loss1, loss2) for key in more1: np.testing.assert_allclose(more1[key], more2[key]) + def test_version_four_defaults_to_non_hessian_loss(self) -> None: + """Version-4 payloads deserialize without inventing Hessian supervision.""" + data = EnergyLoss(starter_learning_rate=1.0).serialize() + + loss_fn = EnergyLoss.deserialize(data) + + self.assertFalse(loss_fn.has_h) + self.assertEqual(loss_fn.start_pref_h, 0.0) + self.assertEqual(loss_fn.limit_pref_h, 0.0) + if __name__ == "__main__": unittest.main() diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py index 242ef0c188..337c325fdc 100644 --- a/source/tests/common/dpmodel/test_loss_padding.py +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -490,6 +490,13 @@ def _padded_force(f_A, f_B): return np.stack([f_A_pad, f_B], axis=0) # [2, NP, 3] +def _padded_hessian(h_A, h_B): + """Pad both Cartesian axes of frame A's Hessian to the batch width.""" + h_A_pad = np.zeros((3 * NP, 3 * NP), dtype=np.float64) + h_A_pad[: 3 * NA, : 3 * NA] = h_A + return np.stack([h_A_pad, h_B], axis=0) + + def _padded_atom(arr_A, arr_B, ncomp): """Pad arr_A from [NA, ncomp] to [NP, ncomp] with zeros, stack with arr_B.""" pad = np.zeros((NP, ncomp), dtype=np.float64) @@ -749,6 +756,299 @@ def test_no_op_for_non_mixed(self): ) +class TestDPModelEnergyLossHessianGradAccum: + """Hessian reductions exclude placeholder rows and columns.""" + + def _make_loss(self): + return EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=0.0, + limit_pref_ae=0.0, + start_pref_pf=0.0, + limit_pref_pf=0.0, + start_pref_h=1.0, + limit_pref_h=1.0, + ) + + @staticmethod + def _set_hessian(model_pred, label, pred_h, label_h): + model_pred["hessian"] = pred_h + label["hessian"] = label_h + label["find_hessian"] = 1.0 + return model_pred, label + + def _loss_fn(self, model_pred, label, natoms): + loss, _ = self._make_loss().call(1.0, natoms, model_pred, label) + return float(loss) + + def test_mse_grad_accum(self): + """Padded Hessian MSE equals the mean of per-frame Hessian MSEs.""" + h_A = _rnd(3 * NA, 3 * NA) + h_A_hat = _rnd(3 * NA, 3 * NA) + h_B = _rnd(3 * NB, 3 * NB) + h_B_hat = _rnd(3 * NB, 3 * NB) + + def make_A(): + pred, label = _full_ener_dicts( + 1, + NA, + np.zeros((1, 1)), + np.zeros((1, 1)), + mask=np.ones((1, NA)), + ) + pred, label = self._set_hessian(pred, label, h_A[None], h_A_hat[None]) + return pred, label, NA + + def make_B(): + pred, label = _full_ener_dicts( + 1, + NB, + np.zeros((1, 1)), + np.zeros((1, 1)), + mask=np.ones((1, NB)), + ) + pred, label = self._set_hessian(pred, label, h_B[None], h_B_hat[None]) + return pred, label, NB + + def make_padded(): + pred, label = _full_ener_dicts( + 2, NP, np.zeros((2, 1)), np.zeros((2, 1)), mask=_MASK_PAD + ) + pred, label = self._set_hessian( + pred, + label, + _padded_hessian(h_A, h_B), + _padded_hessian(h_A_hat, h_B_hat), + ) + return pred, label, NP + + assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + + def test_placeholder_pairs_are_excluded_from_metrics(self): + """Ghost-touching matrix entries do not dilute or pollute MAE/RMSE.""" + pred, label = _full_ener_dicts( + 1, + NP, + np.zeros((1, 1)), + np.zeros((1, 1)), + mask=np.array([[1.0] * NA + [0.0] * (NP - NA)]), + ) + pred_h = np.zeros((1, 3 * NP, 3 * NP), dtype=np.float64) + label_h = np.full_like(pred_h, 100.0) + label_h[:, : 3 * NA, : 3 * NA] = 2.0 + pred, label = self._set_hessian(pred, label, pred_h, label_h) + + loss, more_loss = self._make_loss().call(1.0, NP, pred, label, mae=True) + + np.testing.assert_allclose(loss, 4.0) + np.testing.assert_allclose(more_loss["rmse_h"], 2.0) + np.testing.assert_allclose(more_loss["mae_h"], 2.0) + + def test_pt_placeholder_pairs_are_excluded_from_metrics(self): + """The PyTorch energy loss uses the same Hessian pair mask.""" + import pytest + + torch = pytest.importorskip("torch") + from deepmd.pt.loss.ener import EnergyStdLoss as PTEnergyStdLoss + from deepmd.pt.utils import ( + env, + ) + + pred, label = _full_ener_dicts( + 1, + NP, + np.zeros((1, 1)), + np.zeros((1, 1)), + mask=np.array([[1.0] * NA + [0.0] * (NP - NA)]), + ) + pred_h = np.zeros((1, 3 * NP, 3 * NP), dtype=np.float64) + label_h = np.full_like(pred_h, 100.0) + label_h[:, : 3 * NA, : 3 * NA] = 2.0 + pred, label = self._set_hessian(pred, label, pred_h, label_h) + pt_pred = { + key: torch.as_tensor(value, device=env.DEVICE) + if isinstance(value, np.ndarray) + else value + for key, value in pred.items() + } + pt_label = { + key: torch.as_tensor(value, device=env.DEVICE) + if isinstance(value, np.ndarray) + else value + for key, value in label.items() + } + loss_obj = PTEnergyStdLoss( + starter_learning_rate=1.0, + start_pref_h=1.0, + limit_pref_h=1.0, + ) + + _, loss, more_loss = loss_obj({}, lambda: pt_pred, pt_label, NP, 1.0, mae=True) + + torch.testing.assert_close(loss, loss.new_tensor(4.0)) + torch.testing.assert_close(more_loss["rmse_h"], loss.new_tensor(2.0)) + torch.testing.assert_close(more_loss["mae_h"], loss.new_tensor(2.0)) + + def test_no_op_for_non_mixed(self): + """An all-ones mask preserves the established Hessian reduction.""" + pred_h = _rnd(1, 3 * NP, 3 * NP) + label_h = _rnd(1, 3 * NP, 3 * NP) + pred_mask, label_mask = _full_ener_dicts( + 1, + NP, + np.zeros((1, 1)), + np.zeros((1, 1)), + mask=np.ones((1, NP)), + ) + pred_plain, label_plain = _full_ener_dicts( + 1, NP, np.zeros((1, 1)), np.zeros((1, 1)) + ) + pred_mask, label_mask = self._set_hessian( + pred_mask, label_mask, pred_h, label_h + ) + pred_plain, label_plain = self._set_hessian( + pred_plain, label_plain, pred_h, label_h + ) + + loss_mask, metrics_mask = self._make_loss().call( + 1.0, NP, pred_mask, label_mask, mae=True + ) + loss_plain, metrics_plain = self._make_loss().call( + 1.0, NP, pred_plain, label_plain, mae=True + ) + + np.testing.assert_allclose(loss_mask, loss_plain) + np.testing.assert_allclose(metrics_mask["rmse_h"], metrics_plain["rmse_h"]) + np.testing.assert_allclose(metrics_mask["mae_h"], metrics_plain["mae_h"]) + + def _make_mae_loss(self): + return EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=0.0, + limit_pref_ae=0.0, + start_pref_pf=0.0, + limit_pref_pf=0.0, + start_pref_h=1.0, + limit_pref_h=1.0, + loss_func="mae", + ) + + def _constant_residual_dicts(self, residual): + """One frame whose Hessian residual is *residual* everywhere.""" + pred, label = _full_ener_dicts( + 1, NP, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NP)) + ) + pred_h = np.zeros((1, 3 * NP, 3 * NP), dtype=np.float64) + label_h = np.full_like(pred_h, residual) + return self._set_hessian(pred, label, pred_h, label_h) + + def test_mae_loss_func_uses_l1_for_the_hessian(self): + """``loss_func='mae'`` must not train the Hessian term on squared error.""" + residual = 10.0 + pred, label = self._constant_residual_dicts(residual) + + loss, more_loss = self._make_mae_loss().call(1.0, NP, pred, label, mae=True) + + # L2 would give residual ** 2; the displays stay MSE-derived. + np.testing.assert_allclose(loss, residual) + np.testing.assert_allclose(more_loss["mae_h"], residual) + np.testing.assert_allclose(more_loss["rmse_h"], residual) + + mse_pred, mse_label = self._constant_residual_dicts(residual) + mse_loss, _ = self._make_loss().call(1.0, NP, mse_pred, mse_label) + np.testing.assert_allclose(mse_loss, residual**2) + + def test_unknown_loss_func_is_rejected_for_the_hessian(self): + """Any other loss_func must fail rather than silently train on MSE.""" + import pytest + + loss_obj = self._make_loss() + # loss_func is validated in __init__, so set it directly to reach the + # Hessian branch's dispatch. + loss_obj.loss_func = "huber" + pred, label = self._constant_residual_dicts(1.0) + with pytest.raises(NotImplementedError, match="hessian loss"): + loss_obj.call(1.0, NP, pred, label) + + def test_huber_rejects_hessian_supervision(self): + """use_huber must not silently fall back to a raw MSE Hessian term.""" + import pytest + + with pytest.raises(RuntimeError, match="Huber loss is not implemented"): + EnergyLoss( + starter_learning_rate=1.0, + start_pref_h=1.0, + limit_pref_h=1.0, + use_huber=True, + ) + + def test_pt_mae_loss_func_uses_l1_for_the_hessian(self): + """The PyTorch loss must dispatch on loss_func for the Hessian too.""" + import pytest + + torch = pytest.importorskip("torch") + from deepmd.pt.loss.ener import EnergyStdLoss as PTEnergyStdLoss + from deepmd.pt.utils import ( + env, + ) + + residual = 10.0 + pred, label = self._constant_residual_dicts(residual) + pt_pred = { + key: torch.as_tensor(value, device=env.DEVICE) + if isinstance(value, np.ndarray) + else value + for key, value in pred.items() + } + pt_label = { + key: torch.as_tensor(value, device=env.DEVICE) + if isinstance(value, np.ndarray) + else value + for key, value in label.items() + } + + def make_loss(loss_func): + return PTEnergyStdLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_h=1.0, + limit_pref_h=1.0, + loss_func=loss_func, + ) + + _, mae_loss, more_loss = make_loss("mae")( + {}, lambda: pt_pred, pt_label, NP, 1.0, mae=True + ) + _, mse_loss, _ = make_loss("mse")({}, lambda: pt_pred, pt_label, NP, 1.0) + + torch.testing.assert_close(mae_loss, mae_loss.new_tensor(residual)) + torch.testing.assert_close(mse_loss, mse_loss.new_tensor(residual**2)) + torch.testing.assert_close(more_loss["mae_h"], mae_loss.new_tensor(residual)) + + with pytest.raises(RuntimeError, match="Huber loss is not implemented"): + PTEnergyStdLoss( + starter_learning_rate=1.0, + start_pref_h=1.0, + limit_pref_h=1.0, + use_huber=True, + ) + + class TestDPModelEnergyLossVirialGradAccum: """Idiom 2 (extensive, k=9) for the virial (has_v) term. diff --git a/source/tests/common/test_argcheck_backend_docs.py b/source/tests/common/test_argcheck_backend_docs.py index aa21033d81..ee5d40398a 100644 --- a/source/tests/common/test_argcheck_backend_docs.py +++ b/source/tests/common/test_argcheck_backend_docs.py @@ -71,7 +71,7 @@ def test_representative_declared_support_labels(self) -> None: energy_loss = argcheck.loss_args_plugin.get_argument("ener") self.assertTrue( energy_loss["start_pref_h"].doc.startswith( - "(Supported Backend: PyTorch, PaddlePaddle) " + "(Supported Backend: PyTorch, JAX, PaddlePaddle) " ) ) self.assertTrue( diff --git a/source/tests/common/test_argcheck_loss.py b/source/tests/common/test_argcheck_loss.py new file mode 100644 index 0000000000..b623e1ab4e --- /dev/null +++ b/source/tests/common/test_argcheck_loss.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.utils.argcheck import ( + loss_args, +) + + +def test_energy_hessian_loss_type_is_energy_alias() -> None: + """The legacy Hessian loss type must normalize to the energy schema.""" + canonical = loss_args().normalize_value( + { + "type": "ener", + "start_pref_h": 2.0, + "limit_pref_h": 1.0, + } + ) + legacy = loss_args().normalize_value( + { + "type": "ener_hess", + "start_pref_h": 2.0, + "limit_pref_h": 1.0, + } + ) + + loss_args().check_value(legacy, strict=True) + assert legacy == canonical + assert legacy["type"] == "ener" diff --git a/source/tests/common/test_data_system_hessian.py b/source/tests/common/test_data_system_hessian.py new file mode 100644 index 0000000000..193126e08b --- /dev/null +++ b/source/tests/common/test_data_system_hessian.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Mixed-system batching must pad Hessian labels as square blocks. + +``DeepmdDataSystem.get_batch()`` merges frames from systems with different +atom counts. A Hessian is a ``(3 * natoms, 3 * natoms)`` matrix, so neither +the non-atomic concatenate nor the atomic flat-prefix copy can pad it: the +former raises on ragged systems and the latter would scatter the rows. The +JAX trainer calls ``get_batch()`` directly, so this is reachable from a +``batch_size: "mixed:N"`` run with Hessian supervision. +""" + +import shutil +import tempfile +import unittest +from pathlib import ( + Path, +) + +import numpy as np + +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) + +NATOMS = (2, 3) +NFRAMES = 2 +# One constant per system, so a mis-scattered block is visible as a stray value. +HESSIAN_VALUES = (1.0, 2.0) + + +def _write_system(root: Path, natoms: int, value: float) -> None: + """Write a one-set system whose Hessian is a constant square block.""" + root.mkdir(parents=True) + (root / "type.raw").write_text("\n".join(["0"] * natoms) + "\n") + (root / "type_map.raw").write_text("A\n") + set_dir = root / "set.000" + set_dir.mkdir() + rng = np.random.default_rng(natoms) + np.save(set_dir / "coord.npy", rng.random((NFRAMES, natoms * 3))) + np.save(set_dir / "box.npy", np.tile(np.eye(3).reshape(9) * 20.0, (NFRAMES, 1))) + np.save(set_dir / "energy.npy", rng.random((NFRAMES, 1))) + dof = natoms * 3 + np.save( + set_dir / "hessian.npy", + np.full((NFRAMES, dof * dof), value, dtype=np.float64), + ) + + +class TestMixedBatchHessian(unittest.TestCase): + def setUp(self) -> None: + self.tmpdir = Path(tempfile.mkdtemp()) + systems = [] + for natoms, value in zip(NATOMS, HESSIAN_VALUES, strict=True): + path = self.tmpdir / f"sys_{natoms}" + _write_system(path, natoms, value) + systems.append(str(path)) + self.ds = DeepmdDataSystem(systems, "mixed:2", 1, 2.0) + self.ds.add("energy", 1, atomic=False, must=True, high_prec=True) + self.ds.add("hessian", 1, atomic=False, must=True, special_shape="hessian") + + def tearDown(self) -> None: + shutil.rmtree(self.tmpdir, ignore_errors=True) + + def _ragged_batch(self) -> dict: + """Merge exactly one frame from each system. + + ``get_batch_mixed`` picks systems at random, so it does not reliably + produce the ragged case this covers. + """ + batch_data = [] + for sys_idx in range(self.ds.nsystems): + bb_data = self.ds.data_systems[sys_idx].get_batch(1) + bb_data["natoms_vec"] = self.ds.natoms_vec[sys_idx] + bb_data["default_mesh"] = self.ds.default_mesh[sys_idx] + batch_data.append(bb_data) + return self.ds._merge_batch_data(batch_data) + + def test_mixed_batch_embeds_each_hessian_block(self) -> None: + """Each frame's block sits in the top-left of the padded square.""" + batch = self._ragged_batch() + max_natoms = int(batch["natoms_vec"][0]) + self.assertEqual(max_natoms, max(NATOMS)) + padded_dof = max_natoms * 3 + + self.assertEqual(batch["hessian"].shape, (len(NATOMS), padded_dof * padded_dof)) + hessian = batch["hessian"].reshape(-1, padded_dof, padded_dof) + + for frame, (natoms, value) in enumerate( + zip(NATOMS, HESSIAN_VALUES, strict=True) + ): + frame_dof = natoms * 3 + self.assertEqual(int(batch["real_natoms_vec"][frame, 0]), natoms) + np.testing.assert_array_equal(hessian[frame, :frame_dof, :frame_dof], value) + # Everything outside the block is padding and must stay zero. + np.testing.assert_array_equal(hessian[frame, frame_dof:, :], 0.0) + np.testing.assert_array_equal(hessian[frame, :, frame_dof:], 0.0) + + def test_get_batch_accepts_a_mixed_hessian_batch(self) -> None: + """The public entry point must not raise for any sampled combination.""" + for _ in range(20): + batch = self.ds.get_batch() + padded_dof = int(batch["natoms_vec"][0]) * 3 + self.assertEqual(batch["hessian"].shape[1], padded_dof * padded_dof) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/jax/test_make_hessian_model.py b/source/tests/jax/test_make_hessian_model.py index bb25bd67ca..e9fe727139 100644 --- a/source/tests/jax/test_make_hessian_model.py +++ b/source/tests/jax/test_make_hessian_model.py @@ -174,3 +174,13 @@ def test_output_def(self) -> None: self.model_hess.model_output_def()["energy_derv_r_derv_r"].category, OutputVariableCategory.DERV_R_DERV_R, ) + + def test_enable_hessian_is_idempotent(self) -> None: + """Restored Hessian models may be enabled again by loss requirements.""" + self.model_valu.enable_hessian() + enabled_def = self.model_valu.hess_fitting_def + + self.model_valu.enable_hessian() + + self.assertIs(self.model_valu.hess_fitting_def, enabled_def) + self.assertTrue(self.model_valu.model_output_def()["energy"].r_hessian) diff --git a/source/tests/jax/test_training.py b/source/tests/jax/test_training.py index c175de8da4..0055a35b4b 100644 --- a/source/tests/jax/test_training.py +++ b/source/tests/jax/test_training.py @@ -76,6 +76,9 @@ from deepmd.utils.compat import ( convert_optimizer_v31_to_v32, ) +from deepmd.utils.data import ( + DataRequirementItem, +) from deepmd.utils.env_mat_stat import ( StatItem, ) @@ -205,6 +208,74 @@ def _minimal_jax_multitask_config(model_params: dict) -> dict: } +class _RequirementModel: + """Minimal model double for trainer requirement-driven output tests.""" + + def __init__(self) -> None: + self.hessian_enable_calls = 0 + + def enable_hessian(self) -> None: + self.hessian_enable_calls += 1 + + def get_dim_fparam(self) -> int: + return 0 + + +@patch("deepmd.jax.train.trainer.DPTrainer._build_losses") +@patch("deepmd.jax.train.trainer.get_model") +def test_jax_hessian_mode_follows_loss_data_requirement( + get_model, + build_losses, +) -> None: + """The trainer consumes the loss requirement instead of its prefactors.""" + model = _RequirementModel() + get_model.return_value = model + build_losses.return_value = { + "Default": SimpleNamespace( + label_requirement=[DataRequirementItem("hessian", ndof=1)] + ) + } + model_params = {"type_map": ["O"], "descriptor": {}} + + trainer = DPTrainer(_minimal_jax_config(model_params)) + + assert model.hessian_enable_calls == 1 + assert trainer.model_def_script["hessian_mode"] is True + + +@patch("deepmd.jax.train.trainer.DPTrainer._build_losses") +@patch("deepmd.jax.train.trainer.get_model") +def test_jax_multitask_hessian_only_enables_requesting_branch( + get_model, + build_losses, +) -> None: + """Each multi-task branch follows its own loss data requirements.""" + model_a = _RequirementModel() + model_b = _RequirementModel() + get_model.side_effect = [model_a, model_b] + build_losses.return_value = { + "task_a": SimpleNamespace( + label_requirement=[DataRequirementItem("hessian", ndof=1)] + ), + "task_b": SimpleNamespace( + label_requirement=[DataRequirementItem("energy", ndof=1)] + ), + } + model_params = { + "model_dict": { + "task_a": {"type_map": ["O"], "descriptor": {}}, + "task_b": {"type_map": ["O"], "descriptor": {}}, + } + } + + trainer = DPTrainer(_minimal_jax_multitask_config(model_params)) + + assert model_a.hessian_enable_calls == 1 + assert model_b.hessian_enable_calls == 0 + assert trainer.model_def_script["model_dict"]["task_a"]["hessian_mode"] is True + assert "hessian_mode" not in trainer.model_def_script["model_dict"]["task_b"] + + def _shared_jax_model_config(*, share_fitting: bool = True) -> dict: shared_dict: dict = { "shared_type_map": ["O", "H", "B"], diff --git a/source/tests/pd/test_loss.py b/source/tests/pd/test_loss.py index a7b8109e10..b5b0f503ed 100644 --- a/source/tests/pd/test_loss.py +++ b/source/tests/pd/test_loss.py @@ -12,8 +12,13 @@ ) from deepmd.pd.loss import ( + EnergyHessianStdLoss, EnergyStdLoss, ) +from deepmd.pd.train.training import ( + get_loss, + whether_hessian, +) from deepmd.pd.utils.dataset import ( DeepmdDataSetForLoader, ) @@ -581,5 +586,78 @@ def fake_model(): self.assertTrue(np.isnan(pd_more_loss_absent[f"l2_{key}_loss"].numpy())) +class TestEnergyHessianLossCompatibility(unittest.TestCase): + """Paddle uses the same Hessian activation and data-shape contract.""" + + def test_either_schedule_endpoint_enables_hessian(self) -> None: + """Ramp-up and ramp-down schedules must both activate supervision.""" + for loss_type in ("ener", "ener_hess"): + for start_pref_h, limit_pref_h in ((1.0, 0.0), (0.0, 1.0)): + params = { + "type": loss_type, + "start_pref_h": start_pref_h, + "limit_pref_h": limit_pref_h, + } + + self.assertTrue(whether_hessian(params)) + loss = get_loss(params, start_lr=1.0, _ntypes=0, _model=None) + self.assertIsInstance(loss, EnergyHessianStdLoss) + self.assertTrue(loss.has_h) + hessian_req = next( + item for item in loss.label_requirement if item.key == "hessian" + ) + self.assertFalse(hessian_req.atomic) + self.assertEqual(hessian_req.special_shape, "hessian") + + def test_zero_schedule_keeps_standard_energy_loss(self) -> None: + """Two zero endpoints must leave Hessian supervision disabled.""" + params = { + "type": "ener", + "start_pref_h": 0.0, + "limit_pref_h": 0.0, + } + + self.assertFalse(whether_hessian(params)) + loss = get_loss(params, start_lr=1.0, _ntypes=0, _model=None) + self.assertIs(type(loss), EnergyStdLoss) + + def test_huber_is_rejected_for_hessian_supervision(self) -> None: + """Paddle must not silently optimize Hessians with raw MSE under Huber.""" + with self.assertRaisesRegex( + RuntimeError, "Huber loss is not implemented for hessian" + ): + EnergyHessianStdLoss(start_pref_h=1.0, use_huber=True) + + def test_hessian_padding_is_masked_on_both_atom_axes(self) -> None: + """Only the real-real Cartesian block contributes to Hessian loss.""" + residual = np.full((1, 6, 6), 10.0) + residual[:, :3, :3] = 2.0 + model_pred = { + "hessian": paddle.zeros([1, 6, 6], dtype="float64"), + "mask": paddle.to_tensor([[1.0, 0.0]], dtype="float64"), + } + label = { + "hessian": paddle.to_tensor(residual, dtype="float64"), + "find_hessian": 1.0, + } + loss_module = EnergyHessianStdLoss( + starter_learning_rate=1.0, + start_pref_h=1.0, + limit_pref_h=1.0, + ) + + def fake_model(**kwargs): + return model_pred + + _, loss, more_loss = loss_module( + {}, fake_model, label, natoms=2, learning_rate=1.0, mae=True + ) + + self.assertTrue(np.allclose(loss.numpy(), 4.0)) + self.assertTrue(np.allclose(more_loss["l2_hessian_loss"].numpy(), 4.0)) + self.assertTrue(np.allclose(more_loss["rmse_h"].numpy(), 2.0)) + self.assertTrue(np.allclose(more_loss["mae_h"].numpy(), 2.0)) + + if __name__ == "__main__": unittest.main() diff --git a/source/tests/pt/model/test_make_hessian_model.py b/source/tests/pt/model/test_make_hessian_model.py index f2387ca8b0..2bb1539f71 100644 --- a/source/tests/pt/model/test_make_hessian_model.py +++ b/source/tests/pt/model/test_make_hessian_model.py @@ -182,3 +182,13 @@ def test_output_def(self) -> None: self.model_hess.model_output_def()["energy_derv_r_derv_r"].category, OutputVariableCategory.DERV_R_DERV_R, ) + + def test_enable_hessian_is_idempotent(self) -> None: + """Restored Hessian models may be enabled again by loss requirements.""" + self.model_valu.enable_hessian() + enabled_type = type(self.model_valu) + + self.model_valu.enable_hessian() + + self.assertIs(type(self.model_valu), enabled_type) + self.assertTrue(self.model_valu.model_output_def()["energy"].r_hessian) diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py index 494338a2da..21210a6260 100644 --- a/source/tests/pt/model/test_sezm_model.py +++ b/source/tests/pt/model/test_sezm_model.py @@ -1005,7 +1005,7 @@ def _build_wrapper(use_compile: bool) -> ModelWrapper: "water_1": {"type": "ener"}, "water_2": {"type": "ener"}, } - models = get_model_for_wrapper(mt_cfg, _loss_params=loss_params) + models = get_model_for_wrapper(mt_cfg) prepare_model_for_loss(models, loss_params) wrapper = ModelWrapper(models) wrapper.share_params(shared_links, {"water_1": 0.5, "water_2": 0.5}) diff --git a/source/tests/pt/test_change_bias.py b/source/tests/pt/test_change_bias.py index a883449d22..746656ac8c 100644 --- a/source/tests/pt/test_change_bias.py +++ b/source/tests/pt/test_change_bias.py @@ -96,7 +96,6 @@ def setUp(self) -> None: self.model_path_user_bias = Path(current_path) / ( model_name + "user_bias" + ".pt" ) - self.loss_params = self.config["loss"] def test_change_bias_with_data(self) -> None: run_dp( @@ -108,7 +107,6 @@ def test_change_bias_with_data(self) -> None: model_params = state_dict["model"]["_extra_state"]["model_params"] model_for_wrapper = get_model_for_wrapper( model_params, - _loss_params=self.loss_params, ) wrapper = ModelWrapper(model_for_wrapper) wrapper.load_state_dict(state_dict["model"]) @@ -134,7 +132,6 @@ def test_change_bias_with_data_sys_file(self) -> None: model_params = state_dict["model"]["_extra_state"]["model_params"] model_for_wrapper = get_model_for_wrapper( model_params, - _loss_params=self.loss_params, ) wrapper = ModelWrapper(model_for_wrapper) wrapper.load_state_dict(state_dict["model"]) @@ -158,7 +155,6 @@ def test_change_bias_with_user_defined(self) -> None: model_params = state_dict["model"]["_extra_state"]["model_params"] model_for_wrapper = get_model_for_wrapper( model_params, - _loss_params=self.loss_params, ) wrapper = ModelWrapper(model_for_wrapper) wrapper.load_state_dict(state_dict["model"]) diff --git a/source/tests/pt/test_loss.py b/source/tests/pt/test_loss.py index 2519111357..c5176f1fea 100644 --- a/source/tests/pt/test_loss.py +++ b/source/tests/pt/test_loss.py @@ -708,6 +708,130 @@ def fake_model(): self.assertTrue(np.isnan(pt_more_loss_h_absent[f"l2_{key}_loss"])) +class TestEnergyHessianLossCompatibility(unittest.TestCase): + """Hessian supervision is canonical on EnergyStdLoss in PyTorch.""" + + def test_legacy_class_matches_energy_loss(self) -> None: + kwargs = { + "starter_learning_rate": 1.0, + "start_pref_h": 2.0, + "limit_pref_h": 1.0, + } + canonical = EnergyStdLoss(**kwargs) + legacy = EnergyHessianStdLoss(**kwargs) + + self.assertEqual(canonical.serialize(), legacy.serialize()) + self.assertEqual(canonical.serialize()["@version"], 5) + self.assertEqual(canonical.label_requirement, legacy.label_requirement) + hessian_req = next( + item for item in canonical.label_requirement if item.key == "hessian" + ) + self.assertFalse(hessian_req.atomic) + self.assertEqual(hessian_req.special_shape, "hessian") + + def test_legacy_loss_type_dispatches_to_energy_loss(self) -> None: + from deepmd.pt.train.training import ( + get_loss, + ) + + for loss_type in ("ener", "ener_hess"): + params = { + "type": loss_type, + "start_pref_h": 0.0, + "limit_pref_h": 1.0, + } + loss = get_loss(params, start_lr=1.0, _ntypes=0, _model=None) + self.assertIs(type(loss), EnergyStdLoss) + self.assertTrue(loss.has_h) + self.assertTrue( + any(item.key == "hessian" for item in loss.label_requirement) + ) + + def test_inference_only_does_not_request_hessian(self) -> None: + """The change-bias mock loss must not allocate unused quadratic labels.""" + loss = EnergyStdLoss(starter_learning_rate=1.0, inference=True) + + self.assertFalse(loss.has_h) + self.assertNotIn( + "hessian", + {item.key for item in loss.label_requirement}, + ) + + def test_version_four_defaults_to_non_hessian_loss(self) -> None: + """Older energy-loss payloads retain their non-Hessian behavior.""" + data = EnergyStdLoss(starter_learning_rate=1.0).serialize() + + loss = EnergyStdLoss.deserialize(data) + + self.assertEqual(data["@version"], 4) + self.assertFalse(loss.has_h) + self.assertEqual(loss.start_pref_h, 0.0) + self.assertEqual(loss.limit_pref_h, 0.0) + + def test_model_hessian_mode_follows_loss_data_requirement(self) -> None: + from deepmd.pt.train.training import ( + prepare_model_for_data_requirements, + ) + + class HessianCapableModel: + def __init__(self) -> None: + self.enable_calls = 0 + self.model_def_script = "" + + def enable_hessian(self) -> None: + self.enable_calls += 1 + + model = HessianCapableModel() + model_params: dict = {} + requirements = [ + DataRequirementItem( + "hessian", + ndof=1, + special_shape="hessian", + ) + ] + + prepare_model_for_data_requirements(model, requirements, model_params) + + self.assertEqual(model.enable_calls, 1) + self.assertTrue(model_params["hessian_mode"]) + self.assertEqual(model.model_def_script, '{"hessian_mode": true}') + + def test_multitask_hessian_mode_only_enables_requesting_branch(self) -> None: + from deepmd.pt.train.training import ( + prepare_model_for_data_requirements, + ) + + class HessianCapableModel: + def __init__(self) -> None: + self.enable_calls = 0 + self.model_def_script = "" + + def enable_hessian(self) -> None: + self.enable_calls += 1 + + models = { + "with_hessian": HessianCapableModel(), + "without_hessian": object(), + } + requirements = { + "with_hessian": [DataRequirementItem("hessian", ndof=1)], + "without_hessian": [DataRequirementItem("energy", ndof=1)], + } + model_params = { + "model_dict": { + "with_hessian": {}, + "without_hessian": {}, + } + } + + prepare_model_for_data_requirements(models, requirements, model_params) + + self.assertEqual(models["with_hessian"].enable_calls, 1) + self.assertTrue(model_params["model_dict"]["with_hessian"]["hessian_mode"]) + self.assertNotIn("hessian_mode", model_params["model_dict"]["without_hessian"]) + + class TestEnerSpinLoss(LossCommonTest): def setUp(self) -> None: self.start_lr = 1.1