feat(tf): implement change-bias command

deepmd/main.py

@@ -721,12 +721,13 @@ def main_parser() -> argparse.ArgumentParser:
     parser_change_bias = subparsers.add_parser(
         "change-bias",
         parents=[parser_log],
-        help="(Supported backend: PyTorch) Change model out bias according to the input data.",
+        help="Change model out bias according to the input data.",
         formatter_class=RawTextArgumentDefaultsHelpFormatter,
         epilog=textwrap.dedent(
             """\
         examples:
             dp change-bias model.pt -s data -n 10 -m change
+            dp --tf change-bias checkpoint_dir -s data -n 10 -m change
         """
         ),
     )

deepmd/tf/entrypoints/main.py

@@ -2,6 +2,7 @@
 """DeePMD-Kit entry point module."""
 
 import argparse
+import logging
 from pathlib import (
     Path,
 )
@@ -13,6 +14,9 @@
 from deepmd.backend.suffix import (
     format_model_suffix,
 )
+from deepmd.common import (
+    expand_sys_str,
+)
 from deepmd.main import (
     get_ll,
     main_parser,
@@ -34,9 +38,184 @@
 from deepmd.tf.nvnmd.entrypoints.train import (
     train_nvnmd,
 )
+from deepmd.utils.data_system import (
+    DeepmdDataSystem,
+)
 
 __all__ = ["get_ll", "main", "main_parser", "parse_args"]
 
+log = logging.getLogger(__name__)
+
+
+def change_bias(
+    input_file: str,
+    mode: str = "change",
+    bias_value: Optional[list] = None,
+    datafile: Optional[str] = None,
+    system: str = ".",
+    numb_batch: int = 0,
+    model_branch: Optional[str] = None,
+    output: Optional[str] = None,
+) -> None:
+    """Change model out bias according to the input data.
+
+    Parameters
+    ----------
+    input_file : str
+        The input checkpoint folder or frozen model file
+    mode : str, optional
+        The mode for changing energy bias, by default "change"
+    bias_value : Optional[list], optional
+        The user defined value for each type, by default None
+    datafile : Optional[str], optional
+        The path to the datafile, by default None
+    system : str, optional
+        The system dir, by default "."
+    numb_batch : int, optional
+        The number of frames for bias changing, by default 0
+    model_branch : Optional[str], optional
+        Model branch chosen for changing bias if multi-task model, by default None
+    output : Optional[str], optional
+        The model after changing bias, by default None
+    """
+    import os
+    from pathlib import (
+        Path,
+    )
+
+    from deepmd.tf.train.trainer import (
+        DPTrainer,
+    )
+    from deepmd.tf.utils.argcheck import (
+        normalize,
+    )
+    from deepmd.tf.utils.compat import (
+        update_deepmd_input,
+    )
+
+    input_path = Path(input_file)
+
+    # Check if input is a checkpoint directory or frozen model
+    if input_path.is_dir():
+        # Checkpoint directory
+        checkpoint_folder = str(input_path)
+        # Check for valid checkpoint early
+        if not (input_path / "checkpoint").exists():
+            raise RuntimeError(f"No valid checkpoint found in {checkpoint_folder}")
+    elif input_file.endswith((".pb", ".pbtxt")):
+        # Frozen model - for now, not supported
+        raise NotImplementedError(
+            "Bias changing for frozen models (.pb/.pbtxt) is not yet implemented. "
+            "Please provide a checkpoint directory instead. "
+            "You can train a model to create checkpoints, then use this command "
+            "to modify the bias, and finally freeze the modified model."
+        )
+    else:
+        raise RuntimeError(
+            "The model provided must be a checkpoint directory or frozen model file (.pb/.pbtxt)"
+        )
+
+    bias_adjust_mode = "change-by-statistic" if mode == "change" else "set-by-statistic"
+
+    if bias_value is not None:
+        raise NotImplementedError(
+            "User-defined bias setting is not yet implemented for TensorFlow models. "
+            "Please use the data-based bias adjustment mode."
+        )
+
+    # Load data systems for bias calculation
+    if datafile is not None:
+        with open(datafile) as datalist:
+            all_sys = datalist.read().splitlines()
+    else:
+        all_sys = expand_sys_str(system)
+
+    # Load the data systems
+    data = DeepmdDataSystem(
+        systems=all_sys,
+        batch_size=1,
+        test_size=1,
+        rcut=None,
+        set_prefix="set",
+    )
+
+    # Read the checkpoint to get the model configuration
+    checkpoint_path = Path(checkpoint_folder)
+
+    # Find the input.json file or create a minimal config
+    # We need this to reconstruct the model
+    input_json_path = checkpoint_path / "input.json"
+    if not input_json_path.exists():
+        # Look for input.json in parent directories or common locations
+        for parent in checkpoint_path.parents:
+            potential_input = parent / "input.json"
+            if potential_input.exists():
+                input_json_path = potential_input
+                break
+        else:
+            raise RuntimeError(
+                f"Cannot find input.json configuration file needed to load the model. "
+                f"Please ensure input.json is available in {checkpoint_folder} or its parent directories."
+            )
+
+    # Load the configuration
+    with open(input_json_path) as f:
+        import json
+
+        jdata = json.load(f)
+
+    # Update and normalize the configuration
+    jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json")
+    jdata = normalize(jdata)
+
+    # Determine output path
+    if output is None:
+        output = str(checkpoint_path) + "_bias_updated"
+
+    # Create trainer to access model methods
+    from deepmd.tf.train.run_options import (
+        RunOptions,
+    )
+
+    run_opt = RunOptions(
+        init_model=checkpoint_folder,
+        restart=None,
+        finetune=None,
+        init_frz_model=None,
+        train_data=all_sys,
+        valid_data=None,
+    )
+
+    trainer = DPTrainer(jdata, run_opt)
+
+    # Get the type map from the model
+    type_map = data.get_type_map()
+    if len(type_map) == 0:
+        # If data doesn't have type_map, get from model
+        type_map = trainer.model.get_type_map()
+
+    log.info(f"Changing bias for model with type_map: {type_map}")
+    log.info(f"Using bias adjustment mode: {bias_adjust_mode}")
+
+    # Use the trainer's change energy bias functionality
+    trainer._change_energy_bias(
+        data,
+        checkpoint_folder,  # Use checkpoint as frozen model path for compatibility
+        type_map,
+        bias_adjust_mode=bias_adjust_mode,
+    )
+
+    # Save the updated model
+    import shutil
+
+    shutil.copytree(checkpoint_folder, output, dirs_exist_ok=True)
+    trainer.save_checkpoint(os.path.join(output, "model.ckpt"))
+
+    log.info(f"Bias changing complete. Updated model saved to {output}")
+    log.info(
+        f"You can now freeze this model using: dp freeze -c {output} -o model_updated.pb"
+    )
+
 
 def main(args: Optional[Union[list[str], argparse.Namespace]] = None) -> None:
     """DeePMD-Kit entry point.
@@ -86,6 +265,17 @@ def main(args: Optional[Union[list[str], argparse.Namespace]] = None) -> None:
         compress(**dict_args)
     elif args.command == "convert-from":
         convert(**dict_args)
+    elif args.command == "change-bias":
+        change_bias(
+            input_file=dict_args["INPUT"],
+            mode=dict_args["mode"],
+            bias_value=dict_args["bias_value"],
+            datafile=dict_args["datafile"],
+            system=dict_args["system"],
+            numb_batch=dict_args["numb_batch"],
+            model_branch=dict_args["model_branch"],
+            output=dict_args["output"],
+        )
     elif args.command == "train-nvnmd":  # nvnmd
         train_nvnmd(**dict_args)
     elif args.command is None:

source/tests/tf/test_change_bias.py

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import tempfile
+import unittest
+from pathlib import (
+    Path,
+)
+
+from deepmd.tf.entrypoints.main import (
+    change_bias,
+)
+
+
+class TestChangeBias(unittest.TestCase):
+    def setUp(self):
+        """Set up test fixtures."""
+        self.temp_dir = tempfile.mkdtemp()
+        self.temp_path = Path(self.temp_dir)
+
+    def tearDown(self):
+        """Clean up test fixtures."""
+        import shutil
+
+        shutil.rmtree(self.temp_dir, ignore_errors=True)
+
+    def test_change_bias_frozen_model_not_implemented(self):
+        """Test that frozen model support raises NotImplementedError."""
+        fake_pb = self.temp_path / "model.pb"
+        fake_pb.write_text("fake model content")
+
+        with self.assertRaises(NotImplementedError) as cm:
+            change_bias(
+                input_file=str(fake_pb),
+                mode="change",
+                system=".",
+            )
+
+        self.assertIn("Bias changing for frozen models", str(cm.exception))
+        self.assertIn(".pb/.pbtxt", str(cm.exception))
+
+    def test_change_bias_invalid_model_type(self):
+        """Test that invalid model types raise RuntimeError."""
+        fake_model = self.temp_path / "model.xyz"
+        fake_model.write_text("fake model content")
+
+        with self.assertRaises(RuntimeError) as cm:
+            change_bias(
+                input_file=str(fake_model),
+                mode="change",
+                system=".",
+            )
+
+        self.assertIn("checkpoint directory or frozen model file", str(cm.exception))
+
+    def test_change_bias_no_checkpoint_in_directory(self):
+        """Test that missing checkpoint in directory raises RuntimeError."""
+        fake_dir = self.temp_path / "fake_checkpoint"
+        fake_dir.mkdir()
+
+        # Create a fake data system for the test
+        fake_data_dir = self.temp_path / "fake_data"
+        fake_data_dir.mkdir()
+        fake_set_dir = fake_data_dir / "set.000"
+        fake_set_dir.mkdir()
+
+        with self.assertRaises(RuntimeError) as cm:
+            change_bias(
+                input_file=str(fake_dir),
+                mode="change",
+                system=str(fake_data_dir),
+            )
+
+        self.assertIn("No valid checkpoint found", str(cm.exception))
+
+    def test_change_bias_user_defined_not_implemented(self):
+        """Test that user-defined bias raises NotImplementedError."""
+        fake_dir = self.temp_path / "fake_checkpoint"
+        fake_dir.mkdir()
+        (fake_dir / "checkpoint").write_text("fake checkpoint")
+
+        with self.assertRaises(NotImplementedError) as cm:
+            change_bias(
+                input_file=str(fake_dir),
+                mode="change",
+                bias_value=[1.0, 2.0],
+                system=".",
+            )
+
+        self.assertIn(
+            "User-defined bias setting is not yet implemented", str(cm.exception)
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()