Source code for atomate2.qchem.jobs.base

"""Definition of a base QChem Maker."""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Callable

from emmet.core.qc_tasks import TaskDoc
from jobflow import Maker, Response, job
from monty.serialization import dumpfn
from monty.shutil import gzip_dir
from pymatgen.io.qchem.inputs import QCInput

from atomate2.qchem.files import copy_qchem_outputs
from atomate2.qchem.run import run_qchem, should_stop_children
from atomate2.qchem.sets.base import QCInputGenerator

if TYPE_CHECKING:
    from pymatgen.core.structure import Molecule




def qchem_job(method: Callable) -> job:
    """
    Decorate the ``make`` method of QChem job makers.

    This is a thin wrapper around :obj:`~jobflow.core.job.Job` that configures common
    settings for all QChem jobs. It also configures the output schema to be a QChem
    :obj:`.TaskDoc`.

    Any makers that return QChem jobs (not flows) should decorate the ``make`` method
    with @qchem_job. For example:

    .. code-block:: python

        class MyQChemMaker(BaseQChemMaker):
            @qchem_job
            def make(molecule):
                # code to run QChem job.
                pass

    Parameters
    ----------
    method : callable
        A BaseQChemMaker.make method. This should not be specified directly and is
        implied by the decorator.

    Returns
    -------
    callable
        A decorated version of the make function that will generate QChem jobs.
    """
    return job(method, data=QCInput, output_schema=TaskDoc)





@dataclass
class BaseQCMaker(Maker):
    """
    Base QChem job maker.

    Parameters
    ----------
    name : str
        The job name.
    input_set_generator : .QChemInputGenerator
        A generator used to make the input set.
    write_input_set_kwargs : dict
        Keyword arguments that will get passed to :obj:`.write_qchem_input_set`.
    copy_qchem_kwargs : dict
        Keyword arguments that will get passed to :obj:`.copy_qchem_outputs`.
    run_qchem_kwargs : dict
        Keyword arguments that will get passed to :obj:`.run_qchem`.
    task_document_kwargs : dict
        Keyword arguments that will get passed to :obj:`.TaskDoc.from_directory`.
    stop_children_kwargs : dict
        Keyword arguments that will get passed to :obj:`.should_stop_children`.
    write_additional_data : dict
        Additional data to write to the current directory. Given as a dict of
        {filename: data}. Note that if using FireWorks, dictionary keys cannot contain
        the "." character which is typically used to denote file extensions. To avoid
        this, use the ":" character, which will automatically be converted to ".". E.g.
        ``{"my_file:txt": "contents of the file"}``.
    """

    name: str = "base qchem job"
    input_set_generator: QCInputGenerator = field(
        default_factory=lambda: QCInputGenerator(
            job_type="sp", scf_algorithm="diis", basis_set="def2-qzvppd"
        )
    )
    write_input_set_kwargs: dict = field(default_factory=dict)
    copy_qchem_kwargs: dict = field(default_factory=dict)
    run_qchem_kwargs: dict = field(default_factory=dict)
    task_document_kwargs: dict = field(default_factory=dict)
    stop_children_kwargs: dict = field(default_factory=dict)
    write_additional_data: dict = field(default_factory=dict)



    @qchem_job
    def make(
        self, molecule: Molecule, prev_qchem_dir: str | Path | None = None
    ) -> Response:
        """Run a QChem calculation.

        Parameters
        ----------
        molecule : Molecule
            A pymatgen molecule object.
        prev_qchem_dir : str or Path or None
            A previous QChem calculation directory to copy output files from.
        """
        # copy previous inputs
        from_prev = prev_qchem_dir is not None
        if prev_qchem_dir is not None:
            copy_qchem_outputs(prev_qchem_dir, **self.copy_qchem_kwargs)

        self.write_input_set_kwargs.setdefault("from_prev", from_prev)

        # write qchem input files
        # self.input_set_generator.get_input_set(molecule).write_inputs()
        self.input_set_generator.get_input_set(molecule)

        # write any additional data
        for filename, data in self.write_additional_data.items():
            dumpfn(data, filename.replace(":", "."))

        # run qchem
        run_qchem(**self.run_qchem_kwargs)

        # parse qchem outputs
        task_doc = TaskDoc.from_directory(Path.cwd(), **self.task_document_kwargs)
        # task_doc.task_label = self.name
        task_doc.task_type = self.name

        # decide whether child jobs should proceed
        stop_children = should_stop_children(task_doc, **self.stop_children_kwargs)

        # gzip folder
        gzip_dir(".")

        return Response(
            stop_children=stop_children,
            stored_data={"custodian": task_doc.custodian},
            output=task_doc,
        )