How to develop a new atomate2 workflow

Anatomy of an atomate2 computational workflow (i.e., what do I need to write?)

Every atomate2 workflow is an instance of jobflow’s Flow class, which is a collection of Job and/or other Flow objects. So your end goal is to produce a Flow.

In the context of computational materials science, Flow objects are most easily created by a Maker, which contains a factory method make() that produces a Flow, given certain inputs. Typically, the input to Maker.make() includes atomic coordinate information in the form of a pymatgen Structure or Molecule object. So the basic signature looks like this:

class ExampleMaker(Maker):
    def make(self, structure: Structure) -> Flow:
        # take the input structure and return a `Flow`
        return Flow(...)

The Maker class usually contains most calculation parameters and other settings needed to set up the calculation correctly. Much of this logic can be written within methods of the Maker or independent functions and then turned into a ‘ Job ‘ via the make method and the @job decorator, or into a Flow via the make method and by returning a Flow object.

One common task in almost any materials science calculation is writing calculation input files to disk so the underlying software (e.g., VASP, Q-Chem, CP2K, etc.) can execute them. This is preferably done via a pymatgen InputSet class. InputSet is essentially a dict-like container that specifies which files to write and their contents. It has a write_input() method that writes those files to disk, e.g.,

class ExampleInputSet(InputSet):
    def __init__(self, inputs: dict[PathLike, InputFile]):
        self.inputs = inputs

    def write_input(self, directory: str | Path) -> None:
        # write the files held by this `InputSet` to `directory`
        for path, file in self.inputs.items():
            file.write(directory / path)

Similarly to the way that Maker classes generate Flows and Jobs, InputSets are most easily created by InputGenerator classes. InputGenerator classes have a method get_input_set() that typically takes atomic coordinates (e.g., a Structure or Molecule object) and produces an InputSet, e.g.,

class ExampleInputGenerator(InputGenerator):
    def get_input_set(self, structure: Structure) -> ExampleInputSet:
        # take the input structure, determine appropriate
        # input file contents, and return an `InputSet`
        return ExampleInputSet(...)

pymatgen already contains InputGenerator / InputSet pairs for many common codes, so when developing a workflow Maker, it is convenient to reuse them to prepare your files. The Maker holds an InputGenerator as a class parameter and calls get_input_set() / write_input() from within its make() method.

Finally, most atomate2 workflows return structured output in the form of “Task Documents”. Task documents are instances of emmet’s BaseTaskDocument class (a pydantic.BaseModel subclass, an object similar to a python @dataclass but with additional type checking and validation) that define schemas for storing calculation outputs. emmet already contains calculation schemas for codes utilized by the Materials Project (e.g., VASP, Q-Chem, FEFF) as well as a number of schemas for code-agnostic structural and molecular information (for example, the MaterialsDoc is a schema for solid material calculation data). atomate2 can also interpret output generated by cclib, which is able to parse the output of many additional codes.

Putting this together for the specific case of executing a simulation code, a Maker’s make() method is itself decorated with @job, which turns make() into the Job that actually gets executed, potentially on a remote compute resource — this is why any code that touches the filesystem or launches an external process must live inside make() (or another @job-decorated function it calls), rather than elsewhere in the Maker. A typical make() for such a task will (1) write the input files, (2) run the underlying code, and (3) parse the output directory into a TaskDocument, which is returned as the Job’s output via a Response:

class ExampleMaker(Maker):
    input_set_generator: ExampleInputGenerator = field(
        default_factory=ExampleInputGenerator
    )

    @job
    def make(self, structure: Structure) -> Response:
        # create and write the `InputSet` to the current directory
        input_set = self.input_set_generator.get_input_set(structure)
        input_set.write_input(".")

        # run the underlying code
        run_example_code()

        # parse the output directory into a task document
        task_doc = ExampleTaskDoc.from_directory(".")

        return Response(output=task_doc)

For a real, more elaborate example, see e.g. atomate2.vasp.jobs.base.BaseVaspMaker.make in atomate2/vasp/jobs/base.py, which follows this same pattern using write_vasp_input_set(), run_vasp(), and TaskDoc.from_directory().

In summary, a new atomate2 workflow consists of the following components:

  • A Maker that actually generates the workflow

  • One or more Job and/or Flow classes that define the discrete steps in the workflow

  • (optionally) an InputGenerator that produces a pymatgen InputSet for writing calculation input files

  • (optionally) a TaskDocument that defines a schema for storing the output data

Where do I put my code?

Because of the distributed design of the MP Software Ecosystem, writing a complete new workflow may involve making contributions to more than one GitHub repository. The following guidelines should help you understand where to put your contribution.

  • All workflow code (Job, Flow, Maker) belongs in atomate2

  • InputSet and InputGenerator code belongs in pymatgen. However, if you need to create these classes from scratch (i.e., you are working with a code that is not already supported inpymatgen), then it is recommended to include them in atomate2 at first to facilitate rapid iteration. Once mature, they can be moved to pymatgen or to a pymatgen addon package.

  • TaskDocument schemas should generally be developed in atomate2 alongside the workflow code. We recommend that you first check emmet to see if there is an existing schema that matches what you need. If so, you can import it. If not, check cclib. cclib output can be imported via atomate2.common.schemas.TaskDocument. If neither code has what you need, then new schemas should be developed within atomate2 (or cclib).