From eeca00976fff44586dbda9fa40a81c1d808e8902 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 20 Jan 2026 20:11:39 +0800 Subject: [PATCH 1/5] fix: Huber energy loss uses raw model energy instead of energy_pred that includes atom_ener_coeff weighting, so coefficients are silently dropped when use_huber=True --- deepmd/pt/loss/ener.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py index 1778806343..5626bf8300 100644 --- a/deepmd/pt/loss/ener.py +++ b/deepmd/pt/loss/ener.py @@ -224,8 +224,8 @@ def forward( loss += atom_norm * (pref_e * l2_ener_loss) else: l_huber_loss = custom_huber_loss( - atom_norm * model_pred["energy"], - atom_norm * label["energy"], + atom_norm * energy_pred, + atom_norm * energy_label, delta=self.huber_delta, ) loss += pref_e * l_huber_loss From f08533282752d28432dcc75fa1647d4986e0ad0f Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 20 Jan 2026 22:42:53 +0800 Subject: [PATCH 2/5] fix also pd --- deepmd/pd/loss/ener.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/pd/loss/ener.py b/deepmd/pd/loss/ener.py index fbd806b26d..81d2589844 100644 --- a/deepmd/pd/loss/ener.py +++ b/deepmd/pd/loss/ener.py @@ -211,8 +211,8 @@ def forward(self, input_dict, model, label, natoms, learning_rate, mae=False): loss += atom_norm * (pref_e * l2_ener_loss) else: l_huber_loss = custom_huber_loss( - atom_norm * model_pred["energy"], - atom_norm * label["energy"], + atom_norm * energy_pred, + atom_norm * energy_label, delta=self.huber_delta, ) loss += pref_e * l_huber_loss From 750cbb24c97992dfe83b452920b653ca9634f9c4 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 20 Jan 2026 22:53:02 +0800 Subject: [PATCH 3/5] add ut --- source/tests/consistent/loss/test_ener.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/source/tests/consistent/loss/test_ener.py b/source/tests/consistent/loss/test_ener.py index 681af5608e..47d9b8866f 100644 --- a/source/tests/consistent/loss/test_ener.py +++ b/source/tests/consistent/loss/test_ener.py @@ -57,12 +57,15 @@ @parameterized( - (False, True), # use_huber + (False, False), # huber, enable_atom_ener_coeff + (True, False), + (False, True), + (True, True), ) class TestEner(CommonTest, LossTest, unittest.TestCase): @property def data(self) -> dict: - (use_huber,) = self.param + (use_huber, enable_atom_ener_coeff) = self.param return { "start_pref_e": 0.02, "limit_pref_e": 1.0, @@ -75,6 +78,7 @@ def data(self) -> dict: "start_pref_pf": 1.0 if not use_huber else 0.0, "limit_pref_pf": 1.0 if not use_huber else 0.0, "use_huber": use_huber, + "enable_atom_ener_coeff": enable_atom_ener_coeff, } skip_tf = CommonTest.skip_tf @@ -124,11 +128,13 @@ def setUp(self) -> None: self.natoms, ) ), + "atom_ener_coeff": rng.random((self.nframes, self.natoms)), "atom_pref": np.ones((self.nframes, self.natoms, 3)), "find_energy": 1.0, "find_force": 1.0, "find_virial": 1.0, "find_atom_ener": 1.0, + "find_atom_ener_coeff": 1.0, "find_atom_pref": 1.0, } From 2f98a70da5bb4882bcbe812b3fc13112a8b46edb Mon Sep 17 00:00:00 2001 From: Han Wang Date: Tue, 20 Jan 2026 23:42:07 +0800 Subject: [PATCH 4/5] fix issue in ut --- source/tests/consistent/loss/test_ener.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/source/tests/consistent/loss/test_ener.py b/source/tests/consistent/loss/test_ener.py index 47d9b8866f..36a8fba44a 100644 --- a/source/tests/consistent/loss/test_ener.py +++ b/source/tests/consistent/loss/test_ener.py @@ -57,10 +57,8 @@ @parameterized( - (False, False), # huber, enable_atom_ener_coeff - (True, False), - (False, True), - (True, True), + (False, True), # huber + (False, True), # enable_atom_ener_coeff ) class TestEner(CommonTest, LossTest, unittest.TestCase): @property From 0fea7664bce0ed7b59d024d5aefa314e7871dac9 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Wed, 21 Jan 2026 07:59:43 +0800 Subject: [PATCH 5/5] fix array api issue in ut --- deepmd/dpmodel/loss/ener.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py index 9441c97e32..3bf9695852 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/deepmd/dpmodel/loss/ener.py @@ -131,8 +131,8 @@ def call( # E = - E(A) - E(B) + E(C) + E(D) # A, B, C, D could be put far away from each other atom_ener_coeff = label_dict["atom_ener_coeff"] - atom_ener_coeff = xp.reshape(atom_ener_coeff, xp.shape(atom_ener)) - energy = xp.sum(atom_ener_coeff * atom_ener, 1) + atom_ener_coeff = xp.reshape(atom_ener_coeff, atom_ener.shape) + energy = xp.sum(atom_ener_coeff * atom_ener, axis=1) if self.has_f or self.has_pf or self.relative_f or self.has_gf: force_reshape = xp.reshape(force, (-1,)) force_hat_reshape = xp.reshape(force_hat, (-1,))