Merge branch 'main' of https://github.com/rdguha1995/atomate

rdguha1995 · rdguha1995 · commit 7cebcb8613bc · 2022-07-12T11:37:14.000-07:00
diff --git a/atomate/qchem/firetasks/parse_outputs.py b/atomate/qchem/firetasks/parse_outputs.py
@@ -116,3 +116,101 @@ def run_task(self, fw_spec):
             stored_data={"task_id": task_doc.get("task_id", None)},
             update_spec=update_spec,
         )
+
+
+@explicit_serialize
+class ProtCalcToDb(FiretaskBase):
+    """
+    Enter a QChem run of a proton into the database. Uses current directory unless you
+    specify calc_dir or calc_loc.
+
+    Optional params:
+        calc_dir (str): path to dir (on current filesystem) that contains QChem
+            input and output files for both the H0 and H2+. Default: use current working directory.
+        calc_loc (str OR bool): if True will set most recent calc_loc. If str
+            search for the most recent calc_loc with the matching name
+        input_file_H0 (str): name of the QChem input file for H0 calculation
+        output_file_H0 (str): name of the QChem output file for H0 calculation
+        input_file_H2 (str): name of the QChem input file for H2 calculation
+        output_file_H2 (str): name of the QChem output file for H2 calculation
+        additional_fields (dict): dict of additional fields to add
+        db_file (str): path to file containing the database credentials.
+            Supports env_chk. Default: write data to JSON file.
+        runs (list): Series of file suffixes that the Drone should look for
+            when parsing output.
+    """
+
+    optional_params = [
+        "calc_dir",
+        "calc_loc",
+        "input_file_H0",
+        "output_file_H0",
+        "input_file_H2",
+        "output_file_H2",
+        "additional_fields",
+        "db_file",
+        "runs",
+    ]
+
+    def run_task(self, fw_spec):
+        # get the directory that contains the QChem dir to parse
+        calc_dir = os.getcwd()
+        if "calc_dir" in self:
+            calc_dir = self["calc_dir"]
+        elif self.get("calc_loc"):
+            calc_dir = get_calc_loc(self["calc_loc"], fw_spec["calc_locs"])["path"]
+        input_file_H0 = self.get("input_file_H0", "H0.qin")
+        output_file_H0 = self.get("output_file_H0", "H0.qout")
+        input_file_H2 = self.get("input_file_H2", "H2_plus.qin")
+        output_file_H2 = self.get("output_file_H2", "H2_plus.qout")
+        runs = self.get("runs", None)
+
+        # parse the QChem directory
+        logger.info(f"PARSING DIRECTORY: {calc_dir}")
+
+        additional_fields = self.get("additional_fields", {})
+
+        drone = QChemDrone(runs=runs, additional_fields=additional_fields)
+
+        # assimilate (i.e., parse)
+        task_doc_1 = drone.assimilate(
+            path=calc_dir,
+            input_file=input_file_H0,
+            output_file=output_file_H0,
+            multirun=False,
+        )
+
+        task_doc_2 = drone.assimilate(
+            path=calc_dir,
+            input_file=input_file_H2,
+            output_file=output_file_H2,
+            multirun=False,
+        )
+
+        task_doc_clean = task_doc_1
+        task_doc_clean["calcs_reversed"].append(task_doc_2["calcs_reversed"][0])
+        task_doc_clean["input"]["initial_molecule"]["charge"] = 1
+        task_doc_clean["input"]["initial_molecule"]["spin_multiplicity"] = 1
+        task_doc_clean["orig"]["molecule"]["charge"] = 1
+        task_doc_clean["orig"]["molecule"]["spin_multiplicity"] = 1
+        task_doc_clean["output"]["initial_molecule"]["charge"] = 1
+        task_doc_clean["output"]["initial_molecule"]["spin_multiplicity"] = 1
+        task_doc_clean["output"]["final_energy"] = (
+            task_doc_2["output"]["final_energy"] - task_doc_1["output"]["final_energy"]
+        )
+
+        # get the database connection
+        db_file = env_chk(self.get("db_file"), fw_spec)
+
+        # db insertion or taskdoc dump
+        if not db_file:
+            with open(os.path.join(calc_dir, "task.json"), "w") as f:
+                f.write(json.dumps(task_doc_clean, default=DATETIME_HANDLER))
+        else:
+            mmdb = QChemCalcDb.from_db_file(db_file, admin=True)
+            t_id = mmdb.insert(task_doc_clean)
+            logger.info(f"Finished parsing with task_id: {t_id}")
+
+        return FWAction(
+            stored_data={"task_id": task_doc_clean.get("task_id", None)},
+        )
diff --git a/atomate/qchem/firetasks/run_calc.py b/atomate/qchem/firetasks/run_calc.py
@@ -144,7 +144,9 @@ def run_task(self, fw_spec):
         freq_before_opt = self.get("freq_before_opt", False)
 
         handler_groups = {
-            "default": [QChemErrorHandler(input_file=input_file, output_file=output_file)],
+            "default": [
+                QChemErrorHandler(input_file=input_file, output_file=output_file)
+            ],
             "no_handler": [],
         }
 
