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

Draft
njzjz wants to merge 2 commits into
base: master
Choose a base branch
from
Draft

fix(jax): add hessian energy loss #5607

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ae453b0
fix(jax): add hessian energy loss
njzjz Jun 29, 2026
fcfc2ec
fix(data): declare hessian label shape
njzjz Jul 7, 2026
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
40 changes: 40 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,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",
)
)
Comment on lines +565 to +575

@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 +602,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
95 changes: 55 additions & 40 deletions deepmd/utils/data.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down Expand Up @@ -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_"):
Expand All @@ -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

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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_:
Expand Down Expand Up @@ -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`."
Expand All @@ -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:
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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:
Expand All @@ -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])

Expand Down Expand Up @@ -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__(
Expand All @@ -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
Expand All @@ -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:
Expand All @@ -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:
Expand Down
5 changes: 5 additions & 0 deletions deepmd/utils/data_system.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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.

Expand Down Expand Up @@ -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(
Expand All @@ -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:
Expand Down
Loading