feat: support rglob patterns for dp test

Author: caic99

deepmd/entrypoints/test.py

@@ -14,6 +14,7 @@
 
 from deepmd.common import (
     expand_sys_str,
+    j_loader,
 )
 from deepmd.infer.deep_dipole import (
     DeepDipole,
@@ -41,6 +42,9 @@
 from deepmd.utils.data import (
     DeepmdData,
 )
+from deepmd.utils.data_system import (
+    process_systems,
+)
 from deepmd.utils.weight_avg import (
     weighted_average,
 )
@@ -60,6 +64,7 @@ def test(
     model: str,
     system: str,
     datafile: str,
+    input_json: Optional[str] = None,
     numb_test: int,
     rand_seed: Optional[int],
     shuffle_test: bool,
@@ -78,6 +83,8 @@ def test(
         system directory
     datafile : str
         the path to the list of systems to test
+    input_json : Optional[str]
+        the training input json file. Validation systems in this file will be used.
     numb_test : int
         munber of tests to do. 0 means all data.
     rand_seed : Optional[int]
@@ -101,7 +108,20 @@ 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 input_json is not None:
+        jdata = j_loader(input_json)
+        val_params = jdata.get("training", {}).get("validation_data", {})
+        validation = val_params.get("systems")
+        if not validation:
+            raise RuntimeError("No validation data found in input json")
+        root = Path(input_json).parent
+        if isinstance(validation, str):
+            validation = str((root / Path(validation)).resolve())
+        else:
+            validation = [str((root / Path(ss)).resolve()) for ss in validation]
+        patterns = val_params.get("rglob_patterns", None)
+        all_sys = process_systems(validation, patterns=patterns)
+    elif datafile is not None:
         with open(datafile) as datalist:
             all_sys = datalist.read().splitlines()
     else:

deepmd/main.py

@@ -371,6 +371,13 @@ 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(
+        "-i",
+        "--input-json",
+        default=None,
+        type=str,
+        help="The training input json file. Validation data in the training script will be used for testing.",
+    )
     parser_tst.add_argument(
         "-S",
         "--set-prefix",

source/tests/pt/test_dp_test.py

@@ -30,7 +30,7 @@
 
 
 class DPTest:
-    def test_dp_test_1_frame(self) -> None:
+    def _run_dp_test(self, use_input_json: bool, numb_test: int = 0) -> None:
         trainer = get_trainer(deepcopy(self.config))
         with torch.device("cpu"):
             input_dict, label_dict, _ = trainer.get_data(is_train=False)
@@ -44,12 +44,16 @@ 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=val_sys,
             datafile=None,
+            input_json=self.input_json if use_input_json else None,
             set_prefix="set",
-            numb_test=0,
+            numb_test=numb_test,
             rand_seed=None,
             shuffle_test=False,
             detail_file=self.detail_file,
@@ -93,6 +97,12 @@ 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 tearDown(self) -> None:
         for f in os.listdir("."):
             if f.startswith("model") and f.endswith(".pt"):
@@ -140,6 +150,61 @@ 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=self.config["training"]["validation_data"]["systems"],
+            datafile=None,
+            input_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 TestDPTestPropertySeA(unittest.TestCase):
     def setUp(self) -> None:
         self.detail_file = "test_dp_test_property_detail"