Copying Sam's changes for drone, adding ForceFW

Author: espottesmith

atomate/qchem/drones.py

@@ -198,6 +198,9 @@ def generate_doc(self, dir_name, qcinput_files, qcoutput_files, multirun):
             elif "ESP" in d_calc_final:
                 d["output"]["esp"] = d_calc_final["ESP"][-1]
 
+            if "nbo_data" in d_calc_final:
+                d["output"]["nbo"] = d_calc_final["nbo_data"]
+
             if d["output"]["job_type"] in ["opt", "optimization", "ts"]:
                 if "molecule_from_optimized_geometry" in d_calc_final:
                     d["output"]["optimized_molecule"] = d_calc_final[
@@ -233,6 +236,13 @@ def generate_doc(self, dir_name, qcinput_files, qcoutput_files, multirun):
                 d["output"]["scan_geometries"] = d_calc_final.get("optimized_geometries")
                 d["output"]["scan_molecules"] = d_calc_final.get("molecules_from_optimized_geometries")
 
+            if d["output"]["job_type"] == "force":
+                d["output"]["gradients"] = d_calc_final["gradients"][0]
+                if d_calc_final["pcm_gradients"] is not None:
+                    d["output"]["pcm_gradients"] = d_calc_final["pcm_gradients"][0]
+                if d_calc_final["CDS_gradients"] is not None:
+                    d["output"]["CDS_gradients"] = d_calc_final["CDS_gradients"][0]
+
             opt_trajectory = []
             calcs = copy.deepcopy(d["calcs_reversed"])
             calcs.reverse()

atomate/qchem/fireworks/core.py

@@ -100,6 +100,83 @@ def __init__(self,
             **kwargs)
 
 
+class ForceFW(Firework):
+    def __init__(
+        self,
+        molecule=None,
+        name="force calculation",
+        qchem_cmd=">>qchem_cmd<<",
+        multimode=">>multimode<<",
+        max_cores=">>max_cores<<",
+        qchem_input_params=None,
+        db_file=None,
+        parents=None,
+        **kwargs
+    ):
+        """
+        Converge the electron density and calculate the atomic forces, aka the gradient.
+        Args:
+            molecule (Molecule): Input molecule.
+            name (str): Name for the Firework.
+            qchem_cmd (str): Command to run QChem. Supports env_chk.
+            multimode (str): Parallelization scheme, either openmp or mpi. Defaults to openmp.
+            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.
+                                       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 conjuction 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 {}
+        input_file = "mol.qin"
+        output_file = "mol.qout"
+        t = []
+        t.append(
+            WriteInputFromIOSet(
+                molecule=molecule,
+                qchem_input_set="ForceSet",
+                input_file=input_file,
+                qchem_input_params=qchem_input_params,
+            )
+        )
+        t.append(
+            RunQChemCustodian(
+                qchem_cmd=qchem_cmd,
+                multimode=multimode,
+                input_file=input_file,
+                output_file=output_file,
+                max_cores=max_cores,
+                job_type="normal",
+            )
+        )
+        t.append(
+            QChemToDb(
+                db_file=db_file,
+                input_file=input_file,
+                output_file=output_file,
+                additional_fields={"task_label": name},
+            )
+        )
+        super(ForceFW, self).__init__(t, parents=parents, name=name, **kwargs)
+
+
 class OptimizeFW(Firework):
     def __init__(self,
                  molecule=None,