[WIP] Port Lattice Dynamics workflow into main

atomate/common/powerups.py

@@ -1,14 +1,15 @@
 """
 This module defines general powerups that can be used for all workflows
 """
+from typing import Dict, Any, Optional
 from importlib import import_module
 from typing import List
 
 from atomate.utils.utils import get_fws_and_tasks
 from fireworks import Workflow, FileWriteTask
 from fireworks.utilities.fw_utilities import get_slug
 
-__author__ = "Janine George, Guido Petretto, Ryan Kingsbury"
+__author__ = "Janine George, Guido Petretto, Ryan Kingsbury, Alex Ganose"
 __email__ = (
     "janine.george@uclouvain.be, guido.petretto@uclouvain.be, RKingsbury@lbl.gov"
 )
@@ -190,20 +191,55 @@ def set_execution_options(
     return original_wf
 
 
+def update_firetask(
+    original_wf: Workflow,
+    update_params: Dict[str, Any],
+    task_name_constraint: str,
+    fw_name_constraint: Optional[str] = None,
+) -> Workflow:
+    """
+    General powerup for arbitrary updates to a Firetask.
+
+    Args:
+        original_wf: The original workflow.
+        update_params: A dictionary of the keyword arguments to update.
+        task_name_constraint: Only apply changes to the Firetasks where the
+            Firetask class name contains this string.
+        fw_name_constraint: Only apply changes to Fireworks where the Firework
+            name contains this substring.
+
+    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[idx_t].update(update_params)
+
+    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.
+    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')
+        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