How to use the atomate.utils.utils.get_fws_and_tasks function in atomate

To help you get started, we’ve selected a few atomate examples, based on popular ways it is used in public projects.

Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.

github hackingmaterials / atomate / atomate / vasp / powerups.py View on Github external
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)
github hackingmaterials / atomate / atomate / vasp / powerups.py View on Github external
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)
github hackingmaterials / atomate / atomate / vasp / powerups.py View on Github external
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
github hackingmaterials / atomate / atomate / vasp / powerups.py View on Github external
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)
github hackingmaterials / atomate / atomate / vasp / powerups.py View on Github external
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
github hackingmaterials / atomate / atomate / vasp / powerups.py View on Github external
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)
github hackingmaterials / atomate / atomate / vasp / powerups.py View on Github external
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
github hackingmaterials / atomate / atomate / vasp / powerups.py View on Github external
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
github hackingmaterials / atomate / atomate / vasp / powerups.py View on Github external
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