@@ -206,7 +208,9 @@ def run_task(self, fw_spec):
         # construct handlers
         handlers = handler_groups[self.get("handler_group", "default")]
 
-        c = Custodian(handlers, jobs, max_errors=max_errors, gzipped_output=gzipped_output)
+        c = Custodian(
+            handlers, jobs, max_errors=max_errors, gzipped_output=gzipped_output
+        )
 
         c.run()
 
@@ -248,7 +252,9 @@ def _verify_inputs(self):
         ref_qin = QCInput.from_file(os.path.join(self["ref_dir"], input_file))
 
         np.testing.assert_equal(ref_qin.molecule.species, user_qin.molecule.species)
-        np.testing.assert_allclose(ref_qin.molecule.cart_coords, user_qin.molecule.cart_coords, atol=0.0001)
+        np.testing.assert_allclose(
+            ref_qin.molecule.cart_coords, user_qin.molecule.cart_coords, atol=0.0001
+        )
         for key in ref_qin.rem:
             if user_qin.rem.get(key) != ref_qin.rem.get(key):
                 raise ValueError(f"Rem key {key} is inconsistent!")
diff --git a/atomate/qchem/fireworks/core.py b/atomate/qchem/fireworks/core.py
@@ -5,11 +5,13 @@
 from itertools import chain
 
 from fireworks import Firework
+from pymatgen.core.sites import Site
+from pymatgen.core.structure import Molecule
 
 from atomate.qchem.firetasks.critic2 import ProcessCritic2, RunCritic2
 from atomate.qchem.firetasks.fragmenter import FragmentMolecule
 from atomate.qchem.firetasks.geo_transformations import PerturbGeometry
-from atomate.qchem.firetasks.parse_outputs import QChemToDb
+from atomate.qchem.firetasks.parse_outputs import ProtCalcToDb, QChemToDb
 from atomate.qchem.firetasks.run_calc import RunQChemCustodian
 from atomate.qchem.firetasks.write_inputs import WriteInputFromIOSet
 from pymatgen.core.structure import Molecule
