diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 050b0ee52f..5692b019cd 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -119,13 +119,17 @@ def get_loss( def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: additional_data_requirement: list[DataRequirementItem] = [] if _model.get_dim_fparam() > 0: + has_default_fparam = _model.has_default_fparam() + fparam_default = ( + np.asarray(_model.get_default_fparam()) if has_default_fparam else 0.0 + ) additional_data_requirement.append( DataRequirementItem( "fparam", _model.get_dim_fparam(), atomic=False, - must=False, - default=0.0, + must=not has_default_fparam, + default=fparam_default, ) ) if _model.get_dim_aparam() > 0: diff --git a/source/tests/pt_expt/test_training.py b/source/tests/pt_expt/test_training.py index 16bf20cffa..bb3123b1ed 100644 --- a/source/tests/pt_expt/test_training.py +++ b/source/tests/pt_expt/test_training.py @@ -545,6 +545,74 @@ def test_batch_shapes(self) -> None: shutil.rmtree(tmpdir, ignore_errors=True) +class TestAdditionalDataRequirement(unittest.TestCase): + """Cover both branches of fparam handling in get_additional_data_requirement. + + Regression guard for the bug where pt_expt hard-coded ``default=0.0`` for + fparam, ignoring the model's ``default_fparam``. When the dataset has no + ``fparam.npy``, that caused the data system to feed zeros to the fitting + net while pt fed the model default — producing a constant DC offset + (e.g. ``[0, -100]`` after fparam normalisation) that slowed training. + """ + + def test_has_default_fparam(self) -> None: + """has_default_fparam==True: must=False and default broadcasts the model value.""" + from deepmd.pt_expt.train.training import ( + get_additional_data_requirement, + ) + + class _M: + def get_dim_fparam(self) -> int: + return 2 + + def get_dim_aparam(self) -> int: + return 0 + + def has_default_fparam(self) -> bool: + return True + + def get_default_fparam(self) -> list[float]: + return [0.0, 1.0] + + reqs = get_additional_data_requirement(_M()) + self.assertEqual(len(reqs), 1) + fparam_req = reqs[0] + self.assertEqual(fparam_req.key, "fparam") + self.assertEqual(fparam_req.ndof, 2) + self.assertFalse(fparam_req.must) + # default is the model's default_fparam, not 0.0 + self.assertNotIsInstance(fparam_req.default, float) + import numpy as np + + np.testing.assert_array_equal(np.asarray(fparam_req.default), [0.0, 1.0]) + + def test_no_default_fparam(self) -> None: + """has_default_fparam==False: must=True (data file must exist) and default=0.0.""" + from deepmd.pt_expt.train.training import ( + get_additional_data_requirement, + ) + + class _M: + def get_dim_fparam(self) -> int: + return 2 + + def get_dim_aparam(self) -> int: + return 0 + + def has_default_fparam(self) -> bool: + return False + + def get_default_fparam(self) -> None: + return None + + reqs = get_additional_data_requirement(_M()) + self.assertEqual(len(reqs), 1) + fparam_req = reqs[0] + self.assertEqual(fparam_req.key, "fparam") + self.assertTrue(fparam_req.must) + self.assertEqual(fparam_req.default, 0.0) + + class TestRestart(unittest.TestCase): """Test restart and init_model resume paths."""