fixed bson serialization problem by moving the name from key to value

Author: jmmshn

atomate/common/powerups.py

@@ -224,31 +224,27 @@ def set_queue_adapter(
 
 def powerup_by_kwargs(
     original_wf: Workflow,
-    add_powerup_module: List[str] = None,
-    **powerup_kwargs,
+    powerup_dicts: List[dict],
 ):
     """
-    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},
-    }
+    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( "add_additional_fields_to_taskdocs" : {
-                                                                "update_dict" : {"foo" : "bar"}
-                                                                }
+        power_up_by_kwargs([
+            {"powerup_name" : "add_additional_fields_to_taskdocs",
+             "kwargs: {"update_dict" : {"foo" : "bar"}}}
+             ]
         )
 
     Args:
         original_wf: workflow that will be changed
-        add_powerup_module: user_made modules that contain powerups
-        powerup_kwargs: KWARGS used to apply any power, EXAMPLE:
-                {"add_modify_incar": {
-                                      "modify_incar_params": {
-                                      "incar_update": {"KPAR": 8}}
-                                        }
-                }
+        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)
@@ -257,17 +253,27 @@ def powerup_by_kwargs(
         "atomate.qchem.powerups",
         "atomate.common.powerups",
     ]
-    # user can add new modules containing custom powerups
-    if add_powerup_module is not None:
-        powerup_modules = add_powerup_module + powerup_modules
-
-    for k, v in powerup_kwargs.items():
-        for module_name in powerup_modules:
-            try:
-                module = import_module(module_name)
-                powerup = getattr(module, k)
-                original_wf = powerup(original_wf, **v)
-                break
-            except Exception:
-                pass
+
+    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

@@ -83,7 +83,16 @@ def test_add_tags(self):
 
     def test_powerup_by_kwargs(self):
         my_wf = copy_wf(self.bs_wf)
-        my_wf = powerup_by_kwargs(my_wf, add_tags={"tags_list": ["foo", "bar"]})
+        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):

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
@@ -316,10 +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},
+            }
         )
-    return powerup_by_kwargs(wf, **d_pu)
+    return powerup_by_kwargs(wf, powerup_list)

atomate/vasp/tests/test_vasp_powerups.py

@@ -247,8 +247,13 @@ def test_use_potcar_spec(self):
 
     def test_powerup_by_kwargs(self):
         my_wf = copy_wf(self.bs_wf)
-        my_wf = powerup_by_kwargs(my_wf, add_trackers={})
-        my_wf = powerup_by_kwargs(my_wf, add_tags={"tags_list": ["foo", "bar"]})
+        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"])

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
@@ -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

@@ -41,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(