fix(jax): add hessian energy loss

[njzjz]

Summary

• add Hessian prefactors directly to dpmodel EnergyLoss instead of introducing a separate loss class
• add Hessian label requirements, loss/RMSE reporting, and serialization fields
• enable Hessian outputs in the JAX trainer when the Hessian loss is configured

Tests

• source venv/bin/activate && pytest source/tests/common/dpmodel/test_loss_ener.py -q
• source venv/bin/activate && ruff check .
• source venv/bin/activate && ruff format .

Summary by CodeRabbit

• New Features

  • Added optional Hessian loss support to the energy loss, including configurable scheduling parameters for Hessian weighting.
  • When Hessian targets are enabled, training now requests Hessian labels and computes an additional Hessian error term (with corresponding metrics).

• Tests

  • Added dedicated test coverage for the Hessian-loss computation path.
  • Updated loss serialization/round-trip tests to include the new Hessian weighting parameters and regenerated Hessian test data.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 162d37c7-9fa9-4b6f-a35c-1b8c82486b72

📥 Commits

Reviewing files that changed from the base of the PR and between b56a253 and ae453b0.

📒 Files selected for processing (3)

• deepmd/dpmodel/loss/ener.py
• deepmd/jax/train/trainer.py
• source/tests/common/dpmodel/test_loss_ener.py

🚧 Files skipped from review as they are similar to previous changes (3)

• deepmd/jax/train/trainer.py
• source/tests/common/dpmodel/test_loss_ener.py
• deepmd/dpmodel/loss/ener.py

---

📝 Walkthrough

Walkthrough

EnergyLoss adds optional Hessian loss prefactors and computes a Hessian L2 term when enabled. The JAX trainer enables Hessian support and exposes Hessian tensors to the loss. Tests cover the new path and serialization.

Changes

Hessian Loss Feature

Layer / File(s)	Summary
EnergyLoss Hessian parameters, call logic, label requirement, and serialization deepmd/dpmodel/loss/ener.py	Adds start_pref_h/limit_pref_h params and has_h flag to __init__, computes pref_h and Hessian MSE loss in call(), appends "hessian" DataRequirementItem in label_requirement, and exports new fields in serialize().
JAX trainer Hessian wiring deepmd/jax/train/trainer.py	DPTrainer calls self.model.enable_hessian() when loss.has_h is true; both loss_fn and loss_fn_more_loss populate model_dict["hessian"] from energy_derv_r_derv_r by squeezing axis -3.
Tests source/tests/common/dpmodel/test_loss_ener.py	_make_data extended with hessian/hessian_key options; TestEnergyLossHessian validates loss value, rmse_h metric, and label_requirement; serialize/deserialize round-trip updated to cover Hessian prefactors.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• deepmodeling/deepmd-kit#5325: Modifies EnergyLoss.call return path and more_loss dict in the same file, directly adjacent to the Hessian loss branch added here.

Suggested labels

new feature, Python

Suggested reviewers

• wanghan-iapcm
• anyangml

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: adding Hessian energy loss support in JAX.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 OpenGrep (1.23.0)

source/tests/common/dpmodel/test_loss_ener.py

┌──────────────┐
│ Opengrep CLI │
└──────────────┘

�[32m✔�[39m �[1mOpengrep OSS�[0m
�[32m✔�[39m Basic security coverage for first-party code vulnerabilities.

[00.13][ERROR]: unable to find a config; path .coderabbit-opengrep-fallback.yml does not exist

deepmd/dpmodel/loss/ener.py

┌──────────────┐
│ Opengrep CLI │
└──────────────┘

�[32m✔�[39m �[1mOpengrep OSS�[0m
�[32m✔�[39m Basic security coverage for first-party code vulnerabilities.

[00.15][ERROR]: unable to find a config; path .coderabbit-opengrep-fallback.yml does not exist

deepmd/jax/train/trainer.py

┌──────────────┐
│ Opengrep CLI │
└──────────────┘

�[32m✔�[39m �[1mOpengrep OSS�[0m
�[32m✔�[39m Basic security coverage for first-party code vulnerabilities.

[00.14][ERROR]: unable to find a config; path .coderabbit-opengrep-fallback.yml does not exist

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with 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.

Inline comments:
In `@deepmd/dpmodel/loss/ener.py`:
- Around line 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.

In `@deepmd/jax/train/trainer.py`:
- Around line 120-121: The Hessian enablement in Trainer setup is unguarded, so
`self.model.enable_hessian()` can break for model/loss combinations like the JAX
`zbl` path that do not implement it. Update the `Trainer` logic to check that
the model actually exposes `enable_hessian` before calling it, using the
existing `self.loss.has_h` condition as the trigger and keeping the current
hessian tensor handling unchanged.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 23778160-e75a-4d2b-be1f-2babb0b6710f

📥 Commits

Reviewing files that changed from the base of the PR and between 1550599 and b56a253.

📒 Files selected for processing (3)

• deepmd/dpmodel/loss/ener.py
• deepmd/jax/train/trainer.py
• source/tests/common/dpmodel/test_loss_ener.py

[njzjz]

Possible reviewers based on changed lines, exact file history, and exact-file review history:

• @wanghan-iapcm — 15 commits on changed files; 26 reviews on exact changed files (deepmd/dpmodel/loss/ener.py, deepmd/jax/train/trainer.py, deepmd/utils/data.py, deepmd/utils/data_system.py, source/tests/common/dpmodel/test_loss_ener.py).
• @OutisLi — 18 reviews on exact changed files (deepmd/dpmodel/loss/ener.py, deepmd/jax/train/trainer.py).

No review request was made automatically.

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh