Skip to content

Commit

Permalink
format the index of train and lmp
Browse files Browse the repository at this point in the history
  • Loading branch information
Han Wang committed Feb 14, 2022
1 parent 553d45d commit 41d14ea
Show file tree
Hide file tree
Showing 3 changed files with 34 additions and 17 deletions.
6 changes: 4 additions & 2 deletions dpgen2/constants.py
Original file line number Diff line number Diff line change
@@ -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'
Expand Down
23 changes: 16 additions & 7 deletions dpgen2/flow/prep_run_dp_train.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@
upload_artifact,
download_artifact,
argo_range,
argo_len,
argo_sequence,
)
from dflow.python import(
PythonOPTemplate,
Expand All @@ -20,7 +22,9 @@
Artifact,
Slices,
)

from dpgen2.constants import (
train_index_pattern,
)
import os
from typing import Set, List
from pathlib import Path
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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"],
Expand All @@ -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)
Expand Down
22 changes: 14 additions & 8 deletions dpgen2/flow/prep_run_lmp.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,9 @@
Artifact,
Slices,
)

from dpgen2.constants import (
lmp_index_pattern,
)
import os
from typing import Set, List
from pathlib import Path
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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"],
Expand All @@ -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)
Expand Down

0 comments on commit 41d14ea

Please sign in to comment.