Frequency Flattening Optimizer (#863)

* flows for ffopt

* opt freq works

* ffopt working

* remove vim swap files and add extensions to gitignore

* precommit

* Add prev_dir kwarg to qchem jobs and deprecation warning for prev_qchem_dir

* linting

* cleanup

* fix failing set test

* slight cleanup

* move H2O.xyz file to pytest fixture

* make qchem taskdoc task_type more flexible, add tests for ffopt

* precommit

* remove duplicate frequency_analysis_2 test dir

* precommit

---------

Co-authored-by: Rohith Srinivaas Mohanakrishnan <[email protected]>
Co-authored-by: Rohith Srinivaas Mohanakrishnan <[email protected]>
Co-authored-by: Aaron Kaplan <[email protected]>
Co-authored-by: esoteric-ephemera <[email protected]>

.gitignore

@@ -64,6 +64,10 @@ docs/reference/atomate2.*
 
 .DS_Store
 
+# vim extensions
+*.swp
+*.swo
+
 # see https://github.com/materialsproject/atomate2/issues/345
 *.doctrees*
 

src/atomate2/qchem/flows/__init__.py

@@ -0,0 +1 @@
+"""Flows for running Q-Chem calculations."""

src/atomate2/qchem/flows/core.py

@@ -0,0 +1,225 @@
+"""Define core QChem flows."""
+
+from __future__ import annotations
+
+from copy import deepcopy
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from jobflow import Flow, Maker, Response, job
+
+from atomate2.qchem.jobs.core import FreqMaker, OptMaker
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from jobflow import Job
+    from pymatgen.core.structure import Molecule
+
+    from atomate2.qchem.jobs.base import BaseQCMaker
+
+
+@dataclass
+class DoubleOptMaker(Maker):
+    """
+    Maker to perform a double Qchem relaxation.
+
+    Parameters
+    ----------
+    name : str
+        Name of the flows produced by this maker.
+    relax_maker1 : .BaseVaspMaker
+        Maker to use to generate the first relaxation.
+    relax_maker2 : .BaseVaspMaker
+        Maker to use to generate the second relaxation.
+    """
+
+    name: str = "double opt"
+    opt_maker1: BaseQCMaker | None = field(default_factory=OptMaker)
+    opt_maker2: BaseQCMaker = field(default_factory=OptMaker)
+
+    def make(self, molecule: Molecule, prev_dir: str | Path | None = None) -> Flow:
+        """
+        Create a flow with two chained molecular optimizations.
+
+        Parameters
+        ----------
+        molecule : .Molecule
+            A pymatgen Molecule object.
+        prev_dir : str or Path or None
+            A previous QChem calculation directory to copy output files from.
+
+        Returns
+        -------
+        Flow
+            A flow containing two geometric optimizations.
+        """
+        jobs: list[Job] = []
+        if self.opt_maker1:
+            # Run a pre-relaxation
+            opt1 = self.opt_maker1.make(molecule, prev_dir=prev_dir)
+            opt1.name += " 1"
+            jobs += [opt1]
+            molecule = opt1.output.optimized_molecule
+            prev_dir = opt1.output.dir_name
+
+        opt2 = self.opt_maker2.make(molecule, prev_dir=prev_dir)
+        opt2.name += " 2"
+        jobs += [opt2]
+
+        return Flow(jobs, output=opt2.output, name=self.name)
+
+    @classmethod
+    def from_opt_maker(cls, opt_maker: BaseQCMaker) -> DoubleOptMaker:
+        """
+        Instantiate the DoubleRelaxMaker with two relax makers of the same type.
+
+        Parameters
+        ----------
+        opt_maker : .BaseQCMaker
+            Maker to use to generate the first and second geometric optimizations.
+        """
+        return cls(relax_maker1=deepcopy(opt_maker), relax_maker2=deepcopy(opt_maker))
+
+
+@dataclass
+class FrequencyOptMaker(Maker):
+    """
+    Maker to perform a frequency calculation after an optimization.
+
+    Parameters
+    ----------
+    name : str
+        Name of the flows produced by this maker.
+    opt_maker : .BaseQCMaker
+        Maker to use to generate the opt maker
+    freq_maker : .BaseQCMaker
+        Maker to use to generate the freq maker
+    """
+
+    name: str = "opt frequency"
+    opt_maker: BaseQCMaker = field(default_factory=OptMaker)
+    freq_maker: BaseQCMaker = field(default_factory=FreqMaker)
+
+    def make(self, molecule: Molecule, prev_dir: str | Path | None = None) -> Flow:
+        """
+        Create a flow with optimization followed by frequency calculation.
+
+        Parameters
+        ----------
+        molecule : .Molecule
+            A pymatgen Molecule object.
+        prev_dir : str or Path or None
+            A previous QChem calculation directory to copy output files from.
+
+        Returns
+        -------
+        Flow
+            A flow containing with optimization and frequency calculation.
+        """
+        jobs: list[Job] = []
+        opt = self.opt_maker.make(molecule, prev_dir=prev_dir)
+        opt.name = "Geometry Optimization"
+        jobs += [opt]
+
+        freq = self.freq_maker.make(
+            molecule=opt.output.output.optimized_molecule,
+            prev_dir=opt.output.dir_name,
+        )
+        freq.name = "Frequency Analysis"
+        jobs += [freq]
+
+        return Flow(
+            jobs, output={"opt": opt.output, "freq": freq.output}, name=self.name
+        )
+
+
+@dataclass
+class FrequencyOptFlatteningMaker(Maker):
+    """
+    Maker to perform a frequency calculation after an optimization.
+
+    Parameters
+    ----------
+    name : str
+        Name of the flows produced by this maker.
+    opt_maker : .BaseQCMaker
+        Maker to use to generate the opt maker
+    freq_maker : .BaseQCMaker
+        Maker to use to generate the freq maker
+    """
+
+    name: str = "frequency flattening opt"
+    opt_maker: BaseQCMaker = field(default_factory=OptMaker)
+    freq_maker: BaseQCMaker = field(default_factory=FreqMaker)
+    scale: float = 1.0
+    max_ffopt_runs: int = 5
+
+    @job
+    def make(
+        self,
+        molecule: Molecule,
+        mode: list | None = None,
+        lowest_freq: float = -1.0,
+        ffopt_runs: int = 0,
+        overwrite_inputs: dict | None = None,
+        prev_dir: str | Path | None = None,
+    ) -> Flow:
+        """
+        Optimize geometry and perturb negative frequency modes.
+
+        Parameters
+        ----------
+        molecule : .Molecule
+            A pymatgen Molecule object.
+        prev_dir : str or Path or None
+            A previous QChem calculation directory to copy output files from.
+
+        Returns
+        -------
+        Flow
+            A flow containing with optimization and frequency calculation.
+        """
+        mode = mode or [[0.0, 0.0, 0.0] for _ in range(len(molecule))]
+
+        if overwrite_inputs is not None:
+            self.opt_maker.input_set_generator.overwrite_inputs = overwrite_inputs
+            self.freq_maker.input_set_generator.overwrite_inputs = overwrite_inputs
+
+        new_flow = None
+        new_output = None
+
+        if (lowest_freq < 0) and (ffopt_runs < self.max_ffopt_runs):
+            jobs: list[Job] = []
+
+            for idx in range(len(molecule)):
+                molecule.translate_sites(
+                    indices=[idx], vector=[self.scale * v for v in mode[idx]]
+                )
+
+            opt = self.opt_maker.make(molecule, prev_dir=prev_dir)
+            opt.name = "Geometry Optimization"
+            jobs += [opt]
+            molecule = opt.output.output.optimized_molecule
+
+            freq = self.freq_maker.make(molecule, prev_dir=prev_dir)
+            freq.name = f"Frequency Analysis {ffopt_runs + 1}"
+            jobs += [freq]
+
+            recursive = self.make(
+                molecule,
+                mode=freq.output.output.frequency_modes[0],
+                lowest_freq=freq.output.output.frequencies[0],
+                ffopt_runs=ffopt_runs + 1,
+                prev_dir=prev_dir,
+            )
+            new_flow = Flow([*jobs, recursive], output=recursive.output)
+            new_output = recursive.output
+
+        elif ffopt_runs == 0:
+            freq = self.freq_maker.make(molecule, prev_dir=prev_dir)
+            freq.name = f"Frequency Analysis {ffopt_runs + 1}"
+            new_flow = [freq]
+            new_output = freq.output
+
+        return Response(replace=new_flow, output=new_output)

src/atomate2/qchem/jobs/base.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import logging
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import TYPE_CHECKING, Callable
@@ -19,6 +20,8 @@
 if TYPE_CHECKING:
     from pymatgen.core.structure import Molecule
 
+logger = logging.getLogger(__name__)
+
 
 def qchem_job(method: Callable) -> job:
     """
@@ -94,30 +97,51 @@ class BaseQCMaker(Maker):
     task_document_kwargs: dict = field(default_factory=dict)
     stop_children_kwargs: dict = field(default_factory=dict)
     write_additional_data: dict = field(default_factory=dict)
+    task_type: str | None = None
 
     @qchem_job
     def make(
-        self, molecule: Molecule, prev_qchem_dir: str | Path | None = None
+        self,
+        molecule: Molecule,
+        prev_dir: str | Path | None = None,
+        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
+        prev_dir : str or Path or None
+            A previous calculation directory to copy output files from.
+        prev_qchem_dir (deprecated): 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)
+            logger.warning(
+                "`prev_qchem_dir` will be deprecated in a future release. "
+                "Please use `prev_dir` instead."
+            )
+            if prev_dir is not None:
+                logger.warning(
+                    "You set both `prev_dir` and `prev_qchem_dir`, "
+                    "only `prev_dir` will be used."
+                )
+            else:
+                prev_dir = prev_qchem_dir
+
+        if from_prev := (prev_dir is not None):
+            copy_qchem_outputs(prev_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)
+        self.input_set_generator.get_input_set(molecule).write_input(
+            directory=Path.cwd()
+        )
 
         # write any additional data
         for filename, data in self.write_additional_data.items():
@@ -129,7 +153,7 @@ def make(
         # 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
+        task_doc.task_type = self.name if self.task_type is None else self.task_type
 
         # decide whether child jobs should proceed
         stop_children = should_stop_children(task_doc, **self.stop_children_kwargs)

src/atomate2/qchem/jobs/core.py

@@ -197,6 +197,7 @@ class FreqMaker(BaseQCMaker):
 
     name: str = "frequency"
     input_set_generator: QCInputGenerator = field(default_factory=FreqSetGenerator)
+    task_type: str = "Frequency Analysis"
 
 
 @dataclass

src/atomate2/qchem/run.py

@@ -21,7 +21,7 @@
     from emmet.core.qc_tasks import TaskDoc
 
 
-_DEFAULT_HANDLERS = (QChemErrorHandler,)
+_DEFAULT_HANDLERS = (QChemErrorHandler(),)
 
 logger = logging.getLogger(__name__)
 
@@ -95,7 +95,6 @@ def run_qchem(
         scratch_dir=scratch_dir,
         **custodian_kwargs,
     )
-
     logger.info("Running QChem using custodian.")
     c.run()
 

src/atomate2/qchem/sets/base.py

@@ -66,11 +66,13 @@ def write_input(
         inputs.update(self.optional_files)
 
         for key, val in inputs.items():
-            if val is not None and (overwrite or not (directory / key).exists()):
-                with zopen(directory / key, "wt") as file:
+            inp_key = "mol.qin" if key == "Input_Dict" else f"{key}.qin"
+            if val is not None and (overwrite or not (directory / inp_key).exists()):
+                # should this be open instead of zopen? can QChem inputs be gzipped?
+                with zopen(directory / inp_key, "wt") as file:
                     file.write(str(val))
-            elif not overwrite and (directory / key).exists():
-                raise FileExistsError(f"{directory / key} already exists.")
+            elif not overwrite and (directory / inp_key).exists():
+                raise FileExistsError(f"{directory / inp_key} already exists.")
 
     @staticmethod
     def from_directory(
@@ -91,9 +93,9 @@ def from_directory(
 
         inputs = {}
         for name, obj in objs.items():
-            if (directory / name).exists():
-                inputs[name.lower()] = obj.from_file(directory / name)
-
+            file_path = directory / ("mol.qin" if name == "Input_Dict" else name)
+            if file_path.exists():
+                inputs[name.lower()] = obj.from_file(file_path)
         optional_inputs = {}
         if optional_files is not None:
             for name, obj in optional_files.items():

src/atomate2/settings.py

@@ -220,7 +220,7 @@ class Atomate2Settings(BaseSettings):
     # QChem specific settings
 
     QCHEM_CMD: str = Field(
-        "qchem_std", description="Command to run standard version of qchem."
+        "qchem", description="Command to run standard version of qchem."
     )
 
     QCHEM_CUSTODIAN_MAX_ERRORS: int = Field(