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/pd/loss/ener.py b/deepmd/pd/loss/ener.py index cf093b90d4..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 @@ -592,7 +597,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/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/tf/loss/ener.py b/deepmd/tf/loss/ener.py index 91607245a2..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 @@ -571,7 +576,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) 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()