Skip to content

fix(jax): add hessian energy loss #5607

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
njzjz wants to merge 1 commit into
base: master
Choose a base branch
from
Open

fix(jax): add hessian energy loss #5607

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 39 additions & 0 deletions deepmd/dpmodel/loss/ener.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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 = {}
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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,
)
)
Comment on lines +565 to +574

@coderabbitai coderabbitai Bot Jun 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift

label_requirement advertises the wrong Hessian shape.

DataRequirementItem("hessian", ndof=1, atomic=True) describes an nframes x natoms x 1 label, but the new loss path and tests consume a full nframes x (3 * natoms) x (3 * natoms) tensor. DPTrainer.data_requirements forwards this schema to the dataset loader, so real Hessian training will request/load hessian.npy with an incompatible contract even though these unit tests pass by building label_dict manually. Please make the requirement describe the real on-disk Hessian layout, or add dedicated loader support before enabling this in training.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@deepmd/dpmodel/loss/ener.py` around lines 565 - 574, The Hessian entry in
`DPTrainer.data_requirements` is advertising the wrong tensor shape for
`label_requirement`. Update the `DataRequirementItem("hessian", ...)` definition
in `ener.py` so it matches the real on-disk Hessian layout used by the new loss
path and tests, rather than the current atomic `ndof=1` schema. If the dataset
loader cannot yet consume the full Hessian tensor, add the loader support first
and keep `has_h` gated until the contract is consistent.

return label_requirement

def serialize(self) -> dict:
Expand Down Expand Up @@ -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,
Expand Down
6 changes: 6 additions & 0 deletions deepmd/jax/train/trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Comment thread
coderabbitai[bot] marked this conversation as resolved.
else:
raise RuntimeError("unknown loss type " + loss_type)

Expand Down Expand Up @@ -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],
Expand Down Expand Up @@ -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],
Expand Down
54 changes: 52 additions & 2 deletions source/tests/common/dpmodel/test_loss_ener.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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."""

Expand All @@ -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)
Expand Down
Loading