def use_potcar_spec(original_wf,
fw_name_constraint=None,
vasp_to_db_kwargs=None):
"""
In all WriteVasp tasks, enable the potcar_spec option. In this mode,
POTCAR files will be written as POTCAR.spec files, containing only the
atomic symbols. Furthermore, POTCAR files will not be parsed by the
VaspToDb drone.
The primary use case for this powerup is to enable easier testing of
atomate workflows. Typically, writing VaspInputSets requires having the
VASP pseudopotentials installed. Due to licensing restraints, the
VASP pseudopotentials are not installed in the atomate testing environment.
Use of this powerup therefore enables testing of atomate workflows in the
absence of installed pseudopotentials.
Note: this powerup should also be combined with RunFakeVasp with
check_potcar set to False.
Args:
original_wf (Workflow)
fw_name_constraint (str): Only apply changes to FWs where fw_name
contains this substring.
vasp_to_db_kwargs (dict): Additional kwargs to pass to VaspToDb.
Returns:
Workflow
"""
idx_list = get_fws_and_tasks(
original_wf,
fw_name_constraint=fw_name_constraint,
task_name_constraint="WriteVasp",
)
idx_list.extend(
get_fws_and_tasks(
original_wf,
fw_name_constraint=fw_name_constraint,
task_name_constraint="UpdateScanRelaxBandgap",
))
for idx_fw, idx_t in idx_list:
original_wf.fws[idx_fw].tasks[idx_t]["potcar_spec"] = True
idx_list = get_fws_and_tasks(
original_wf,
fw_name_constraint=fw_name_constraint,
task_name_constraint="VaspToDb",
)
vasp_to_db_kwargs = vasp_to_db_kwargs if vasp_to_db_kwargs else {}
for idx_fw, idx_t in idx_list:
original_wf.fws[idx_fw].tasks[idx_t]["parse_potcar_file"] = False
original_wf.fws[idx_fw].tasks[idx_t].update(vasp_to_db_kwargs)
return original_wf
def set_queue_adapter(
original_wf: Workflow,
queueadapter: dict = None,
fw_name_constraint: str = None,
task_name_constraint: str = None,
) -> Workflow:
"""
set _queueadapter spec of Fireworker(s) of a Workflow. It can be used to change the overall queueadapter during the run.
Args:
original_wf (Workflow): workflow that will be changed
queueadapter (dict): dict to change _queueadapter
fw_name_constraint (str): name of the Fireworks to be tagged (all if None is passed)
task_name_constraint (str): name of the Firetasks to be tagged (e.g. None or 'RunVasp')
Returns:
Workflow: modified workflow with specified Fireworkers tagged
"""
for idx_fw, idx_t in get_fws_and_tasks(
original_wf,
fw_name_constraint=fw_name_constraint,
task_name_constraint=task_name_constraint,
):
q = original_wf.fws[idx_fw].spec.get("_queueadapter", {})
q.update(queueadapter)
original_wf.fws[idx_fw].spec["_queueadapter"] = q
return 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
Returns:
Workflow
"""
if tracked_files is None:
tracked_files = ["OUTCAR", "OSZICAR"]
trackers = [
Tracker(f, nlines=nlines, allow_zipped=True) for f in tracked_files
]
idx_list = get_fws_and_tasks(original_wf, task_name_constraint="RunVasp")
for idx_fw, idx_t in idx_list:
if "_trackers" in original_wf.fws[idx_fw].spec:
original_wf.fws[idx_fw].spec["_trackers"].extend(trackers)
else:
original_wf.fws[idx_fw].spec["_trackers"] = trackers
return original_wf
def add_modify_kpoints(
original_wf, modify_kpoints_params=None, fw_name_constraint=None
):
"""
Every FireWork that runs VASP has a ModifyKpoints task just beforehand. For
example, allows you to modify the KPOINTS based on the Worker using env_chk
or using hard-coded changes.
Args:
original_wf (Workflow)
modify_kpoints_params (dict): dict of parameters for ModifyKpoints.
fw_name_constraint (str): Only apply changes to FWs where fw_name
contains this substring.
Returns:
Workflow
"""
modify_kpoints_params = modify_kpoints_params or {
"kpoints_update": ">>kpoints_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, ModifyKpoints(**modify_kpoints_params)
)
return original_wf
def clean_up_files(
original_wf,
files=("WAVECAR*",),
fw_name_constraint=None,
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
def add_bandgap_check(original_wf,
check_bandgap_params=None,
fw_name_constraint=None):
"""
Every FireWork that enters into the Db has a band gap check afterwards,
e.g. min_gap and max_gap
Args:
original_wf (Workflow)
check_bandgap_params (dict): a **kwargs** style dict of params, e.g.
min_gap or max_gap
fw_name_constraint (str) - Only apply changes to FWs where fw_name
contains this substring.
Returns:
Workflow
"""
check_bandgap_params = check_bandgap_params or {}
idx_list = get_fws_and_tasks(
original_wf,
fw_name_constraint=fw_name_constraint,
task_name_constraint="VaspToDb",
)
for idx_fw, idx_t in idx_list:
original_wf.fws[idx_fw].tasks.append(
CheckBandgap(**check_bandgap_params))
return original_wf
def add_stability_check(original_wf,
check_stability_params=None,
fw_name_constraint=None):
"""
Every FireWork that enters into the Db 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.
Returns:
Workflow
"""
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="VaspToDb"):
original_wf.fws[idx_fw].tasks.append(
CheckStability(**check_stability_params))
return original_wf
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 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.
Useful for increasing the k-point mesh for metallic or small gap systems.
Note that this powerup only works on FireWorks with the appropriate WriteVasp* tasks that
accept the small_gap_multiply argument...
Args:
original_wf (Workflow)
gap_cutoff (float): Only multiply k-points for materials with gap < gap_cutoff (eV)
density_multiplier (float): Multiply k-point density by this amount
fw_name_constraint (str): Only apply changes to FWs where fw_name contains this substring.
Returns:
Workflow
"""
for idx_fw, idx_t in get_fws_and_tasks(
original_wf,
fw_name_constraint=fw_name_constraint,
task_name_constraint="WriteVasp"):
original_wf.fws[idx_fw].tasks[idx_t]["small_gap_multiply"] = [
gap_cutoff, density_multiplier
]
return original_wf
def set_execution_options(original_wf,
fworker_name=None,
category=None,
fw_name_constraint=None,
task_name_constraint=None):
"""
set _fworker spec of Fireworker(s) of a Workflow. It can be used to specify a queue;
e.g. run large-memory jobs on a separate queue.
Args:
original_wf (Workflow):
fworker_name (str): user-defined tag to be added under fw.spec._fworker
e.g. "large memory", "big", etc
category (str): category of FWorker that should pul job
fw_name_constraint (str): name of the Fireworks to be tagged (all if None is passed)
task_name_constraint (str): name of the Firetasks to be tagged (e.g. None or 'RunVasp')
Returns:
Workflow: modified workflow with specified Fireworkers tagged
"""
for idx_fw, idx_t in get_fws_and_tasks(
original_wf,
fw_name_constraint=fw_name_constraint,
task_name_constraint=task_name_constraint):
if fworker_name:
original_wf.fws[idx_fw].spec["_fworker"] = fworker_name
if category:
original_wf.fws[idx_fw].spec["_category"] = category
return original_wf
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".
Returns:
Workflow
"""
custodian_params = custodian_params if custodian_params else {}
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"] = original_wf.fws[idx_fw].tasks[
idx_t]["vasp_cmd"]
original_wf.fws[idx_fw].tasks[idx_t] = \
RunVaspCustodian(**custodian_params)
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.
Returns:
Workflow
"""
modify_incar_params = modify_incar_params or {
"incar_update": ">>incar_update<<"
}
for idx_fw, idx_t in get_fws_and_tasks(
original_wf,
fw_name_constraint=fw_name_constraint,
task_name_constraint="RunVasp"):
original_wf.fws[idx_fw].tasks.insert(
idx_t, ModifyIncar(**modify_incar_params))
return original_wf
def get_powerups_wf(original_wf):
"""
get user powerups setting.
"""
idx_list = get_fws_and_tasks(original_wf)
for idx_fw, idx_t in idx_list:
f0 = original_wf.fws[idx_fw].tasks[idx_t]
if not isinstance(f0, Iterable) or isinstance(f0, str) : continue
for k0 in f0:
if debug: print("level 0", k0, type(f0))
if k0=='powerups' :
if debug: print("level 0", f0[k0])
return f0[k0]
else:
try:
f1 = f0[k0]
except:
f1 = k0
if not isinstance(f1, Iterable) or isinstance(f1, str) : continue
for k1 in f1:
if debug: print("level 1", k1, type(f1))
if str(k1)=='powerups' :
if debug: print("level 1", f1[k1])
return f1[k1]
else:
try:
f2 = f1[k1]
except:
f2 = k1
if not isinstance(f2, Iterable) or isinstance(f2, str) : continue
for k2 in f2:
if debug: print("level 2", k2, type(f2))
if str(k2)=='powerups' :
if debug: print("level 2", f2[k2])
return f2[k2]
else:
try:
f3 = f2[k2]
except:
f3=k2
if not isinstance(f3, Iterable) or isinstance(f3, str) : continue
for k3 in f3:
if debug: print("level 3", k3, type(f3))
if str(k3)=='powerups' :
if debug: print(type(f0),type(f1),type(f2),type(f3))
if debug: print("level 3", f3[k3])
return f3[k3]
else:
try:
f4 = f3[k3]
except:
f4=k3
if not isinstance(f4, Iterable) or isinstance(f4, str) : continue
for k4 in f4:
if debug: print("level 4", k4, type(f4))
if str(k4)=='powerups' :
if debug: print("level 4", f4[k4])
return f4[k4]
return {}
def generate_elastic_workflow(structure, tags=None):
"""
Generates a standard production workflow.
Notes:
Uses a primitive structure transformed into
the conventional basis (for equivalent deformations).
Adds the "minimal" category to the minimal portion
of the workflow necessary to generate the elastic tensor,
and the "minimal_full_stencil" category to the portion that
includes all of the strain stencil, but is symmetrically complete
"""
if tags == None:
tags = []
# transform the structure
ieee_rot = Tensor.get_ieee_rotation(structure)
if not SquareTensor(ieee_rot).is_rotation(tol=0.005):
raise ValueError(
"Rotation matrix does not satisfy rotation conditions")
symm_op = SymmOp.from_rotation_and_translation(ieee_rot)
ieee_structure = structure.copy()
ieee_structure.apply_operation(symm_op)
# construct workflow
wf = wf_elastic_constant(ieee_structure)
# Set categories, starting with optimization
opt_fws = get_fws_and_tasks(wf, fw_name_constraint="optimization")
wf.fws[opt_fws[0][0]].spec['elastic_category'] = "minimal"
# find minimal set of fireworks using symmetry reduction
fws_by_strain = {
Strain(fw.tasks[-1]['pass_dict']['strain']): n
for n, fw in enumerate(wf.fws) if 'deformation' in fw.name
}
unique_tensors = symmetry_reduce(list(fws_by_strain.keys()),
ieee_structure)
for unique_tensor in unique_tensors:
fw_index = get_tkd_value(fws_by_strain, unique_tensor)
if np.isclose(unique_tensor, 0.005).any():
wf.fws[fw_index].spec['elastic_category'] = "minimal"
else:
wf.fws[fw_index].spec['elastic_category'] = "minimal_full_stencil"
# Add tags
if tags:
wf = add_tags(wf, tags)
wf = add_modify_incar(wf)
priority = 500 - structure.num_sites
wf = add_priority(wf, priority)
for fw in wf.fws:
if fw.spec.get('elastic_category') == 'minimal':
fw.spec['_priority'] += 2000
elif fw.spec.get('elastic_category') == 'minimal_full_stencil':
fw.spec['_priority'] += 1000
return wf
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:
Workflow: modified with SOC
"""
if structure is None:
try:
sid = get_fws_and_tasks(original_wf, fw_name_constraint="structure optimization",
task_name_constraint="WriteVasp")
fw_id = sid[0][0]
task_id = sid[0][1]
structure = original_wf.fws[fw_id].tasks[task_id]["vasp_input_set"].structure
except:
raise ValueError("modify_to_soc powerup requires the structure in vasp_input_set")
magmom = ""
for _ 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"):
original_wf.fws[idx_fw].tasks[idx_t]["vasp_cmd"] = ">>vasp_ncl<<"
original_wf.fws[idx_fw].tasks.insert(idx_t, ModifyIncar(**modify_incar_params))
original_wf.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"):
original_wf.fws[idx_fw].name += " soc"
return original_wf
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 test_use_potcar_spec(self):
wf = copy_wf(self.bs_wf)
wf = use_potcar_spec(wf)
idx_list = get_fws_and_tasks(wf, task_name_constraint="WriteVasp")
self.assertTrue(len(idx_list) > 0)
for idx_fw, idx_t in idx_list:
task = wf.fws[idx_fw].tasks[idx_t]
self.assertTrue(task["potcar_spec"])
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 test_add_metadata(self):
my_wf = copy_wf(self.bs_wf)
my_wf.metadata = {"what": "ever"}
meta_dict = {"foo": "bar", "baz": 42}
my_wf = add_metadata(my_wf, meta_dict, fw_name_constraint="NonSCFFW")
self.assertEqual(my_wf.metadata, {"what": "ever", "foo": "bar", "baz": 42})
for [fw, _] in get_fws_and_tasks(my_wf, fw_name_constraint="NonSCFFW"):
for key, val in meta_dict.items():
self.assertEqual(fw.spec[key], val)
def modify_gzip_vasp(original_wf, gzip_output):
"""
For all RunVaspCustodian tasks, modify gzip_output boolean
Args:
original_wf (Workflow)
gzip_output (bool): Value to set gzip_output to for RunVaspCustodian
Returns:
Workflow
"""
for idx_fw, idx_t in get_fws_and_tasks(original_wf, task_name_constraint="RunVaspCustodian"):
original_wf.fws[idx_fw].tasks[idx_t]["gzip_output"] = gzip_output
return original_wf
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 set_queue_options(
original_wf,
walltime=None,
time_min=None,
qos=None,
pmem=None,
fw_name_constraint=None,
task_name_constraint=None,
):
"""
Modify queue submission parameters of Fireworks in a Workflow.
This powerup overrides paramters in the qadapter file by setting values in
the 'queueadapter' key of a Firework spec. For example, the walltime
requested from a queue can be modified on a per-workflow basis.
Args:
original_wf (Workflow):
walltime (str): Total walltime to request for the job in HH:MM:SS
format e.g., "00:10:00" for 10 minutes.
time_min (str): Minimum walltime to request in HH:MM:SS format.
Specifying both `walltime` and `time_min` can improve throughput on
some queues.
qos (str): QoS level to request. Typical examples include "regular",
"flex", and "scavenger". For Cori KNL "flex" QoS, it is necessary
to specify a `time_min` of no more than 2 hours.
fw_name_constraint (str): name of the Fireworks to be tagged (all if
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})
if pmem:
qsettings.update({"pmem": pmem})
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 use_gamma_vasp(original_wf, gamma_vasp_cmd):
"""
For all RunVaspCustodian tasks, add the desired scratch dir.
Args:
original_wf (Workflow)
gamma_vasp_cmd (str): path to gamma_vasp_cmd. Supports env_chk
Returns:
Workflow
"""
for idx_fw, idx_t in get_fws_and_tasks(original_wf, task_name_constraint="RunVaspCustodian"):
original_wf.fws[idx_fw].tasks[idx_t]["gamma_vasp_cmd"] = gamma_vasp_cmd
return original_wf
def use_scratch_dir(original_wf, scratch_dir):
"""
For all RunVaspCustodian tasks, add the desired scratch dir.
Args:
original_wf (Workflow)
scratch_dir (path): Path to the scratch dir to use. Supports env_chk
Returns:
Workflow
"""
for idx_fw, idx_t in get_fws_and_tasks(original_wf, task_name_constraint="RunVaspCustodian"):
original_wf.fws[idx_fw].tasks[idx_t]["scratch_dir"] = scratch_dir
return original_wf
def clear_modify(original_wf, fw_name_constraint=None):
"""
Simple powerup that clears the modifications to a workflow.
Args:
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
Returns:
Workflow
"""
for idx_fw, idx_t in get_fws_and_tasks(original_wf, task_name_constraint=task_name_constraint):
original_wf.fws[idx_fw].tasks[idx_t]["additional_fields"].update(update_dict)
return original_wf
def add_tags(original_wf, tags_list):
"""
Adds tags to all Fireworks in the Workflow, WF metadata, as well as
additional_fields for the Drone to track them later (e.g. tag all fireworks
and tasks related to a specific research project).
Tags are written to the "_spec" key of each Firework in the workflow and
to the "metadata.tags" key of each Workflow. If the workflow contains any
Firetasks ending in "ToDb", e.g. VaspToDb, QChemToDb, etc., then the tags
are also passed as "additional_fields" to these tasks and included in the
resulting task documents.
Args:
original_wf (Workflow)
tags_list: list of tags parameters (list of strings)
Returns:
Workflow
"""
# WF metadata
if "tags" in original_wf.metadata:
original_wf.metadata["tags"].extend(tags_list)
else:
original_wf.metadata["tags"] = tags_list
# FW metadata
for idx_fw in range(len(original_wf.fws)):
if "tags" in original_wf.fws[idx_fw].spec:
original_wf.fws[idx_fw].spec["tags"].extend(tags_list)
else:
original_wf.fws[idx_fw].spec["tags"] = tags_list
# DB insertion tasks
idxs = get_fws_and_tasks(original_wf, task_name_constraint="ToDb")
for idx_fw, idx_t in idxs:
if "additional_fields" in original_wf.fws[idx_fw].tasks[
idx_t].optional_params:
if "tags" in original_wf.fws[idx_fw].tasks[idx_t][
"additional_fields"]:
original_wf.fws[idx_fw].tasks[idx_t]["additional_fields"][
"tags"].extend(tags_list)
else:
original_wf.fws[idx_fw].tasks[idx_t]["additional_fields"][
"tags"] = tags_list
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 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.
Returns:
Workflow
"""
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 = original_wf.fws[idx_fw].tasks[idx_t]["vasp_cmd"]
original_wf.fws[idx_fw].tasks[idx_t] = RunVaspDirect(vasp_cmd=vasp_cmd)
return original_wf
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_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 preserve_fworker(original_wf, fw_name_constraint=None):
"""
set _preserve_fworker spec of Fireworker(s) of a Workflow. Can be used to pin a workflow to
the first fworker it is run with. Very useful when running on multiple machines that can't
share files. fw_name_constraint can be used to only preserve fworker after a certain point
where file passing becomes important
Args:
original_wf (Workflow):
fw_name_constraint (str): name of the Fireworks to be tagged (all if None is passed)
Returns:
Workflow: modified workflow with specified Fireworkers tagged
"""
for idx_fw, idx_t in get_fws_and_tasks(original_wf, fw_name_constraint=fw_name_constraint):
original_wf.fws[idx_fw].spec["_preserve_fworker"] = True
return original_wf
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 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_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 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)
Returns:
Workflow
"""
# WF metadata
if "tags" in original_wf.metadata:
original_wf.metadata["tags"].extend(tags_list)
else:
original_wf.metadata["tags"] = tags_list
# FW metadata
for idx_fw in range(len(original_wf.fws)):
if "tags" in original_wf.fws[idx_fw].spec:
original_wf.fws[idx_fw].spec["tags"].extend(tags_list)
else:
original_wf.fws[idx_fw].spec["tags"] = tags_list
# DB insertion tasks
for constraint in ["VaspToDb", "BoltztrapToDb"]:
idxs = get_fws_and_tasks(original_wf, task_name_constraint=constraint)
for idx_fw, idx_t in idxs:
if "tags" in original_wf.fws[idx_fw].tasks[idx_t]["additional_fields"]:
original_wf.fws[idx_fw].tasks[idx_t]["additional_fields"][
"tags"
].extend(tags_list)
else:
original_wf.fws[idx_fw].tasks[idx_t]["additional_fields"][
"tags"
] = tags_list
return original_wf
def add_additional_fields_to_taskdocs(original_wf,
update_dict=None,
task_name_constraint="ToDb"):
"""
For all XXToDbTasks 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
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 get_wf_elastic_constant(structure,
metadata,
strain_states=None,
stencils=None,
db_file=None,
conventional=False,
order=2,
vasp_input_set=None,
analysis=True,
sym_reduce=False,
tag='elastic',
copy_vasp_outputs=False,
**kwargs):
"""
Returns a workflow to calculate elastic constants.
Firework 1 : write vasp input set for structural relaxation,
run vasp,
pass run location,
database insertion.
Firework 2 - number of total deformations: Static runs on the deformed structures
last Firework : Analyze Stress/Strain data and fit the elastic tensor
Args:
structure (Structure): input structure to be optimized and run.
strain_states (list of Voigt-notation strains): list of ratios of nonzero elements
of Voigt-notation strain, e. g. [(1, 0, 0, 0, 0, 0), (0, 1, 0, 0, 0, 0), etc.].
stencils (list of floats, or list of list of floats): values of strain to multiply
by for each strain state, i. e. stencil for the perturbation along the strain
state direction, e. g. [-0.01, -0.005, 0.005, 0.01]. If a list of lists,
stencils must correspond to each strain state provided.
db_file (str): path to file containing the database credentials.
conventional (bool): flag to convert input structure to conventional structure,
defaults to False.
order (int): order of the tensor expansion to be determined. Defaults to 2 and
currently supports up to 3.
vasp_input_set (VaspInputSet): vasp input set to be used. Defaults to static
set with ionic relaxation parameters set. Take care if replacing this,
default ensures that ionic relaxation is done and that stress is calculated
for each vasp run.
analysis (bool): flag to indicate whether analysis task should be added
and stresses and strains passed to that task
sym_reduce (bool): Whether or not to apply symmetry reductions
tag (str):
copy_vasp_outputs (bool): whether or not to copy previous vasp outputs.
kwargs (keyword arguments): additional kwargs to be passed to get_wf_deformations
Returns:
Workflow
"""
# Convert to conventional if specified
if conventional:
structure = SpacegroupAnalyzer(
structure).get_conventional_standard_structure()
uis_elastic = {
"IBRION": 2,
"NSW": 99,
"ISIF": 2,
"ISTART": 1,
"PREC": "High"
}
vis = vasp_input_set or MPStaticSet(structure,
user_incar_settings=uis_elastic)
strains = []
if strain_states is None:
strain_states = get_default_strain_states(order)
if stencils is None:
stencils = [np.linspace(-0.01, 0.01, 5 +
(order - 2) * 2)] * len(strain_states)
if np.array(stencils).ndim == 1:
stencils = [stencils] * len(strain_states)
for state, stencil in zip(strain_states, stencils):
strains.extend(
[Strain.from_voigt(s * np.array(state)) for s in stencil])
# Remove zero strains
strains = [strain for strain in strains if not (abs(strain) < 1e-10).all()]
# Adding the zero strains for the purpose of calculating at finite pressure or thermal expansion
_strains = [Strain.from_deformation([[1, 0, 0], [0, 1, 0], [0, 0, 1]])]
strains.extend(_strains)
"""
"""
vstrains = [strain.voigt for strain in strains]
if np.linalg.matrix_rank(vstrains) < 6:
# TODO: check for sufficiency of input for nth order
raise ValueError(
"Strain list is insufficient to fit an elastic tensor")
deformations = [s.get_deformation_matrix() for s in strains]
"""
print(strains)
print(deformations)
"""
if sym_reduce:
# Note this casts deformations to a TensorMapping
# with unique deformations as keys to symmops
deformations = symmetry_reduce(deformations, structure)
wf_elastic = get_wf_deformations(structure,
deformations,
tag=tag,
db_file=db_file,
vasp_input_set=vis,
copy_vasp_outputs=copy_vasp_outputs,
**kwargs)
if analysis:
defo_fws_and_tasks = get_fws_and_tasks(
wf_elastic,
fw_name_constraint="deformation",
task_name_constraint="Transmuted")
for idx_fw, idx_t in defo_fws_and_tasks:
defo = \
wf_elastic.fws[idx_fw].tasks[idx_t]['transformation_params'][0][
'deformation']
pass_dict = {
'strain': Deformation(defo).green_lagrange_strain.tolist(),
'stress': '>>output.ionic_steps.-1.stress',
'deformation_matrix': defo
}
if sym_reduce:
pass_dict.update({'symmops': deformations[defo]})
mod_spec_key = "deformation_tasks->{}".format(idx_fw)
pass_task = pass_vasp_result(pass_dict=pass_dict,
mod_spec_key=mod_spec_key)
wf_elastic.fws[idx_fw].tasks.append(pass_task)
fw_analysis = Firework(ElasticTensorToDb(structure=structure,
db_file=db_file,
order=order,
fw_spec_field='tags',
metadata=metadata,
vasp_input_set=vis),
name="Analyze Elastic Data",
spec={"_allow_fizzled_parents": True})
wf_elastic.append_wf(Workflow.from_Firework(fw_analysis),
wf_elastic.leaf_fw_ids)
wf_elastic.name = "{}:{}".format(structure.composition.reduced_formula,
"elastic constants")
return wf_elastic
def get_wf_elastic_constant(structure, strain_states=None, stencils=None,
db_file=None,
conventional=False, order=2, vasp_input_set=None,
analysis=True,
sym_reduce=False, tag='elastic',
copy_vasp_outputs=False, **kwargs):
"""
Returns a workflow to calculate elastic constants.
Firework 1 : write vasp input set for structural relaxation,
run vasp,
pass run location,
database insertion.
Firework 2 - number of total deformations: Static runs on the deformed structures
last Firework : Analyze Stress/Strain data and fit the elastic tensor
Args:
structure (Structure): input structure to be optimized and run.
strain_states (list of Voigt-notation strains): list of ratios of nonzero elements
of Voigt-notation strain, e. g. [(1, 0, 0, 0, 0, 0), (0, 1, 0, 0, 0, 0), etc.].
stencils (list of floats, or list of list of floats): values of strain to multiply
by for each strain state, i. e. stencil for the perturbation along the strain
state direction, e. g. [-0.01, -0.005, 0.005, 0.01]. If a list of lists,
stencils must correspond to each strain state provided.
db_file (str): path to file containing the database credentials.
conventional (bool): flag to convert input structure to conventional structure,
defaults to False.
order (int): order of the tensor expansion to be determined. Defaults to 2 and
currently supports up to 3.
vasp_input_set (VaspInputSet): vasp input set to be used. Defaults to static
set with ionic relaxation parameters set. Take care if replacing this,
default ensures that ionic relaxation is done and that stress is calculated
for each vasp run.
analysis (bool): flag to indicate whether analysis task should be added
and stresses and strains passed to that task
sym_reduce (bool): Whether or not to apply symmetry reductions
tag (str):
copy_vasp_outputs (bool): whether or not to copy previous vasp outputs.
kwargs (keyword arguments): additional kwargs to be passed to get_wf_deformations
Returns:
Workflow
"""
# Convert to conventional if specified
if conventional:
structure = SpacegroupAnalyzer(
structure).get_conventional_standard_structure()
uis_elastic = {"IBRION": 2, "NSW": 99, "ISIF": 2, "ISTART": 1,
"PREC": "High"}
vis = vasp_input_set or MPStaticSet(structure,
user_incar_settings=uis_elastic)
strains = []
if strain_states is None:
strain_states = get_default_strain_states(order)
if stencils is None:
stencils = [np.linspace(-0.01, 0.01, 5 + (order - 2) * 2)] * len(
strain_states)
if np.array(stencils).ndim == 1:
stencils = [stencils] * len(strain_states)
for state, stencil in zip(strain_states, stencils):
strains.extend(
[Strain.from_voigt(s * np.array(state)) for s in stencil])
# Remove zero strains
strains = [strain for strain in strains if not (abs(strain) < 1e-10).all()]
vstrains = [strain.voigt for strain in strains]
if np.linalg.matrix_rank(vstrains) < 6:
# TODO: check for sufficiency of input for nth order
raise ValueError("Strain list is insufficient to fit an elastic tensor")
deformations = [s.get_deformation_matrix() for s in strains]
if sym_reduce:
# Note this casts deformations to a TensorMapping
# with unique deformations as keys to symmops
deformations = symmetry_reduce(deformations, structure)
wf_elastic = get_wf_deformations(structure, deformations, tag=tag,
db_file=db_file,
vasp_input_set=vis,
copy_vasp_outputs=copy_vasp_outputs,
**kwargs)
if analysis:
defo_fws_and_tasks = get_fws_and_tasks(wf_elastic,
fw_name_constraint="deformation",
task_name_constraint="Transmuted")
for idx_fw, idx_t in defo_fws_and_tasks:
defo = \
wf_elastic.fws[idx_fw].tasks[idx_t]['transformation_params'][0][
'deformation']
pass_dict = {
'strain': Deformation(defo).green_lagrange_strain.tolist(),
'stress': '>>output.ionic_steps.-1.stress',
'deformation_matrix': defo}
if sym_reduce:
pass_dict.update({'symmops': deformations[defo]})
mod_spec_key = "deformation_tasks->{}".format(idx_fw)
pass_task = pass_vasp_result(pass_dict=pass_dict,
mod_spec_key=mod_spec_key)
wf_elastic.fws[idx_fw].tasks.append(pass_task)
fw_analysis = Firework(
ElasticTensorToDb(structure=structure, db_file=db_file,
order=order, fw_spec_field='tags'),
name="Analyze Elastic Data", spec={"_allow_fizzled_parents": True})
wf_elastic.append_wf(Workflow.from_Firework(fw_analysis),
wf_elastic.leaf_fw_ids)
wf_elastic.name = "{}:{}".format(structure.composition.reduced_formula,
"elastic constants")
return wf_elastic