Repository: deepmodeling/dpgen2 (language: python)

Issue:
Implement the op RunVasp
Implement the op RunVasp. This OP run a VASP DFT task prepared by PrepVasp, and outputs the labeled configuration (coordinate, simulation cell, dft energy, force and virial) in the `deepmd/npy` data format provided by [`dpdata`](https://github.com/deepmodeling/dpdata/).

We have to implement the execute method of the class. What the op does is explained in the [docstr of the class](https://github.com/wanghan-iapcm/dpgen2/blob/ef08f44f7a39c81fe469885ffaa4cd8be9d4d0a3/dpgen2/op/run_vasp.py#L12-L20) and the interface of the execute method is provided in [the docstr of the method](https://github.com/wanghan-iapcm/dpgen2/blob/ef08f44f7a39c81fe469885ffaa4cd8be9d4d0a3/dpgen2/op/run_vasp.py#L42-L63).

One can also take the [implementation of RunDPTrain](https://github.com/wanghan-iapcm/dpgen2/blob/master/dpgen2/op/run_dp_train.py) as an example for the RunVasp


Discussion / hints:

Fix patch:
diff --git a/dpgen2/constants.py b/dpgen2/constants.py
index 90b03f32..e36adc43 100644
--- a/dpgen2/constants.py
+++ b/dpgen2/constants.py
@@ -16,4 +16,6 @@
 vasp_input_name = 'INCAR'
 vasp_pot_name = 'POTCAR'
 vasp_kp_name = 'KPOINTS'
+vasp_default_log_name = 'vasp.log'
+vasp_default_out_data_name = 'data'
 
diff --git a/dpgen2/op/run_vasp.py b/dpgen2/op/run_vasp.py
index fa385f27..830baa0e 100644
--- a/dpgen2/op/run_vasp.py
+++ b/dpgen2/op/run_vasp.py
@@ -2,11 +2,33 @@
     OP,
     OPIO,
     OPIOSign,
-    Artifact
+    Artifact,
+    TransientError,
+    FatalError,
 )
-import os, json
-from typing import Tuple, List, Set
+import os, json, dpdata
 from pathlib import Path
+from typing import (
+    Tuple, 
+    List, 
+    Set,
+)
+from dpgen2.utils.run_command import run_command
+from dpgen2.utils.chdir import set_directory
+from dpgen2.constants import(
+    vasp_conf_name,
+    vasp_input_name,
+    vasp_pot_name,
+    vasp_kp_name,
+    vasp_default_log_name,
+    vasp_default_out_data_name,
+)
+from dargs import (
+    dargs, 
+    Argument, 
+    Variant, 
+    ArgumentEncoder,
+)
 
 class RunVasp(OP):
     r"""Execute a VASP task.
@@ -22,9 +44,9 @@ class RunVasp(OP):
     @classmethod
     def get_input_sign(cls):
         return OPIOSign({
+            "config" : dict,
             "task_name": str,
             "task_path" : Artifact(Path),
-            "config" : dict,
         })
 
     @classmethod
@@ -46,6 +68,7 @@ def execute(
         ip : dict
             Input dict with components:
         
+            - `config`: (`dict`) The config of vasp task. Check `RunVasp.vasp_args` for definitions.
             - `task_name`: (`str`) The name of task.
             - `task_path`: (`Artifact(Path)`) The path that contains all input files prepareed by `PrepVasp`.
 
@@ -61,5 +84,54 @@ def execute(
         TransientError
             On the failure of VASP execution. 
         """
-        raise NotImplementedError
+        config = ip['config'] if ip['config'] is not None else {}
+        config = RunVasp.normalize_config(config)
+        command = config['command']
+        log_name = config['log']
+        out_name = config['out']
+        task_name = ip['task_name']
+        task_path = ip['task_path']
+        input_files = [vasp_conf_name, vasp_input_name, vasp_pot_name, vasp_kp_name]
+        input_files = [(Path(task_path)/ii).resolve() for ii in input_files]
+        work_dir = Path(task_name)
+
+        with set_directory(work_dir):
+            # link input files
+            for ii in input_files:
+                iname = ii.name
+                Path(iname).symlink_to(ii)
+            # run vasp
+            command = [command, '>', log_name]
+            ret, out, err = run_command(command)
+            if ret != 0:
+                raise TransientError('vasp failed')
+            # convert the output to deepmd/npy format
+            sys = dpdata.LabeledSystem('OUTCAR')
+            sys.to('deepmd/npy', out_name)
+
+        return OPIO({
+            "log": work_dir / log_name,
+            "labeled_data": work_dir / out_name,
+        })
+
+
+    @staticmethod
+    def vasp_args():
+        doc_vasp_cmd = "The command of VASP"
+        doc_vasp_log = "The log file name of VASP"
+        doc_vasp_out = "The output dir name of labeled data. In `deepmd/npy` format provided by `dpdata`."
+        return [
+            Argument("command", str, optional=True, default='vasp', doc=doc_vasp_cmd),
+            Argument("log", str, optional=True, default=vasp_default_log_name, doc=doc_vasp_log),
+            Argument("out", str, optional=True, default=vasp_default_out_data_name, doc=doc_vasp_out),
+        ]
+
+    @staticmethod
+    def normalize_config(data = {}):
+        ta = RunVasp.vasp_args()
+        base = Argument("base", dict, ta)
+        data = base.normalize_value(data, trim_pattern="_*")
+        base.check_value(data, strict=True)
+        return data
 
+    


Test patch:
diff --git a/tests/op/test_run_vasp.py b/tests/op/test_run_vasp.py
new file mode 100644
index 00000000..17d99992
--- /dev/null
+++ b/tests/op/test_run_vasp.py
@@ -0,0 +1,143 @@
+from op.context import dpgen2
+import numpy as np
+import unittest, json, shutil
+from mock import mock, patch, call
+from dflow.python import (
+    OP,
+    OPIO,
+    OPIOSign,
+    Artifact,
+    TransientError,
+)
+from pathlib import Path
+from dpgen2.constants import (
+    vasp_conf_name,
+    vasp_input_name,
+    vasp_pot_name,
+    vasp_kp_name,
+    vasp_default_log_name,
+    vasp_default_out_data_name,
+)
+from dpgen2.op.run_vasp import RunVasp
+
+
+class TestRunVasp(unittest.TestCase):
+    def setUp(self):
+        self.task_path = Path('task/path')
+        self.task_path.mkdir(parents=True, exist_ok=True)
+        self.model_path = Path('models/path')
+        self.model_path.mkdir(parents=True, exist_ok=True)
+        (self.task_path/vasp_conf_name).write_text('foo')
+        (self.task_path/vasp_input_name).write_text('bar')
+        (self.task_path/vasp_pot_name).write_text('dee')
+        (self.task_path/vasp_kp_name).write_text('por')
+        self.task_name = 'task_000'
+        
+    def tearDown(self):
+        if Path('task').is_dir():
+            shutil.rmtree('task')
+        if Path(self.task_name).is_dir():
+            shutil.rmtree(self.task_name)
+
+    @patch('dpgen2.op.run_vasp.run_command')
+    def test_success(self, mocked_run):
+        mocked_run.side_effect = [ (0, 'foo\n', '') ]
+        op = RunVasp()
+        def new_to(obj, foo, bar):
+            data_path = Path('data')
+            data_path.mkdir()
+            (data_path/'foo').write_text('bar')
+        def new_init(obj, foo):
+            pass
+        with mock.patch.object(dpgen2.op.run_vasp.dpdata.LabeledSystem, 
+                               'to', 
+                               new=new_to):
+            with mock.patch.object(dpgen2.op.run_vasp.dpdata.LabeledSystem, 
+                                   '__init__', 
+                                   new=new_init):
+                out = op.execute(
+                    OPIO({
+                        'config' : {'command' : 'myvasp',
+                                    'log': 'foo.log', 
+                                    'out':'data',
+                                    },
+                        'task_name' : self.task_name,
+                        'task_path' : self.task_path,
+                    }))
+        work_dir = Path(self.task_name)
+        # check output
+        self.assertEqual(out['log'], work_dir/'foo.log')
+        self.assertEqual(out['labeled_data'], work_dir/'data')
+        # check call
+        calls = [
+            call(['myvasp', '>', 'foo.log']),
+        ]
+        mocked_run.assert_has_calls(calls)
+        # check input files are correctly linked
+        self.assertEqual((work_dir/vasp_conf_name).read_text(), 'foo')
+        self.assertEqual((work_dir/vasp_input_name).read_text(), 'bar')
+        self.assertEqual((work_dir/vasp_pot_name).read_text(), 'dee')
+        self.assertEqual((work_dir/vasp_kp_name).read_text(), 'por')
+        # check output
+        self.assertEqual((Path(self.task_name)/'data'/'foo').read_text(), 'bar')
+
+
+    @patch('dpgen2.op.run_vasp.run_command')
+    def test_success(self, mocked_run):
+        mocked_run.side_effect = [ (0, 'foo\n', '') ]
+        op = RunVasp()
+        def new_to(obj, foo, bar):
+            data_path = Path('data')
+            data_path.mkdir()
+            (data_path/'foo').write_text('bar')
+        def new_init(obj, foo):
+            pass
+        with mock.patch.object(dpgen2.op.run_vasp.dpdata.LabeledSystem, 
+                               'to', 
+                               new=new_to):
+            with mock.patch.object(dpgen2.op.run_vasp.dpdata.LabeledSystem, 
+                                   '__init__', 
+                                   new=new_init):
+                out = op.execute(
+                    OPIO({
+                        'config' : {'command' : 'myvasp',
+                                    },
+                        'task_name' : self.task_name,
+                        'task_path' : self.task_path,
+                    }))
+        work_dir = Path(self.task_name)
+        # check output
+        self.assertEqual(out['log'], work_dir/vasp_default_log_name)
+        self.assertEqual(out['labeled_data'], work_dir/vasp_default_out_data_name)
+        # check call
+        calls = [
+            call(['myvasp', '>', vasp_default_log_name]),
+        ]
+        mocked_run.assert_has_calls(calls)
+        # check input files are correctly linked
+        self.assertEqual((work_dir/vasp_conf_name).read_text(), 'foo')
+        self.assertEqual((work_dir/vasp_input_name).read_text(), 'bar')
+        self.assertEqual((work_dir/vasp_pot_name).read_text(), 'dee')
+        self.assertEqual((work_dir/vasp_kp_name).read_text(), 'por')
+        # check output
+        self.assertEqual((Path(self.task_name)/'data'/'foo').read_text(), 'bar')
+    
+
+    @patch('dpgen2.op.run_vasp.run_command')
+    def test_error(self, mocked_run):
+        mocked_run.side_effect = [ (1, 'foo\n', '') ]
+        op = RunVasp()
+        with self.assertRaises(TransientError) as ee:
+            out = op.execute(
+                OPIO({
+                    'config' : {'command' : 'myvasp'},
+                    'task_name' : self.task_name,
+                    'task_path' : self.task_path,
+                }))
+        # check call
+        calls = [
+            call(['myvasp', '>', vasp_default_log_name]),
+        ]
+        mocked_run.assert_has_calls(calls)
+                        
+