From ee9142694adb33c88a5d1c31567022d588330a7f Mon Sep 17 00:00:00 2001
From: Duo <50307526+iProzd@users.noreply.github.com>
Date: Fri, 26 Jan 2024 16:16:17 +0800
Subject: [PATCH] fix path in UTs

---
 source/tests/pt/test_LKF.py           |  5 ++++-
 source/tests/pt/test_calculator.py    |  9 +++++++--
 source/tests/pt/test_dp_test.py       |  7 +++++--
 source/tests/pt/test_jit.py           | 13 ++++++++-----
 source/tests/pt/test_permutation.py   |  5 ++++-
 source/tests/pt/test_saveload_dpa1.py |  8 +++++++-
 source/tests/pt/test_training.py      | 23 +++++++++++++++++++----
 source/tests/pt/water/lkf.json        |  4 ++--
 source/tests/pt/water/se_atten.json   |  4 ++--
 source/tests/pt/water/se_e2_a.json    |  4 ++--
 10 files changed, 60 insertions(+), 22 deletions(-)

diff --git a/source/tests/pt/test_LKF.py b/source/tests/pt/test_LKF.py
index b3b59f690f..667a536fb2 100644
--- a/source/tests/pt/test_LKF.py
+++ b/source/tests/pt/test_LKF.py
@@ -4,11 +4,14 @@
 from deepmd.pt.entrypoints.main import (
     main,
 )
+from pathlib import (
+    Path,
+)
 
 
 class TestLKF(unittest.TestCase):
     def test_lkf(self):
-        main(["train", "tests/water/lkf.json"])
+        main(["train", str(Path(__file__).parent / "water/lkf.json")])
 
 
 if __name__ == "__main__":
diff --git a/source/tests/pt/test_calculator.py b/source/tests/pt/test_calculator.py
index 69bb43366f..e2709d3143 100644
--- a/source/tests/pt/test_calculator.py
+++ b/source/tests/pt/test_calculator.py
@@ -13,19 +13,24 @@
 from deepmd.pt.utils.ase_calc import (
     DPCalculator,
 )
+from pathlib import (
+    Path,
+)
 
 dtype = torch.float64
 
 
 class TestCalculator(unittest.TestCase):
     def setUp(self):
-        input_json = "tests/water/se_atten.json"
+        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/data_0")
+        self.config["training"]["training_data"]["systems"] = data_file
         self.config["training"]["validation_data"]["systems"] = [
-            "tests/water/data/single"
+            str(Path(__file__).parent / "water/data/single")
         ]
         self.input_json = "test_dp_test.json"
         with open(self.input_json, "w") as fp:
diff --git a/source/tests/pt/test_dp_test.py b/source/tests/pt/test_dp_test.py
index a14dad72c3..524a6b1f57 100644
--- a/source/tests/pt/test_dp_test.py
+++ b/source/tests/pt/test_dp_test.py
@@ -6,6 +6,9 @@
 from copy import (
     deepcopy,
 )
+from pathlib import (
+    Path,
+)
 
 import numpy as np
 
@@ -19,13 +22,13 @@
 
 class TestDPTest(unittest.TestCase):
     def setUp(self):
-        input_json = "tests/water/se_atten.json"
+        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
         self.config["training"]["validation_data"]["systems"] = [
-            "tests/water/data/single"
+            str(Path(__file__).parent / "water/data/single")
         ]
         self.input_json = "test_dp_test.json"
         with open(self.input_json, "w") as fp:
diff --git a/source/tests/pt/test_jit.py b/source/tests/pt/test_jit.py
index d4c13da438..d951802826 100644
--- a/source/tests/pt/test_jit.py
+++ b/source/tests/pt/test_jit.py
@@ -22,6 +22,9 @@
     model_hybrid,
     model_se_e2_a,
 )
+from pathlib import (
+    Path,
+)
 
 
 class TestJIT:
@@ -43,7 +46,7 @@ def tearDown(self):
 
 class TestEnergyModelSeA(unittest.TestCase, TestJIT):
     def setUp(self):
-        input_json = "tests/water/se_atten.json"
+        input_json = str(Path(__file__).parent / "water/data/single")
         with open(input_json) as f:
             self.config = json.load(f)
         self.config["model"] = deepcopy(model_se_e2_a)
@@ -56,7 +59,7 @@ def tearDown(self):
 
 class TestEnergyModelDPA1(unittest.TestCase, TestJIT):
     def setUp(self):
