fix: use prepared CALYPSO model by default

dpgen2/op/run_caly_dp_optim.py

@@ -0,0 +1,176 @@
+import json
+import logging
+import pickle
+import shutil
+from pathlib import (
+    Path,
+)
+from typing import (
+    List,
+    Tuple,
+)
+
+from dflow.python import (
+    OP,
+    OPIO,
+    Artifact,
+    BigParameter,
+    OPIOSign,
+    Parameter,
+    TransientError,
+)
+
+from dpgen2.constants import (
+    calypso_check_opt_file,
+    calypso_run_opt_file,
+)
+from dpgen2.exploration.task import (
+    ExplorationTaskGroup,
+)
+from dpgen2.utils import (
+    BinaryFileInput,
+    set_directory,
+)
+from dpgen2.utils.run_command import (
+    run_command,
+)
+
+
+class RunCalyDPOptim(OP):
+    r"""Perform structure optimization with DP in `ip["work_path"]`.
+
+    The `optim_results_dir` and `traj_results` will be returned as `op["optim_results_dir"]`
+    and `op["traj_results"]`.
+    """
+
+    @classmethod
+    def get_input_sign(cls):
+        return OPIOSign(
+            {
+                "config": BigParameter(dict),
+                "task_name": Parameter(str),  # calypso_task.idx
+                "finished": Parameter(str),
+                "cnt_num": Parameter(int),
+                "task_dir": Artifact(Path),  # ready to run structure optimization
+            }
+        )
+
+    @classmethod
+    def get_output_sign(cls):
+        return OPIOSign(
+            {
+                "task_name": Parameter(str),
+                "optim_results_dir": Artifact(Path),
+                "traj_results": Artifact(Path),
+            }
+        )
+
+    @OP.exec_sign_check
+    def execute(
+        self,
+        ip: OPIO,
+    ) -> OPIO:
+        r"""Execute the OP.
+
+        Parameters
+        ----------
+        ip : dict
+            Input dict with components:
+            - `config`: (`dict`) The config of calypso task to obtain the command of calypso.
+            - `task_name` : (`str`)
+            - `finished` : (`str`)
+            - `cnt_num` : (`int`)
+            - `task_dir` : (`Path`)
+
+        Returns
+        -------
+        op : dict
+            Output dict with components:
+
+            - `task_name`: (`str`)
+            - `optim_results_dir`: (`List[str]`)
+            - `traj_results`: (`Artifact(List[Path])`)
+        """
+        finished = ip["finished"]
+        cnt_num = ip["cnt_num"]
+
+        task_path = ip["task_dir"]
+        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 {}
+        # PrepCalyDPOptim keeps the backend-specific model filename. Build the
+        # default command from the model that was actually prepared.
+        model_name = next(
+            (
+                candidate
+                for candidate in ("frozen_model.pb", "model.ckpt.pt")
+                if any(path.name == candidate for path in input_files)
+            ),
+            "model.ckpt.pt",
+        )
+        command = config.get(
+            "run_opt_command", f"python -u calypso_run_opt.py {model_name}"
+        )
+
+        work_dir = Path(ip["task_name"])
+
+        with set_directory(work_dir):
+            # link input files
+            for ii in input_files:
+                iname = ii.name
+                Path(iname).symlink_to(ii)
+
+            if finished == "false":
+                ret, out, err = run_command(command, shell=True)
+                if ret != 0:
+                    logging.error(
+                        "".join(
+                            (
+                                "opt failed\n",
+                                "\ncommand was: ",
+                                command,
+                                "\nout msg: ",
+                                out,
+                                "\n",
+                                "\nerr msg: ",
+                                err,
+                                "\n",
+                            )
+                        )
+                    )
+                    raise TransientError("opt failed")
+
+                optim_results_dir = Path("optim_results_dir")
+                optim_results_dir.mkdir(parents=True, exist_ok=True)
+                for poscar in Path().glob("POSCAR_*"):
+                    target = optim_results_dir.joinpath(poscar.name)
+                    shutil.copyfile(poscar, target)
+                for contcar in Path().glob("CONTCAR_*"):
+                    target = optim_results_dir.joinpath(contcar.name)
+                    shutil.copyfile(contcar, target)
+                for outcar in Path().glob("OUTCAR_*"):
+                    target = optim_results_dir.joinpath(outcar.name)
+                    shutil.copyfile(outcar, target)
+
+                traj_results_dir = Path("traj_results")
+                traj_results_dir.mkdir(parents=True, exist_ok=True)
+                for traj in Path().glob("*.traj"):
+                    target = traj_results_dir.joinpath(str(cnt_num) + "." + traj.name)
+                    shutil.copyfile(traj, target)
+
+            else:
+                optim_results_dir = Path("optim_results_dir")
+                optim_results_dir.mkdir(parents=True, exist_ok=True)
+                traj_results_dir = Path("traj_results")
+                traj_results_dir.mkdir(parents=True, exist_ok=True)
+
+        return OPIO(
+            {
+                "task_name": str(work_dir),
+                "optim_results_dir": work_dir / optim_results_dir,
+                "traj_results": work_dir / traj_results_dir,
+            }
+        )

tests/op/test_run_caly_dp_optim.py

