From 41d14eac2d55a05aadbab8f7d9d7e0cdf2e9a951 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Mon, 14 Feb 2022 16:36:20 +0800
Subject: [PATCH] format the index of train and lmp

---
 dpgen2/constants.py              |  6 ++++--
 dpgen2/flow/prep_run_dp_train.py | 23 ++++++++++++++++-------
 dpgen2/flow/prep_run_lmp.py      | 22 ++++++++++++++--------
 3 files changed, 34 insertions(+), 17 deletions(-)

diff --git a/dpgen2/constants.py b/dpgen2/constants.py
index e102ff51..3792687e 100644
--- a/dpgen2/constants.py
+++ b/dpgen2/constants.py
@@ -1,8 +1,10 @@
-train_task_pattern = 'task.%04d'
+train_index_pattern = "%04d"
+train_task_pattern = 'task.' + train_index_pattern
 train_script_name = 'input.json'
 train_log_name = 'train.log'
 model_name_pattern = 'model.%03d.pb'
-lmp_task_pattern = 'task.%06d'
+lmp_index_pattern = "%06d"
+lmp_task_pattern = 'task.' + lmp_index_pattern
 lmp_conf_name = 'conf.lmp'
 lmp_input_name = 'in.lammps'
 lmp_traj_name = 'traj.dump'
diff --git a/dpgen2/flow/prep_run_dp_train.py b/dpgen2/flow/prep_run_dp_train.py
index 532e80ab..384efa2e 100644
--- a/dpgen2/flow/prep_run_dp_train.py
+++ b/dpgen2/flow/prep_run_dp_train.py
@@ -11,6 +11,8 @@
     upload_artifact,
     download_artifact,
     argo_range,
+    argo_len,
+    argo_sequence,
 )
 from dflow.python import(
     PythonOPTemplate,
@@ -20,7 +22,9 @@
     Artifact,
     Slices,
 )
-
+from dpgen2.constants import (
+    train_index_pattern,
+)
 import os
 from typing import Set, List
 from pathlib import Path
@@ -67,10 +71,14 @@ def __init__(
         
         self._keys = ['prep-train', 'run-train']
         self.step_keys = {}
-        for ii in self._keys:
-            self.step_keys[ii] = '--'.join(
-                ["%s"%self.inputs.parameters["block_id"], ii]
-            )
+        ii = 'prep-train'
+        self.step_keys[ii] = '--'.join(
+            ["%s"%self.inputs.parameters["block_id"], ii]
+        )
+        ii = 'run-train'
+        self.step_keys[ii] = '--'.join(
+            ["%s"%self.inputs.parameters["block_id"], ii + "-{{item}}"]
+        )
 
         self = _prep_run_dp_train(
             self, 
@@ -138,7 +146,7 @@ def _prep_run_dp_train(
             run_train_op,
             image=run_train_image,
             slices = Slices(
-                "{{item}}",
+                "int('{{item}}')",
                 input_parameter = ["task_name"],
                 input_artifact = ["task_path", "init_model"],
                 output_artifact = ["model", "lcurve", "log", "script"],
@@ -155,7 +163,8 @@ def _prep_run_dp_train(
             "init_data": train_steps.inputs.artifacts['init_data'],
             "iter_data": train_steps.inputs.artifacts['iter_data'],
         },
-        with_param=argo_range(train_steps.inputs.parameters["numb_models"]),
+        with_sequence=argo_sequence(argo_len(prep_train.outputs.parameters["task_names"]), format=train_index_pattern),
+        # with_param=argo_range(train_steps.inputs.parameters["numb_models"]),
         key = step_keys['run-train'],
     )
     train_steps.add(run_train)
diff --git a/dpgen2/flow/prep_run_lmp.py b/dpgen2/flow/prep_run_lmp.py
index c4ebb4a4..27accccd 100644
--- a/dpgen2/flow/prep_run_lmp.py
+++ b/dpgen2/flow/prep_run_lmp.py
@@ -22,7 +22,9 @@
     Artifact,
     Slices,
 )
-
+from dpgen2.constants import (
+    lmp_index_pattern,
+)
 import os
 from typing import Set, List
 from pathlib import Path
@@ -69,10 +71,14 @@ def __init__(
         
         self._keys = ['prep-lmp', 'run-lmp']
         self.step_keys = {}
-        for ii in self._keys:
-            self.step_keys[ii] = '--'.join(
-                ["%s"%self.inputs.parameters["block_id"], ii]
-            )
+        ii = 'prep-lmp'
+        self.step_keys[ii] = '--'.join(
+            ["%s"%self.inputs.parameters["block_id"], ii]
+        )
+        ii = 'run-lmp'
+        self.step_keys[ii] = '--'.join(
+            ["%s"%self.inputs.parameters["block_id"], ii + "-{{item}}"]
+        )
 
         self = _prep_run_lmp(
             self, 
@@ -140,7 +146,7 @@ def _prep_run_lmp(
             run_op,
             image=run_image,
             slices = Slices(
-                "{{item}}",
+                "int('{{item}}')",
                 input_parameter = ["task_name"],
                 input_artifact = ["task_path"],
                 output_artifact = ["log", "traj", "model_devi"],
@@ -155,8 +161,8 @@ def _prep_run_lmp(
             'task_path' : prep_lmp.outputs.artifacts['task_paths'],
             "models" : prep_run_steps.inputs.artifacts['models'],
         },
-        # with_sequence=argo_sequence(argo_len(prep_lmp.outputs.parameters["task_names"])),
-        with_param=argo_range(argo_len(prep_lmp.outputs.parameters["task_names"])),
+        with_sequence=argo_sequence(argo_len(prep_lmp.outputs.parameters["task_names"]), format=lmp_index_pattern),
+        # with_param=argo_range(argo_len(prep_lmp.outputs.parameters["task_names"])),
         key = step_keys['run-lmp'],
     )
     prep_run_steps.add(run_lmp)