feat: allow testing with training data

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,8 @@ def test(
     model: str,
     system: str,
     datafile: str,
+    input_json: Optional[str] = None,
+    use_train: bool = False,
     numb_test: int,
     rand_seed: Optional[int],
     shuffle_test: bool,
@@ -78,6 +84,10 @@ 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.
+    use_train : bool
+        use training systems in the input json file instead of validation systems
     numb_test : int
         munber of tests to do. 0 means all data.
     rand_seed : Optional[int]
@@ -101,7 +111,25 @@ 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)
+        data_key = "training_data" if use_train else "validation_data"
+        data_params = jdata.get("training", {}).get(data_key, {})
+        systems = data_params.get("systems")
+        if not systems:
+            raise RuntimeError(
+                f"No {'training' if use_train else 'validation'} data found in input json"
+            )
+        root = Path(input_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 use_train:
+        raise RuntimeError("--train-data requires --input-json")
+    elif datafile is not None:
         with open(datafile) as datalist:
             all_sys = datalist.read().splitlines()
     else:

deepmd/main.py

@@ -356,7 +356,7 @@ def main_parser() -> argparse.ArgumentParser:
         type=str,
         help="Frozen model file (prefix) to import. TensorFlow backend: suffix is .pb; PyTorch backend: suffix is .pth.",
     )
-    parser_tst_subgroup = parser_tst.add_mutually_exclusive_group()
+    parser_tst_subgroup = parser_tst.add_mutually_exclusive_group(required=True)
     parser_tst_subgroup.add_argument(
         "-s",
         "--system",
@@ -371,6 +371,19 @@ 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(
+        "--train-data",
+        action="store_true",
+        dest="use_train",
+        help="Use training data in the input json file instead of validation data.",
+    )
     parser_tst.add_argument(
         "-S",
         "--set-prefix",

source/tests/pt/test_dp_test.py

@@ -30,7 +30,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)
@@ -44,12 +46,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=val_sys,
             datafile=None,
+            input_json=self.input_json if use_input_json else None,
+            use_train=use_train,
             set_prefix="set",
-            numb_test=0,
+            numb_test=numb_test,
             rand_seed=None,
             shuffle_test=False,
             detail_file=self.detail_file,
@@ -93,6 +100,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"):
@@ -140,6 +161,117 @@ 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 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=self.config["training"]["validation_data"]["systems"],
+            datafile=None,
+            input_json=self.input_json,
+            use_train=True,
+            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"