QCEngine Support Overhaul

docs/source/conf.py

@@ -176,7 +176,7 @@
 
 
 # This skips classes that derived from ModuleBase, because those classes will
-# have Module API documentation producable by PluginPlay
+# have Module API documentation producible by PluginPlay
 def skip_pluginplay_modules(app, what, name, obj, skip, options):
     bases = obj.obj['bases'] if 'bases' in obj.obj.keys() else []
     if 'pluginplay.ModuleBase' in bases:

docs/source/nitpick_exceptions

@@ -5,3 +5,5 @@ py:class chemist.ChemicalSystem
 py:class Varies depending on the requested property
 
 py:class qcelemental.models.Molecule
+
+py:obj QCEngineEnergy

src/python/friendzone/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .nwx2nwchem import load_nwchem_modules
+from .nwx2qcengine import load_qcengine_modules
 from .nwx2qcelemental import load_qcelemental_modules
 
 
@@ -25,5 +25,5 @@ def load_modules(mm):
     :param mm: The ModuleManager that the all Modules will be loaded into.
     :type mm: pluginplay.ModuleManager
     """
-    load_nwchem_modules(mm)
+    load_qcengine_modules(mm)
     load_qcelemental_modules(mm)

src/python/friendzone/nwx2qcelemental/__init__.py

@@ -19,14 +19,14 @@
 
 
 class SystemViaMolSSI(pp.ModuleBase):
+    """Creates an NWChemEx ChemicalSystem by going through MolSSI's string 
+       parser.
+    """
 
     def __init__(self):
         pp.ModuleBase.__init__(self)
         self.satisfies_property_type(MoleculeFromString())
-        self.description("""
-            Creates an NWChemEx ChemicalSystem by going through MolSSI's
-            string parser.
-            """)
+        self.description(SystemViaMolSSI.__doc__)
 
     def run_(self, inputs, submods):
         pt = MoleculeFromString()

src/python/friendzone/nwx2qcelemental/chemical_system_conversions.py

@@ -43,7 +43,10 @@ def chemical_system2qc_mol(chem_sys):
         y = str(atom_i.y * au2ang)
         z = str(atom_i.z * au2ang)
         out += symbol + " " + x + " " + y + " " + z + "\n"
-    return qcel.models.Molecule.from_data(out)
+    return qcel.models.Molecule.from_data(out,
+                                          fix_com=True,
+                                          fix_orientation=True,
+                                          fix_symmetry="C1")
 
 
 def qc_mol2molecule(qc_mol):

src/python/friendzone/nwx2qcengine/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023 NWChemEx-Project
+# Copyright 2024 NWChemEx-Project
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,3 +11,124 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+from ..friends import is_friend_enabled
+import pluginplay as pp
+from simde import TotalEnergy, EnergyNuclearGradientStdVectorD
+from .call_qcengine import call_qcengine
+
+
+def _run_impl(driver, inputs, rv, runtime):
+    """
+    Our strategy here is to use the fact that the inputs to the TotalEnergy 
+    PT are a subset of those to other PTs
+    """
+    # Step 0: Figure out the PT we're being run as
+    egy_pt = TotalEnergy()
+    grad_pt = EnergyNuclearGradientStdVectorD()
+
+    # Step 1: Unwrap the inputs
+    mol = None
+    if driver == 'energy':
+        mol, = egy_pt.unwrap_inputs(inputs)
+    elif driver == 'gradient':
+        mol, _ = grad_pt.unwrap_inputs(inputs)
+        #TODO: verify ignored second input (the point at which to take the
+        #derivative) is equal to the geometry of mol.
+    else:
+        raise RuntimeError('Unexpected driver type')
+
+    program = inputs['program'].value()
+    method = inputs['method'].value()
+    basis = inputs['basis set'].value()
+
+    # Step 2: Call QCEngine
+    model = {'method': method, 'basis': basis}
+    keywords = {}
+    outputs = call_qcengine(driver,
+                            mol,
+                            program,
+                            runtime,
+                            model=model,
+                            keywords=keywords)
+
+    # Step 3: Prepare results
+    if driver == 'gradient':
+        grad = outputs['gradient'].flatten().tolist()
+        rv = grad_pt.wrap_results(rv, grad)
+
+    return egy_pt.wrap_results(rv, outputs['energy'])
+
+
+class QCEngineEnergy(pp.ModuleBase):
+    """ Driver module for computing energies with QCEngine.
+
+        This class relies on _run_impl to actually implement run_.
+    """
+
+    def __init__(self):
+        pp.ModuleBase.__init__(self)
+        self.satisfies_property_type(TotalEnergy())
+        self.description(QCEngineEnergy.__doc__)
+        self.add_input('program').set_description('Friend to call')
+        self.add_input('method').set_description('Level of theory')
+        self.add_input('basis set').set_description('Name of AO basis set')
+
+    def run_(self, inputs, submods):
+        return _run_impl('energy', inputs, self.results(), self.get_runtime())
+
+
+class QCEngineGradient(QCEngineEnergy):
+    """ Driver module for computing gradients with QCEngine.
+
+        This class extends QCEngineEnergy (QCEngine always computes the energy
+        when computing the gradient thus this module will also compute the
+        energy). Relative to QCEngineEnergy the main differences are:
+
+        - Addition of gradient property type
+        - Invocation of _run_impl with 'gradient' instead of 'energy'
+    """
+
+    def __init__(self):
+        QCEngineEnergy.__init__(self)
+        self.satisfies_property_type(EnergyNuclearGradientStdVectorD())
+
+    def run_(self, inputs, submods):
+        return _run_impl('gradient', inputs, self.results(),
+                         self.get_runtime())
+
+
+def load_qcengine_modules(mm):
+    """Loads the collection of modules that wrap QCElemental calls.
+
+    Currently, the friends exported by this function are:
+
+    #. NWChem
+
+    the levels of theory are: 
+
+    #. SCF
+    #. B3LYP
+    #. MP2
+    #. CCSD
+    #. CCSD(T)
+
+    and we have 0-th and 1-st derivatives.
+
+    The final set of modules is the Cartesian product of all of the above.
+
+    :param mm: The ModuleManager that the NWChem Modules will be loaded into.
+    :type mm: pluginplay.ModuleManager
+    """
+
+    for program in ['nwchem']:
+        if is_friend_enabled(program):
+            for method in ['SCF', 'B3LYP', 'MP2', 'CCSD', 'CCSD(T)']:
+                egy_key = program + ' : ' + method
+                grad_key = egy_key + ' Gradient'
+                mm.add_module(egy_key, QCEngineEnergy())
+                mm.add_module(grad_key, QCEngineGradient())
+
+                for key in [egy_key, grad_key]:
+                    mm.change_input(key, 'program', program)
+                    mm.change_input(key, 'method', method)

src/python/friendzone/nwx2qcengine/call_qcengine.py

@@ -15,10 +15,10 @@
 import qcengine as qcng
 import qcelemental as qcel
 from ..nwx2qcelemental.chemical_system_conversions import chemical_system2qc_mol
-from .pt2driver import pt2driver
+from qcengine.config import TaskConfig
 
 
-def call_qcengine(pt, mol, program, **kwargs):
+def call_qcengine(driver, mol, program, runtime, **kwargs):
     """ Wraps calling a program through the QCEngine API.
 
         .. note::
