Repository: deepmodeling/dpgen2 (language: python)

Issue:
Implement the op RunLmp
Implement the op `RunLmp`. This OP run a LAMMPS MD task prepared by `PrepLmp`, and outputs a MD trajectory and the model deviation of each configuration on the trajectory.

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_lmp.py#L12-L18) 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_lmp.py#L44-L66).

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 `RunLmp`


Discussion / hints:

Fix patch:
diff --git a/dpgen2/op/run_lmp.py b/dpgen2/op/run_lmp.py
index e88bc1a0..6c354d5d 100644
--- a/dpgen2/op/run_lmp.py
+++ b/dpgen2/op/run_lmp.py
@@ -1,12 +1,33 @@
+import os, json
+from pathlib import Path
 from dflow.python import (
     OP,
     OPIO,
     OPIOSign,
-    Artifact
+    Artifact,
+    TransientError,
+    FatalError,
+)
+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 (
+    lmp_conf_name,
+    lmp_input_name,
+    lmp_log_name,
+    lmp_traj_name,
+    lmp_model_devi_name,
+)
+from dargs import (
+    dargs, 
+    Argument, 
+    Variant, 
+    ArgumentEncoder,
 )
-import os, json
-from typing import Tuple, List, Set
-from pathlib import Path
 
 class RunLmp(OP):
     r"""Execute a LAMMPS task.
@@ -48,6 +69,7 @@ def execute(
         ip : dict
             Input dict with components:
         
+            - `config`: (`dict`) The config of lmp task. Check `RunLmp.lmp_args` for definitions.
             - `task_name`: (`str`) The name of the task.
             - `task_path`: (`Artifact(Path)`) The path that contains all input files prepareed by `PrepLmp`.
             - `models`: (`Artifact(List[Path])`) The frozen model to estimate the model deviation. The first model with be used to drive molecular dynamics simulation.
@@ -65,6 +87,52 @@ def execute(
         TransientError
             On the failure of LAMMPS execution. Handle different failure cases? e.g. loss atoms.
         """
-        raise NotImplementedError
+        config = ip['config'] if ip['config'] is not None else {}
+        config = RunLmp.normalize_config(config)
+        command = config['command']
+        task_name = ip['task_name']
+        task_path = ip['task_path']
+        models = ip['models']
+        input_files = [lmp_conf_name, lmp_input_name]
+        input_files = [(Path(task_path)/ii).resolve() for ii in input_files]
+        model_files = [Path(ii).resolve() for ii in models]
+        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)
+            # link models
+            for mm in model_files:
+                mname = mm.name
+                Path(mname).symlink_to(mm)
+            # run lmp
+            command = [command, '-i', lmp_input_name, '-log', lmp_log_name]
+            ret, out, err = run_command(command)
+            if ret != 0:
+                raise TransientError('lmp failed')
+            
+        return OPIO({
+            "log" : work_dir / lmp_log_name,
+            "traj" : work_dir / lmp_traj_name,
+            "model_devi" : work_dir / lmp_model_devi_name,
+        })        
+
+            
+    @staticmethod
+    def lmp_args():
+        doc_lmp_cmd = "The command of LAMMPS"
+        return [
+            Argument("command", str, optional=True, default='lmp', doc=doc_lmp_cmd),
+        ]
 
+    @staticmethod
+    def normalize_config(data = {}):
+        ta = RunLmp.lmp_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_lmp.py b/tests/op/test_run_lmp.py
new file mode 100644
index 00000000..4a944b30
--- /dev/null
+++ b/tests/op/test_run_lmp.py
@@ -0,0 +1,90 @@
+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 (
+    lmp_conf_name,
+    lmp_input_name,
+    lmp_log_name,
+    lmp_traj_name,
+    lmp_model_devi_name,
+)
+from dpgen2.op.run_lmp import RunLmp
+
+
+class TestRunLmp(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/lmp_conf_name).write_text('foo')
+        (self.task_path/lmp_input_name).write_text('bar')
+        self.task_name = 'task_000'
+        self.models = [self.model_path/Path(f'model_{ii}.pb') for ii in range(4)]
+        for idx,ii in enumerate(self.models):
+            ii.write_text(f'model{idx}')
+        
+    def tearDown(self):
+        if Path('task').is_dir():
+            shutil.rmtree('task')
+        if Path('models').is_dir():
+            shutil.rmtree('models')
+        if Path(self.task_name).is_dir():
+            shutil.rmtree(self.task_name)
+
+    @patch('dpgen2.op.run_lmp.run_command')
+    def test_success(self, mocked_run):
+        mocked_run.side_effect = [ (0, 'foo\n', '') ]
+        op = RunLmp()
+        out = op.execute(
+            OPIO({
+                'config' : {'command' : 'mylmp'},
+                'task_name' : self.task_name,
+                'task_path' : self.task_path,
+                'models' : self.models,
+            }))
+        work_dir = Path(self.task_name)
+        # check output
+        self.assertEqual(out['log'], work_dir/lmp_log_name)
+        self.assertEqual(out['traj'], work_dir/lmp_traj_name)
+        self.assertEqual(out['model_devi'], work_dir/lmp_model_devi_name)
+        # check call
+        calls = [
+            call(['mylmp', '-i', lmp_input_name, '-log', lmp_log_name]),
+        ]
+        mocked_run.assert_has_calls(calls)
+        # check input files are correctly linked
+        self.assertEqual((work_dir/lmp_conf_name).read_text(), 'foo')
+        self.assertEqual((work_dir/lmp_input_name).read_text(), 'bar')
+        for ii in range(4):
+            self.assertEqual((work_dir/f'model_{ii}.pb').read_text(), f'model{ii}')
+
+    
+    @patch('dpgen2.op.run_lmp.run_command')
+    def test_error(self, mocked_run):
+        mocked_run.side_effect = [ (1, 'foo\n', '') ]
+        op = RunLmp()
+        with self.assertRaises(TransientError) as ee:
+            out = op.execute(
+                OPIO({
+                    'config' : {'command' : 'mylmp'},
+                    'task_name' : self.task_name,
+                    'task_path' : self.task_path,
+                    'models' : self.models,
+                }))
+        # check call
+        calls = [
+            call(['mylmp', '-i', lmp_input_name, '-log', lmp_log_name]),
+        ]
+        mocked_run.assert_has_calls(calls)
+                        
+