Merge branch 'main' of https://www.github.com/hackingmaterials/atomate into main

Author: computron

.pre-commit-config.yaml

@@ -1,19 +1,21 @@
 repos:
-  - repo: git@github.com:humitos/mirrors-autoflake.git
-    rev: v1.1
+  - repo: https://github.com/myint/autoflake
+    rev: v1.4
     hooks:
     - id: autoflake
-      args: [ '--in-place', '--remove-all-unused-imports', '--remove-unused-variable' ]
-  - repo: https://github.com/ambv/black
-    rev: stable
+      args: [--in-place, --remove-all-unused-imports, --remove-unused-variable, --ignore-init-module-imports]
+  - repo: https://github.com/psf/black
+    rev: 21.4b2
     hooks:
     - id: black
-      language_version: python3
+  - repo: https://gitlab.com/pycqa/flake8
+    rev: 3.8.4
+    hooks:
+      - id: flake8
+        args: [--max-line-length=125]
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v2.3.0
+    rev: v3.4.0
     hooks:
       - id: check-yaml
       - id: end-of-file-fixer
       - id: trailing-whitespace
-      - id: flake8
-        args: ["--max-line-length=125"]

CONTRIBUTING.md

@@ -1,15 +1,19 @@
 # Contributing to atomate
+
 We love your input! We want to make contributing to atomate as easy and transparent as possible, whether it's:
+
 * Reporting a bug
 * Discussing the current state of the code
 * Submitting a fix
 * Proposing or implementing new features
 * Becoming a maintainer
 
 ## Reporting bugs, getting help, and discussion