-        input_json = "tests/water/se_atten.json"
+        input_json = str(Path(__file__).parent / "water/data/single")
         with open(input_json) as f:
             self.config = json.load(f)
         self.config["model"] = deepcopy(model_dpa1)
@@ -69,7 +72,7 @@ def tearDown(self):
 
 class TestEnergyModelDPA2(unittest.TestCase, TestJIT):
     def setUp(self):
-        input_json = "tests/water/se_atten.json"
+        input_json = str(Path(__file__).parent / "water/data/single")
         with open(input_json) as f:
             self.config = json.load(f)
         self.config["model"] = deepcopy(model_dpa2)
@@ -89,7 +92,7 @@ def setUp(self):
 @unittest.skip("hybrid not supported at the moment")
 class TestEnergyModelHybrid(unittest.TestCase, TestJIT):
     def setUp(self):
-        input_json = "tests/water/se_atten.json"
+        input_json = str(Path(__file__).parent / "water/data/single")
         with open(input_json) as f:
             self.config = json.load(f)
         self.config["model"] = deepcopy(model_hybrid)
@@ -100,7 +103,7 @@ def setUp(self):
 @unittest.skip("hybrid not supported at the moment")
 class TestEnergyModelHybrid2(unittest.TestCase, TestJIT):
     def setUp(self):
-        input_json = "tests/water/se_atten.json"
+        input_json = str(Path(__file__).parent / "water/data/single")
         with open(input_json) as f:
             self.config = json.load(f)
         self.config["model"] = deepcopy(model_hybrid)
diff --git a/source/tests/pt/test_permutation.py b/source/tests/pt/test_permutation.py
index 6a7adc5857..0910b8d136 100644
--- a/source/tests/pt/test_permutation.py
+++ b/source/tests/pt/test_permutation.py
@@ -19,6 +19,9 @@
 from deepmd.pt.utils.stat import (
     make_stat_input,
 )
+from pathlib import (
+    Path,
+)
 
 dtype = torch.float64
 
