fix: calypso interface is compatible with dpv3 (#201)

dpgen2/exploration/task/calypso/caly_input.py

@@ -18,6 +18,7 @@
 calypso_run_opt_str = """#!/usr/bin/env python3
 
 import os
+import sys
 import time
 import glob
 import numpy as np
@@ -100,7 +101,7 @@ def Write_Outcar(outcar, element, ele, volume, lat, pos, ene, force, stress, pst
 def run_opt(fmax, stress):
     # Using the ASE&DP to Optimize Configures
 
-    calc = DP(model='frozen_model.pb')    # init the model before iteration
+    calc = DP(model=sys.argv[1])    # init the model before iteration
 
     Opt_Step = 1000
     start = time.time()

dpgen2/op/prep_caly_input.py

@@ -43,6 +43,7 @@
 calypso_run_opt_str = """#!/usr/bin/env python3
 
 import os
+import sys
 import time
 import glob
 import numpy as np
@@ -126,7 +127,7 @@ def Write_Outcar(outcar, element, ele, volume, lat, pos, ene, force, stress, pst
 def run_opt(fmax, stress):
     '''Using the ASE&DP to Optimize Configures'''
 
-    calc = DP(model='frozen_model.pb')    # init the model before iteration
+    calc = DP(model=sys.argv[1])    # init the model before iteration
 
     Opt_Step = 1000
     start = time.time()

dpgen2/op/prep_run_dp_optim.py

@@ -117,17 +117,25 @@ def execute(
         caly_run_opt_file = _caly_run_opt_file.resolve()
         caly_check_opt_file = _caly_check_opt_file.resolve()
         poscar_list = [poscar.resolve() for poscar in poscar_dir.rglob("POSCAR_*")]
-        model_list = [model.resolve() for model in models_dir.rglob("*model*pb")]
+
+        model_name = "frozen_model.pb"
+        model_list = [model.resolve() for model in models_dir.rglob(model_name)]
+        if len(model_list) == 0:
+            model_name = "model.ckpt.pt"
+            model_list = [model.resolve() for model in models_dir.rglob(model_name)]
+
         model_list = sorted(model_list, key=lambda x: str(x).split(".")[1])
         model_file = model_list[0]
 
         config = ip["config"] if ip["config"] is not None else {}
-        command = config.get("run_opt_command", "python -u calypso_run_opt.py")
+        command = config.get(
+            "run_opt_command", f"python -u calypso_run_opt.py {model_name}"
+        )
 
         with set_directory(work_dir):
             for idx, poscar in enumerate(poscar_list):
                 Path(poscar.name).symlink_to(poscar)
-            Path("frozen_model.pb").symlink_to(model_file)
+            Path(model_name).symlink_to(model_file)
             Path(caly_run_opt_file.name).symlink_to(caly_run_opt_file)
             Path(caly_check_opt_file.name).symlink_to(caly_check_opt_file)
 

dpgen2/op/run_caly_model_devi.py

@@ -273,14 +273,15 @@ def parse_traj(traj_file):
 
 
 def write_model_devi_out(devi: np.ndarray, fname: Union[str, Path], header: str = ""):
-    assert devi.shape[1] == 7
+    assert devi.shape[1] == 8
     header = "%s\n%10s" % (header, "step")
     for item in "vf":
         header += "%19s%19s%19s" % (
             f"max_devi_{item}",
             f"min_devi_{item}",
             f"avg_devi_{item}",
         )
+        header += "%19s" % "devi_e"
     with open(fname, "ab") as fp:
         np.savetxt(
             fp,

tests/op/test_prep_run_dp_optim.py

@@ -50,9 +50,9 @@ def setUp(self):
         self.models_dir = Path("models_dir")
         self.models_dir.mkdir(parents=True, exist_ok=True)
         for i in range(4):
-            self.models_dir.joinpath(model_name_pattern % i).write_text(
-                f"model.{str(i)}.pb"
-            )
+            path_name = self.models_dir.joinpath(f"task.{i}")
+            path_name.mkdir(exist_ok=True, parents=True)
+            path_name.joinpath("frozen_model.pb").write_text("frozen_model.pb")
 
         self.task_name = calypso_task_pattern % 0
 
@@ -85,6 +85,7 @@ def tearDown(self):
         shutil.rmtree(self.ref_traj_results_dir)
         shutil.rmtree(Path(self.task_name))
         shutil.rmtree(self.poscar_dir)
+        shutil.rmtree(self.models_dir)
 
     @patch("dpgen2.op.prep_run_dp_optim.run_command")
     def test_success_00(self, mocked_run):

tests/op/test_run_caly_model_devi.py

@@ -143,7 +143,7 @@ def side_effect_1(*args, **kwargs):
         mocked_run_1.side_effect = side_effect_1
 
         def side_effect_2(*args, **kwargs):
-            return [[1, 1, 1, 1, 1, 1, 1]]
+            return [[1, 1, 1, 1, 1, 1, 1, 1]]
 
         mocked_run_2.side_effect = side_effect_2