From ae453b010110c98eb1813deb819c1f0f8d6a4fa8 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 30 Jun 2026 02:09:21 +0800 Subject: [PATCH 1/2] fix(jax): add hessian energy loss --- deepmd/dpmodel/loss/ener.py | 39 ++++++++++++++ deepmd/jax/train/trainer.py | 6 +++ source/tests/common/dpmodel/test_loss_ener.py | 54 ++++++++++++++++++- 3 files changed, 97 insertions(+), 2 deletions(-) diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py index 7515f19b9a..0a8d57fc8c 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/deepmd/dpmodel/loss/ener.py @@ -71,6 +71,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. @@ -123,6 +127,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", @@ -155,12 +161,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" @@ -270,6 +279,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 = {} @@ -457,6 +467,23 @@ 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", None) + ) + if self.has_h and hessian is not None and "hessian" in label_dict: + find_hessian = label_dict.get("find_hessian", 0.0) + diff_h = xp.reshape(label_dict["hessian"], (-1,)) - xp.reshape( + hessian, + (-1,), + ) + l2_hessian_loss = xp.mean(xp.square(diff_h)) + loss += pref_h * find_hessian * l2_hessian_loss + more_loss["rmse_h"] = self.display_if_exist( + xp.sqrt(l2_hessian_loss), find_hessian + ) + if mae: + mae_h = xp.mean(xp.abs(diff_h)) + more_loss["mae_h"] = self.display_if_exist(mae_h, find_hessian) self.l2_l = loss more_loss["rmse"] = xp.sqrt(loss) @@ -535,6 +562,16 @@ def label_requirement(self) -> list[DataRequirementItem]: default=1.0, ) ) + if self.has_h: + label_requirement.append( + DataRequirementItem( + "hessian", + ndof=1, + atomic=True, + must=False, + high_prec=False, + ) + ) return label_requirement def serialize(self) -> dict: @@ -564,6 +601,8 @@ def serialize(self) -> dict: "start_pref_gf": self.start_pref_gf, "limit_pref_gf": self.limit_pref_gf, "numb_generalized_coord": self.numb_generalized_coord, + "start_pref_h": self.start_pref_h, + "limit_pref_h": self.limit_pref_h, "use_huber": self.use_huber, "huber_delta": self.huber_delta, "loss_func": self.loss_func, diff --git a/deepmd/jax/train/trainer.py b/deepmd/jax/train/trainer.py index 180249eaef..46d8fc581e 100644 --- a/deepmd/jax/train/trainer.py +++ b/deepmd/jax/train/trainer.py @@ -117,6 +117,8 @@ def get_lr_and_coef(lr_param: dict) -> LearningRateExp: loss_type = loss_param.get("type", "ener") if loss_type == "ener": self.loss = EnergyLoss.get_loss(loss_param) + if getattr(self.loss, "has_h", False): + self.model.enable_hessian() else: raise RuntimeError("unknown loss type " + loss_type) @@ -211,6 +213,8 @@ def loss_fn( 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) loss, more_loss = self.loss( learning_rate=lr, natoms=label_dict["type"].shape[1], @@ -249,6 +253,8 @@ def loss_fn_more_loss( 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) loss, more_loss = self.loss( learning_rate=lr, natoms=label_dict["type"].shape[1], diff --git a/source/tests/common/dpmodel/test_loss_ener.py b/source/tests/common/dpmodel/test_loss_ener.py index ebf9ba0a64..cd6803a4ab 100644 --- a/source/tests/common/dpmodel/test_loss_ener.py +++ b/source/tests/common/dpmodel/test_loss_ener.py @@ -15,7 +15,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 +50,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 +156,38 @@ def test_forward(self) -> None: self.assertIsNotNone(loss) +class TestEnergyLossHessian(TestEnergyLossBase): + """Test Hessian loss inside the dpmodel energy loss.""" + + 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}, + ) + + class TestEnergyLossSerialize(TestEnergyLossBase): """Test serialize/deserialize round-trip.""" @@ -160,10 +203,17 @@ 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["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) From fcfc2ece7482a70c5d3d68bf30c9afa93a3ca841 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 7 Jul 2026 23:40:37 +0800 Subject: [PATCH 2/2] fix(data): declare hessian label shape --- deepmd/dpmodel/loss/ener.py | 3 +- deepmd/utils/data.py | 95 +++++++++++-------- deepmd/utils/data_system.py | 5 + source/tests/common/dpmodel/test_loss_ener.py | 35 +++++++ 4 files changed, 97 insertions(+), 41 deletions(-) diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py index 0a8d57fc8c..996787a043 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/deepmd/dpmodel/loss/ener.py @@ -567,9 +567,10 @@ def label_requirement(self) -> list[DataRequirementItem]: DataRequirementItem( "hessian", ndof=1, - atomic=True, + atomic=False, must=False, high_prec=False, + special_shape="hessian", ) ) return label_requirement diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py index 38db47809d..ddc29ce8cf 100644 --- a/deepmd/utils/data.py +++ b/deepmd/utils/data.py @@ -159,6 +159,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,6 +188,9 @@ 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_"): @@ -202,6 +206,7 @@ def add( "default": default, "dtype": dtype, "output_natoms_for_type_sel": output_natoms_for_type_sel, + "special_shape": special_shape, } return self @@ -732,6 +737,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: @@ -808,7 +814,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 @@ -836,7 +844,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_: @@ -868,24 +888,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`." @@ -898,7 +903,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: @@ -925,8 +932,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 @@ -959,7 +970,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) @@ -983,7 +996,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: @@ -1008,23 +1031,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]) @@ -1137,6 +1146,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__( @@ -1151,6 +1163,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 @@ -1162,6 +1175,7 @@ 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: @@ -1176,6 +1190,7 @@ def to_dict(self) -> dict: "default": self.default, "dtype": self.dtype, "output_natoms_for_type_sel": self.output_natoms_for_type_sel, + "special_shape": self.special_shape, } def __getitem__(self, key: str) -> np.ndarray: diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 9d13cb4699..f020c5cf38 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: diff --git a/source/tests/common/dpmodel/test_loss_ener.py b/source/tests/common/dpmodel/test_loss_ener.py index cd6803a4ab..154688e5b0 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, @@ -187,6 +193,35 @@ def test_forward_hessian(self) -> None: {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."""