@@ -179,7 +182,7 @@
 
 def make_sample(model_params):
     training_systems = [
-        "tests/water/data/data_0",
+        str(Path(__file__).parent / "water/data/data_0"),
     ]
     data_stat_nbatch = model_params.get("data_stat_nbatch", 10)
     train_data = DpLoaderSet(
diff --git a/source/tests/pt/test_saveload_dpa1.py b/source/tests/pt/test_saveload_dpa1.py
index d7d311edbd..e14fa97022 100644
--- a/source/tests/pt/test_saveload_dpa1.py
+++ b/source/tests/pt/test_saveload_dpa1.py
@@ -31,6 +31,9 @@
 from deepmd.tf.common import (
     expand_sys_str,
 )
+from pathlib import (
+    Path,
+)
 
 
 def get_dataset(config):
@@ -62,12 +65,15 @@ def get_dataset(config):
 
 class TestSaveLoadDPA1(unittest.TestCase):
     def setUp(self):
-        input_json = "tests/water/se_atten.json"
+        input_json = str(Path(__file__).parent / "water/se_atten.json")
         with open(input_json) as fin:
             self.config = json.load(fin)
         self.config["loss"]["starter_learning_rate"] = self.config["learning_rate"][
             "start_lr"
         ]
+        data_file = str(Path(__file__).parent / "water/data/data_0")
+        self.config["training"]["training_data"]["systems"] = data_file
+        self.config["training"]["validation_data"]["systems"] = data_file
         self.dataset, self.sampled = get_dataset(self.config)
         self.training_dataloader = DataLoader(
             self.dataset,
diff --git a/source/tests/pt/test_training.py b/source/tests/pt/test_training.py
index d7ca6a4829..7549a29367 100644
--- a/source/tests/pt/test_training.py
+++ b/source/tests/pt/test_training.py
@@ -17,6 +17,9 @@
     model_hybrid,
     model_se_e2_a,
 )
+from pathlib import (
+    Path,
+)
 
 
 class TestDPTrain:
@@ -37,9 +40,12 @@ def tearDown(self):
 
 class TestEnergyModelSeA(unittest.TestCase, TestDPTrain):
     def setUp(self):
-        input_json = "tests/water/se_atten.json"
+        input_json = str(Path(__file__).parent / "water/se_atten.json")
         with open(input_json) as f:
             self.config = json.load(f)
+        data_file = str(Path(__file__).parent / "water/data/data_0")
+        self.config["training"]["training_data"]["systems"] = data_file
+        self.config["training"]["validation_data"]["systems"] = data_file
         self.config["model"] = deepcopy(model_se_e2_a)
         self.config["training"]["numb_steps"] = 1
         self.config["training"]["save_freq"] = 1
@@ -47,9 +53,12 @@ def setUp(self):
 
 class TestEnergyModelDPA1(unittest.TestCase, TestDPTrain):
     def setUp(self):
-        input_json = "tests/water/se_atten.json"
+        input_json = str(Path(__file__).parent / "water/se_atten.json")
         with open(input_json) as f:
             self.config = json.load(f)
+        data_file = str(Path(__file__).parent / "water/data/data_0")
+        self.config["training"]["training_data"]["systems"] = data_file
+        self.config["training"]["validation_data"]["systems"] = data_file
         self.config["model"] = deepcopy(model_dpa1)
         self.config["training"]["numb_steps"] = 1
         self.config["training"]["save_freq"] = 1
@@ -57,9 +66,12 @@ def setUp(self):
 
 class TestEnergyModelDPA2(unittest.TestCase, TestDPTrain):
     def setUp(self):
-        input_json = "tests/water/se_atten.json"
+        input_json = str(Path(__file__).parent / "water/se_atten.json")
         with open(input_json) as f:
             self.config = json.load(f)
+        data_file = str(Path(__file__).parent / "water/data/data_0")
+        self.config["training"]["training_data"]["systems"] = data_file
+        self.config["training"]["validation_data"]["systems"] = data_file
         self.config["model"] = deepcopy(model_dpa2)
         self.config["model"]["descriptor"]["rcut"] = self.config["model"]["descriptor"][
             "repinit_rcut"
@@ -77,9 +89,12 @@ def setUp(self):
 @unittest.skip("hybrid not supported at the moment")
 class TestEnergyModelHybrid(unittest.TestCase, TestDPTrain):
     def setUp(self):
-        input_json = "tests/water/se_atten.json"
+        input_json = str(Path(__file__).parent / "water/se_atten.json")
         with open(input_json) as f:
             self.config = json.load(f)
+        data_file = str(Path(__file__).parent / "water/data/data_0")
+        self.config["training"]["training_data"]["systems"] = data_file
+        self.config["training"]["validation_data"]["systems"] = data_file
         self.config["model"] = deepcopy(model_hybrid)
         self.config["training"]["numb_steps"] = 1
         self.config["training"]["save_freq"] = 1
diff --git a/source/tests/pt/water/lkf.json b/source/tests/pt/water/lkf.json
index db5714d2cb..4385d02136 100644
--- a/source/tests/pt/water/lkf.json
+++ b/source/tests/pt/water/lkf.json
@@ -53,14 +53,14 @@
   "training": {
     "training_data": {
       "systems": [
-        "tests/water/data/data_0"
+        "pt/water/data/data_0"
       ],
       "batch_size": 3,
       "_comment": "that's all"
     },
     "validation_data": {
       "systems": [
-        "tests/water/data/data_0"
+        "pt/water/data/data_0"
       ],
       "batch_size": 1,
       "numb_btch": 3,
diff --git a/source/tests/pt/water/se_atten.json b/source/tests/pt/water/se_atten.json
index 8b6611cc7b..8867e0db41 100644
--- a/source/tests/pt/water/se_atten.json
+++ b/source/tests/pt/water/se_atten.json
@@ -61,14 +61,14 @@
   "training": {
     "training_data": {
       "systems": [
-        "tests/water/data/data_0"
+        "pt/water/data/data_0"
       ],
       "batch_size": 1,
       "_comment": "that's all"
     },
     "validation_data": {
       "systems": [
-        "tests/water/data/data_0"
+        "pt/water/data/data_0"
       ],
       "batch_size": 1,
       "numb_btch": 1,
diff --git a/source/tests/pt/water/se_e2_a.json b/source/tests/pt/water/se_e2_a.json
index b729adf440..425ca3cbf5 100644
--- a/source/tests/pt/water/se_e2_a.json
+++ b/source/tests/pt/water/se_e2_a.json
@@ -53,14 +53,14 @@
   "training": {
     "training_data": {
       "systems": [
-        "tests/water/data/data_0"
+        "pt/water/data/data_0"
       ],
       "batch_size": 1,
       "_comment": "that's all"
     },
     "validation_data": {
       "systems": [
-        "tests/water/data/data_0"
+        "pt/water/data/data_0"
       ],
       "batch_size": 1,
       "numb_btch": 3,