Merge pull request #35 from qchempku2017/generic_makers

Generic makers

WFacer/jobs.py

@@ -1,20 +1,16 @@
 """Unitary jobs used by an atomate2 workflow."""
+import importlib
 import logging
 from copy import deepcopy
 from warnings import warn
 
 import numpy as np
-from atomate2.vasp.jobs.core import RelaxMaker, StaticMaker, TightRelaxMaker
-from atomate2.vasp.sets.core import (
-    RelaxSetGenerator,
-    StaticSetGenerator,
-    TightRelaxSetGenerator,
-)
 from emmet.core.vasp.task_valid import TaskState  # atomate2 >= 0.0.11
 from jobflow import Flow, OnMissing, Response, job
 from pymatgen.analysis.elasticity.strain import Deformation
 from smol.cofe import ClusterExpansion
 from smol.moca import CompositionSpace
+from smol.utils.class_utils import class_name_from_str
 
 from .enumeration import (
     enumerate_compositions_as_counts,
@@ -110,33 +106,90 @@ def _enumerate_structures(
     return new_structures, new_sc_matrices, new_features
 
 
-def _get_vasp_makers(options):
-    """Get the required VASP makers."""
+def _get_structure_job_maker(maker_name, generator_kwargs=None, maker_kwargs=None):
+    """Get a single job maker from a structure job."""
+    # Format for a maker name: module:name-of-maker. module name must be in full path.
+    maker_module, maker_name = maker_name.split(":")[:2]
+    maker_module = maker_module.lower()
+
+    # Only cp2k, vasp and forcefields are supported.
+    if "amset" in maker_module or "common" in maker_module or "lobster" in maker_module:
+        raise NotImplementedError(
+            f"Makers in {maker_module} are not supported by WFacer!"
+        )
+
+    # Get maker from name. Note that only vasp supports tight relax (atomate2==0.0.10).
+    maker_name = class_name_from_str(maker_name)
+    try:
+        maker_class = getattr(importlib.import_module(maker_module), maker_name)
+    except AttributeError:
+        warn(
+            f"Maker {maker_name} is not supported by atomate module {maker_module}!"
+            f" Skipped."
+        )
+        return None
+    except ModuleNotFoundError:
+        warn(f"Atomate module {maker_module} not found! Skipped.")
+        return None
+
+    generator_kwargs = generator_kwargs or {}
+    maker_kwargs = maker_kwargs or {}
+
+    # vasp and cp2k makers require input set generators, and allows stopping children.
+    if "forcefield" not in maker_module:
+        # replace jobs.core with sets.core; replace Maker with SetGenerator
+        generator_module = ".".join(maker_module.split(".")[:-2]) + ".sets.core"
+        generator_name = maker_name[:-5] + "SetGenerator"
+        generator_class = getattr(
+            importlib.import_module(generator_module), generator_name
+        )
+        generator = generator_class(**generator_kwargs)
+
+        maker_kwargs["stop_children_kwargs"] = {"handle_unsuccessful": "error"}
+        return maker_class(input_set_generator=generator, **maker_kwargs)
+
+    # Force field makers does not require input set. Up to atomate2==0.0.10,
+    # they also won't support children kwargs as well as tight relaxation.
+    else:
+        return maker_class(**maker_kwargs)
+
+
+def _get_structure_calculation_makers(options):
+    """Get required calculation makers for a single structure."""
+    # Here, the maker names contain module specification.
+    relax_maker_name = options["relax_maker_name"]
     relax_gen_kwargs = options["relax_generator_kwargs"]
-    relax_generator = RelaxSetGenerator(**relax_gen_kwargs)
     relax_maker_kwargs = options["relax_maker_kwargs"]
-    # Force throwing out an error instead of defusing children, as parsing and
-    # fitting jobs are children of structure jobs and will be defused
-    # as well if not taken care of!
-
-    # Error handling will be taken care of by lost run detection and
-    # fixing functionalities in WFacer.fireworks_patches.
-    relax_maker_kwargs["stop_children_kwargs"] = {"handle_unsuccessful": "error"}
-    relax_maker = RelaxMaker(input_set_generator=relax_generator, **relax_maker_kwargs)
+    relax_maker = _get_structure_job_maker(
+        relax_maker_name,
+        generator_kwargs=relax_gen_kwargs,
+        maker_kwargs=relax_maker_kwargs,
+    )
+
+    static_maker_name = options["static_maker_name"]
     static_gen_kwargs = options["static_generator_kwargs"]
-    static_generator = StaticSetGenerator(**static_gen_kwargs)
     static_maker_kwargs = options["static_maker_kwargs"]
-    static_maker_kwargs["stop_children_kwargs"] = {"handle_unsuccessful": "error"}
-    static_maker = StaticMaker(
-        input_set_generator=static_generator, **static_maker_kwargs
-    )
-    tight_gen_kwargs = options["tight_generator_kwargs"]
-    tight_generator = TightRelaxSetGenerator(**tight_gen_kwargs)
-    tight_maker_kwargs = options["tight_maker_kwargs"]
-    tight_maker_kwargs["stop_children_kwargs"] = {"handle_unsuccessful": "error"}
-    tight_maker = TightRelaxMaker(
-        input_set_generator=tight_generator, **tight_maker_kwargs
+    static_maker = _get_structure_job_maker(
+        static_maker_name,
+        generator_kwargs=static_gen_kwargs,
+        maker_kwargs=static_maker_kwargs,
     )
+
+    # Note that only vasp supports tight relax (atomate2==0.0.10).
+    # You can also specify tight_maker_name as None to avoid using it.
+    # No extra argument is needed in options to achieve that.
+    tight_maker_name = options["tight_maker_name"]
+    if tight_maker_name is not None:
+        tight_gen_kwargs = options["tight_generator_kwargs"]
+        tight_maker_kwargs = options["tight_maker_kwargs"]
+        tight_maker = _get_structure_job_maker(
+            tight_maker_name,
+            generator_kwargs=tight_gen_kwargs,
+            maker_kwargs=tight_maker_kwargs,
+        )
+    else:
+        tight_maker = None
+
     return relax_maker, static_maker, tight_maker
 
 
@@ -281,7 +334,7 @@ def get_structure_calculation_flows(enum_output, last_ce_document):
     else:
         struct_id = len(last_ce_document.enumerated_structures)
     options = last_ce_document.ce_options
-    relax_maker, static_maker, tight_maker = _get_vasp_makers(options)
+    relax_maker, static_maker, tight_maker = _get_structure_calculation_makers(options)
 
     log.info("Performing ab-initio calculations.")
     new_structures = enum_output["new_structures"]
@@ -298,7 +351,7 @@ def get_structure_calculation_flows(enum_output, last_ce_document):
 
         jobs = list()
         jobs.append(relax_maker.make(def_structure))
-        if options["add_tight_relax"]:
+        if tight_maker is not None:
             tight_job = tight_maker.make(
                 jobs[-1].output.structure, prev_vasp_dir=jobs[-1].output.dir_name
             )

WFacer/preprocessing.py

@@ -403,6 +403,13 @@ def process_calculation_options(d):
                 Default is [1.03, 1.02, 1.01], which means to
                 stretch the structure by 3%, 2% and 1% along a, b, and c
                 directions, respectively.
+             relax_maker_name(str):
+                Name of relax maker class.
+                Default to "atomate2.vasp.jobs.core:RelaxMaker",
+                while **cp2k** and **force field** makers are also supported.
+                The module to which the Maker belongs must be specified first,
+                and separated from the maker class name with a quotation mark ":".
+                For example, write *atomate2.vasp.jobs.core:RelaxMaker*.
              relax_generator_kwargs(dict):
                 Additional arguments to pass into an :mod:`atomate2`
                 class :class:`VaspInputGenerator` that is used to specify
@@ -411,11 +418,13 @@ class :class:`VaspInputGenerator` that is used to specify
              relax_maker_kwargs(dict):
                 Additional arguments to initialize an :mod:`atomate2`
                 class :class:`RelaxMaker`. **Not frequently used**.
-             add_tight_relax(bool):
-                Whether to add a tight relaxation job after a coarse
-                relaxation. Default to True.
-                You may want to disable this if your system has
-                difficulty converging forces or energies.
+             tight_maker_name(str or None):
+                Name of the tight relax maker class.
+                Default to :class:`atomate2.vasp.jobs.core:TightRelaxMaker`,
+                setting to None will disable tight relaxation.
+
+                .. note:: Tight relax can only be supported by VASP.
+
              tight_generator_kwargs(dict):
                 Additional arguments to pass into an :mod:`atomate2`
                 class :class:`VaspInputGenerator`
@@ -425,6 +434,10 @@ class :class:`VaspInputGenerator`
                 Additional arguments to pass into an :mod:`atomate2`
                 class :class:`TightRelaxMaker`. A tight relax is performed after
                 relaxation, if add_tight_relax is True. **Not frequently used**.
+             static_maker_name(str):
+                 Name of static maker.
+                 Default to "atomate2.vasp.jobs.core:StaticMaker",
+                 while **cp2k** and **force fields** are also supported.
              static_generator_kwargs(dict):
                 Additional arguments to pass into an :mod:`atomate2`
                 class :class:`VaspInputGenerator` that is used to initialize
@@ -461,11 +474,19 @@ class :class:`VaspInputGenerator` that is used to initialize
 
     return {
         "apply_strain": strain_before_relax.tolist(),
+        "relax_maker_name": d.get(
+            "relax_maker_name", "atomate2.vasp.jobs.core:RelaxMaker"
+        ),
         "relax_generator_kwargs": d.get("relax_generator_kwargs", {}),
         "relax_maker_kwargs": d.get("relax_maker_kwargs", {}),
-        "add_tight_relax": d.get("add_tight_relax", True),
+        "tight_maker_name": d.get(
+            "tight_maker_name", "atomate2.vasp.jobs.core:TightRelaxMaker"
+        ),
         "tight_generator_kwargs": d.get("tight_generator_kwargs", {}),
         "tight_maker_kwargs": d.get("tight_maker_kwargs", {}),
+        "static_maker_name": d.get(
+            "static_maker_name", "atomate2.vasp.jobs.core:StaticMaker"
+        ),
         "static_generator_kwargs": d.get("static_generator_kwargs", {}),
         "static_maker_kwargs": d.get("static_maker_kwargs", {}),
         "other_properties": d.get("other_properties"),

WFacer/specie_decorators/charge.py

@@ -177,6 +177,8 @@ class MagneticChargeDecorator(GpOptimizedDecorator, ChargeDecorator):
     """Assign charges from magnitudes of total magentic moments on sites.
 
     Uses Gaussian process to optimize charge assignment.
+
+    .. note:: Does not support cp2k or force field!
     """
 
     decorated_prop_name = "oxi_state"

WFacer/utils/task_document.py

@@ -1,10 +1,18 @@
-"""Handle atomate2 taskdocument output.
+"""Handle atomate2 task document output.
 
 Extend the methods in this file if more site-wise properties other than magmom should
 be extracted.
+
+.. note:
+ Currently only supports reading from class :class:`emmet.core.tasks.TaskDoc`,
+ class :class:`atomate2.cp2k.schemas.task.TaskDocument` and
+ class :class:`atomate2.forcefields.schemas.ForceFieldTaskDocument`.
 """
 
-from pymatgen.entries.computed_entries import ComputedStructureEntry
+from atomate2.cp2k.schemas.task import TaskDocument
+from atomate2.forcefields.schemas import ForceFieldTaskDocument
+from emmet.core.tasks import TaskDoc
+from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry
 
 from ..specie_decorators.base import get_site_property_query_names_from_decorator
 from .query import get_property_from_object
@@ -40,9 +48,10 @@ def get_entry_from_taskdoc(taskdoc, property_and_queries=None, decorator_names=N
     """Get the computed structure entry from :class:`TaskDoc`.
 
     Args:
-        taskdoc(TaskDoc):
-            A task document generated as vasp task output by emmet-core.
-        property_and_queries(list of (str, str) or str): optional
+        taskdoc(StructureMetadata):
+            A task document generated as vasp task output by emmet-core, CP2K
+            or force fields.
+        property_and_queries(list of (str, str) or list of str): optional
             A list of property names to be retrieved from taskdoc,
             and the query string to retrieve them, paired in tuples.
             If only strings are given, will also query with the given
@@ -61,10 +70,21 @@ def get_entry_from_taskdoc(taskdoc, property_and_queries=None, decorator_names=N
             property required by decorator, and the properties
             dict ready to be inserted into a :class:`CeDataWangler`.
     """
+    if not isinstance(taskdoc, (TaskDoc, TaskDocument, ForceFieldTaskDocument)):
+        raise ValueError(f"Document type {type(taskdoc)} not supported!")
     # Final optimized structure.
     structure = taskdoc.structure
     # The computed entry, not including the structure.
-    computed_entry = taskdoc.entry
+    if isinstance(taskdoc, (TaskDoc, TaskDocument)):
+        computed_entry = taskdoc.entry
+    # In ForcefieldTaskdocument, Need to retrieve the entry from ionic steps.
+    else:
+        computed_entry = ComputedEntry(
+            composition=taskdoc.structure.composition,
+            energy=taskdoc.output.energy,
+            correction=0.0,
+        )
+
     prop_dict = {}
     if property_and_queries is not None:
         for p in property_and_queries:
@@ -84,7 +104,8 @@ def get_entry_from_taskdoc(taskdoc, property_and_queries=None, decorator_names=N
             site_property_query_names = get_site_property_query_names_from_decorator(d)
             for sp, query in site_property_query_names:
                 # Total magnetization on each site is already read and added to
-                # structure by atomate2. It should be overwritten.
+                # structure by atomate2 for all kinds of makers (VASP, CP2K, force fields).
+                # There is no need to do it again.
                 if sp != "magmom":
                     site_props[sp] = get_property_from_object(taskdoc, query)
     for sp, prop in site_props.items():

tests/conftest.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import pytest
+from atomate2.forcefields.schemas import ForceFieldTaskDocument
 from emmet.core.tasks import TaskDoc  # emmet-core >= 0.60.0.
 from monty.serialization import loadfn
 from pydantic import parse_file_as
@@ -180,6 +181,11 @@ def single_wrangler_sin(single_ensemble_sin):
     return gen_random_wrangler(single_ensemble_sin, 30)
 
 
-@pytest.fixture(scope="package", params=["zns_taskdoc.json"])
+@pytest.fixture(scope="package", params=["zns_taskdoc.json", "zns_ff_taskdoc.json"])
 def single_taskdoc(request):
-    return parse_file_as(TaskDoc, os.path.join(DATA_DIR, request.param))
+    if "ff" not in request.param:
+        return parse_file_as(TaskDoc, os.path.join(DATA_DIR, request.param))
+    else:
+        return parse_file_as(
+            ForceFieldTaskDocument, os.path.join(DATA_DIR, request.param)
+        )

tests/data/default_options.yaml

@@ -17,11 +17,13 @@ duplicacy_criteria: "correlations"
 n_parallel: null
 keep_ground_states: true
 apply_strain: [[1.03, 0, 0], [0, 1.02, 0], [0, 0, 1.01]]
+relax_maker_name: "atomate2.vasp.jobs.core:RelaxMaker"
 relax_generator_kwargs: {}
 relax_maker_kwargs: {}
-add_tight_relax: true
+tight_maker_name: "atomate2.vasp.jobs.core:TightRelaxMaker"
 tight_generator_kwargs: {}
 tight_maker_kwargs: {}
+static_maker_name: "atomate2.vasp.jobs.core:StaticMaker"
 static_generator_kwargs: {}
 static_maker_kwargs: {}
 other_properties: null