diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py index 7ea35d5596..7515f19b9a 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/deepmd/dpmodel/loss/ener.py @@ -71,6 +71,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: @@ -124,6 +127,7 @@ def __init__( huber_delta: float | list[float] = 0.01, loss_func: str = "mse", f_use_norm: bool = False, + use_default_pf: bool = False, intensive_ener_virial: bool = False, **kwargs: Any, ) -> None: @@ -164,6 +168,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 self.intensive_ener_virial = intensive_ener_virial if self.f_use_norm and not (self.use_huber or self.loss_func == "mae"): raise RuntimeError( @@ -203,7 +208,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, @@ -504,6 +511,7 @@ def label_requirement(self) -> list[DataRequirementItem]: must=False, high_prec=False, repeat=3, + default=1.0, ) ) if self.has_gf > 0: @@ -539,7 +547,7 @@ def serialize(self) -> dict: """ return { "@class": "EnergyLoss", - "@version": 3, + "@version": 4, "starter_learning_rate": self.starter_learning_rate, "start_pref_e": self.start_pref_e, "limit_pref_e": self.limit_pref_e, @@ -560,6 +568,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, "intensive_ener_virial": self.intensive_ener_virial, } @@ -579,7 +588,7 @@ def deserialize(cls, data: dict) -> "Loss": """ data = data.copy() version = data.pop("@version") - check_version_compatibility(version, 3, 1) + check_version_compatibility(version, 4, 1) data.pop("@class") # Backward compatibility: version 1-2 used legacy normalization if version < 3: diff --git a/deepmd/pd/loss/ener.py b/deepmd/pd/loss/ener.py index d4ea5ce587..4ab9d71ab9 100644 --- a/deepmd/pd/loss/ener.py +++ b/deepmd/pd/loss/ener.py @@ -138,6 +138,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 @@ -577,7 +581,7 @@ def serialize(self) -> dict: """ return { "@class": "EnergyLoss", - "@version": 3, + "@version": 4, "starter_learning_rate": self.starter_learning_rate, "start_pref_e": self.start_pref_e, "limit_pref_e": self.limit_pref_e, @@ -598,6 +602,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), "intensive_ener_virial": self.intensive_ener_virial, } @@ -617,7 +622,7 @@ def deserialize(cls, data: dict) -> "TaskLoss": """ data = data.copy() version = data.pop("@version") - check_version_compatibility(version, 3, 1) + check_version_compatibility(version, 4, 1) data.pop("@class") # Handle backward compatibility for older versions without intensive_ener_virial if version < 3: diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py index f1dffee218..50d83a4ac9 100644 --- a/deepmd/pt/loss/ener.py +++ b/deepmd/pt/loss/ener.py @@ -59,6 +59,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 | list[float] = 0.01, intensive_ener_virial: bool = False, @@ -107,6 +108,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: @@ -161,6 +165,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 @@ -381,7 +386,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) @@ -538,7 +545,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", @@ -577,6 +584,7 @@ def label_requirement(self) -> list[DataRequirementItem]: must=False, high_prec=False, repeat=3, + default=1.0, ) ) if self.has_gf > 0: @@ -612,7 +620,7 @@ def serialize(self) -> dict: """ return { "@class": "EnergyLoss", - "@version": 3, + "@version": 4, "starter_learning_rate": self.starter_learning_rate, "start_pref_e": self.start_pref_e, "limit_pref_e": self.limit_pref_e, @@ -633,6 +641,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, "intensive_ener_virial": self.intensive_ener_virial, } @@ -652,7 +661,7 @@ def deserialize(cls, data: dict) -> "TaskLoss": """ data = data.copy() version = data.pop("@version") - check_version_compatibility(version, 3, 1) + check_version_compatibility(version, 4, 1) data.pop("@class") # Handle backward compatibility for older versions without intensive_ener_virial if version < 3: diff --git a/deepmd/tf/loss/ener.py b/deepmd/tf/loss/ener.py index da7e5d4462..fbaee24fb2 100644 --- a/deepmd/tf/loss/ener.py +++ b/deepmd/tf/loss/ener.py @@ -147,6 +147,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 @@ -557,7 +561,7 @@ def serialize(self, suffix: str = "") -> dict: """ return { "@class": "EnergyLoss", - "@version": 3, + "@version": 4, "starter_learning_rate": self.starter_learning_rate, "start_pref_e": self.start_pref_e, "limit_pref_e": self.limit_pref_e, @@ -578,6 +582,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), "intensive_ener_virial": self.intensive_ener_virial, } @@ -599,7 +604,7 @@ def deserialize(cls, data: dict, suffix: str = "") -> "Loss": """ data = data.copy() version = data.pop("@version") - check_version_compatibility(version, 3, 1) + check_version_compatibility(version, 4, 1) data.pop("@class") # Handle backward compatibility for older versions without intensive_ener_virial if version < 3: diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index a4291a2d31..67226a51de 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -3217,6 +3217,13 @@ 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. " + "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") doc_numb_generalized_coord = "The dimension of generalized coordinates. Required when generalized force loss is used." @@ -3339,6 +3346,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-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/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 new file mode 100644 index 0000000000..c7b837123d --- /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.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"])) + + 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"], 4) + + loss_fn2 = EnergyStdLoss.deserialize(data) + self.assertTrue(loss_fn2.use_default_pf) + + +if __name__ == "__main__": + unittest.main()