@@ -0,0 +1,172 @@
+import os
+import shutil
+import unittest
+from pathlib import (
+    Path,
+)
+
+import numpy as np
+from dflow.python import (
+    OP,
+    OPIO,
+    Artifact,
+    OPIOSign,
+    TransientError,
+)
+from mock import (
+    call,
+    mock,
+    patch,
+)
+
+# isort: off
+from .context import (
+    dpgen2,
+)
+from dpgen2.constants import (
+    calypso_task_pattern,
+    model_name_pattern,
+    calypso_run_opt_file,
+    calypso_check_opt_file,
+)
+from dpgen2.op import RunCalyDPOptim
+from dpgen2.utils import (
+    BinaryFileInput,
+)
+
+# isort: on
+
+
+class TestRunDPOptim(unittest.TestCase):
+    def setUp(self):
+        self.task_name = "task_dir"
+        self.task_dir = Path("pre_run_optim")
+
+        self.task_dir.mkdir(parents=True, exist_ok=True)
+        self.task_dir.joinpath("frozen_model.pb").write_text("1")
+        self.task_dir.joinpath("calypso_run_opt.py").write_text("1")
+        self.task_dir.joinpath("calypso_check_opt.py").write_text("1")
+        for i in range(1, 11):
+            self.task_dir.joinpath(f"POSCAR_{i}").write_text("1")
+
+    def tearDown(self):
+        shutil.rmtree(self.task_name)
+        shutil.rmtree(self.task_dir)
+
+    @patch("dpgen2.op.run_caly_dp_optim.run_command")
+    def test_00_success(self, mocked_run):
+        def side_effect(*args, **kwargs):
+            for i in range(1, 11):
+                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")
+            return (0, "foo\n", "")
+
+        mocked_run.side_effect = side_effect
+        op = RunCalyDPOptim()
+        out = op.execute(
+            OPIO(
+                {
+                    "config": {"run_calypso_command": "echo 1"},
+                    "task_name": self.task_name,
+                    "finished": "false",
+                    "cnt_num": 0,
+                    "task_dir": self.task_dir,
+                }
+            )
+        )
+        # check output
+        self.assertEqual(out["task_name"], self.task_name)
+
+        optim_results_dir = out["optim_results_dir"]
+        list_optim_results_dir = list(optim_results_dir.iterdir())
+        counts_optim_results_dir = len(list_optim_results_dir)
+        counts_outcar_in_optim_results_dir = len(
+            list(optim_results_dir.rglob("OUTCAR_*"))
+        )
+
+        self.assertTrue(optim_results_dir, Path(self.task_name) / "optim_results_dir")
+        self.assertEqual(counts_optim_results_dir, 30)
+        self.assertEqual(counts_outcar_in_optim_results_dir, 10)
+        self.assertTrue(
+            Path(self.task_name) / "optim_results_dir" / "CONTCAR_4"
+            in list_optim_results_dir
+        )
+
+        traj_results_dir = out["traj_results"]
+        list_traj_results_dir = list(traj_results_dir.glob("*.traj"))
+        counts_traj = len(list_traj_results_dir)
+        self.assertEqual(traj_results_dir, Path(self.task_name) / "traj_results")
+        self.assertEqual(counts_traj, 10)
+        self.assertTrue(
+            Path(self.task_name) / "traj_results" / "0.3.traj" in list_traj_results_dir
+        )
+        mocked_run.assert_called_once_with(
+            "python -u calypso_run_opt.py frozen_model.pb", shell=True
+        )
+
+    @patch("dpgen2.op.run_caly_dp_optim.run_command")
+    def test_01_error(self, mocked_run):
+        def side_effect(*args, **kwargs):
+            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")
+            return (1, "foo\n", "")
+
+        mocked_run.side_effect = side_effect
+        op = RunCalyDPOptim()
+        self.assertRaises(
+            TransientError,
+            op.execute,
+            OPIO(
+                {
+                    "config": {"run_calypso_command": "echo 1"},
+                    "task_name": self.task_name,
+                    "finished": "false",
+                    "cnt_num": 0,
+                    "task_dir": self.task_dir,
+                }
+            ),
+        )
+
+    def test_02_success(self):
+        op = RunCalyDPOptim()
+        out = op.execute(
+            OPIO(
+                {
+                    "config": {"run_calypso_command": "echo 1"},
+                    "task_name": self.task_name,
+                    "finished": "true",
+                    "cnt_num": 0,
+                    "task_dir": self.task_dir,
+                }
+            )
+        )
+        # check output
+        self.assertEqual(out["task_name"], self.task_name)
+
+        optim_results_dir = out["optim_results_dir"]
+        list_optim_results_dir = list(optim_results_dir.iterdir())
+        counts_optim_results_dir = len(list_optim_results_dir)
+        counts_outcar_in_optim_results_dir = len(
+            list(optim_results_dir.rglob("OUTCAR_*"))
+        )
+
+        self.assertTrue(optim_results_dir, Path(self.task_name) / "optim_results_dir")
+        self.assertEqual(counts_optim_results_dir, 0)
+        self.assertEqual(counts_outcar_in_optim_results_dir, 0)
+        self.assertTrue(
+            Path(self.task_name) / "optim_results_dir" / "CONTCAR_4"
+            not in list_optim_results_dir
+        )
+
+        traj_results_dir = out["traj_results"]
+        list_traj_results_dir = list(traj_results_dir.glob("*.traj"))
+        counts_traj = len(list_traj_results_dir)
+        self.assertEqual(traj_results_dir, Path(self.task_name) / "traj_results")
+        self.assertEqual(counts_traj, 0)
+        self.assertTrue(
+            Path(self.task_name) / "traj_results" / "0.3.traj"
+            not in list_traj_results_dir
+        )