diff --git a/dpgen2/op/run_caly_dp_optim.py b/dpgen2/op/run_caly_dp_optim.py new file mode 100644 index 00000000..ff3facdd --- /dev/null +++ b/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, + } + ) diff --git a/tests/op/test_run_caly_dp_optim.py b/tests/op/test_run_caly_dp_optim.py new file mode 100644 index 00000000..a49abd52 --- /dev/null +++ b/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 + )