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