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/6] 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/6] 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/6] 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/6] 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/6] 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) From a4be6080f186ccd13ac3c0155236fd75c9f2303b Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Sun, 17 May 2026 21:26:17 +0800 Subject: [PATCH 6/6] update doc --- doc/model/train-energy.md | 26 ++++++++++++++++++++++++++ doc/model/train-se-a-mask.md | 16 ---------------- 2 files changed, 26 insertions(+), 16 deletions(-) diff --git a/doc/model/train-energy.md b/doc/model/train-energy.md index d936ce9e16..3c053ac535 100644 --- a/doc/model/train-energy.md +++ b/doc/model/train-energy.md @@ -150,4 +150,30 @@ The {ref}`intensive_ener_virial ` option (defa If one does not want to train with virial, then he/she may set the virial prefactors {ref}`start_pref_v ` and {ref}`limit_pref_v ` to 0. +### Prefactor force loss with default atom preference + +:::{note} +**Supported backends**: PyTorch {{ pytorch_icon }}, DP {{ dpmodel_icon }} +::: + +When using the prefactor force loss (controlled by {ref}`start_pref_pf ` and {ref}`limit_pref_pf `), the training data typically requires an `atom_pref.npy` file in each system directory to specify per-atom prefactors $q_k$. If `atom_pref.npy` is not provided, the {ref}`use_default_pf ` option can be set to `true` to use a default atom preference of 1.0 for all atoms: + +```json + "loss" : { + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0, + "start_pref_pf": 1.0, + "limit_pref_pf": 1.0, + "use_default_pf": true + } +``` + +This allows using the prefactor force loss without requiring `atom_pref.npy` files. When `atom_pref.npy` is provided in the training data, it will be used as-is regardless of the `use_default_pf` setting. + +Note that `use_default_pf` is only effective for the PyTorch and DP (NumPy reference) backends. The TensorFlow and Paddle backends raise `NotImplementedError` when `use_default_pf` is set to `true`. + [^1]: This section is built upon Jinzhe Zeng, Duo Zhang, Denghui Lu, Pinghui Mo, Zeyu Li, Yixiao Chen, Marián Rynik, Li'ang Huang, Ziyao Li, Shaochen Shi, Yingze Wang, Haotian Ye, Ping Tuo, Jiabin Yang, Ye Ding, Yifan Li, Davide Tisi, Qiyu Zeng, Han Bao, Yu Xia, Jiameng Huang, Koki Muraoka, Yibo Wang, Junhan Chang, Fengbo Yuan, Sigbjørn Løland Bore, Chun Cai, Yinnian Lin, Bo Wang, Jiayan Xu, Jia-Xin Zhu, Chenxing Luo, Yuzhi Zhang, Rhys E. A. Goodall, Wenshuo Liang, Anurag Kumar Singh, Sikai Yao, Jingchao Zhang, Renata Wentzcovitch, Jiequn Han, Jie Liu, Weile Jia, Darrin M. York, Weinan E, Roberto Car, Linfeng Zhang, Han Wang, [J. Chem. Phys. 159, 054801 (2023)](https://doi.org/10.1063/5.0155600) licensed under a [Creative Commons Attribution (CC BY) license](http://creativecommons.org/licenses/by/4.0/). diff --git a/doc/model/train-se-a-mask.md b/doc/model/train-se-a-mask.md index 98a70c483d..1356cdd566 100644 --- a/doc/model/train-se-a-mask.md +++ b/doc/model/train-se-a-mask.md @@ -85,22 +85,6 @@ 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).