@@ -210,6 +212,121 @@ def __init__(
         super().__init__(t, parents=parents, name=name, **kwargs)
 
 
+class ProtonEnergyFW(Firework):
+    def __init__(
+        self,
+        name="proton electronic energy",
+        qchem_cmd=">>qchem_cmd<<",
+        multimode=">>multimode<<",
+        max_cores=">>max_cores<<",
+        qchem_input_params=None,
+        db_file=None,
+        parents=None,
+        max_errors=5,
+        **kwargs
+    ):
+        """
+        For this custom Firework the electronic energy of a proton in a specific solvent environment is approximated.
+        Since a proton has 0 electrons,running a QChem job would yield an error. The energy can be approximated by
+        calculating the electronic energy of a proton and a hydrogen atom at infinite separation and then subtracting
+        the electronic energy of a hydrogen atom. This Firework combines these two calculations and adds a task doc to the
+        DB with the separate calculation details and the effective energy after subtraction.
+
+        Args:
+            name (str): Name for the Firework.
+            qchem_cmd (str): Command to run QChem. Supports env_chk.
+            multimode (str): Parallelization scheme, either openmp or mpi. Supports env_chk.
+            max_cores (int): Maximum number of cores to parallelize over. Supports env_chk.
+            qchem_input_params (dict): Specify kwargs for instantiating the input set parameters
+                calculating the electronic energy of the proton in a specific solvent environment.
+                The energy of the proton will be effectively zero in vacuum. Use either pcm_dieletric
+                or some smd_solvent.Basic uses would be to modify the default inputs of the set,
+                such as dft_rung, basis_set, pcm_dielectric, scf_algorithm, or max_scf_cycles.
+                See pymatgen/io/qchem/sets.py for default values of all input parameters. For
+                instance, if a user wanted to use a more advanced DFT functional, include a pcm
+                with a dielectric of 30, and use a larger basis, the user would set
+                qchem_input_params = {"dft_rung": 5, "pcm_dielectric": 30, "basis_set":
+                "6-311++g**"}. However, more advanced customization of the input is also
+                possible through the overwrite_inputs key which allows the user to directly
+                modify the rem, pcm, smd, and solvent dictionaries that QChemDictSet passes to
+                inputs.py to print an actual input file. For instance, if a user wanted to set
+                the sym_ignore flag in the rem section of the input file to true, then they
+                would set qchem_input_params = {"overwrite_inputs": "rem": {"sym_ignore":
+                "true"}}. Of course, overwrite_inputs could be used in conjunction with more
+                typical modifications, as seen in the test_double_FF_opt workflow test.
+            db_file (str): Path to file specifying db credentials to place output parsing.
+            parents ([Firework]): Parents of this particular Firework.
+            **kwargs: Other kwargs that are passed to Firework.__init__.
+        """
+
+        qchem_input_params = qchem_input_params or {}
+
+        H_site = Site("H", [0.0, 0.0, 0.0])
+        H_site_inf = Site("H", [100000.0, 0.0, 0.0])
+        H0_atom = Molecule.from_sites([H_site])
+        H2_plus_mol = Molecule.from_sites([H_site, H_site_inf])
+        H0_atom.set_charge_and_spin(0, 2)
+        H2_plus_mol.set_charge_and_spin(1, 2)
+        input_file_1 = "H0.qin"
+        output_file_1 = "H0.qout"
+        input_file_2 = "H2_plus.qin"
+        output_file_2 = "H2_plus.qout"
+        t = []
+        t.append(
+            WriteInputFromIOSet(
+                molecule=H0_atom,
+                qchem_input_set="SinglePointSet",
+                input_file=input_file_1,
+                qchem_input_params=qchem_input_params,
+            )
+        )
+        t.append(
+            RunQChemCustodian(
+                qchem_cmd=qchem_cmd,
+                multimode=multimode,
+                input_file=input_file_1,
+                output_file=output_file_1,
+                max_cores=max_cores,
+                max_errors=max_errors,
+                job_type="normal",
+                gzipped_output=False,
+            )
+        )
+
+        t.append(
+            WriteInputFromIOSet(
+                molecule=H2_plus_mol,
+                qchem_input_set="SinglePointSet",
+                input_file=input_file_2,
+                qchem_input_params=qchem_input_params,
+            )
+        )
+        t.append(
+            RunQChemCustodian(
+                qchem_cmd=qchem_cmd,
+                multimode=multimode,
+                input_file=input_file_2,
+                output_file=output_file_2,
+                max_cores=max_cores,
+                max_errors=max_errors,
+                job_type="normal",
+                gzipped_output=False,
+            )
+        )
+        t.append(
+            ProtCalcToDb(
+                db_file=db_file,
+                input_file_H0=input_file_1,
+                output_file_H0=output_file_1,
+                input_file_H2=input_file_2,
+                output_file_H2=output_file_2,
+                additional_fields={"task_label": name},
+            )
+        )
+
+        super().__init__(t, parents=parents, name=name, **kwargs)
+
+
 class ForceFW(Firework):
     def __init__(
         self,
diff --git a/atomate/qchem/fireworks/tests/test_core.py b/atomate/qchem/fireworks/tests/test_core.py
@@ -4,11 +4,14 @@
 
 import numpy as np
 from pymatgen.io.qchem.outputs import QCOutput
+from pymatgen.core.structure import Molecule
+from pymatgen.core.sites import Site
 
 from atomate.qchem.firetasks.critic2 import ProcessCritic2, RunCritic2
 from atomate.qchem.firetasks.fragmenter import FragmentMolecule
 from atomate.qchem.firetasks.geo_transformations import PerturbGeometry
 from atomate.qchem.firetasks.parse_outputs import QChemToDb
+from atomate.qchem.firetasks.parse_outputs import ProtCalcToDb
 from atomate.qchem.firetasks.run_calc import RunQChemCustodian
 from atomate.qchem.firetasks.write_inputs import WriteInputFromIOSet
 from atomate.qchem.fireworks.core import (
@@ -19,6 +22,7 @@
     FrequencyFW,
     OptimizeFW,
     PESScanFW,
+    ProtonEnergyFW,
     SinglePointFW,
     TransitionStateFW,
 )
@@ -132,6 +136,72 @@ def test_SinglePointFW_not_defaults(self):
         self.assertEqual(firework.parents, [])
         self.assertEqual(firework.name, "special single point")
 
+    def test_ProtonEnergyFW(self):
+        H_site = Site("H", [0.0, 0.0, 0.0])
+        H_site_inf = Site("H", [100000.0, 0.0, 0.0])
+        H0_atom = Molecule.from_sites([H_site])
+        H2_plus_mol = Molecule.from_sites([H_site, H_site_inf])
+        H0_atom.set_charge_and_spin(0, 2)
+        H2_plus_mol.set_charge_and_spin(1, 2)
+        firework = ProtonEnergyFW(qchem_input_params={"smd_solvent": "water"})
+        self.assertEqual(
+            firework.tasks[0].as_dict(),
+            WriteInputFromIOSet(
+                molecule=H0_atom,
+                qchem_input_set="SinglePointSet",
+                input_file="H0.qin",
+                qchem_input_params={"smd_solvent": "water"},
+            ).as_dict(),
+        )
+        self.assertEqual(
+            firework.tasks[1].as_dict(),
+            RunQChemCustodian(
+                qchem_cmd=">>qchem_cmd<<",
+                multimode=">>multimode<<",
+                input_file="H0.qin",
+                output_file="H0.qout",
+                max_cores=">>max_cores<<",
+                max_errors=5,
+                job_type="normal",
+                gzipped_output=False,
+            ).as_dict(),
+        )
+        self.assertEqual(
+            firework.tasks[2].as_dict(),
+            WriteInputFromIOSet(
+                molecule=H2_plus_mol,
+                qchem_input_set="SinglePointSet",
+                input_file="H2_plus.qin",
+                qchem_input_params={"smd_solvent": "water"},
+            ).as_dict(),
+        )
+        self.assertEqual(
+            firework.tasks[3].as_dict(),
+            RunQChemCustodian(
+                qchem_cmd=">>qchem_cmd<<",
+                multimode=">>multimode<<",
+                input_file="H2_plus.qin",
+                output_file="H2_plus.qout",
+                max_cores=">>max_cores<<",
+                max_errors=5,
+                job_type="normal",
+                gzipped_output=False,
+            ).as_dict(),
+        )
+        self.assertEqual(
+            firework.tasks[4].as_dict(),
+            ProtCalcToDb(
+                db_file=None,
+                input_file_H0="H0.qin",
+                output_file_H0="H0.qout",
+                input_file_H2="H2_plus.qin",
+                output_file_H2="H2_plus.qout",
+                additional_fields={"task_label": "proton electronic energy"},
+            ).as_dict(),
+        )
+        self.assertEqual(firework.parents, [])
+        self.assertEqual(firework.name, "proton electronic energy")
+
     def test_OptimizeFW_defaults(self):
         firework = OptimizeFW(molecule=self.act_mol)
         self.assertEqual(
diff --git a/atomate/utils/tests/test_database.py b/atomate/utils/tests/test_database.py
