Merge branch 'main' of https://github.com/hackingmaterials/atomate into qcinput-vdw

Author: rkingsbury

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/powerups.py

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

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())

atomate/vasp/config.py

@@ -38,6 +38,10 @@
 # useful for storing duplicate of FW.json
 STORE_ADDITIONAL_JSON = False
 
+# store Bader analysis when parsing VASP directories
+# (also requires "bader" executable to be in path)
+STORE_BADER = False
+
 # vasp output files that will be copied to lobster run
 VASP_OUTPUT_FILES = [
     "OUTCAR",

atomate/vasp/drones.py

@@ -35,7 +35,7 @@
 
 from atomate.utils.utils import get_logger
 from atomate import __version__ as atomate_version
-from atomate.vasp.config import STORE_VOLUMETRIC_DATA, STORE_ADDITIONAL_JSON
+from atomate.vasp.config import STORE_VOLUMETRIC_DATA, STORE_ADDITIONAL_JSON, STORE_BADER
 
 __author__ = "Kiran Mathew, Shyue Ping Ong, Shyam Dwaraknath, Anubhav Jain"
 __email__ = "[email protected]"
@@ -44,7 +44,9 @@
 
 logger = get_logger(__name__)
 
-BADER_EXE_EXISTS = which("bader") or which("bader.exe")
+BADER_EXE_EXISTS = (which("bader") or which("bader.exe"))
+STORE_BADER = STORE_BADER and BADER_EXE_EXISTS
+
 
 
 class VaspDrone(AbstractDrone):
@@ -135,7 +137,7 @@ def __init__(
         parse_locpot=True,
         additional_fields=None,
         use_full_uri=True,
-        parse_bader=BADER_EXE_EXISTS,
+        parse_bader=STORE_BADER,
         parse_chgcar=False,
         parse_aeccar=False,
         parse_potcar_file=True,

atomate/vasp/firetasks/approx_neb_dynamic_tasks.py

@@ -0,0 +1,205 @@
+from fireworks import FiretaskBase, FWAction, explicit_serialize, Workflow
+from atomate.utils.utils import env_chk
+from atomate.vasp.database import VaspCalcDb
+from atomate.vasp.fireworks.approx_neb import ImageFW
+from atomate.common.powerups import powerup_by_kwargs
+
+__author__ = "Ann Rutt"
+__email__ = "[email protected]"
+
+
+@explicit_serialize
+class GetImageFireworks(FiretaskBase):
+    """
+    Adds ImageFWs to the workflow for the provided images_key
+    according to the scheme specified by launch_mode. Optional
+    parameters such as "handler_group", "add_additional_fields",
+    and "add_tags" can be used to modify the resulting ImageFWs.
+
+    Args:
+        db_file (str): path to file containing the database
+            credentials.
+        approx_neb_wf_uuid (str): unique id for approx neb workflow
+            record keeping.
+        images_key (str): specifies a key corresponding the images
+            field of the approx_neb collection which specifies the
+            desired combination of end points to interpolate images
+            between. images_key should be a string of format "0+1",
+            "0+2", etc. matching end_points_combo input of
+            PathfinderToDb Firetask or pathfinder_key input of
+            AddSelectiveDynamics Firetask. If images_key is not
+            provided images will be launched for all paths/keys in
+            the approx_neb collection images field.
+        launch_mode (str): "all" or "screening"
+        vasp_cmd (str): the name of the full executable for running
+            VASP.
+    Optional Params:
+        vasp_input_set (VaspInputSet class): can use to
+            define VASP input parameters.
+            See pymatgen.io.vasp.sets module for more
+            information. MPRelaxSet() and
+            override_default_vasp_params are used if
+            vasp_input_set = None.
+        override_default_vasp_params (dict): if provided,
+            vasp_input_set is disregarded and the Vasp Input
+            Set is created by passing
+            override_default_vasp_params to MPRelaxSet().
+            Allows for easy modification of MPRelaxSet().
+            For example, to set ISIF=2 in the INCAR use:
+            {"user_incar_settings":{"ISIF":2}}
+        handler_group (str or [ErrorHandler]): group of handlers to
+            use for RunVaspCustodian firetask. See handler_groups
+            dict in the code for the groups and complete list of
+            handlers in each group. Alternatively, you can specify a
+            list of ErrorHandler objects.
+        add_additional_fields (dict): dict of additional fields to
+            add to task docs (by additional_fields of VaspToDb).
+        add_tags (list of strings): added to the "tags" field of the
+            task docs.
+    """
+
+    required_params = [
+        "db_file",
+        "approx_neb_wf_uuid",
+        "images_key",
+        "launch_mode",
+        "vasp_cmd",
+    ]
+    optional_params = [
+        "vasp_input_set",
+        "override_default_vasp_params",
+        "handler_group",
+        "add_additional_fields",
+        "add_tags",
+    ]
+
+    def run_task(self, fw_spec):
+        # get the database connection
+        db_file = env_chk(self["db_file"], fw_spec)
+        mmdb = VaspCalcDb.from_db_file(db_file, admin=True)
+        mmdb.collection = mmdb.db["approx_neb"]
+        wf_uuid = self["approx_neb_wf_uuid"]
+        launch_mode = self["launch_mode"]
+        images_key = self["images_key"]
+
+        approx_neb_doc = mmdb.collection.find_one({"wf_uuid": wf_uuid}, {"images": 1})
+        all_images = approx_neb_doc["images"]
+
+        # get structure_path of desired images and sort into structure_paths
+        if images_key and isinstance(all_images, (dict)):
+            images = all_images[images_key]
+            max_n = len(images)
+            if launch_mode == "all":
+                structure_paths = [
+                    "images." + images_key + "." + str(n) + ".input_structure"
+                    for n in range(0, max_n)
+                ]
+            elif launch_mode == "screening":
+                structure_paths = self.get_and_sort_paths(
+                    max_n=max_n, images_key=images_key
+                )
+        elif isinstance(all_images, (dict)):
+            structure_paths = dict()
+            if launch_mode == "all":
+                for key, images in all_images.items():
+                    max_n = len(images)
+                    structure_paths[key] = [
+                        "images." + key + "." + str(n) + ".input_structure"
+                        for n in range(0, max_n)
+                    ]
+            elif launch_mode == "screening":
+                for key, images in all_images.items():
+                    structure_paths[key] = self.get_and_sort_paths(
+                        max_n=len(images), images_key=key
+                    )
+
+        # get list of fireworks to launch
+        if isinstance(structure_paths, (list)):
+            if isinstance(structure_paths[0], (str)):
+                relax_image_fws = []
+                for path in structure_paths:
+                    relax_image_fws.append(self.get_fw(structure_path=path))
+            else:
+                relax_image_fws = self.get_screening_fws(sorted_paths=structure_paths)
+        elif isinstance(structure_paths, (dict)):
+            relax_image_fws = []
+            if launch_mode == "all":
+                for key in structure_paths.keys():
+                    for path in structure_paths[key]:
+                        relax_image_fws.append(self.get_fw(structure_path=path))
+            elif launch_mode == "screening":
+                for key in structure_paths.keys():
+                    sorted_paths = structure_paths[key]
+                    relax_image_fws.extend(
+                        self.get_screening_fws(sorted_paths=sorted_paths)
+                    )
+
+        # place fws in temporary wf in order to use powerup_by_kwargs
+        # to apply powerups to image fireworks
+        if "vasp_powerups" in fw_spec.keys():
+            temp_wf = Workflow(relax_image_fws)
+            powerup_dicts = fw_spec["vasp_powerups"]
+            temp_wf = powerup_by_kwargs(temp_wf, powerup_dicts)
+            relax_image_fws = temp_wf.fws
+
+        return FWAction(additions=relax_image_fws)
+
+    def get_and_sort_paths(self, max_n, images_key=""):
+        sorted_paths = [[], [], []]
+        mid_n = int(max_n / 2)
+        q1 = int((max_n - mid_n) / 2)  # for second screening pass
+        q3 = int((max_n + mid_n) / 2)  # for second screening pass
+
+        for n in range(0, max_n):
+            path = "images." + images_key + "." + str(n) + ".input_structure"
+            if n == mid_n:  # path for first screening pass (center image index)
+                sorted_paths[0].append(path)
+            elif n in [q1, q3]:
+                sorted_paths[1].append(path)
+            else:
+                sorted_paths[-1].append(path)
+
+        return sorted_paths
+
+    def get_fw(self, structure_path, parents=None):
+        add_tags = self.get("add_tags")
+        fw = ImageFW(
+            approx_neb_wf_uuid=self["approx_neb_wf_uuid"],
+            structure_path=structure_path,
+            db_file=self["db_file"],
+            vasp_input_set=self.get("vasp_input_set"),
+            vasp_cmd=self["vasp_cmd"],
+            override_default_vasp_params=self.get("override_default_vasp_params"),
+            handler_group=self.get("handler_group"),
+            parents=parents,
+            add_additional_fields=self.get("add_additional_fields"),
+            add_tags=add_tags,
+        )
+        if isinstance(add_tags, (list)):
+            if "tags" in fw.spec.keys():
+                fw.spec["tags"].extend(add_tags)
+            else:
+                fw.spec["tags"] = add_tags
+        return fw
+
+    def get_screening_fws(self, sorted_paths):
+        if isinstance(sorted_paths, (list)) != True:
+            if (
+                any([isinstance(i, (list)) for i in sorted_paths]) != True
+                or len(sorted_paths) != 3
+            ):
+                raise TypeError("sorted_paths must be a list containing 3 lists")
+
+        s1_fw = self.get_fw(structure_path=sorted_paths[0][0])
+        # ToDo: modify this firework to add firetask that checks whether to run/defuse children
+
+        s2_fws = []
+        for path in sorted_paths[1]:
+            s2_fws.append(self.get_fw(structure_path=path, parents=s1_fw))
+            # ToDo: modify this firework to add firetask that checks whether to run/defuse children
+
+        remaining_fws = []
+        for path in sorted_paths[-1]:
+            remaining_fws.append(self.get_fw(structure_path=path, parents=s2_fws))
+
+        return [s1_fw] + s2_fws + remaining_fws