Run lammps extra outputs

dpgen2/op/run_lmp.py

@@ -83,6 +83,7 @@ def get_output_sign(cls):
                 "model_devi": Artifact(Path),
                 "plm_output": Artifact(Path, optional=True),
                 "optional_output": Artifact(Path, optional=True),
+                "extra_outputs": Artifact(List[Path]),
             }
         )
 
@@ -213,6 +214,10 @@ def execute(
         if ele_temp is not None:
             ret_dict["optional_output"] = work_dir / "job.json"
 
+        extra_outputs = []
+        for fname in config["extra_output_files"]:
+            extra_outputs += list(work_dir.glob(fname))
+        ret_dict["extra_outputs"] = extra_outputs  # type: ignore
         return OPIO(ret_dict)
 
     def get_model_devi(self, model_devi_file):
@@ -226,6 +231,7 @@ def lmp_args():
         doc_head = "Select a head from multitask"
         doc_use_ele_temp = "Whether to use electronic temperature, 0 for no, 1 for frame temperature, and 2 for atomic temperature"
         doc_use_hdf5 = "Use HDF5 to store trajs and model_devis"
+        doc_extra_output_files = "Extra output file names, support wildcards"
         return [
             Argument("command", str, optional=True, default="lmp", doc=doc_lmp_cmd),
             Argument(
@@ -256,6 +262,13 @@ def lmp_args():
                 default=False,
                 doc=doc_use_hdf5,
             ),
+            Argument(
+                "extra_output_files",
+                list,
+                optional=True,
+                default=[],
+                doc=doc_extra_output_files,
+            ),
         ]
 
     @staticmethod

dpgen2/superop/prep_run_lmp.py

@@ -76,6 +76,7 @@ def __init__(
             "model_devis": OutputArtifact(),
             "plm_output": OutputArtifact(),
             "optional_outputs": OutputArtifact(),
+            "extra_outputs": OutputArtifact(),
         }
 
         super().__init__(
@@ -179,6 +180,7 @@ def _prep_run_lmp(
                     "model_devi",
                     "plm_output",
                     "optional_output",
+                    "extra_outputs",
                 ],
                 **template_slice_config,
             ),
@@ -217,5 +219,8 @@ def _prep_run_lmp(
     prep_run_steps.outputs.artifacts[
         "optional_outputs"
     ]._from = run_lmp.outputs.artifacts["optional_output"]
+    prep_run_steps.outputs.artifacts["extra_outputs"]._from = run_lmp.outputs.artifacts[
+        "extra_outputs"
+    ]
 
     return prep_run_steps

dpgen2/utils/download_dpgen2_artifacts.py

@@ -63,7 +63,8 @@ def add_output(
     "prep-run-explore": DownloadDefinition()
     .add_output("logs")
     .add_output("trajs")
-    .add_output("model_devis"),
+    .add_output("model_devis")
+    .add_output("extra_outputs"),
     "prep-run-fp": DownloadDefinition()
     .add_input("confs")
     .add_output("logs")

tests/mocked_ops.py

@@ -424,6 +424,7 @@ def execute(
                 "log": work_dir / log,
                 "traj": work_dir / traj,
                 "model_devi": work_dir / model_devi,
+                "extra_outputs": [],
             }
         )
 

tests/op/test_run_lmp.py

@@ -126,6 +126,29 @@ def test_error(self, mocked_run):
         ]
         mocked_run.assert_has_calls(calls)
 
+    def test_extra_outputs(self):
+        op = RunLmp()
+        out = op.execute(
+            OPIO(
+                {
+                    "config": {
+                        "command": "echo Hello > foo.txt",
+                        "extra_output_files": ["foo.txt"],
+                    },
+                    "task_name": self.task_name,
+                    "task_path": self.task_path,
+                    "models": self.models,
+                }
+            )
+        )
+        work_dir = Path(self.task_name)
+        # check output
+        self.assertEqual(out["extra_outputs"], [work_dir / "foo.txt"])
+        self.assertEqual(
+            (work_dir / "foo.txt").read_text().strip(),
+            "Hello -i in.lammps -log log.lammps",
+        )
+
 
 class TestRunLmpDist(unittest.TestCase):
     lmp_config = """variable        NSTEPS          equal 1000

tests/utils/test_dl_dpgen2_arti.py

@@ -123,6 +123,11 @@ def test_lmp_download(self, mocked_dl):
                 path=Path("iter-000001/prep-run-explore/outputs"),
                 skip_exists=True,
             ),
+            mock.call(
+                "arti-extra_outputs",
+                path=Path("iter-000001/prep-run-explore/outputs"),
+                skip_exists=True,
+            ),
         ]
         self.assertEqual(len(mocked_dl.call_args_list), len(expected))
         for ii, jj in zip(mocked_dl.call_args_list, expected):
@@ -253,6 +258,11 @@ def test_update_finished_steps_exist_steps(self, mocked_dl):
                 path=Path("iter-000001/prep-run-explore/outputs"),
                 skip_exists=True,
             ),
+            mock.call(
+                "arti-extra_outputs",
+                path=Path("iter-000001/prep-run-explore/outputs"),
+                skip_exists=True,
+            ),
         ]
         self.assertEqual(len(mocked_dl.call_args_list), len(expected))
         for ii, jj in zip(mocked_dl.call_args_list, expected):
@@ -315,6 +325,11 @@ def test_update_finished_steps_none_steps(self, mocked_dl):
                 path=Path("iter-000001/prep-run-explore/outputs"),
                 skip_exists=True,
             ),
+            mock.call(
+                "arti-extra_outputs",
+                path=Path("iter-000001/prep-run-explore/outputs"),
+                skip_exists=True,
+            ),
         ]
         for ii, jj in zip(mocked_dl.call_args_list, expected):
             self.assertEqual(ii, jj)