Merge pull request #623 from jmmshn/dynamic_poweups

better handling of dynamic powerups

Author: utf

atomate/common/powerups.py

@@ -3,12 +3,13 @@
 """
 This module defines general powerups that can be used for all workflows
 """
+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"
 __email__ = (
     "[email protected], [email protected], [email protected]"
@@ -219,3 +220,60 @@ def set_queue_adapter(
         original_wf.fws[idx_fw].spec["_queueadapter"] = q
 
     return original_wf
+
+
+def powerup_by_kwargs(
+    original_wf: Workflow,
+    powerup_dicts: List[dict],
+):
+    """
+    apply powerups in the form using a list of dictionaries
+    [
+        {"powerup_name" : powerup_function1, "kwargs": {parameter1 : value1, parameter2: value2}},
+        {"powerup_name" : powerup_function2, "kwargs": {parameter1 : value1, parameter2: value2}},
+    ]
+
+    As an example:
+        power_up_by_kwargs([
+            {"powerup_name" : "add_additional_fields_to_taskdocs",
+             "kwargs: {"update_dict" : {"foo" : "bar"}}}
+             ]
+        )
+
+    Args:
+        original_wf: workflow that will be changed
+        powerup_dicts: dictionary containing the powerup_name and kwarg.
+            if "." is present in the name it will be imported as a full path
+            if not we will use standard atomate modules where the powerups are kept
+
+    """
+    # a list of possible powerups in atomate (most specific first)
+    powerup_modules = [
+        "atomate.vasp.powerups",
+        "atomate.qchem.powerups",
+        "atomate.common.powerups",
+    ]
+
+    for pd in powerup_dicts:
+        name = pd["powerup_name"]
+        kwargs = pd["kwargs"]
+        found = False
+        if "." in name:
+            module_name, method_name = name.rsplit(".", 1)
+            module = import_module(module_name)
+            powerup = getattr(module, method_name)
+            original_wf = powerup(original_wf, **kwargs)
+            found = True
+        else:
+            for module_name in powerup_modules:
+                try:
+                    module = import_module(module_name)
+                    powerup = getattr(module, name)
+                    original_wf = powerup(original_wf, **kwargs)
+                    found = True
+                    break
+                except Exception:
+                    pass
+        if not found:
+            raise RuntimeError("Could not find powerup {}.".format(name))
+    return original_wf

atomate/common/tests/test_powerups.py

@@ -5,7 +5,12 @@
 from pymatgen.io.vasp.sets import MPRelaxSet
 from pymatgen.util.testing import PymatgenTest
 
-from atomate.common.powerups import set_queue_adapter, add_priority, add_tags
+from atomate.common.powerups import (
+    set_queue_adapter,
+    add_priority,
+    add_tags,
+    powerup_by_kwargs,
+)
 from fireworks import Firework, ScriptTask, Workflow
 
 __author__ = "Janine George, Guido Petretto"
@@ -76,6 +81,20 @@ def test_add_tags(self):
         self.assertEqual(b_found, 1)
         self.assertEqual(v_found, 4)
 
+    def test_powerup_by_kwargs(self):
+        my_wf = copy_wf(self.bs_wf)
+        my_wf = powerup_by_kwargs(
+            my_wf,
+            [
+                {"powerup_name": "add_tags", "kwargs": {"tags_list": ["foo", "bar"]}},
+                {
+                    "powerup_name": "atomate.common.powerups.add_priority",
+                    "kwargs": {"root_priority": 123},
+                },
+            ],
+        )
+        self.assertEqual(my_wf.metadata["tags"], ["foo", "bar"])
+
     def test_set_queue_adapter(self):
         # test fw_name_constraint
         fw1 = Firework([ScriptTask(script=None)], fw_id=-1, name="Firsttask")

atomate/vasp/firetasks/electrode_tasks.py

@@ -1,5 +1,4 @@
 import math
-from collections import defaultdict
 
 from fireworks import FiretaskBase, explicit_serialize, FWAction, Firework, Workflow
 from pymatgen.core import Structure
@@ -13,7 +12,7 @@
 
 from atomate.vasp.database import VaspCalcDb
 from atomate.vasp.firetasks import pass_vasp_result
-from atomate.vasp.powerups import powerup_by_kwargs, POWERUP_NAMES
+from atomate.common.powerups import powerup_by_kwargs
 
 __author__ = "Jimmy Shen"
 __email__ = "[email protected]"
@@ -316,11 +315,13 @@ def get_powerup_wf(wf, fw_spec, additional_fields=None):
     Returns:
         Updated workflow
     """
-    d_pu = defaultdict(dict)
-    d_pu.update(fw_spec.get("vasp_powerups", {}))
+    powerup_list = []
+    powerup_list.extend(fw_spec.get("vasp_powerups", []))
     if additional_fields is not None:
-        d_pu["add_additional_fields_to_taskdocs"].update(
-            {"update_dict": additional_fields}
+        powerup_list.append(
+            {
+                "powerup_name": "add_additional_fields_to_taskdocs",
+                "kwargs": {"update_dict": additional_fields},
+            }
         )
-    p_kwargs = {k: d_pu[k] for k in POWERUP_NAMES if k in d_pu}
-    return powerup_by_kwargs(wf, **p_kwargs)
+    return powerup_by_kwargs(wf, powerup_list)

atomate/vasp/powerups.py

@@ -33,8 +33,6 @@
 __author__ = "Anubhav Jain, Kiran Mathew, Alex Ganose"
 __email__ = "[email protected], [email protected]"
 
-POWERUP_NAMES = []
-
 
 @deprecated(replacement=common_add_priority)
 def add_priority(original_wf, root_priority, child_priority=None):
@@ -829,35 +827,3 @@ def use_fake_lobster(original_wf, ref_dirs, params_to_check=None):
                         )
 
     return original_wf
-
-
-local_names = dict(locals())  # locals() changes as the program runs, make a copy here
-
-for k, v in local_names.items():
-    if (
-        hasattr(v, "__module__")
-        and v.__module__ == "atomate.vasp.powerups"
-        and k != "power_up_by_kwargs"
-    ):
-        POWERUP_NAMES.append(k)
-
-local_names = {k: v for k, v in local_names.items() if k in POWERUP_NAMES}
-
-
-def powerup_by_kwargs(wf, **kwargs):
-    """
-    apply powerups in the form using a kwargs dictionary of the form:
-    {
-        powerup_function_name1 : {parameter1 : value1, parameter2: value2},
-        powerup_function_name2 : {parameter1 : value1, parameter2: value2},
-    }
-
-    As an example:
-        power_up_by_kwargs( "add_additional_fields_to_taskdocs" : {
-                                                                "update_dict" : {"foo" : "bar"}
-                                                                }
-        )
-    """
-    for k, v in kwargs.items():
-        wf = local_names[k](wf, **v)
-    return wf

atomate/vasp/tests/test_vasp_powerups.py

@@ -19,6 +19,7 @@
     set_queue_options,
     use_potcar_spec,
 )
+from atomate.common.powerups import powerup_by_kwargs
 from atomate.vasp.workflows.base.core import get_wf
 
 from pymatgen.io.vasp.sets import MPRelaxSet
@@ -244,6 +245,19 @@ def test_use_potcar_spec(self):
             task = wf.fws[idx_fw].tasks[idx_t]
             self.assertTrue(task["potcar_spec"])
 
+    def test_powerup_by_kwargs(self):
+        my_wf = copy_wf(self.bs_wf)
+        my_wf = powerup_by_kwargs(
+            my_wf, [{"powerup_name": "add_trackers", "kwargs": {}}]
+        )
+        my_wf = powerup_by_kwargs(
+            my_wf,
+            [{"powerup_name": "add_tags", "kwargs": {"tags_list": ["foo", "bar"]}}],
+        )
+        for fw in my_wf.fws:
+            self.assertEqual(len(fw.spec["_trackers"]), 2)
+        self.assertEqual(my_wf.metadata["tags"], ["foo", "bar"])
+
 
 def copy_wf(wf):
     return Workflow.from_dict(wf.to_dict())

atomate/vasp/workflows/base/electrode.py

@@ -1,3 +1,5 @@
+from typing import List
+
 from fireworks import Workflow
 from pymatgen.core import Structure
 from pymatgen.analysis.structure_matcher import StructureMatcher
@@ -9,7 +11,7 @@
 __email__ = "[email protected]"
 
 from atomate.vasp.fireworks import Firework, OptimizeFW, StaticFW, pass_vasp_result
-from atomate.vasp.powerups import powerup_by_kwargs
+from atomate.common.powerups import powerup_by_kwargs
 
 """
 Define workflow related to battery material simulation --- they all have a working ion
@@ -23,7 +25,7 @@ def get_ion_insertion_wf(
     db_file: str = DB_FILE,
     vasptodb_kwargs: dict = None,
     volumetric_data_type: str = "CHGCAR",
-    vasp_powerups: dict = None,
+    vasp_powerups: List[dict] = None,
     max_insertions: int = 4,
     allow_fizzled_parents: bool = True,
     optimizefw_kwargs: dict = None,
@@ -106,7 +108,7 @@ def get_ion_insertion_wf(
 
     # Apply the vasp powerup if present
     if vasp_powerups is not None:
-        wf = powerup_by_kwargs(wf, **vasp_powerups)
+        wf = powerup_by_kwargs(wf, vasp_powerups)
         for fw in wf.fws:
             fw.spec["vasp_powerups"] = vasp_powerups
 

atomate/vasp/workflows/tests/test_insertion_workflow.py

@@ -21,6 +21,9 @@
 wf_dir = ref_dir / "insertion_wf"
 
 VASP_CMD = None  # for fake VASP
+DEBUG_MODE = (
+    False  # If true, retains the database and output dirs at the end of the test
+)
 
 
 class TestInsertionWorkflow(AtomateTest):
@@ -38,19 +41,23 @@ def setUp(self):
             working_ion="Mg",
             volumetric_data_type="AECCAR",
             db_file=db_dir / "db.json",
-            vasp_powerups={
-                "add_modify_incar": {
-                    "modify_incar_params": {"incar_update": {"KPAR": 8}}
+            vasp_powerups=[
+                {
+                    "powerup_name": "add_modify_incar",
+                    "kwargs": {"modify_incar_params": {"incar_update": {"KPAR": 8}}},
                 },
-                "use_fake_vasp": {
-                    "ref_dirs": calc_dirs,
-                    "check_incar": False,
-                    "check_kpoints": False,
-                    "check_poscar": False,
-                    "check_potcar": False,
+                {
+                    "powerup_name": "use_fake_vasp",
+                    "kwargs": {
+                        "ref_dirs": calc_dirs,
+                        "check_incar": False,
+                        "check_kpoints": False,
+                        "check_poscar": False,
+                        "check_potcar": False,
+                    },
                 },
-                "use_potcar_spec": {},
-            },
+                {"powerup_name": "use_potcar_spec", "kwargs": {}},
+            ],
             optimizefw_kwargs={"ediffg": -0.05},
         )
         wf = use_fake_vasp(