Merge branch 'main' of github.com:janosh/atomate

Author: JaGeo

atomate/vasp/powerups.py

@@ -83,17 +83,19 @@ def use_custodian(original_wf, fw_name_constraint=None, custodian_params=None):
     Returns:
        Workflow
     """
-    custodian_params = custodian_params if custodian_params else {}
+    custodian_params = custodian_params or {}
     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"
-            ]
+        for param in (
+            RunVaspCustodian.optional_params + RunVaspCustodian.required_params
+        ):
+            value = original_wf.fws[idx_fw].tasks[idx_t].get(param, None)
+            if value is not None and param not in custodian_params:
+                custodian_params[param] = value
         original_wf.fws[idx_fw].tasks[idx_t] = RunVaspCustodian(**custodian_params)
     return original_wf
 
@@ -723,12 +725,14 @@ def use_gamma_vasp(original_wf, gamma_vasp_cmd):
 
 def modify_gzip_vasp(original_wf, gzip_output):
     """
-    For all RunVaspCustodian tasks, modify gzip_output boolean
+    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
+       Workflow: Same workflow as passed in with in-place modified gzip_output.
     """
     idx_list = get_fws_and_tasks(original_wf, task_name_constraint="RunVaspCustodian")
     for idx_fw, idx_t in idx_list:

atomate/vasp/tests/test_vasp_powerups.py

@@ -5,6 +5,7 @@
 
 from atomate.vasp.powerups import (
     add_priority,
+    modify_gzip_vasp,
     use_custodian,
     add_trackers,
     add_modify_incar,
@@ -64,21 +65,32 @@ def test_custodian_powerups(self):
             self.assertTrue("RunVaspDirect" in fw.tasks[task_idx]._fw_name)
             self.assertEqual(fw.tasks[task_idx]["vasp_cmd"], "test_VASP")
 
-        my_wf_double_relax = remove_custodian(copy_wf(self.bs_wf))
-        my_wf_double_relax = use_custodian(
-            my_wf_double_relax,
+        wf_double_relax = remove_custodian(copy_wf(self.bs_wf))
+
+        wf_double_relax = use_custodian(
+            wf_double_relax,
             fw_name_constraint="structure optimization",
             custodian_params={"job_type": "double_relaxation_run"},
         )
 
-        for fw in my_wf_double_relax.fws:
+        for fw in wf_double_relax.fws:
             if "structure optimization" in fw.name:
                 self.assertTrue("RunVaspCustodian" in fw.tasks[1]._fw_name)
                 self.assertEqual(fw.tasks[1]["job_type"], "double_relaxation_run")
             else:
                 self.assertTrue("RunVaspDirect" in fw.tasks[2]._fw_name)
                 self.assertFalse("job_type" in fw.tasks[2])
 
+        # test use_custodian() does not overwrite existing custodian params
+        wf_double_relax = modify_gzip_vasp(wf_double_relax, False)
+        wf_double_relax = use_custodian(wf_double_relax)
+        idx_list = get_fws_and_tasks(
+            wf_double_relax, task_name_constraint="RunVaspCustodian"
+        )
+        for idx_fw, idx_t in idx_list:
+            task = wf_double_relax.fws[idx_fw].tasks[idx_t]
+            assert task["gzip_output"] is False
+
     def test_modify_incar(self):
         my_wf = add_modify_incar(
             copy_wf(self.bs_wf),