+
 At any time, feel free to start a thread on our [Discourse forum](https://discuss.matsci.org/c/atomate).
 
 If you are making a bug report, incorporate as many elements of the following as possible to ensure a timely response and avoid the need for followups:
+
 * A quick summary and/or background
 * Steps to reproduce - be specific! **Provide sample code.**
 * What you expected would happen, compared to what actually happens
@@ -19,16 +23,19 @@ If you are making a bug report, incorporate as many elements of the following as
 We love thorough bug reports as this means the development team can make quick and meaningful fixes. When we confirm your bug report, we'll move it to the GitHub issues where its progress can be further tracked.
 
 ## Contributing code modifications or additions through Github
-We use github to host code, to track issues and feature requests, as well as accept pull requests. We maintain a list of all contributors to atomate [here.](https://atomate.org/contributors.html)
+
+We use github to host code, to track issues and feature requests, as well as accept pull requests. We maintain a list of all contributors to atomate [here](https://atomate.org/contributors.html).
 
 Pull requests are the best way to propose changes to the codebase. Follow the [Github flow](https://www.atlassian.com/git/tutorials/comparing-workflows/forking-workflow) for more information on this procedure.
 
 The basic procedure for making a PR is:
+
 * Fork the repo and create your branch from main.
 * Commit your improvements to your branch and push to your Github fork (repo).
 * When you're finished, go to your fork and make a Pull Request. It will automatically update if you need to make further changes.
 
 ### How to Make a **Great** Pull Request
+
 We have a few tips for writing good PRs that are accepted into the main repo:
 
 * Use the Google Code style for all of your code. Find an example [here.](https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html)
@@ -41,4 +48,5 @@ When you submit your PR, our CI service will automatically run your tests.
 We welcome good discussion on the best ways to write your code, and the comments on your PR are an excellent area for discussion.
 
 #### References
+
 This document was adapted from the open-source contribution guidelines for Facebook's Draft, as well as briandk's [contribution template](https://gist.github.com/briandk/3d2e8b3ec8daf5a27a62).

atomate/common/firetasks/glue_tasks.py

@@ -1,6 +1,3 @@
-# coding: utf-8
-
-
 import os
 import monty
 import shutil
@@ -72,7 +69,7 @@ def get_calc_loc(target_name, calc_locs):
         for doc in reversed(calc_locs):
             if doc["name"] == target_name:
                 return doc
-        raise ValueError("Could not find the target_name: {}".format(target_name))
+        raise ValueError(f"Could not find the target_name: {target_name}")
     else:
         return calc_locs[-1]
 

atomate/common/firetasks/parse_outputs.py

@@ -1,6 +1,3 @@
-# coding: utf-8
-
-
 import json
 import os
 
@@ -10,7 +7,7 @@
 from atomate.utils.utils import env_chk, get_logger, load_class
 from atomate.common.firetasks.glue_tasks import get_calc_loc
 
-__author__ = 'Shyam Dwaraknath <[email protected]>, Anubhav Jain <[email protected]>'
+__author__ = "Shyam Dwaraknath <[email protected]>, Anubhav Jain <[email protected]>"
 
 logger = get_logger(__name__)
 
@@ -37,7 +34,14 @@ class ToDbTask(FiretaskBase):
     """
 
     required_params = ["drone"]
-    optional_params = ["mmdb", "db_file", "calc_dir", "calc_loc", "additional_fields", "options"]
+    optional_params = [
+        "mmdb",
+        "db_file",
+        "calc_dir",
+        "calc_loc",
+        "additional_fields",
+        "options",
+    ]
 
     def run_task(self, fw_spec):
         # get the directory that contains the dir to parse
@@ -48,12 +52,15 @@ def run_task(self, fw_spec):
             calc_dir = get_calc_loc(self["calc_loc"], fw_spec["calc_locs"])["path"]
 
         # parse the calc directory
-        logger.info("PARSING DIRECTORY: {} USING DRONE: {}".format(
-            calc_dir, self['drone'].__class__.__name__))
+        logger.info(
+            "PARSING DIRECTORY: {} USING DRONE: {}".format(
+                calc_dir, self["drone"].__class__.__name__
+            )
+        )
         # get the database connection
-        db_file = env_chk(self.get('db_file'), fw_spec)
+        db_file = env_chk(self.get("db_file"), fw_spec)
 
-        drone = self['drone'].__class__()
+        drone = self["drone"].__class__()
         task_doc = drone.assimilate(calc_dir)
         if not db_file:
             with open("task.json", "w") as f:
@@ -66,7 +73,9 @@ def run_task(self, fw_spec):
 
             # insert the task document
             t_id = db.insert(task_doc)
-            logger.info("Finished parsing with task_id: {}".format(t_id))
+            logger.info(f"Finished parsing with task_id: {t_id}")
 
-        return FWAction(stored_data={"task_id": task_doc.get("task_id", None)},
-                        defuse_children=(task_doc["state"] != "successful"))
+        return FWAction(
+            stored_data={"task_id": task_doc.get("task_id", None)},
+            defuse_children=(task_doc["state"] != "successful"),
+        )

atomate/common/firetasks/run_calc.py

@@ -1,6 +1,3 @@
-# coding: utf-8
-
-
 import os
 import subprocess
 
@@ -10,7 +7,7 @@
 
 from atomate.utils.utils import env_chk, get_logger
 
-__author__ = 'Anubhav Jain <[email protected]>'
+__author__ = "Anubhav Jain <[email protected]>"
 
 logger = get_logger(__name__)
 
@@ -34,9 +31,9 @@ def run_task(self, fw_spec):
         if self.get("expand_vars", False):
             cmd = os.path.expandvars(cmd)
 
-        logger.info("Running command: {}".format(cmd))
+        logger.info(f"Running command: {cmd}")
         return_code = subprocess.call(cmd, shell=True)
-        logger.info("Command {} finished running with returncode: {}".format(cmd, return_code))
+        logger.info(f"Command {cmd} finished running with returncode: {return_code}")
 
 
 @explicit_serialize
@@ -58,6 +55,10 @@ class RunCustodianFromObjects(FiretaskBase):
     optional_params = ["validators", "custodian_params"]
 
     def run_task(self, fw_spec):
-        c = Custodian(self["handlers"], self["jobs"], self.get("validators"),
-                      **self.get("custodian_params", {}))
+        c = Custodian(
+            self["handlers"],
+            self["jobs"],
+            self.get("validators"),
+            **self.get("custodian_params", {}),
+        )
         c.run()

atomate/common/firetasks/tests/test_parse_outputs.py

@@ -1,55 +1,56 @@
-
-
 import os
 import unittest
-import shutil
 
-from fireworks import LaunchPad, Firework, Workflow
+from fireworks import Firework, Workflow
 from fireworks.core.rocket_launcher import rapidfire
 
 from pymatgen.apps.borg.hive import AbstractDrone
 
 from atomate.common.firetasks.parse_outputs import ToDbTask
 from atomate.utils.testing import AtomateTest
 
-__author__ = 'Shyam Dwaraknath <[email protected]>'
+__author__ = "Shyam Dwaraknath <[email protected]>"
 
 module_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)))
 db_dir = os.path.join(module_dir, "..", "..", "test_files")
 
 
 class TestDrone(AbstractDrone):
-
     def __init__(self):
         pass
 
     def assimilate(self, path):
-        return {"drone": "Test Drone",
-                "dir_name": "/test",
-                'state': "successful"}
+        return {"drone": "Test Drone", "dir_name": "/test", "state": "successful"}
 
     def get_valid_paths(self, path):
         return path
 
 
 class TestToDbTask(AtomateTest):
-
     def test_ToDbTask(self):
         d = TestDrone()
 
-        fw1 = Firework([ToDbTask(drone=d,
-                                 mmdb="atomate.vasp.database.VaspCalcDb",
-                                 db_file=os.path.join(db_dir, "db.json"),
-                                 calc_dir=db_dir)], name="fw1")
+        fw1 = Firework(
+            [
+                ToDbTask(
+                    drone=d,
+                    mmdb="atomate.vasp.database.VaspCalcDb",
+                    db_file=os.path.join(db_dir, "db.json"),
+                    calc_dir=db_dir,
+                )
+            ],
+            name="fw1",
+        )
 
         wf = Workflow([fw1])
         self.lp.add_wf(wf)
         rapidfire(self.lp)
 
-        task1 = self.lp.db.tasks.find_one({"task_id":1})
-        self.assertEqual(task1['task_id'],1)
-        self.assertEqual(task1['dir_name'],"/test" )
-        self.assertEqual(task1['drone'],'Test Drone')
+        task1 = self.lp.db.tasks.find_one({"task_id": 1})
+        self.assertEqual(task1["task_id"], 1)
+        self.assertEqual(task1["dir_name"], "/test")
+        self.assertEqual(task1["drone"], "Test Drone")
+
 
 if __name__ == "__main__":
     unittest.main()

atomate/common/powerups.py

@@ -1,5 +1,3 @@
-# coding: utf-8
-
 """
 This module defines general powerups that can be used for all workflows
 """
@@ -102,7 +100,7 @@ def add_namefile(original_wf, use_slug=True):
        Workflow
     """
     for idx, fw in enumerate(original_wf.fws):
-        fname = "FW--{}".format(fw.name)
+        fname = f"FW--{fw.name}"
         if use_slug:
             fname = get_slug(fname)
 
@@ -132,6 +130,34 @@ def add_additional_fields_to_taskdocs(
     return original_wf
 
 
+def add_metadata(wf, meta_dict, fw_name_constraint=None):
+    """
+    Add a metadata dictionary to a Workflow and all its Fireworks. The dictionary
+    is merged into the "metadata" key of the Workflow and into the "_spec" key of
+    each Firework in the workflow.
+
+    Can be used in combination with add_additional_fields_to_taskdocs to add the
+    same set of key-value pairs to Workflows, Fireworks and Tasks collections.
+
+    Args:
+        wf (Workflow)
+        meta_dict: dictionary of custom metadata
+
+    Returns:
+       Workflow
+    """
+
+    # add metadata to Workflow metadata
+    wf.metadata.update(meta_dict)
+
+    # add metadata to Firework metadata
+    for fw in wf.fws:
+        if fw_name_constraint is None or fw_name_constraint in fw.name:
+            fw.spec.update(meta_dict)
+
+    return wf
+
+
 def preserve_fworker(original_wf, fw_name_constraint=None):
     """
     set _preserve_fworker spec of Fireworker(s) of a Workflow. Can be used to
@@ -275,5 +301,5 @@ def powerup_by_kwargs(
                 except Exception:
                     pass
         if not found:
-            raise RuntimeError("Could not find powerup {}.".format(name))
+            raise RuntimeError(f"Could not find powerup {name}.")
     return original_wf

atomate/common/tests/test_powerups.py

@@ -1,5 +1,6 @@
 import unittest
 
+from atomate.utils.utils import get_fws_and_tasks
 from atomate.vasp.workflows.base.core import get_wf
 
 from pymatgen.io.vasp.sets import MPRelaxSet
@@ -10,6 +11,8 @@
     add_priority,
     add_tags,
     powerup_by_kwargs,
+    add_additional_fields_to_taskdocs,
+    add_metadata,
 )
 from fireworks import Firework, ScriptTask, Workflow
 
@@ -137,6 +140,36 @@ def test_set_queue_adapter(self):
         )
         self.assertDictEqual(wf.id_fw[-3].spec, {})
 
+    def test_add_additional_fields_to_taskdocs(self):
+
+        my_wf = copy_wf(self.bsboltz_wf)
+        meta_dict = {"foo": "bar", "baz": 42}
+        my_wf = add_additional_fields_to_taskdocs(my_wf, meta_dict)
+
+        found = 0
+
+        for fw in my_wf.fws:
+            for task in fw.tasks:
+                if "ToDb" in str(task):
+                    for key, val in meta_dict.items():
+                        self.assertEqual(task["additional_fields"][key], val)
+
+                    found += 1
+
+        self.assertEqual(found, 5)
+
+    def test_add_metadata(self):
+        my_wf = copy_wf(self.bs_wf)
+        my_wf.metadata = {"what": "ever"}
+        meta_dict = {"foo": "bar", "baz": 42}
+        my_wf = add_metadata(my_wf, meta_dict, fw_name_constraint="NonSCFFW")
+
+        self.assertEqual(my_wf.metadata, {"what": "ever", "foo": "bar", "baz": 42})
+
+        for [fw, _] in get_fws_and_tasks(my_wf, fw_name_constraint="NonSCFFW"):
+            for key, val in meta_dict.items():
+                self.assertEqual(fw.spec[key], val)
+
 
 def copy_wf(wf):
     return Workflow.from_dict(wf.to_dict())