Merge pull request #12 from MolSSI/mv_qcng

Move older QCEngine function to tests

qcmanybody/builder.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import itertools
-from typing import Iterable, Union, Literal, Optional, Dict, Set
+from typing import Iterable, Union, Literal, Optional, Dict, Set, List
 
 from qcmanybody.models import BsseEnum, FragBasIndex
 
@@ -29,7 +29,10 @@ def build_nbody_compute_list(
         Whether the total data (True; energy/gradient/Hessian) of the molecular system has been requested,
         as opposed to interaction data (False).
     supersystem_ie_only
-        ????
+        Target the supersystem total/interaction energy (IE) data over the many-body expansion (MBE) "
+        analysis, thereby omitting intermediate-body calculations.
+    supersystem_max_nbody
+        Maximum n-body to use for a supersystem calculation. Must be specified if "supersystem" is in `nbodies`
 
     Returns
     -------
@@ -62,8 +65,8 @@ def build_nbody_compute_list(
             raise ValueError("supersystem_max_nbody must be provided if 'supersystem' contains nbodies")
 
         include_supersystem = True
-        nbodies = list(nbodies)
-        nbodies.remove("supersystem")
+
+    nbodies: List[int] = [x for x in nbodies if x != "supersystem"]
 
     # What levels do we need?
     fragment_range = range(1, nfragments + 1)

qcmanybody/manybody.py

@@ -452,7 +452,7 @@ def analyze(
 
         # All properties that were passed to us
         # * seed with "energy" so free/no-op jobs can process
-        available_properties = set(["energy"])
+        available_properties: Set[str] = {"energy"}
         for property_data in component_results.values():
             available_properties.update(property_data.keys())
 

qcmanybody/tests/utils.py

@@ -1,10 +1,14 @@
 import json
 import math
 import os
+from typing import Mapping, Union, Literal, Any, Iterable
 
 import numpy
+import qcengine as qcng
 import zstandard
+from qcelemental.models import Molecule, AtomicInput
 
+from qcmanybody import ManyBodyCalculator, delabeler
 from qcmanybody.models import BsseEnum
 
 _my_dir = os.path.dirname(os.path.realpath(__file__))
@@ -61,10 +65,8 @@ def load_component_data(file_base):
         return unjsonify(json.load(f))
 
 
-def generate_component_data(mol, levels, specifications, bsse_type, return_total_data, out_filename):
-    from qcmanybody.qcengine_helper import run_qcengine_base
-
-    mc, component_results = run_qcengine_base(mol, levels, specifications, bsse_type, return_total_data)
+def generate_component_data(mol, levels, specifications, bsse_type, return_total_data, out_filename, supsersytem_ie_only=False):
+    mc, component_results = run_qcengine(mol, levels, specifications, bsse_type, return_total_data, supsersytem_ie_only)
 
     component_results = jsonify(component_results)
     filepath = os.path.join(_my_dir, "component_data", out_filename + ".json.zst")
@@ -154,3 +156,46 @@ def compare_results(qcmb_results, ref_results, levels):
                     res[f"{bstr}-CORRECTED {l}-BODY CONTRIBUTION TO ENERGY"],
                     ref_results[f"{bstr}-CORRECTED {l}-BODY CONTRIBUTION TO ENERGY"],
                 )
+
+
+def run_qcengine(
+    molecule: Molecule,
+    levels: Mapping[Union[int, Literal["supersystem"]], str],
+    specifications: Mapping[str, Mapping[str, Any]],
+    bsse_type: Iterable[BsseEnum],
+    return_total_data: bool,
+    supersystem_ie_only: bool
+):
+
+    mc = ManyBodyCalculator(molecule, bsse_type, levels, return_total_data, supersystem_ie_only)
+
+    component_results = {}
+
+    computation_count = {}
+    for chem, label, imol in mc.iterate_molecules():
+        print(label)
+        inp = AtomicInput(molecule=imol, **specifications[chem]["specification"])
+
+        _, real, bas = delabeler(label)
+        computation_count.setdefault(len(real), 0)
+        computation_count[len(real)] += 1
+
+        result = qcng.compute(inp, specifications[chem]["program"])
+
+        if not result.success:
+            print(result.error.error_message)
+            raise RuntimeError("Calculation did not succeed! Error:\n" + result.error.error_message)
+
+        # pull out stuff
+        props = {"energy", "gradient", "hessian"}
+
+        component_results[label] = {}
+
+        for p in props:
+            if hasattr(result.properties, f"return_{p}"):
+                v = getattr(result.properties, f"return_{p}")
+                # print(f"  {label} {p}: {v}")
+                if v is not None:
+                    component_results[label][p] = v
+
+    return mc, component_results