Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
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)
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 "structure" not in wf_dict["fws"][idx_fw]["name"] and "static" not 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)
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.
Returns:
Workflow
"""
modify_incar_params = modify_incar_params or {
"incar_update": ">>incar_update<<"
}
idx_list = get_fws_and_tasks(
original_wf,
fw_name_constraint=fw_name_constraint,
task_name_constraint="RunVasp",
)
for idx_fw, idx_t in idx_list:
original_wf.fws[idx_fw].tasks.insert(
idx_t, ModifyIncar(**modify_incar_params)
)
return original_wf
def add_stability_check(original_wf, check_stability_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_stability_params (dict): a **kwargs** style dict of params
fw_name_constraint (str) - Only apply changes to FWs where fw_name contains this substring.
"""
check_stability_params = check_stability_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(CheckStability(**check_stability_params).to_dict())
return update_wf(original_wf)
None is passed)
task_name_constraint (str): name of the Firetasks to be tagged (e.g.
None or 'RunVasp')
Returns:
Workflow: workflow with modified queue options
"""
qsettings = {}
if walltime:
qsettings.update({"walltime": walltime})
if time_min:
qsettings.update({"time_min": time_min})
if qos:
qsettings.update({"qos": qos})
idx_list = get_fws_and_tasks(
original_wf,
fw_name_constraint=fw_name_constraint,
task_name_constraint=task_name_constraint,
)
for idx_fw, idx_t in idx_list:
original_wf.fws[idx_fw].spec.update({"_queueadapter": qsettings})
return original_wf
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 clear_modify(original_wf, fw_name_constraint=None):
"""
Simple powerup that clears the modifications to a workflow.
Args:
original_wf (Workflow): The original workflow.
fw_name_constraint (str): name constraint for fireworks to
have their modification tasks removed
"""
idx_list = get_fws_and_tasks(
original_wf,
fw_name_constraint=fw_name_constraint,
task_name_constraint="Modify",
)
idx_list.reverse()
for idx_fw, idx_t in idx_list:
original_wf.fws[idx_fw].tasks.pop(idx_t)
return original_wf
def add_additional_fields_to_taskdocs(
original_wf, update_dict=None, task_name_constraint="VaspToDb"
):
"""
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
task_name_constraint (str): name of the Firetasks to be modified.
Returns:
Workflow
"""
idx_list = get_fws_and_tasks(
original_wf, task_name_constraint=task_name_constraint
)
for idx_fw, idx_t in idx_list:
original_wf.fws[idx_fw].tasks[idx_t]["additional_fields"].update(
update_dict
)
return original_wf
task_name_constraint="RunVasp",
):
"""
Cleans up files after another fireworks. Default behavior is to remove
WAVECAR after running VASP.
Args:
original_wf (Workflow)
files (list): list of patterns to match for files to clean up
fw_name_constraint (str): pattern for fireworks to clean up files after
task_name_constraint (str): pattern for firetask to clean up files
Returns:
Workflow
"""
idx_list = get_fws_and_tasks(
original_wf,
fw_name_constraint=fw_name_constraint,
task_name_constraint=task_name_constraint,
)
for idx_fw, idx_t in idx_list:
original_wf.fws[idx_fw].tasks.insert(
idx_t + 1, DeleteFiles(files=files)
)
return original_wf