fix: fix caly_evo_step superop

dpgen2/op/prep_dp_optim.py

@@ -111,13 +111,17 @@ def execute(
         finished = ip["finished"]
 
         work_dir = Path(ip["task_name"])
+        # print("-----in PrepDPOptim====work_dir(task_name)", work_dir)
         poscar_dir = ip["poscar_dir"]
         models_dir = ip["models_dir"]
+        # print("-----in PrepDPOptim====poscar_dir", poscar_dir)
+        # print("-----in PrepDPOptim====models_dir", models_dir)
         _caly_run_opt_file = ip["caly_run_opt_file"]
         _caly_check_opt_file = ip["caly_check_opt_file"]
         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_*")]
+        # print("-----in PrepDPOptim====len(poscar_dir.rglob(POSCAR*))", len(poscar_list))
         poscar_list = sorted(poscar_list, key=lambda x: int(x.name.strip("POSCAR_")))
         model_name = "frozen_model.pb"
         model_list = [model.resolve() for model in models_dir.rglob(model_name)]
@@ -137,11 +141,12 @@ def execute(
                 ]
 
                 task_dirs = []
-                for idx, poscar_list in enumerate(grouped_poscar_list):
+                for idx, _poscar_list in enumerate(grouped_poscar_list):
                     opt_path = Path(f"opt_path_{idx}")
+                    # print("-----in PrepDPOptim====opt_path, _poscar_list)", opt_path, _poscar_list)
                     task_dirs.append(work_dir / opt_path)
                     with set_directory(opt_path):
-                        for poscar in poscar_list:
+                        for poscar in _poscar_list:
                             Path(poscar.name).symlink_to(poscar)
                         Path(model_name).symlink_to(model_file)
                         Path(caly_run_opt_file.name).symlink_to(caly_run_opt_file)
@@ -157,7 +162,7 @@ def execute(
             {
                 "task_names": task_names,
                 "task_dirs": task_dirs,
-                "caly_run_opt_file": work_dir / calypso_run_opt_file,
-                "caly_check_opt_file": work_dir / calypso_check_opt_file,
+                "caly_run_opt_file": work_dir / caly_run_opt_file.name,
+                "caly_check_opt_file": work_dir / caly_check_opt_file.name,
             }
         )

dpgen2/op/run_dp_optim.py

@@ -97,7 +97,10 @@ def execute(
         cnt_num = ip["cnt_num"]
 
         task_path = ip["task_dir"]
-        input_files = [ii.resolve() for ii in Path(task_path).iterdir()]
+        if task_path is not None:
+            input_files = [ii.resolve() for ii in Path(task_path).iterdir()]
+        else:
+            input_files = []
 
         config = ip["config"] if ip["config"] is not None else {}
         command = config.get(
@@ -161,7 +164,5 @@ def execute(
                 "task_name": str(work_dir),
                 "optim_results_dir": work_dir / optim_results_dir,
                 "traj_results": work_dir / traj_results_dir,
-                # "caly_run_opt_file": work_dir / calypso_run_opt_file,
-                # "caly_check_opt_file": work_dir / calypso_check_opt_file,
             }
         )

dpgen2/superop/caly_evo_step.py

@@ -97,18 +97,19 @@ def __init__(
 
         self.collect_run_calypso_keys = [
             "%s--collect-run-calypso-%s-%s" % (self.inputs.parameters["block_id"], i, j)
-            for i in range(50)
-            for j in range(50)
+            for i in range(20)
+            for j in range(20)
         ]
         self.prep_dp_optim_keys = [
             "%s--prep-dp-optim-%s-%s" % (self.inputs.parameters["block_id"], i, j)
-            for i in range(50)
-            for j in range(50)
+            for i in range(20)
+            for j in range(20)
         ]
         self.run_dp_optim_keys = [
-            "%s--run-dp-optim-%s-%s" % (self.inputs.parameters["block_id"], i, j)
-            for i in range(50)
-            for j in range(50)
+            "%s--run-dp-optim-%s-%s-%s" % (self.inputs.parameters["block_id"], i, j, k)
+            for i in range(20)
+            for j in range(20)
+            for k in range(50)
         ]
         self._keys = (
             self.collect_run_calypso_keys
@@ -243,7 +244,7 @@ def _caly_evo_step(
             run_dp_optim_op,
             slices=Slices(
                 input_parameter=["task_name"],
-                input_artifact=["task_path"],
+                input_artifact=["task_dir"],
                 output_artifact=["traj_results", "optim_results_dir"],
             ),
             python_packages=upload_python_packages,
@@ -264,7 +265,7 @@ def _caly_evo_step(
             caly_evo_step_steps.inputs.parameters["iter_num"],
             caly_evo_step_steps.inputs.parameters["cnt_num"],
         ),
-        executor=prep_executor,  # cpu is enough to run calypso.x, default step config is c2m4
+        executor=run_executor,
         **run_config,
     )
     caly_evo_step_steps.add(run_dp_optim)

tests/mocked_ops.py

@@ -972,6 +972,7 @@ def execute(
 
         # work_dir name: calypso_task.idx
         work_dir = Path(ip["task_name"])
+        # print("==============----------work_dir in collruncaly:", work_dir)
         work_dir.mkdir(exist_ok=True, parents=True)
 
         qhull_input = (
@@ -1121,20 +1122,38 @@ def execute(
         finished = ip["finished"]
         work_dir = Path(ip["task_name"])
         cnt_num = ip["cnt_num"]
-        # print(f"--------=---------task_name: {work_dir}")
+        # print(f"----in RunDpOptim----task_name: {work_dir}")
+        # print(f"----cwd in RunDPOptim------cwd: {cwd}")
         work_dir.mkdir(parents=True, exist_ok=True)
 
         config = ip["config"] if ip["config"] is not None else {}
         command = config.get("run_opt_command", "python -u calypso_run_opt.py")
 
-        os.chdir(work_dir)
+        task_path = ip["task_dir"]
+        # print(f"----in RunDpOptim----task_dir: {task_path}")
+        if task_path is not None:
+            input_files = [ii.resolve() for ii in Path(task_path).iterdir()]
+        else:
+            pass
 
-        for i in range(1, 6):
-            Path().joinpath(f"CONTCAR_{str(i)}").write_text(f"CONTCAR_{str(i)}")
-            Path().joinpath(f"OUTCAR_{str(i)}").write_text(f"OUTCAR_{str(i)}")
-            Path().joinpath(f"{str(i)}.traj").write_text(f"{str(i)}.traj")
+        os.chdir(work_dir)
 
         if finished == "false":
+            for ii in input_files:
+                iname = ii.name
+                Path(iname).symlink_to(ii)
+
+            poscar_list = sorted(Path().rglob("POSCAR_*"))
+            # print(f"-----in RunDpOptim----poscar_list: {poscar_list}")
+            cnt = 0
+            for poscar in poscar_list:
+                cnt += 1
+                poscar_name = poscar.name
+                num = poscar_name.strip("POSCAR_")
+                Path(poscar_name.replace("POSCAR", "CONTCAR")).write_text("")
+                Path(poscar_name.replace("POSCAR", "OUTCAR")).write_text("")
+                Path(f"{num}.traj").write_text("")
+
             optim_results_dir = Path("optim_results_dir")
             optim_results_dir.mkdir(parents=True, exist_ok=True)
             for poscar in Path().glob("POSCAR_*"):
@@ -1165,8 +1184,6 @@ def execute(
                 "task_name": str(work_dir),
                 "optim_results_dir": work_dir / optim_results_dir,
                 "traj_results": work_dir / traj_results_dir,
-                "caly_run_opt_file": work_dir / calypso_run_opt_file,
-                "caly_check_opt_file": work_dir / calypso_check_opt_file,
             }
         )
 

tests/test_caly_evo_step.py

@@ -276,12 +276,11 @@ def setUp(self):
         )
 
     def tearDown(self):
-        pass
-        # shutil.rmtree(self.work_dir, ignore_errors=True)
-        # for i in Path().glob("caly-evo-step-*"):
-        #     shutil.rmtree(i, ignore_errors=True)
-        # for i in Path().glob("caly_task*"):
-        #     shutil.rmtree(i, ignore_errors=True)
+        shutil.rmtree(self.work_dir, ignore_errors=True)
+        for i in Path().glob("caly-evo-step-*"):
+            shutil.rmtree(i, ignore_errors=True)
+        for i in Path().glob("caly_task*"):
+            shutil.rmtree(i, ignore_errors=True)
 
     @unittest.skip("only need to run test_01")
     def test_00(self):
@@ -324,7 +323,7 @@ def test_00(self):
         self.assertEqual(step.phase, "Succeeded")
 
     # @unittest.skip("temp skit")
-    def test_01(self):
+    def test_caly_evo_step(self):
         steps = CalyEvoStep(
             "caly-evo-run",
             MockedCollRunCaly,
@@ -381,6 +380,8 @@ def test_01(self):
 
         download_artifact(step.outputs.artifacts["traj_results"])
 
+        cwd = Path().cwd()
+        print(cwd, self.task_name_list)
         for idx, name in enumerate(self.task_name_list):
             cwd = Path().cwd()
             os.chdir(Path(name))