From e5f7ce5f5ceb9a6613648a2d4de2d9fb333b1967 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Mon, 30 Mar 2026 20:02:53 +0800 Subject: [PATCH 1/5] feat(pt/dpmodel): add use_default_pf --- deepmd/dpmodel/loss/ener.py | 15 +- deepmd/pt/loss/ener.py | 17 +- deepmd/utils/argcheck.py | 12 + doc/model/train-se-a-mask.md | 16 ++ source/tests/pt/test_loss_default_pf.py | 308 ++++++++++++++++++++++++ 5 files changed, 361 insertions(+), 7 deletions(-) create mode 100644 source/tests/pt/test_loss_default_pf.py diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py index 1a99183a79..8c88de59a0 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/deepmd/dpmodel/loss/ener.py @@ -68,6 +68,9 @@ class EnergyLoss(Loss): The prefactor of generalized force loss at the end of the training. numb_generalized_coord : int The dimension of generalized coordinates. + 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. use_huber : bool Enables Huber loss calculation for energy/force/virial terms with user-defined threshold delta (D). The loss function smoothly transitions between L2 and L1 loss: @@ -110,6 +113,7 @@ def __init__( huber_delta: float = 0.01, loss_func: str = "mse", f_use_norm: bool = False, + use_default_pf: bool = False, **kwargs: Any, ) -> None: # Validate loss_func @@ -149,6 +153,7 @@ def __init__( self.use_huber = use_huber self.huber_delta = huber_delta self.f_use_norm = f_use_norm + self.use_default_pf = use_default_pf if self.f_use_norm and not (self.use_huber or self.loss_func == "mae"): raise RuntimeError( "f_use_norm can only be True when use_huber or loss_func='mae'." @@ -182,7 +187,9 @@ def call( find_force = label_dict["find_force"] find_virial = label_dict["find_virial"] find_atom_ener = label_dict["find_atom_ener"] - find_atom_pref = label_dict["find_atom_pref"] + find_atom_pref = ( + label_dict["find_atom_pref"] if not self.use_default_pf else 1.0 + ) xp = array_api_compat.array_namespace( energy, force, @@ -477,6 +484,7 @@ def label_requirement(self) -> list[DataRequirementItem]: must=False, high_prec=False, repeat=3, + default=1.0, ) ) if self.has_gf > 0: @@ -512,7 +520,7 @@ def serialize(self) -> dict: """ return { "@class": "EnergyLoss", - "@version": 2, + "@version": 3, "starter_learning_rate": self.starter_learning_rate, "start_pref_e": self.start_pref_e, "limit_pref_e": self.limit_pref_e, @@ -533,6 +541,7 @@ def serialize(self) -> dict: "huber_delta": self.huber_delta, "loss_func": self.loss_func, "f_use_norm": self.f_use_norm, + "use_default_pf": self.use_default_pf, } @classmethod @@ -550,6 +559,6 @@ def deserialize(cls, data: dict) -> "Loss": The deserialized loss module """ data = data.copy() - check_version_compatibility(data.pop("@version"), 2, 1) + check_version_compatibility(data.pop("@version"), 3, 1) data.pop("@class") return cls(**data) diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py index 66d60aacec..ae74c14164 100644 --- a/deepmd/pt/loss/ener.py +++ b/deepmd/pt/loss/ener.py @@ -56,6 +56,7 @@ def __init__( loss_func: str = "mse", inference: bool = False, use_huber: bool = False, + use_default_pf: bool = False, f_use_norm: bool = False, huber_delta: float = 0.01, **kwargs: Any, @@ -103,6 +104,9 @@ def __init__( MAE loss is less sensitive to outliers compared to MSE loss. inference : bool If true, it will output all losses found in output, ignoring the pre-factors. + 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. use_huber : bool Enables Huber loss calculation for energy/force/virial terms with user-defined threshold delta (D). The loss function smoothly transitions between L2 and L1 loss: @@ -147,6 +151,7 @@ def __init__( self.limit_pref_pf = limit_pref_pf self.start_pref_gf = start_pref_gf self.limit_pref_gf = limit_pref_gf + self.use_default_pf = use_default_pf self.relative_f = relative_f self.enable_atom_ener_coeff = enable_atom_ener_coeff self.numb_generalized_coord = numb_generalized_coord @@ -357,7 +362,9 @@ def forward( if self.has_pf and "atom_pref" in label: atom_pref = label["atom_pref"] - find_atom_pref = label.get("find_atom_pref", 0.0) + find_atom_pref = ( + label.get("find_atom_pref", 0.0) if not self.use_default_pf else 1.0 + ) pref_pf = pref_pf * find_atom_pref atom_pref_reshape = atom_pref.reshape(-1) @@ -514,7 +521,7 @@ def label_requirement(self) -> list[DataRequirementItem]: high_prec=True, ) ) - if self.has_f: + if self.has_f or self.has_pf or self.relative_f is not None or self.has_gf: label_requirement.append( DataRequirementItem( "force", @@ -553,6 +560,7 @@ def label_requirement(self) -> list[DataRequirementItem]: must=False, high_prec=False, repeat=3, + default=1.0, ) ) if self.has_gf > 0: @@ -588,7 +596,7 @@ def serialize(self) -> dict: """ return { "@class": "EnergyLoss", - "@version": 2, + "@version": 3, "starter_learning_rate": self.starter_learning_rate, "start_pref_e": self.start_pref_e, "limit_pref_e": self.limit_pref_e, @@ -609,6 +617,7 @@ def serialize(self) -> dict: "huber_delta": self.huber_delta, "loss_func": self.loss_func, "f_use_norm": self.f_use_norm, + "use_default_pf": self.use_default_pf, } @classmethod @@ -626,7 +635,7 @@ def deserialize(cls, data: dict) -> "TaskLoss": The deserialized loss module """ data = data.copy() - check_version_compatibility(data.pop("@version"), 2, 1) + check_version_compatibility(data.pop("@version"), 3, 1) data.pop("@class") return cls(**data) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 2aca936e6c..0cdf4580ab 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -3189,6 +3189,11 @@ def loss_ener() -> list[Argument]: "atomic prefactor force", label="atom_pref", abbr="pf" ) doc_limit_pref_pf = limit_pref("atomic prefactor force") + doc_use_default_pf = ( + "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 in training data. " + "When atom_pref.npy is provided, it will be used as-is regardless of this setting." + ) doc_start_pref_gf = start_pref("generalized force", label="drdq", abbr="gf") doc_limit_pref_gf = limit_pref("generalized force") doc_numb_generalized_coord = "The dimension of generalized coordinates. Required when generalized force loss is used." @@ -3299,6 +3304,13 @@ def loss_ener() -> list[Argument]: default=0.00, doc=doc_limit_pref_pf, ), + Argument( + "use_default_pf", + bool, + optional=True, + default=False, + doc=doc_use_default_pf, + ), Argument("relative_f", [float, None], optional=True, doc=doc_relative_f), Argument( "enable_atom_ener_coeff", diff --git a/doc/model/train-se-a-mask.md b/doc/model/train-se-a-mask.md index 1356cdd566..98a70c483d 100644 --- a/doc/model/train-se-a-mask.md +++ b/doc/model/train-se-a-mask.md @@ -85,6 +85,22 @@ And the `loss` section in the training input script should be set as follows. } ``` +If `atom_pref.npy` is not provided in the training data, one can set `use_default_pf` to `true` to use a default atom preference of 1.0 for all atoms. This allows using the prefactor force loss (`pf` loss) without requiring `atom_pref.npy` files. When `atom_pref.npy` is provided, it will be used as-is regardless of this setting. + +```json +"loss": { + "type": "ener", + "start_pref_e": 0.0, + "limit_pref_e": 0.0, + "start_pref_f": 0.0, + "limit_pref_f": 0.0, + "start_pref_pf": 1.0, + "limit_pref_pf": 1.0, + "use_default_pf": true, + "_comment": " that's all" + } +``` + ## Type embedding Same as [`se_e2_a`](./train-se-e2-a.md). diff --git a/source/tests/pt/test_loss_default_pf.py b/source/tests/pt/test_loss_default_pf.py new file mode 100644 index 0000000000..0a1b47ad78 --- /dev/null +++ b/source/tests/pt/test_loss_default_pf.py @@ -0,0 +1,308 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for the use_default_pf feature in EnergyStdLoss (PT-only, no TF dependency).""" + +import unittest +from pathlib import ( + Path, +) + +import numpy as np +import torch + +from deepmd.pt.loss import ( + EnergyStdLoss, +) +from deepmd.pt.utils import ( + dp_random, + env, +) +from deepmd.pt.utils.dataset import ( + DeepmdDataSetForLoader, +) +from deepmd.utils.data import ( + DataRequirementItem, +) + +from ..seed import ( + GLOBAL_SEED, +) + +energy_data_requirement = [ + DataRequirementItem( + "energy", + ndof=1, + atomic=False, + must=False, + high_prec=True, + ), + DataRequirementItem( + "force", + ndof=3, + atomic=True, + must=False, + high_prec=False, + ), + DataRequirementItem( + "virial", + ndof=9, + atomic=False, + must=False, + high_prec=False, + ), +] + + +def get_single_batch(dataset, index=None): + if index is None: + index = dp_random.choice(np.arange(len(dataset))) + np_batch = dataset[index] + pt_batch = {} + for key in ["coord", "box", "force", "energy", "virial", "atype", "natoms"]: + if key in np_batch.keys(): + np_batch[key] = np.expand_dims(np_batch[key], axis=0) + pt_batch[key] = torch.as_tensor(np_batch[key], device=env.DEVICE) + if key in ["coord", "force"]: + np_batch[key] = np_batch[key].reshape(1, -1) + return np_batch, pt_batch + + +def get_batch(system, type_map, data_requirement): + dataset = DeepmdDataSetForLoader(system, type_map) + dataset.add_data_requirement(data_requirement) + np_batch, pt_batch = get_single_batch(dataset) + return np_batch, pt_batch + + +class TestEnerStdLossDefaultPf(unittest.TestCase): + """Test use_default_pf feature in EnergyStdLoss.""" + + def setUp(self) -> None: + self.start_lr = 1.1 + self.cur_lr = 1.2 + self.start_pref_e = 0.02 + self.limit_pref_e = 1.0 + self.start_pref_f = 0.0 + self.limit_pref_f = 0.0 + self.start_pref_v = 0.0 + self.limit_pref_v = 0.0 + self.start_pref_pf = 1.0 + self.limit_pref_pf = 1.0 + + self.system = str(Path(__file__).parent / "water/data/data_0") + self.type_map = ["H", "O"] + + np_batch, pt_batch = get_batch( + self.system, self.type_map, energy_data_requirement + ) + natoms = np_batch["natoms"] + self.nloc = int(natoms[0][0]) + rng = np.random.default_rng(GLOBAL_SEED) + + l_energy, l_force, l_virial = ( + np_batch["energy"], + np_batch["force"], + np_batch["virial"], + ) + p_energy, p_force, p_virial = ( + np.ones_like(l_energy), + np.ones_like(l_force), + np.ones_like(l_virial), + ) + nloc = self.nloc + batch_size = pt_batch["coord"].shape[0] + p_atom_energy = rng.random(size=[batch_size, nloc]) + atom_pref = np.ones([batch_size, nloc * 3]) + + self.model_pred = { + "energy": torch.from_numpy(p_energy), + "force": torch.from_numpy(p_force), + "virial": torch.from_numpy(p_virial), + "atom_energy": torch.from_numpy(p_atom_energy), + } + # label WITH find_atom_pref (simulates data with atom_pref.npy) + self.label_with_pref = { + "energy": torch.from_numpy(l_energy), + "find_energy": 1.0, + "force": torch.from_numpy(l_force), + "find_force": 1.0, + "virial": torch.from_numpy(l_virial), + "find_virial": 0.0, + "atom_pref": torch.from_numpy(atom_pref), + "find_atom_pref": 1.0, + } + # label WITHOUT find_atom_pref (simulates data without atom_pref.npy) + self.label_without_pref = { + "energy": torch.from_numpy(l_energy), + "find_energy": 1.0, + "force": torch.from_numpy(l_force), + "find_force": 1.0, + "virial": torch.from_numpy(l_virial), + "find_virial": 0.0, + "atom_pref": torch.from_numpy(atom_pref), + "find_atom_pref": 0.0, + } + self.natoms = pt_batch["natoms"] + + def test_default_pf_enabled(self) -> None: + """With use_default_pf=True, pf loss should be computed even without find_atom_pref.""" + loss_fn = EnergyStdLoss( + self.start_lr, + self.start_pref_e, + self.limit_pref_e, + self.start_pref_f, + self.limit_pref_f, + self.start_pref_v, + self.limit_pref_v, + start_pref_pf=self.start_pref_pf, + limit_pref_pf=self.limit_pref_pf, + use_default_pf=True, + ) + + def fake_model(): + return self.model_pred + + # With find_atom_pref=0.0 but use_default_pf=True, pf loss should still be computed + _, pt_loss, pt_more_loss = loss_fn( + {}, + fake_model, + self.label_without_pref, + self.nloc, + self.cur_lr, + ) + pt_loss_val = pt_loss.detach().cpu().numpy() + # loss should be non-zero because pf loss is activated via use_default_pf + self.assertTrue(pt_loss_val != 0.0) + self.assertIn("rmse_pf", pt_more_loss) + # The pref_force_loss should be a valid number (not NaN) + self.assertFalse(np.isnan(pt_more_loss["l2_pref_force_loss"])) + + def test_default_pf_disabled(self) -> None: + """With use_default_pf=False (default), pf loss should NOT be computed without find_atom_pref.""" + loss_fn = EnergyStdLoss( + self.start_lr, + self.start_pref_e, + self.limit_pref_e, + self.start_pref_f, + self.limit_pref_f, + self.start_pref_v, + self.limit_pref_v, + start_pref_pf=self.start_pref_pf, + limit_pref_pf=self.limit_pref_pf, + use_default_pf=False, + ) + + def fake_model(): + return self.model_pred + + # With find_atom_pref=0.0 and use_default_pf=False, pf loss contribution is zero + _, pt_loss_without, pt_more_loss_without = loss_fn( + {}, + fake_model, + self.label_without_pref, + self.nloc, + self.cur_lr, + ) + # With find_atom_pref=1.0, pf loss should be computed + _, pt_loss_with, pt_more_loss_with = loss_fn( + {}, + fake_model, + self.label_with_pref, + self.nloc, + self.cur_lr, + ) + # without find_atom_pref, the pf part contributes nothing + self.assertTrue(np.isnan(pt_more_loss_without["l2_pref_force_loss"])) + # with find_atom_pref, pf loss should be computed + self.assertFalse(np.isnan(pt_more_loss_with["l2_pref_force_loss"])) + + def test_default_pf_consistency(self) -> None: + """With use_default_pf=True and atom_pref=1.0, results should match explicit find_atom_pref=1.0.""" + loss_fn_default = EnergyStdLoss( + self.start_lr, + self.start_pref_e, + self.limit_pref_e, + self.start_pref_f, + self.limit_pref_f, + self.start_pref_v, + self.limit_pref_v, + start_pref_pf=self.start_pref_pf, + limit_pref_pf=self.limit_pref_pf, + use_default_pf=True, + ) + + def fake_model(): + return self.model_pred + + # use_default_pf=True with find_atom_pref=0.0 + _, pt_loss_default, _ = loss_fn_default( + {}, + fake_model, + self.label_without_pref, + self.nloc, + self.cur_lr, + ) + # use_default_pf=True with find_atom_pref=1.0 (should also give same result) + _, pt_loss_explicit, _ = loss_fn_default( + {}, + fake_model, + self.label_with_pref, + self.nloc, + self.cur_lr, + ) + # Both should be the same since use_default_pf overrides find_atom_pref + self.assertTrue( + np.allclose( + pt_loss_default.detach().cpu().numpy(), + pt_loss_explicit.detach().cpu().numpy(), + ) + ) + + def test_label_requirement_force_included(self) -> None: + """When has_pf=True but has_f=False, force should still be in label_requirement.""" + loss_fn = EnergyStdLoss( + self.start_lr, + 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_pf=self.start_pref_pf, + limit_pref_pf=self.limit_pref_pf, + use_default_pf=True, + ) + label_req = loss_fn.label_requirement + keys = [r.key for r in label_req] + self.assertIn("force", keys) + self.assertIn("atom_pref", keys) + + def test_label_requirement_atom_pref_default(self) -> None: + """atom_pref DataRequirementItem should have default=1.0.""" + loss_fn = EnergyStdLoss( + self.start_lr, + start_pref_pf=self.start_pref_pf, + limit_pref_pf=self.limit_pref_pf, + use_default_pf=True, + ) + label_req = loss_fn.label_requirement + atom_pref_req = next(r for r in label_req if r.key == "atom_pref") + self.assertEqual(atom_pref_req.default, 1.0) + + def test_serialize_deserialize(self) -> None: + """Serialization round-trip should preserve use_default_pf.""" + loss_fn = EnergyStdLoss( + self.start_lr, + start_pref_pf=self.start_pref_pf, + limit_pref_pf=self.limit_pref_pf, + use_default_pf=True, + ) + data = loss_fn.serialize() + self.assertTrue(data["use_default_pf"]) + self.assertEqual(data["@version"], 3) + + loss_fn2 = EnergyStdLoss.deserialize(data) + self.assertTrue(loss_fn2.use_default_pf) + + +if __name__ == "__main__": + unittest.main() From 28e6fdd91f33c2bea8672f4fef2f3321b263b564 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Tue, 31 Mar 2026 01:29:29 +0800 Subject: [PATCH 2/5] fix ut --- deepmd/pd/loss/ener.py | 2 +- deepmd/tf/loss/ener.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/pd/loss/ener.py b/deepmd/pd/loss/ener.py index cf093b90d4..738e03bcb2 100644 --- a/deepmd/pd/loss/ener.py +++ b/deepmd/pd/loss/ener.py @@ -592,7 +592,7 @@ def deserialize(cls, data: dict) -> "TaskLoss": The deserialized loss module """ data = data.copy() - check_version_compatibility(data.pop("@version"), 2, 1) + check_version_compatibility(data.pop("@version"), 3, 1) data.pop("@class") return cls(**data) diff --git a/deepmd/tf/loss/ener.py b/deepmd/tf/loss/ener.py index 91607245a2..89f787cedb 100644 --- a/deepmd/tf/loss/ener.py +++ b/deepmd/tf/loss/ener.py @@ -571,7 +571,7 @@ def deserialize(cls, data: dict, suffix: str = "") -> "Loss": The deserialized loss module """ data = data.copy() - check_version_compatibility(data.pop("@version"), 2, 1) + check_version_compatibility(data.pop("@version"), 3, 1) data.pop("@class") return cls(**data) From 7d8a3074f5ae6a4a54e6d38ba01a73d5343e4986 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Tue, 31 Mar 2026 15:41:50 +0800 Subject: [PATCH 3/5] fix ut --- deepmd/pd/loss/ener.py | 7 ++++++- deepmd/tf/loss/ener.py | 7 ++++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/deepmd/pd/loss/ener.py b/deepmd/pd/loss/ener.py index 738e03bcb2..d646b2fbd6 100644 --- a/deepmd/pd/loss/ener.py +++ b/deepmd/pd/loss/ener.py @@ -125,6 +125,10 @@ def __init__( raise NotImplementedError( "Paddle backend does not support f_use_norm=True." ) + if kwargs.get("use_default_pf", False): + raise NotImplementedError( + "Paddle backend does not support use_default_pf=True." + ) self.starter_learning_rate = starter_learning_rate self.has_e = (start_pref_e != 0.0 and limit_pref_e != 0.0) or inference @@ -554,7 +558,7 @@ def serialize(self) -> dict: """ return { "@class": "EnergyLoss", - "@version": 2, + "@version": 3, "starter_learning_rate": self.starter_learning_rate, "start_pref_e": self.start_pref_e, "limit_pref_e": self.limit_pref_e, @@ -575,6 +579,7 @@ def serialize(self) -> dict: "huber_delta": self.huber_delta, "loss_func": self.loss_func, "f_use_norm": self.f_use_norm, + "use_default_pf": getattr(self, "use_default_pf", False), } @classmethod diff --git a/deepmd/tf/loss/ener.py b/deepmd/tf/loss/ener.py index 89f787cedb..3308c5ae50 100644 --- a/deepmd/tf/loss/ener.py +++ b/deepmd/tf/loss/ener.py @@ -133,6 +133,10 @@ def __init__( raise NotImplementedError( "TensorFlow backend does not support f_use_norm=True." ) + if kwargs.get("use_default_pf", False): + raise NotImplementedError( + "TensorFlow backend does not support use_default_pf=True." + ) self.starter_learning_rate = starter_learning_rate self.start_pref_e = start_pref_e @@ -531,7 +535,7 @@ def serialize(self, suffix: str = "") -> dict: """ return { "@class": "EnergyLoss", - "@version": 2, + "@version": 3, "starter_learning_rate": self.starter_learning_rate, "start_pref_e": self.start_pref_e, "limit_pref_e": self.limit_pref_e, @@ -552,6 +556,7 @@ def serialize(self, suffix: str = "") -> dict: "huber_delta": self.huber_delta, "loss_func": self.loss_func, "f_use_norm": self.f_use_norm, + "use_default_pf": getattr(self, "use_default_pf", False), } @classmethod From 604fff9c7c03e29ba2ac6548a49e83602e9b4a9a Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Sun, 10 May 2026 21:54:48 +0800 Subject: [PATCH 4/5] fix comment --- deepmd/utils/argcheck.py | 4 +- source/tests/consistent/loss/test_ener.py | 178 ++++++++++++++++++++++ source/tests/pt/test_loss_default_pf.py | 6 +- 3 files changed, 184 insertions(+), 4 deletions(-) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 2c9b7b8d50..67226a51de 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -3220,7 +3220,9 @@ def loss_ener() -> list[Argument]: doc_use_default_pf = ( "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 in training data. " - "When atom_pref.npy is provided, it will be used as-is regardless of this setting." + "When atom_pref.npy is provided, it will be used as-is regardless of this setting. " + "Note: this option is only effective for the PyTorch/DPModel backends; " + "the TensorFlow and Paddle backends raise NotImplementedError when set to true." ) doc_start_pref_gf = start_pref("generalized force", label="drdq", abbr="gf") doc_limit_pref_gf = limit_pref("generalized force") diff --git a/source/tests/consistent/loss/test_ener.py b/source/tests/consistent/loss/test_ener.py index 08229606e2..c87e7409cb 100644 --- a/source/tests/consistent/loss/test_ener.py +++ b/source/tests/consistent/loss/test_ener.py @@ -812,3 +812,181 @@ def test_intensive_vs_legacy_scaling_difference(self) -> None: places=5, msg=f"Expected intensive/legacy ratio ~{expected_ratio:.6f}, got {actual_ratio:.6f}", ) + + +class TestEnerDefaultPf(CommonTest, LossTest, unittest.TestCase): + """Test energy loss with use_default_pf=True. + + The pf term is activated through the default atom_pref of 1.0 even though + `find_atom_pref` is 0.0 in the label. This exercises the cross-backend + consistency between PT and DP for the new option. TF and Paddle backends + raise NotImplementedError when use_default_pf=True and are skipped. + """ + + @property + def data(self) -> dict: + return { + "start_pref_e": 0.02, + "limit_pref_e": 1.0, + "start_pref_f": 1000.0, + "limit_pref_f": 1.0, + "start_pref_v": 1.0, + "limit_pref_v": 1.0, + "start_pref_ae": 1.0, + "limit_pref_ae": 1.0, + "start_pref_pf": 1.0, + "limit_pref_pf": 1.0, + "use_default_pf": True, + } + + skip_tf = True + skip_pd = True + skip_pt = CommonTest.skip_pt + skip_pt_expt = not INSTALLED_PT_EXPT + skip_jax = not INSTALLED_JAX + skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT + + tf_class = EnerLossTF + dp_class = EnerLossDP + pt_class = EnerLossPT + pt_expt_class = EnerLossPTExpt + jax_class = EnerLossDP + pd_class = EnerLossPD + array_api_strict_class = EnerLossDP + args = loss_ener() + + def setUp(self) -> None: + CommonTest.setUp(self) + self.learning_rate = 1e-3 + rng = np.random.default_rng(20250105) + self.nframes = 2 + self.natoms = 6 + self.predict = { + "energy": rng.random((self.nframes,)), + "force": rng.random((self.nframes, self.natoms, 3)), + "virial": rng.random((self.nframes, 9)), + "atom_ener": rng.random((self.nframes, self.natoms)), + } + self.predict_dpmodel_style = { + "energy": self.predict["energy"], + "force": self.predict["force"], + "virial": self.predict["virial"], + "atom_energy": self.predict["atom_ener"], + } + # find_atom_pref=0.0 simulates the case where atom_pref.npy is missing; + # use_default_pf=True must override this and still compute the pf loss. + self.label = { + "energy": rng.random((self.nframes,)), + "force": rng.random((self.nframes, self.natoms, 3)), + "virial": rng.random((self.nframes, 9)), + "atom_ener": 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_pref": 0.0, + } + + @property + def additional_data(self) -> dict: + return { + "starter_learning_rate": 1e-3, + } + + def build_tf(self, obj: Any, suffix: str) -> tuple[list, dict]: + # use_default_pf=True is not supported by TensorFlow; skip_tf is True so + # this method is never invoked, but the abstract base requires it. + raise NotImplementedError + + def eval_pt(self, pt_obj: Any) -> Any: + predict = {kk: numpy_to_torch(vv) for kk, vv in self.predict.items()} + label = {kk: numpy_to_torch(vv) for kk, vv in self.label.items()} + predict["atom_energy"] = predict.pop("atom_ener") + _, loss, more_loss = pt_obj( + {}, + lambda: predict, + label, + self.natoms, + self.learning_rate, + mae=False, + ) + loss = torch_to_numpy(loss) + more_loss = {kk: torch_to_numpy(vv) for kk, vv in more_loss.items()} + return loss, more_loss + + def eval_dp(self, dp_obj: Any) -> Any: + return dp_obj( + self.learning_rate, + self.natoms, + self.predict_dpmodel_style, + self.label, + mae=False, + ) + + def eval_pt_expt(self, pt_expt_obj: Any) -> Any: + predict = { + kk: numpy_to_torch(vv) for kk, vv in self.predict_dpmodel_style.items() + } + label = {kk: numpy_to_torch(vv) for kk, vv in self.label.items()} + loss, more_loss = pt_expt_obj( + self.learning_rate, + self.natoms, + predict, + label, + mae=False, + ) + loss = torch_to_numpy(loss) + more_loss = {kk: torch_to_numpy(vv) for kk, vv in more_loss.items()} + return loss, more_loss + + def eval_jax(self, jax_obj: Any) -> Any: + predict = {kk: jnp.asarray(vv) for kk, vv in self.predict_dpmodel_style.items()} + label = {kk: jnp.asarray(vv) for kk, vv in self.label.items()} + loss, more_loss = jax_obj( + self.learning_rate, + self.natoms, + predict, + label, + mae=False, + ) + loss = to_numpy_array(loss) + more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} + return loss, more_loss + + def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: + predict = { + kk: array_api_strict.asarray(vv) + for kk, vv in self.predict_dpmodel_style.items() + } + label = {kk: array_api_strict.asarray(vv) for kk, vv in self.label.items()} + loss, more_loss = array_api_strict_obj( + self.learning_rate, + self.natoms, + predict, + label, + mae=False, + ) + loss = to_numpy_array(loss) + more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} + return loss, more_loss + + def extract_ret(self, ret: Any, backend) -> dict[str, np.ndarray]: + loss = ret[0] + result = {"loss": np.atleast_1d(np.asarray(loss, dtype=np.float64))} + if len(ret) > 1: + more_loss = ret[1] + for k in sorted(more_loss): + if k.startswith("rmse_") or k.startswith("mae_"): + result[k] = np.atleast_1d( + np.asarray(more_loss[k], dtype=np.float64) + ) + return result + + @property + def rtol(self) -> float: + return 1e-10 + + @property + def atol(self) -> float: + return 1e-10 diff --git a/source/tests/pt/test_loss_default_pf.py b/source/tests/pt/test_loss_default_pf.py index 0a1b47ad78..299524cdbc 100644 --- a/source/tests/pt/test_loss_default_pf.py +++ b/source/tests/pt/test_loss_default_pf.py @@ -171,7 +171,7 @@ def fake_model(): ) pt_loss_val = pt_loss.detach().cpu().numpy() # loss should be non-zero because pf loss is activated via use_default_pf - self.assertTrue(pt_loss_val != 0.0) + self.assertNotEqual(float(pt_loss_val), 0.0) self.assertIn("rmse_pf", pt_more_loss) # The pref_force_loss should be a valid number (not NaN) self.assertFalse(np.isnan(pt_more_loss["l2_pref_force_loss"])) @@ -195,7 +195,7 @@ def fake_model(): return self.model_pred # With find_atom_pref=0.0 and use_default_pf=False, pf loss contribution is zero - _, pt_loss_without, pt_more_loss_without = loss_fn( + _, _pt_loss_without, pt_more_loss_without = loss_fn( {}, fake_model, self.label_without_pref, @@ -203,7 +203,7 @@ def fake_model(): self.cur_lr, ) # With find_atom_pref=1.0, pf loss should be computed - _, pt_loss_with, pt_more_loss_with = loss_fn( + _, _pt_loss_with, pt_more_loss_with = loss_fn( {}, fake_model, self.label_with_pref, From f4fda39a41cc37145fcf09ada2e54af50422598c Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Wed, 13 May 2026 18:25:43 +0800 Subject: [PATCH 5/5] Update test_loss_default_pf.py --- source/tests/pt/test_loss_default_pf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/tests/pt/test_loss_default_pf.py b/source/tests/pt/test_loss_default_pf.py index 299524cdbc..c7b837123d 100644 --- a/source/tests/pt/test_loss_default_pf.py +++ b/source/tests/pt/test_loss_default_pf.py @@ -298,7 +298,7 @@ def test_serialize_deserialize(self) -> None: ) data = loss_fn.serialize() self.assertTrue(data["use_default_pf"]) - self.assertEqual(data["@version"], 3) + self.assertEqual(data["@version"], 4) loss_fn2 = EnergyStdLoss.deserialize(data) self.assertTrue(loss_fn2.use_default_pf)