feat: support using train/valid data from input.json for dp test (#4859)

This pull request extends the testing functionality in DeepMD by
allowing users to specify training and validation data directly via
input JSON files, in addition to existing system and datafile options.
It updates the command-line interface, the main test logic, and adds
comprehensive tests to cover these new features, including support for
recursive glob patterns when selecting systems from JSON files.

### Feature enhancements to testing data sources

* The `test` function in `deepmd/entrypoints/test.py` now accepts
`train_json` and `valid_json` arguments, allowing users to specify
training or validation systems for testing via input JSON files. It
processes these files to extract system paths, including support for
recursive glob patterns. The function also raises an error if no valid
data source is specified.
[[1]](diffhunk://#diff-299c01ed4ee7d0b3f636fe4cb4f0d660a5012b7e95ca0740098b3ace617ab16eL61-R71)
[[2]](diffhunk://#diff-299c01ed4ee7d0b3f636fe4cb4f0d660a5012b7e95ca0740098b3ace617ab16eL104-R151)
* **The command-line interface in `deepmd/main.py` is updated to add
`--train-data` and `--valid-data` arguments for the test subparser,
enabling direct specification of input JSON files for training and
validation data.**

### Test coverage improvements

* New and updated tests in `source/tests/pt/test_dp_test.py` verify the
ability to run tests using input JSON files for both training and
validation data, including cases with recursive glob patterns. This
ensures robust handling of various data source configurations.
[[1]](diffhunk://#diff-ce70e95ffdb1996c7887ea3f63b54d1ae0fef98059572ad03875ca36cfef3c34L33-R35)
[[2]](diffhunk://#diff-ce70e95ffdb1996c7887ea3f63b54d1ae0fef98059572ad03875ca36cfef3c34R49-R59)
[[3]](diffhunk://#diff-ce70e95ffdb1996c7887ea3f63b54d1ae0fef98059572ad03875ca36cfef3c34R103-R116)
[[4]](diffhunk://#diff-ce70e95ffdb1996c7887ea3f63b54d1ae0fef98059572ad03875ca36cfef3c34R164-R273)
* Additional argument parser tests in
`source/tests/common/test_argument_parser.py` confirm correct parsing of
the new `--train-data` and `--valid-data` options.

### Internal code improvements

* Refactored imports and type annotations in
`deepmd/entrypoints/test.py` to support the new functionality and
improve code clarity.
[[1]](diffhunk://#diff-299c01ed4ee7d0b3f636fe4cb4f0d660a5012b7e95ca0740098b3ace617ab16eR17)
[[2]](diffhunk://#diff-299c01ed4ee7d0b3f636fe4cb4f0d660a5012b7e95ca0740098b3ace617ab16eR42-R50)
[[3]](diffhunk://#diff-299c01ed4ee7d0b3f636fe4cb4f0d660a5012b7e95ca0740098b3ace617ab16eL77-R95)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- New Features
- Added support for supplying test systems via JSON files, including
selecting training or validation data.
- Introduced CLI options --train-data and --valid-data for the test
command.
- Supports resolving relative paths from JSON and optional recursive
glob patterns.
- Changes
- Test command now requires at least one data source (JSON, data file,
or system); clearer errors when none or no systems found.
- Tests
- Expanded test coverage for JSON-driven inputs and recursive glob
patterns; refactored helpers for improved readability.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Chun Cai <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: caic99

deepmd/entrypoints/test.py

@@ -15,6 +15,7 @@
 
 from deepmd.common import (
     expand_sys_str,
+    j_loader,
 )
 from deepmd.infer.deep_dipole import (
     DeepDipole,
@@ -39,9 +40,15 @@
     DeepWFC,
 )
 from deepmd.utils import random as dp_random
+from deepmd.utils.compat import (
+    update_deepmd_input,
+)
 from deepmd.utils.data import (
     DeepmdData,
 )
+from deepmd.utils.data_system import (
+    process_systems,
+)
 from deepmd.utils.weight_avg import (
     weighted_average,
 )
@@ -59,8 +66,10 @@
 def test(
     *,
     model: str,
-    system: str,
-    datafile: str,
+    system: Optional[str],
+    datafile: Optional[str],
+    train_json: Optional[str] = None,
+    valid_json: Optional[str] = None,
     numb_test: int,
     rand_seed: Optional[int],
     shuffle_test: bool,
@@ -75,12 +84,16 @@ def test(
     ----------
     model : str
         path where model is stored
-    system : str
+    system : str, optional
         system directory
-    datafile : str
+    datafile : str, optional
         the path to the list of systems to test
+    train_json : Optional[str]
+        Path to the input.json file provided via ``--train-data``. Training systems will be used for testing.
+    valid_json : Optional[str]
+        Path to the input.json file provided via ``--valid-data``. Validation systems will be used for testing.
     numb_test : int
-        munber of tests to do. 0 means all data.
+        number of tests to do. 0 means all data.
     rand_seed : Optional[int]
         seed for random generator
     shuffle_test : bool
@@ -102,11 +115,41 @@ def test(
     if numb_test == 0:
         # only float has inf, but should work for min
         numb_test = float("inf")
-    if datafile is not None:
+    if train_json is not None:
+        jdata = j_loader(train_json)
+        jdata = update_deepmd_input(jdata)
+        data_params = jdata.get("training", {}).get("training_data", {})
+        systems = data_params.get("systems")
+        if not systems:
+            raise RuntimeError("No training data found in input json")
+        root = Path(train_json).parent
+        if isinstance(systems, str):
+            systems = str((root / Path(systems)).resolve())
+        else:
+            systems = [str((root / Path(ss)).resolve()) for ss in systems]
+        patterns = data_params.get("rglob_patterns", None)
+        all_sys = process_systems(systems, patterns=patterns)
+    elif valid_json is not None:
+        jdata = j_loader(valid_json)
+        jdata = update_deepmd_input(jdata)
+        data_params = jdata.get("training", {}).get("validation_data", {})
+        systems = data_params.get("systems")
+        if not systems:
+            raise RuntimeError("No validation data found in input json")
+        root = Path(valid_json).parent
+        if isinstance(systems, str):
+            systems = str((root / Path(systems)).resolve())
+        else:
+            systems = [str((root / Path(ss)).resolve()) for ss in systems]
+        patterns = data_params.get("rglob_patterns", None)
+        all_sys = process_systems(systems, patterns=patterns)
+    elif datafile is not None:
         with open(datafile) as datalist:
             all_sys = datalist.read().splitlines()
-    else:
+    elif system is not None:
         all_sys = expand_sys_str(system)
+    else:
+        raise RuntimeError("No data source specified for testing")
 
     if len(all_sys) == 0:
         raise RuntimeError("Did not find valid system")

deepmd/main.py

@@ -384,6 +384,24 @@ def main_parser() -> argparse.ArgumentParser:
         type=str,
         help="The path to the datafile, each line of which is a path to one data system.",
     )
+    parser_tst_subgroup.add_argument(
+        "--train-data",
+        dest="train_json",
+        default=None,
+        type=str,
+        help=(
+            "The input json file. Training data in the file will be used for testing."
+        ),
+    )
+    parser_tst_subgroup.add_argument(
+        "--valid-data",
+        dest="valid_json",
+        default=None,
+        type=str,
+        help=(
+            "The input json file. Validation data in the file will be used for testing."
+        ),
+    )
     parser_tst.add_argument(
         "-S",
         "--set-prefix",

source/tests/common/test_argument_parser.py

@@ -322,6 +322,32 @@ def test_parser_test(self) -> None:
 
         self.run_test(command="test", mapping=ARGS)
 
+    def test_parser_test_train_data(self) -> None:
+        """Test test subparser with train-data."""
+        ARGS = {
+            "--model": {"type": str, "value": "MODEL.PB"},
+            "--train-data": {
+                "type": (str, type(None)),
+                "value": "INPUT.JSON",
+                "dest": "train_json",
+            },
+        }
+
+        self.run_test(command="test", mapping=ARGS)
+
+    def test_parser_test_valid_data(self) -> None:
+        """Test test subparser with valid-data."""
+        ARGS = {
+            "--model": {"type": str, "value": "MODEL.PB"},
+            "--valid-data": {
+                "type": (str, type(None)),
+                "value": "INPUT.JSON",
+                "dest": "valid_json",
+            },
+        }
+
+        self.run_test(command="test", mapping=ARGS)
+
     def test_parser_compress(self) -> None:
         """Test compress subparser."""
         ARGS = {

source/tests/pt/test_dp_test.py

@@ -37,7 +37,9 @@
 
 
 class DPTest:
-    def test_dp_test_1_frame(self) -> None:
+    def _run_dp_test(
+        self, use_input_json: bool, numb_test: int = 0, use_train: bool = False
+    ) -> None:
         trainer = get_trainer(deepcopy(self.config))
         with torch.device("cpu"):
             input_dict, label_dict, _ = trainer.get_data(is_train=False)
@@ -51,12 +53,17 @@ def test_dp_test_1_frame(self) -> None:
         model = torch.jit.script(trainer.model)
         tmp_model = tempfile.NamedTemporaryFile(delete=False, suffix=".pth")
         torch.jit.save(model, tmp_model.name)
+        val_sys = self.config["training"]["validation_data"]["systems"]
+        if isinstance(val_sys, list):
+            val_sys = val_sys[0]
         dp_test(
             model=tmp_model.name,
-            system=self.config["training"]["validation_data"]["systems"][0],
+            system=None if use_input_json else val_sys,
             datafile=None,
+            train_json=self.input_json if use_input_json and use_train else None,
+            valid_json=self.input_json if use_input_json and not use_train else None,
             set_prefix="set",
-            numb_test=0,
+            numb_test=numb_test,
             rand_seed=None,
             shuffle_test=False,
             detail_file=self.detail_file,
@@ -100,6 +107,20 @@ def test_dp_test_1_frame(self) -> None:
                 ).reshape(-1, 3),
             )
 
+    def test_dp_test_1_frame(self) -> None:
+        self._run_dp_test(False)
+
+    def test_dp_test_input_json(self) -> None:
+        self._run_dp_test(True)
+
+    def test_dp_test_input_json_train(self) -> None:
+        with open(self.input_json) as f:
+            cfg = json.load(f)
+        cfg["training"]["validation_data"]["systems"] = ["non-existent"]
+        with open(self.input_json, "w") as f:
+            json.dump(cfg, f, indent=4)
+        self._run_dp_test(True, use_train=True)
+
     def tearDown(self) -> None:
         for f in os.listdir("."):
             if f.startswith("model") and f.endswith(".pt"):
@@ -147,6 +168,116 @@ def setUp(self) -> None:
             json.dump(self.config, fp, indent=4)
 
 
+class TestDPTestSeARglob(unittest.TestCase):
+    def setUp(self) -> None:
+        self.detail_file = "test_dp_test_ener_rglob_detail"
+        input_json = str(Path(__file__).parent / "water/se_atten.json")
+        with open(input_json) as f:
+            self.config = json.load(f)
+        self.config["training"]["numb_steps"] = 1
+        self.config["training"]["save_freq"] = 1
+        data_file = [str(Path(__file__).parent / "water/data/single")]
+        self.config["training"]["training_data"]["systems"] = data_file
+        root_dir = str(Path(__file__).parent)
+        self.config["training"]["validation_data"]["systems"] = root_dir
+        self.config["training"]["validation_data"]["rglob_patterns"] = [
+            "water/data/single"
+        ]
+        self.config["model"] = deepcopy(model_se_e2_a)
+        self.input_json = "test_dp_test_rglob.json"
+        with open(self.input_json, "w") as fp:
+            json.dump(self.config, fp, indent=4)
+
+    def test_dp_test_input_json_rglob(self) -> None:
+        trainer = get_trainer(deepcopy(self.config))
+        with torch.device("cpu"):
+            input_dict, _, _ = trainer.get_data(is_train=False)
+        input_dict.pop("spin", None)
+        model = torch.jit.script(trainer.model)
+        tmp_model = tempfile.NamedTemporaryFile(delete=False, suffix=".pth")
+        torch.jit.save(model, tmp_model.name)
+        dp_test(
+            model=tmp_model.name,
+            system=None,
+            datafile=None,
+            valid_json=self.input_json,
+            set_prefix="set",
+            numb_test=1,
+            rand_seed=None,
+            shuffle_test=False,
+            detail_file=self.detail_file,
+            atomic=False,
+        )
+        os.unlink(tmp_model.name)
+        self.assertTrue(os.path.exists(self.detail_file + ".e.out"))
+
+    def tearDown(self) -> None:
+        for f in os.listdir("."):
+            if f.startswith("model") and f.endswith(".pt"):
+                os.remove(f)
+            if f.startswith(self.detail_file):
+                os.remove(f)
+            if f in ["lcurve.out", self.input_json]:
+                os.remove(f)
+            if f in ["stat_files"]:
+                shutil.rmtree(f)
+
+
+class TestDPTestSeARglobTrain(unittest.TestCase):
+    def setUp(self) -> None:
+        self.detail_file = "test_dp_test_ener_rglob_train_detail"
+        input_json = str(Path(__file__).parent / "water/se_atten.json")
+        with open(input_json) as f:
+            self.config = json.load(f)
+        self.config["training"]["numb_steps"] = 1
+        self.config["training"]["save_freq"] = 1
+        root_dir = str(Path(__file__).parent)
+        self.config["training"]["training_data"]["systems"] = root_dir
+        self.config["training"]["training_data"]["rglob_patterns"] = [
+            "water/data/single"
+        ]
+        data_file = [str(Path(__file__).parent / "water/data/single")]
+        self.config["training"]["validation_data"]["systems"] = data_file
+        self.config["model"] = deepcopy(model_se_e2_a)
+        self.input_json = "test_dp_test_rglob_train.json"
+        with open(self.input_json, "w") as fp:
+            json.dump(self.config, fp, indent=4)
+
+    def test_dp_test_input_json_rglob_train(self) -> None:
+        trainer = get_trainer(deepcopy(self.config))
+        with torch.device("cpu"):
+            input_dict, _, _ = trainer.get_data(is_train=False)
+        input_dict.pop("spin", None)
+        model = torch.jit.script(trainer.model)
+        tmp_model = tempfile.NamedTemporaryFile(delete=False, suffix=".pth")
+        torch.jit.save(model, tmp_model.name)
+        dp_test(
+            model=tmp_model.name,
+            system=None,
+            datafile=None,
+            train_json=self.input_json,
+            set_prefix="set",
+            numb_test=1,
+            rand_seed=None,
+            shuffle_test=False,
+            detail_file=self.detail_file,
+            atomic=False,
+        )
+        os.unlink(tmp_model.name)
+        self.assertTrue(os.path.exists(self.detail_file + ".e.out"))
+
+    def tearDown(self) -> None:
+        for f in os.listdir("."):
+            if f.startswith("model") and f.endswith(".pt"):
+                os.remove(f)
+            if f.startswith(self.detail_file):
+                os.remove(f)
+            if f in ["lcurve.out", self.input_json]:
+                os.remove(f)
+            if f in ["stat_files"]:
+                shutil.rmtree(f)
+
+
 class TestDPTestForceWeight(DPTest, unittest.TestCase):
     def setUp(self) -> None:
         self.detail_file = "test_dp_test_force_weight_detail"