def add_modify_incar(original_wf,
modify_incar_params=None,
fw_name_constraint=None):
"""
Every FireWork that runs VASP has a ModifyIncar task just beforehand. For example, allows
you to modify the INCAR based on the Worker using env_chk or using hard-coded changes.
Args:
original_wf (Workflow)
modify_incar_params (dict) - dict of parameters for ModifyIncar.
fw_name_constraint (str) - Only apply changes to FWs where fw_name contains this substring.
"""
modify_incar_params = modify_incar_params or {
"incar_update": ">>incar_update<<"
}
wf_dict = original_wf.to_dict()
for idx_fw, idx_t in get_fws_and_tasks(
original_wf,
fw_name_constraint=fw_name_constraint,
task_name_constraint="RunVasp"):
wf_dict["fws"][idx_fw]["spec"]["_tasks"].insert(
idx_t,
ModifyIncar(**modify_incar_params).to_dict())
return Workflow.from_dict(wf_dict)
def add_tags(original_wf, tags_list):
"""
Adds tags to all Fireworks in the Workflow, WF metadata,
as well as additional_fields for the VaspDrone to track them later
(e.g. all fireworks and vasp tasks related to a research project)
Args:
original_wf (Workflow)
tags_list: list of tags parameters (list of strings)
"""
wf_dict = original_wf.to_dict()
# WF metadata
if "tags" in wf_dict["metadata"]:
wf_dict["metadata"]["tags"].extend(tags_list)
else:
wf_dict["metadata"]["tags"] = tags_list
# FW metadata
for idx_fw in range(len(original_wf.fws)):
if "tags" in wf_dict["fws"][idx_fw]["spec"]:
wf_dict["fws"][idx_fw]["spec"]["tags"].extend(tags_list)
else:
wf_dict["fws"][idx_fw]["spec"]["tags"] = tags_list
# Drone
for idx_fw, idx_t in get_fws_and_tasks(original_wf, task_name_constraint="VaspToDbTask"):
if "tags" in wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t]["additional_fields"]:
wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t]["additional_fields"]["tags"].extend(tags_list)
else:
wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t]["additional_fields"]["tags"] = tags_list
return Workflow.from_dict(wf_dict)
def use_custodian(original_wf, fw_name_constraint=None, custodian_params=None):
"""
Replaces all tasks with "RunVasp*" (e.g. RunVaspDirect) to be
RunVaspCustodian. Thus, this powerup adds error correction into VASP
runs if not originally present and/or modifies the correction behavior.
Args:
original_wf (Workflow): original workflow
fw_name_constraint (str): Only apply changes to FWs where fw_name contains this substring.
For example, use custodian only for certain runs, or set job_type to
"double_relaxation_run" only for structure optimization run, or set different
handler_group for different runs.
custodian_params (dict): A dict of parameters for RunVaspCustodian. e.g., use it to set
a "scratch_dir" or "handler_group".
"""
custodian_params = custodian_params if custodian_params else {}
wf_dict = original_wf.to_dict()
vasp_fws_and_tasks = get_fws_and_tasks(
original_wf,
fw_name_constraint=fw_name_constraint,
task_name_constraint="RunVasp")
for idx_fw, idx_t in vasp_fws_and_tasks:
if "vasp_cmd" not in custodian_params:
custodian_params["vasp_cmd"] = wf_dict["fws"][idx_fw]["spec"][
"_tasks"][idx_t]["vasp_cmd"]
wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t] = RunVaspCustodian(
**custodian_params).to_dict()
return Workflow.from_dict(wf_dict)
def to_fireworks(self, method='from dict'):
""" Returns a fireworks workflow object """
from fireworks import Firework, Workflow
if method == 'from dict':
fws = []
for step in self.get_steps():
spec = step.copy()
for key in ['id', 'data', 'name', 'inputs', 'outputs']:
del spec[key]
spec = translate_keys(spec, 'tasks', '_tasks')
spec = translate_keys(spec, 'name', '_fw_name')
fws.append({
'name': step['name'],
'fw_id': step['id'],
'spec': spec
})
dct = {
'fws': fws,
'links': self.get_ctrlflow_links_dict(),
'name': self['name'],
'metadata': {}
}
return Workflow.from_dict(dct)
if method == 'from object':
return Workflow(
fireworks=[Firework(step) for step in self.get_steps()],
links_dict=self.get_ctrlflow_links_dict(),
name=self['name']
)
def use_fake_vasp(original_wf, ref_dirs, params_to_check=None):
"""
Replaces all tasks with "RunVasp" (e.g. RunVaspDirect) to be
RunVaspFake. Thus, we do not actually run VASP but copy
pre-determined inputs and outputs.
Args:
original_wf (Workflow)
ref_dirs (dict): key=firework name, value=path to the reference vasp calculation directory
params_to_check (list): optional list of incar parameters to check.
"""
if not params_to_check:
params_to_check = [
"ISPIN", "ENCUT", "ISMEAR", "SIGMA", "IBRION", "LORBIT", "NBANDS",
"LMAXMIX"
]
wf_dict = original_wf.to_dict()
for idx_fw, fw in enumerate(original_wf.fws):
for job_type in ref_dirs.keys():
if job_type in fw.name:
for idx_t, t in enumerate(fw.tasks):
if "RunVasp" in str(t):
wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t] = \
RunVaspFake(ref_dir=ref_dirs[job_type], params_to_check=params_to_check).to_dict()
return Workflow.from_dict(wf_dict)
def use_custodian(original_wf, fw_name_constraint=None, custodian_params=None):
"""
Replaces all tasks with "RunVasp*" (e.g. RunVaspDirect) to be
RunVaspCustodian. Thus, this powerup adds error correction into VASP
runs if not originally present and/or modifies the correction behavior.
Args:
original_wf (Workflow): original workflow
fw_name_constraint (str): Only apply changes to FWs where fw_name contains this substring.
For example, use custodian only for certain runs, or set job_type to
"double_relaxation_run" only for structure optimization run, or set different
handler_group for different runs.
custodian_params (dict): A dict of parameters for RunVaspCustodian. e.g., use it to set
a "scratch_dir" or "handler_group".
"""
custodian_params = custodian_params if custodian_params else {}
wf_dict = original_wf.to_dict()
vasp_fws_and_tasks = get_fws_and_tasks(original_wf, fw_name_constraint=fw_name_constraint,
task_name_constraint="RunVasp")
for idx_fw, idx_t in vasp_fws_and_tasks:
if "vasp_cmd" not in custodian_params:
custodian_params["vasp_cmd"] = \
wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t]["vasp_cmd"]
wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t] = \
RunVaspCustodian(**custodian_params).to_dict()
return Workflow.from_dict(wf_dict)
def use_scratch_dir(original_wf, scratch_dir):
"""
For all RunVaspCustodian tasks, add the desired scratch dir.
:param original_wf:
:param scratch_dir: The scratch dir to use. Supports env_chk
"""
wf_dict = original_wf.to_dict()
for idx_fw, idx_t in get_fws_and_tasks(original_wf, task_name_constraint="RunVaspCustodian"):
wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t]["scratch_dir"] = scratch_dir
return Workflow.from_dict(wf_dict)
def add_additional_fields_to_taskdocs(original_wf, update_dict=None):
"""
For all VaspToDbTasks in a given workflow, add information
to "additional_fields" to be placed in the task doc.
Args:
original_wf (Workflow)
update_dict (Dict): dictionary to add additional_fields
"""
wf_dict = original_wf.to_dict()
for idx_fw, idx_t in get_fws_and_tasks(original_wf, task_name_constraint="VaspToDbTask"):
wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t]["additional_fields"].update(update_dict)
return Workflow.from_dict(wf_dict)
def use_scratch_dir(original_wf, scratch_dir):
"""
For all RunVaspCustodian tasks, add the desired scratch dir.
:param original_wf:
:param scratch_dir: The scratch dir to use. Supports env_chk
"""
wf_dict = original_wf.to_dict()
for idx_fw, idx_t in get_fws_and_tasks(
original_wf, task_name_constraint="RunVaspCustodian"):
wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t][
"scratch_dir"] = scratch_dir
return Workflow.from_dict(wf_dict)
def add_additional_fields_to_taskdocs(original_wf, update_dict=None):
"""
For all VaspToDbTasks in a given workflow, add information
to "additional_fields" to be placed in the task doc.
Args:
original_wf (Workflow)
update_dict (Dict): dictionary to add additional_fields
"""
wf_dict = original_wf.to_dict()
for idx_fw, idx_t in get_fws_and_tasks(
original_wf, task_name_constraint="VaspToDbTask"):
wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t][
"additional_fields"].update(update_dict)
return Workflow.from_dict(wf_dict)
def update_wf(wf):
"""
Simple helper to ensure that the powerup updates to the workflow dict has taken effect.
This is needed because all the powerups that modify workflow do so on the dict representation
of the workflow(or mix thereof eg: add tasks as dict to the fireworks spec etc) and for
inspection the powerups rely on a mix of object and dict representations of workflow object(
along with the constituent fireworks and firetasks) that are not in one to one correspondence
with the updated dict representation.
Args:
wf (Workflow)
Returns:
Workflow
"""
return Workflow.from_dict(wf.as_dict())
def add_trackers(original_wf):
"""
Every FireWork that runs VASP also tracks the OUTCAR and OSZICAR using FWS Trackers.
Args:
original_wf (Workflow)
"""
tracker1 = Tracker('OUTCAR', nlines=25, allow_zipped=True)
tracker2 = Tracker('OSZICAR', nlines=25, allow_zipped=True)
wf_dict = original_wf.to_dict()
for idx_fw, idx_t in get_fws_and_tasks(original_wf, task_name_constraint="RunVasp"):
if "_trackers" in wf_dict["fws"][idx_fw]["spec"]:
wf_dict["fws"][idx_fw]["spec"]["_trackers"].extend([tracker1, tracker2])
else:
wf_dict["fws"][idx_fw]["spec"]["_trackers"] = [tracker1, tracker2]
return Workflow.from_dict(wf_dict)
def add_modify_incar(original_wf, modify_incar_params=None, fw_name_constraint=None):
"""
Every FireWork that runs VASP has a ModifyIncar task just beforehand. For example, allows
you to modify the INCAR based on the Worker using env_chk or using hard-coded changes.
Args:
original_wf (Workflow)
modify_incar_params (dict) - dict of parameters for ModifyIncar.
fw_name_constraint (str) - Only apply changes to FWs where fw_name contains this substring.
"""
modify_incar_params = modify_incar_params or {"incar_update": ">>incar_update<<"}
wf_dict = original_wf.to_dict()
for idx_fw, idx_t in get_fws_and_tasks(original_wf, fw_name_constraint=fw_name_constraint,
task_name_constraint="RunVasp"):
wf_dict["fws"][idx_fw]["spec"]["_tasks"].insert(idx_t, ModifyIncar(**modify_incar_params).to_dict())
return Workflow.from_dict(wf_dict)
def remove_custodian(original_wf, fw_name_constraint=None):
"""
Replaces all tasks with "RunVasp*" (e.g. RunVaspCustodian) to be
RunVaspDirect.
Args:
original_wf (Workflow): original workflow
fw_name_constraint (str): Only apply changes to FWs where fw_name
contains this substring.
"""
wf_dict = original_wf.to_dict()
vasp_fws_and_tasks = get_fws_and_tasks(original_wf, fw_name_constraint=fw_name_constraint,
task_name_constraint="RunVasp")
for idx_fw, idx_t in vasp_fws_and_tasks:
vasp_cmd = wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t]["vasp_cmd"]
wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t] = RunVaspDirect(vasp_cmd=vasp_cmd).to_dict()
return Workflow.from_dict(wf_dict)
def add_small_gap_multiply(original_wf, gap_cutoff, density_multiplier, fw_name_constraint=None):
"""
In all FWs with specified name constraints, add a 'small_gap_multiply' parameter that
multiplies the k-mesh density of compounds with gap < gap_cutoff by density multiplier.
Note that this powerup only works on FireWorks with the appropriate WriteVasp* tasks that
accept the small_gap_multiply argument...
:param original_wf:
:param gap_cutoff:
:param density_multiplier:
:param fw_name_constraint:
"""
wf_dict = original_wf.to_dict()
for idx_fw, idx_t in get_fws_and_tasks(original_wf, fw_name_constraint=fw_name_constraint,
task_name_constraint="WriteVasp"):
wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t]["small_gap_multiply"] = [gap_cutoff, density_multiplier]
return Workflow.from_dict(wf_dict)
def tag_fws(original_wf, tag, fw_name_constraint=None):
"""
Tags VASP Fworker(s) of a Workflow; e.g. it can be used to run large-memory jobs on a separate queue
Args:
original_wf (Workflow):
tag (string): user-defined tag to be added under fw.spec._fworker (e.g. "large memory", "big", etc)
fw_name_constraint (string): name of the fireworks to be modified (all if None is passed)
Returns:
modified workflow with tagged Fworkers
"""
wf_dict = original_wf.to_dict()
for idx_fw, idx_t in get_fws_and_tasks(original_wf, fw_name_constraint=fw_name_constraint,
task_name_constraint="RunVasp"):
wf_dict["fws"][idx_fw]["spec"]["_fworker"] = tag
return Workflow.from_dict(wf_dict)
def add_namefile(original_wf, use_slug=True):
"""
Every FireWork begins by writing an empty file with the name
"FW--<fw.name>". This makes it easy to figure out what jobs are in what
launcher directories, e.g. "ls -l launch*/FW--*" from within a "block" dir.
Args:
original_wf (Workflow)
use_slug (bool): whether to replace whitespace-type chars with a slug
"""
wf_dict = original_wf.to_dict()
for idx, fw in enumerate(wf_dict["fws"]):
fname = "FW--{}".format(fw["name"])
if use_slug:
fname = get_slug(fname)
wf_dict["fws"][idx]["spec"]["_tasks"].insert(0, FileWriteTask(
files_to_write=[{"filename": fname, "contents": ""}]).to_dict())
return Workflow.from_dict(wf_dict)
def modify_to_soc(original_wf, nbands, structure=None, modify_incar_params=None, fw_name_constraint=None):
"""
Takes a regular workflow and transforms its VASP fireworkers that are specified with
fw_name_constraints to non-collinear calculations taking spin orbit coupling into account.
Args:
original_wf (Workflow)
nbands (int): number of bands selected by the user (for now)
structure (Structure)
modify_incar_params ({}): a dictionary containing the setting for modyfining the INCAR (e.g. {"ICHARG": 11})
fw_name_constraint (string): name of the fireworks to be modified (all if None is passed)
Returns:
modified Workflow with SOC
"""
wf_dict = original_wf.to_dict()
if structure is None:
try:
sid = get_fws_and_tasks(original_wf, fw_name_constraint="structure optimization",
task_name_constraint="RunVasp")[0][0]
structure = Structure.from_dict(wf_dict["fws"][sid]["spec"]["_tasks"][1]["vasp_input_set"]["structure"])
except:
raise ValueError("For this workflow, the structure must be provided as an input")
magmom = ""
for i in structure:
magmom += "0 0 0.6 "
# TODO: add saxis as an input parameter with default being (0 0 1)
modify_incar_params = modify_incar_params or {"incar_update": {"LSORBIT": "T", "NBANDS": nbands, "MAGMOM": magmom,
"ISPIN": 1, "LMAXMIX": 4, "ISYM": 0}}
for idx_fw, idx_t in get_fws_and_tasks(original_wf, fw_name_constraint=fw_name_constraint,
task_name_constraint="RunVasp"):
if "nscf" in wf_dict["fws"][idx_fw]["name"]:
wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t]["vasp_cmd"] = ">>vasp_ncl<<"
wf_dict["fws"][idx_fw]["spec"]["_tasks"].insert(idx_t, ModifyIncar(**modify_incar_params).to_dict())
wf_dict["fws"][idx_fw]["name"] += " soc"
for idx_fw, idx_t in get_fws_and_tasks(original_wf, fw_name_constraint=fw_name_constraint,
task_name_constraint="RunBoltztrap"):
wf_dict["fws"][idx_fw]["name"] += " soc"
return Workflow.from_dict(wf_dict)
def add_trackers(original_wf, tracked_files=None, nlines=25):
"""
Every FireWork that runs VASP also tracks the OUTCAR, OSZICAR, etc using FWS Trackers.
Args:
original_wf (Workflow)
tracked_files (list) : list of files to be tracked
nlines (int): number of lines at the end of files to be tracked
"""
if tracked_files is None:
tracked_files = ["OUTCAR", "OSZICAR"]
trackers = [Tracker(f, nlines=nlines, allow_zipped=True) for f in tracked_files]
wf_dict = original_wf.to_dict()
for idx_fw, idx_t in get_fws_and_tasks(original_wf, task_name_constraint="RunVasp"):
if "_trackers" in wf_dict["fws"][idx_fw]["spec"]:
wf_dict["fws"][idx_fw]["spec"]["_trackers"].extend(trackers)
else:
wf_dict["fws"][idx_fw]["spec"]["_trackers"] = trackers
return Workflow.from_dict(wf_dict)
def use_fake_vasp(original_wf):
"""
Replaces all tasks with "RunVasp" (e.g. RunVaspDirect) to be
RunVaspFake. Thus, we do not actually run VASP but copy
pre-determined inputs and outputs.
Args:
original_wf (Workflow)
"""
wf_dict = original_wf.to_dict()
for idx_fw, fw in enumerate(original_wf.fws):
for job_type in fake_dirs.keys():
if job_type in fw.name:
for idx_t, t in enumerate(fw.tasks):
if "RunVasp" in str(t):
wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t] = \
RunVaspFake(fake_dir=fake_dirs[job_type]).to_dict()
return Workflow.from_dict(wf_dict)
def remove_custodian(original_wf, fw_name_constraint=None):
"""
Replaces all tasks with "RunVasp*" (e.g. RunVaspCustodian) to be
RunVaspDirect.
Args:
original_wf (Workflow): original workflow
fw_name_constraint (str): Only apply changes to FWs where fw_name
contains this substring.
"""
wf_dict = original_wf.to_dict()
vasp_fws_and_tasks = get_fws_and_tasks(
original_wf,
fw_name_constraint=fw_name_constraint,
task_name_constraint="RunVasp")
for idx_fw, idx_t in vasp_fws_and_tasks:
vasp_cmd = wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t]["vasp_cmd"]
wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t] = RunVaspDirect(
vasp_cmd=vasp_cmd).to_dict()
return Workflow.from_dict(wf_dict)
def add_trackers(original_wf, tracked_files=None, nlines=25):
"""
Every FireWork that runs VASP also tracks the OUTCAR, OSZICAR, etc using FWS Trackers.
Args:
original_wf (Workflow)
tracked_files (list) : list of files to be tracked
nlines (int): number of lines at the end of files to be tracked
"""
if tracked_files == None:
tracked_files = ["OUTCAR", "OSZICAR"]
trackers = [
Tracker(f, nlines=nlines, allow_zipped=True) for f in tracked_files
]
wf_dict = original_wf.to_dict()
for idx_fw, idx_t in get_fws_and_tasks(original_wf,
task_name_constraint="RunVasp"):
if "_trackers" in wf_dict["fws"][idx_fw]["spec"]:
wf_dict["fws"][idx_fw]["spec"]["_trackers"].extend(trackers)
else:
wf_dict["fws"][idx_fw]["spec"]["_trackers"] = trackers
return Workflow.from_dict(wf_dict)
def add_namefile(original_wf, use_slug=True):
"""
Every FireWork begins by writing an empty file with the name
"FW--<fw.name>". This makes it easy to figure out what jobs are in what
launcher directories, e.g. "ls -l launch*/FW--*" from within a "block" dir.
Args:
original_wf (Workflow)
use_slug (bool): whether to replace whitespace-type chars with a slug
"""
wf_dict = original_wf.to_dict()
for idx, fw in enumerate(wf_dict["fws"]):
fname = "FW--{}".format(fw["name"])
if use_slug:
fname = get_slug(fname)
wf_dict["fws"][idx]["spec"]["_tasks"].insert(
0,
FileWriteTask(files_to_write=[{
"filename": fname,
"contents": ""
}]).to_dict())
return Workflow.from_dict(wf_dict)
def remove_leaf_fws(orig_wf):
"""
Remove the end nodes(last fireworks) from the given workflow.
Args:
orig_wf (Workflow): The original workflow object.
Returns:
Workflow : the new updated workflow.
"""
wf_dict = orig_wf.as_dict()
all_parents = []
for i, f in enumerate(orig_wf.as_dict()["fws"]):
if f["fw_id"] in orig_wf.leaf_fw_ids:
parents = orig_wf.links.parent_links[int(f["fw_id"])]
all_parents.extend(parents)
del wf_dict["links"][str(f["fw_id"])]
del wf_dict["fws"][i]
for p in parents:
wf_dict["links"][str(p)] = []
new_wf = Workflow.from_dict(wf_dict)
return update_wf(new_wf)
def use_fake_vasp(original_wf, ref_dirs, params_to_check=None):
"""
Replaces all tasks with "RunVasp" (e.g. RunVaspDirect) to be
RunVaspFake. Thus, we do not actually run VASP but copy
pre-determined inputs and outputs.
Args:
original_wf (Workflow)
ref_dirs (dict): key=firework name, value=path to the reference vasp calculation directory
params_to_check (list): optional list of incar parameters to check.
"""
if not params_to_check:
params_to_check = ["ISPIN", "ENCUT", "ISMEAR", "SIGMA", "IBRION", "LORBIT", "NBANDS", "LMAXMIX"]
wf_dict = original_wf.to_dict()
for idx_fw, fw in enumerate(original_wf.fws):
for job_type in ref_dirs.keys():
if job_type in fw.name:
for idx_t, t in enumerate(fw.tasks):
if "RunVasp" in str(t):
wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t] = \
RunVaspFake(ref_dir=ref_dirs[job_type], params_to_check=params_to_check).to_dict()
return Workflow.from_dict(wf_dict)
def add_small_gap_multiply(original_wf,
gap_cutoff,
density_multiplier,
fw_name_constraint=None):
"""
In all FWs with specified name constraints, add a 'small_gap_multiply' parameter that
multiplies the k-mesh density of compounds with gap < gap_cutoff by density multiplier.
Note that this powerup only works on FireWorks with the appropriate WriteVasp* tasks that
accept the small_gap_multiply argument...
:param original_wf:
:param gap_cutoff:
:param density_multiplier:
:param fw_name_constraint:
"""
wf_dict = original_wf.to_dict()
for idx_fw, idx_t in get_fws_and_tasks(
original_wf,
fw_name_constraint=fw_name_constraint,
task_name_constraint="WriteVasp"):
wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t][
"small_gap_multiply"] = [gap_cutoff, density_multiplier]
return Workflow.from_dict(wf_dict)
def add_tags(original_wf, tags_list):
"""
Adds tags to all Fireworks in the Workflow, WF metadata,
as well as additional_fields for the VaspDrone to track them later
(e.g. all fireworks and vasp tasks related to a research project)
Args:
original_wf (Workflow)
tags_list: list of tags parameters (list of strings)
"""
wf_dict = original_wf.to_dict()
# WF metadata
if "tags" in wf_dict["metadata"]:
wf_dict["metadata"]["tags"].extend(tags_list)
else:
wf_dict["metadata"]["tags"] = tags_list
# FW metadata
for idx_fw in range(len(original_wf.fws)):
if "tags" in wf_dict["fws"][idx_fw]["spec"]:
wf_dict["fws"][idx_fw]["spec"]["tags"].extend(tags_list)
else:
wf_dict["fws"][idx_fw]["spec"]["tags"] = tags_list
# Drone
for idx_fw, idx_t in get_fws_and_tasks(
original_wf, task_name_constraint="VaspToDbTask"):
if "tags" in wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t][
"additional_fields"]:
wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t][
"additional_fields"]["tags"].extend(tags_list)
else:
wf_dict["fws"][idx_fw]["spec"]["_tasks"][idx_t][
"additional_fields"]["tags"] = tags_list
return Workflow.from_dict(wf_dict)
def run_task(self, fw_spec):
import gromacs
from ..general import make_md_workflow
# bit of an ad-hoc way to grab the checkpoint file
cpt = [
f for f in self['files'] if (('cpt' in f) and ('prev' not in f))
]
if len(cpt) > 1:
raise ValueError("Multiple CPT files in 'files'; include "
"only one.")
elif len(cpt) < 1:
raise ValueError("No CPT file in 'files'; "
"cannot do continue check.")
else:
cpt = os.path.join(self['archive'], cpt[0])
# bit of an ad-hoc way to grab the tpr file
tpr = [f for f in self['files'] if ('tpr' in f)]
if len(tpr) > 1:
raise ValueError("Multiple TPR files in 'files'; include "
"only one.")
elif len(tpr) < 1:
raise ValueError("No TPR file in 'files'; "
"cannot do continue check.")
else:
tpr = os.path.join(self['archive'], tpr[0])
# let's extract the current frame and place it in the archive, since
# this is useful for starting runs up at any point from the current end
gromacs.trjconv(f=cpt,
s=tpr,
o=os.path.join(
self['archive'], '{}.gro'.format(
os.path.splitext(os.path.basename(tpr))[0])),
input=('0', ))
# extract step number from CPT file
out = gromacs.dump(cp=cpt, stdout=False)
step = int([
line.split(' ')[-1] for line in out[1].split('\n')
if 'step = ' in line
][0])
# extract nsteps from TPR file
out = gromacs.dump(s=tpr, stdout=False)
nsteps = int([
line.split(' ')[-1] for line in out[1].split('\n')
if 'nsteps' in line
][0])
# if step < nsteps, we submit a new workflow
if step < nsteps:
wf = make_md_workflow(sim=self['sim'],
archive=self['archive'],
stages=self['stages'],
files=self['files'],
md_engine=self['md_engine'],
md_category=self['md_category'],
local_category=self['local_category'],
postrun_wf=self['postrun_wf'],
post_wf=self['post_wf'])
return FWAction(additions=[wf])
else:
sim = mds.Sim(self['sim'])
sim.categories['md_status'] = 'finished'
# if given, we submit the post workflow
post_wf = self.get('post_wf')
if post_wf:
if isinstance(post_wf, dict):
post_wf = Workflow.from_dict(post_wf)
# this makes a fresh copy without already-used fw_ids
post_wf = Workflow.from_wflow(post_wf)
return FWAction(additions=[post_wf])
def _copy_wf(self, wf):
return Workflow.from_dict(wf.to_dict())
def copy_wf(wf):
return Workflow.from_dict(wf.to_dict())
def _copy_wf(self, wf):
return Workflow.from_dict(wf.to_dict())
def make_md_workflow(sim,
archive,
stages,
md_engine='gromacs',
md_category='md',
local_category='local',
postrun_wf=None,
post_wf=None,
files=None):
"""Construct a general, single MD simulation workflow.
Assumptions
-----------
Queue launcher submission script must define and export the following
environment variables:
1. STAGING : absolute path on resource to staging directory
2. SCRATCH : absolute path on resource to scratch directory
The staging directory must already exist on all resources specified in
``stages``.
The script ``run_md.sh`` must be somewhere on your path, and must take
a single argument giving the directory to execute MD out of. It should
create and change the working directory to that directory before anything
else.
Parameters
----------
sim : str
MDSynthesis Sim.
archive : str
Absolute path to directory to launch from, which holds all required
files for running MD.
stages : list, str
Dicts giving for each of the following keys:
- 'server': server host to transfer to
- 'user': username to authenticate with
- 'staging': absolute path to staging area on remote resource
alternatively, a path to a yaml file giving a list of dictionaries
with the same information.
md_engine : {'gromacs'}
MD engine name; needed to determine continuation mechanism to use.
md_category : str
Category to use for the MD Firework. Used to target to correct rockets.
local_category : str
Category to use for non-MD Fireworks, which should be run by rockets
where the ``archive`` directory is accessible.
postrun_wf : Workflow
Workflow to perform after each copyback; performed in parallel to continuation run.
post_wf : Workflow
Workflow to perform after completed MD (no continuation); use for final
postprocessing.
files : list
Names of files (not paths) needed for each leg of the simulation. Need
not exist, but if they do they will get staged before each run.
Returns
-------
workflow
MD workflow; can be submitted to LaunchPad of choice.
"""
sim = mds.Sim(sim)
#TODO: perhaps move to its own FireTask?
sim.categories['md_status'] = 'running'
#TODO: the trouble with this is that if this workflow is created with the intent
# of being attached to another, these files may not exist at all yet
f_exist = [f for f in files if os.path.exists(os.path.join(archive, f))]
if isinstance(stages, string_types):
with open(stages, 'r') as f:
stages = yaml.load(f)
## Stage files on all resources where MD may run; takes place locally
fts_stage = list()
for stage in stages:
fts_stage.append(
FileTransferTask(mode='rtransfer',
server=stage['server'],
user=stage['user'],
files=[os.path.join(archive, i) for i in files],
dest=os.path.join(stage['staging'], sim.uuid),
max_retry=5,
shell_interpret=True))
fw_stage = Firework(fts_stage,
spec={
'_launch_dir': archive,
'_category': local_category
},
name='staging')
## MD execution; takes place in queue context of compute resource
# make rundir
ft_mkdir = MkRunDirTask(uuid=sim.uuid)
# copy input files to scratch space
ft_copy = FileTransferTask(
mode='copy',
files=[os.path.join('${STAGING}/', sim.uuid, i) for i in files],
dest=os.path.join('${SCRATCHDIR}/', sim.uuid),
ignore_missing=True,
shell_interpret=True)
# next, run MD
ft_md = ScriptTask(script='run_md.sh {}'.format(
os.path.join('${SCRATCHDIR}/', sim.uuid)),
use_shell=True,
fizzle_bad_rc=True)
# send info on where files live to pull firework
ft_info = BeaconTask(uuid=sim.uuid)
fw_md = Firework([ft_mkdir, ft_copy, ft_md, ft_info],
spec={'_category': md_category},
name='md')
## Pull files back to archive; takes place locally
ft_copyback = FilePullTask(dest=archive)
fw_copyback = Firework([ft_copyback],
spec={
'_launch_dir': archive,
'_category': local_category
},
name='pull')
## Decide if we need to continue and submit new workflow if so; takes place
## locally
if md_engine == 'gromacs':
ft_continue = GromacsContinueTask(sim=sim,
archive=archive,
stages=stages,
md_engine=md_engine,
md_category=md_category,
local_category=local_category,
postrun_wf=postrun_wf,
post_wf=post_wf,
files=files)
else:
raise ValueError("No known md engine `{}`.".format(md_engine))
fw_continue = Firework([ft_continue],
spec={
'_launch_dir': archive,
'_category': local_category
},
name='continue')
wf = Workflow([fw_stage, fw_md, fw_copyback, fw_continue],
links_dict={
fw_stage: [fw_md],
fw_md: [fw_copyback],
fw_copyback: [fw_continue]
},
name='{} | md'.format(sim.name),
metadata=dict(sim.categories))
## Mix in postrun workflow, if given
if postrun_wf:
if isinstance(postrun_wf, dict):
postrun_wf = Workflow.from_dict(postrun_wf)
wf.append_wf(Workflow.from_wflow(postrun_wf), [fw_copyback.fw_id])
return wf
