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 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 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_bandgap_check(original_wf,
check_bandgap_params=None,
fw_name_constraint=None):
"""
Every FireWork that runs VASP has a CheckStability task afterward. This
allows defusing jobs that are not stable. In practice, you might want
to set the fw_name_constraint so that the stability is only checked at the
beginning of the workflow
Args:
original_wf (Workflow)
check_bandgap_params (dict): a **kwargs** style dict of params
fw_name_constraint (str) - Only apply changes to FWs where fw_name contains this substring.
"""
check_bandgap_params = check_bandgap_params or {}
for idx_fw, idx_t in get_fws_and_tasks(
original_wf,
fw_name_constraint=fw_name_constraint,
task_name_constraint="DbTask"):
original_wf.fws[idx_fw].spec["_tasks"].append(
CheckBandgap(**check_bandgap_params).to_dict())
return update_wf(original_wf)
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_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)