@@ -34,8 +34,8 @@ def call_qcengine(pt, mol, program, **kwargs):
         objects to their QCElemental equivalents. Right now those mappings
         include:
 
-        - property_type -> driver type
         - ChemicalSystem -> qcel.models.Molecule
+        - RuntimeView -> qcng.TaskConfig
 
         While not supported at the moment, similar conversions for the AO basis
         set are possible.
@@ -56,14 +56,30 @@ def call_qcengine(pt, mol, program, **kwargs):
                         backend?
         :type program: str
         :param kwargs: Key-value pairs which will be forwarded to QCElemental's
-                       ``AtomicInput`` class via the ``model`` key.
+                       ``AtomicInput`` class as kwargs.
 
-        :return: The requested property.
+        :return: A dictionary containing the requested property and any other
+                 property of potential interest.
         :rtype: Varies depending on the requested property
     """
 
-    driver = pt2driver(pt)
+    # Step 1: Prepare the chemistry-related input
     qc_mol = chemical_system2qc_mol(mol)
-    inp = qcel.models.AtomicInput(molecule=qc_mol, driver=driver, model=kwargs)
-    results = qcng.compute(inp, program)
-    return results.return_result
+    inp = qcel.models.AtomicInput(molecule=qc_mol, driver=driver, **kwargs)
+
+    # Step 2: Prepare the runtime-related input
+    # I *think* ncores is supposed to be the number of threads per MPI rank
+    task_config = {'nnodes': runtime.size(), 'ncores': 1, 'retries': 0}
+
+    # Step 3: Run QCEngine
+    results = qcng.compute(inp, program, task_config=task_config)
+
+    # Step 4: Verify the computation ran correctly
+    if type(results) == qcel.models.common_models.FailedOperation:
+        raise RuntimeError(results.error.error_message)
+
+    # Step 5: Prepare the results
+    rv = {driver: results.return_result}
+    if (driver == "gradient" and "qcvars" in results.extras):
+        rv['energy'] = float(results.extras["qcvars"]["CURRENT ENERGY"])
+    return rv

tests/python/unit_tests/nwx2nwchem/test_nwchem.py → tests/python/unit_tests/nwx2qcengine/test_nwchem.py

@@ -14,7 +14,8 @@
 
 from pluginplay import ModuleManager
 from friendzone import friends, load_modules
-from simde import TotalEnergy
+from simde import TotalEnergy, EnergyNuclearGradientStdVectorD
+from chemist import PointSetD
 from molecules import make_h2
 import unittest
 
@@ -28,6 +29,17 @@ def test_scf(self):
         egy = self.mm.run_as(TotalEnergy(), key, mol)
         self.assertAlmostEqual(egy, -1.094184522864, places=5)
 
+    def test_scf_gradient(self):
+        mol = make_h2()
+        key = 'NWChem : SCF Gradient'
+        self.mm.change_input(key, 'basis set', 'sto-3g')
+        grad = self.mm.run_as(EnergyNuclearGradientStdVectorD(), key, mol,
+                              PointSetD())
+
+        corr = [0.0, 0.0, -0.11827177600466043, 0.0, 0.0, 0.11827177600466043]
+        for g, c in zip(grad, corr):
+            self.assertAlmostEqual(g, c, places=4)
+
     def test_mp2(self):
         mol = make_h2()
         key = 'NWChem : MP2'