Adapt new v0.31.0 features for QCSchema v2 work

.github/workflows/CI.yml

@@ -5,10 +5,12 @@ on:
     branches:
       - master
       - next2024
+      - next2025
   pull_request:
     branches:
       - master
       - next2024
+      - next2025
   schedule:
     - cron: "9 16 * * 1"
 

qcengine/procedures/qcmanybody.py

@@ -1,13 +1,15 @@
-from typing import TYPE_CHECKING, Any, ClassVar, Dict, Union
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, Tuple, Union
 
 from qcelemental.util import safe_version, which_import
 
+from ..util import model_wrapper
 from .model import ProcedureHarness
 
 if TYPE_CHECKING:
-    from ..config import TaskConfig
     from qcmanybody.models import ManyBodyInput, ManyBodyResult
 
+    from ..config import TaskConfig
+
 
 class QCManyBodyProcedure(ProcedureHarness):
 
@@ -23,10 +25,49 @@ def found(self, raise_error: bool = False) -> bool:
             raise_msg="Please install via `conda install qcmanybody -c conda-forge`.",
         )
 
-    def build_input_model(self, data: Union[Dict[str, Any], "ManyBodyInput"]) -> "ManyBodyInput":
+    def build_input_model(
+        self, data: Union[Dict[str, Any], "ManyBodyInput"], *, return_input_schema_version: bool = False
+    ) -> "ManyBodyInput":
         from qcmanybody.models import ManyBodyInput
 
-        return self._build_model(data, ManyBodyInput)
+        return self._build_model(data, "ManyBodyInput", return_input_schema_version=return_input_schema_version)
+
+    # temporary can't use procedure/model.py b/c different import location than qcel.models
+    def _build_model(
+        self, data: Dict[str, Any], model: "BaseModel", /, *, return_input_schema_version: bool = False
+    ) -> Union["BaseModel", Tuple["BaseModel", int]]:
+        """
+        Quick wrapper around util.model_wrapper for inherited classes
+        """
+        import qcmanybody
+
+        v1_model = getattr(qcmanybody.models, model)
+        v2_model = None
+
+        if isinstance(data, v1_model):
+            mdl = model_wrapper(data, v1_model)
+        elif isinstance(data, v2_model):
+            mdl = model_wrapper(data, v2_model)
+        elif isinstance(data, dict):
+            # remember these are user-provided dictionaries, so they'll have the mandatory fields,
+            #   like driver, not the helpful discriminator fields like schema_version.
+            # so long as versions distinguishable by a *required* field, id by dict is reliable.
+
+            # if data.get("specification", False) or data.get("schema_version") == 2:
+            #     mdl = model_wrapper(data, v2_model)
+            # else:
+            mdl = model_wrapper(data, v1_model)
+
+        input_schema_version = mdl.schema_version
+        if input_schema_version != 1:
+            raise InputError("Can't use v2 ManyBody")
+
+        if return_input_schema_version:
+            return mdl, input_schema_version
+            # return mdl.convert_v(2), input_schema_version
+        else:
+            return mdl
+            # return mdl.convert_v(2)
 
     def get_version(self) -> str:
         self.found(raise_error=True)

qcengine/testing.py

@@ -251,6 +251,11 @@ def schema_versions2(request):
         return (qcel.models.v2, 2, qcel.models.v2)
 
 
+@pytest.fixture(scope="function", params=[None])
+def schema_versions0(request):
+    return qcel.models, -1, qcel.models
+
+
 def checkver_and_convert(mdl, tnm, prepost, vercheck: bool = True, cast_dict_as=None):
     import json
 

qcengine/tests/test_procedures.py

@@ -625,34 +625,62 @@ def test_torsiondrive_generic(schema_versions, request):
 
 @using("mace")
 @using("torsiondrive")
-def test_torsiondrive_extra_constraints():
-
-    input_data = TorsionDriveInput(
-        keywords=TDKeywords(dihedrals=[(3, 0, 1, 2)], grid_spacing=[180]),
-        input_specification=QCInputSpecification(driver=DriverEnum.gradient, model=Model(method="small", basis=None)),
-        initial_molecule=[qcng.get_molecule("propane")],
-        optimization_spec=OptimizationSpecification(
-            procedure="geomeTRIC",
-            keywords={
-                "coordsys": "dlc",
-                # use mace as it does not have convergence issues like UFF
-                "program": "mace",
-                "constraints": {
-                    "set": [
-                        {
-                            "type": "dihedral",  # hold a dihedral through the other C-C bond fixed
-                            "indices": (0, 1, 2, 10),
-                            "value": 0.0,
-                        }
-                    ]
+def test_torsiondrive_extra_constraints(schema_versions, request):
+    models, retver, _ = schema_versions
+
+    keywords = {
+        "coordsys": "dlc",
+        "constraints": {
+            "set": [
+                {
+                    "type": "dihedral",  # hold a dihedral through the other C-C bond fixed
+                    "indices": (0, 1, 2, 10),
+                    "value": 0.0,
+                }
+            ]
+        },
+    }
+
+    if from_v2(request.node.name):
+        input_data = models.TorsionDriveInput(
+            initial_molecules=[models.Molecule(**qcng.get_molecule("propane", return_dict=True))],
+            specification=models.TorsionDriveSpecification(
+                keywords=models.TDKeywords(dihedrals=[(3, 0, 1, 2)], grid_spacing=[180]),
+                specification=models.OptimizationSpecification(
+                    program="geomeTRIC",
+                    keywords=keywords,
+                    specification=models.AtomicSpecification(
+                        # use mace as it does not have convergence issues like UFF
+                        program="mace",
+                        driver=models.DriverEnum.gradient,
+                        model=models.Model(method="small", basis=None),
+                    ),
+                ),
+            ),
+        )
+    else:
+        input_data = models.TorsionDriveInput(
+            keywords=models.TDKeywords(dihedrals=[(3, 0, 1, 2)], grid_spacing=[180]),
+            input_specification=models.QCInputSpecification(
+                driver=models.DriverEnum.gradient, model=models.Model(method="small", basis=None)
+            ),
+            initial_molecule=[models.Molecule(**qcng.get_molecule("propane", return_dict=True))],
+            optimization_spec=models.OptimizationSpecification(
+                procedure="geomeTRIC",
+                keywords={
+                    **keywords,
+                    # use mace as it does not have convergence issues like UFF
+                    "program": "mace",
                 },
-            },
-        ),
-    )
+            ),
+        )
 
-    ret = qcng.compute_procedure(input_data, "torsiondrive", raise_error=True)
+    input_data = checkver_and_convert(input_data, request.node.name, "pre")
+    ret = qcng.compute(input_data, "torsiondrive", raise_error=True, return_version=retver)
+    ret = checkver_and_convert(ret, request.node.name, "post")
 
-    assert ret.error is None
+    if "_v2" not in request.node.name:
+        assert ret.error is None
     assert ret.success
 
     expected_grid_ids = {"180", "0"}
@@ -671,7 +699,10 @@ def test_torsiondrive_extra_constraints():
 
     assert ret.provenance.creator.lower() == "torsiondrive"
     assert ret.optimization_history["180"][0].provenance.creator.lower() == "geometric"
-    assert ret.optimization_history["180"][0].trajectory[0].provenance.creator.lower() == "mace"
+    if "v2" in request.node.name:
+        assert ret.optimization_history["180"][0].trajectory_results[0].provenance.creator.lower() == "mace"
+    else:
+        assert ret.optimization_history["180"][0].trajectory[0].provenance.creator.lower() == "mace"
 
     assert "Using MACE-OFF23 MODEL for MACECalculator" in ret.stdout
     assert "All optimizations converged at lowest energy. Job Finished!\n" in ret.stdout

qcengine/tests/test_qcmanybody.py

@@ -1,19 +1,23 @@
 import pprint
 
 import pytest
-
 from qcelemental import constants
-from qcelemental.models import Molecule
 from qcelemental.testing import compare, compare_values
 
 import qcengine as qcng
+from qcengine.testing import checkver_and_convert, from_v2
+from qcengine.testing import schema_versions0 as schema_versions
 from qcengine.testing import using
 
+# TODO full schema_versions when manybody v2 schema ready
+
 
 @pytest.fixture
-def he_tetramer():
+def he_tetramer(schema_versions, request):
+    models, _, _ = schema_versions
+
     a2 = 2 / constants.bohr2angstroms
-    return Molecule(
+    return models.Molecule(
         symbols=["He", "He", "He", "He"],
         fragments=[[0], [1], [2], [3]],
         geometry=[0, 0, 0, 0, 0, a2, 0, a2, 0, 0, a2, a2],
@@ -55,7 +59,11 @@ def he_tetramer():
         ),
     ],
 )
-def test_nbody_he4_single(program, basis, keywords, mbe_keywords, anskey, calcinfo_nmbe, he_tetramer, request):
+def test_nbody_he4_single(
+    schema_versions, program, basis, keywords, mbe_keywords, anskey, calcinfo_nmbe, he_tetramer, request
+):
+    models, retver, _ = schema_versions
+
     from qcmanybody.models import AtomicSpecification, ManyBodyInput
 
     atomic_spec = AtomicSpecification(
@@ -75,7 +83,9 @@ def test_nbody_he4_single(program, basis, keywords, mbe_keywords, anskey, calcin
         molecule=he_tetramer,
     )
 
-    ret = qcng.compute_procedure(mbe_model, "qcmanybody", raise_error=True)
+    mbe_model = checkver_and_convert(mbe_model, request.node.name, "pre")
+    ret = qcng.compute(mbe_model, "qcmanybody", raise_error=True, return_version=retver)
+    ret = checkver_and_convert(ret, request.node.name, "post")
     pprint.pprint(ret.dict(), width=200)
 
     assert ret.extras == {}, f"[w] extras wrongly present: {ret.extras.keys()}"
@@ -139,11 +149,13 @@ def test_nbody_he4_single(program, basis, keywords, mbe_keywords, anskey, calcin
         pytest.param("psi4", marks=using("psi4")),
     ],
 )
-def test_bsse_ene_tu6_cp_ne2(qcprog):
+def test_bsse_ene_tu6_cp_ne2(schema_versions, request, qcprog):
     """
     from https://github.com/psi4/psi4/blob/master/tests/tu6-cp-ne2/input.dat
     Example potential energy surface scan and CP-correction for Ne2
     """
+    models, retver, _ = schema_versions
+
     from qcmanybody.models import ManyBodyInput
 
     tu6_ie_scan = {2.5: 0.757717, 3.0: 0.015685, 4.0: -0.016266}  # Ang: kcal/mol IE
@@ -181,19 +193,21 @@ def test_bsse_ene_tu6_cp_ne2(qcprog):
     }
 
     for R in tu6_ie_scan:
-        nene = Molecule(
+        nene = models.Molecule(
             symbols=["Ne", "Ne"], fragments=[[0], [1]], geometry=[0, 0, 0, 0, 0, R / constants.bohr2angstroms]
         )
         mbe_data["molecule"] = nene
 
         mbe_model = ManyBodyInput(**mbe_data)
         if qcprog == "gamess":
             with pytest.raises(RuntimeError) as exe:
-                qcng.compute_procedure(mbe_model, "qcmanybody", raise_error=True)
+                qcng.compute(mbe_model, "qcmanybody", raise_error=True, return_version=retver)
             assert "GAMESS+QCEngine can't handle ghost atoms yet" in str(exe.value)
             pytest.xfail("GAMESS can't do ghosts")
 
-        ret = qcng.compute_procedure(mbe_model, "qcmanybody", raise_error=True)
+        mbe_model = checkver_and_convert(mbe_model, request.node.name, "pre")
+        ret = qcng.compute(mbe_model, "qcmanybody", raise_error=True, return_version=retver)
+        ret = checkver_and_convert(ret, request.node.name, "post")
         pprint.pprint(ret.dict(), width=200)
 
         assert compare_values(
@@ -206,7 +220,9 @@ def test_bsse_ene_tu6_cp_ne2(qcprog):
 
 
 @using("qcmanybody")
-def test_mbe_error():
+def test_mbe_error(schema_versions, request):
+    models, retver, _ = schema_versions
+
     from qcmanybody.models import ManyBodyInput
 
     mbe_data = {
@@ -225,7 +241,7 @@ def test_mbe_error():
         "molecule": None,
     }
 
-    nene = Molecule(
+    nene = models.Molecule(
         symbols=["Ne", "Ne"], fragments=[[0], [1]], geometry=[0, 0, 0, 0, 0, 3.0 / constants.bohr2angstroms]
     )
     mbe_data["molecule"] = nene
@@ -234,12 +250,12 @@ def test_mbe_error():
 
     # test 1
     with pytest.raises(RuntimeError) as exc:
-        qcng.compute_procedure(mbe_model, "qcmanybody", raise_error=True)
+        qcng.compute(mbe_model, "qcmanybody", raise_error=True, return_version=retver)
 
     assert "Program cms is not registered to QCEngine" in str(exc.value)
 
     # test 2
-    ret = qcng.compute_procedure(mbe_model, "qcmanybody")
+    ret = qcng.compute(mbe_model, "qcmanybody", return_version=retver)
     assert ret.success is False
     assert "Program cms is not registered to QCEngine" in ret.error.error_message
 
@@ -259,9 +275,10 @@ def test_mbe_error():
         # pytest.param("gengeometric", "cp", True, marks=using("geometric_genopt")),
     ],
 )
-def test_bsse_opt_hf_trimer(optimizer, bsse_type, sio):
+def test_bsse_opt_hf_trimer(schema_versions, request, optimizer, bsse_type, sio):
+    models, retver, _ = schema_versions
 
-    initial_molecule = Molecule.from_data(
+    initial_molecule = models.Molecule.from_data(
         """
 F         -0.04288        2.78905        0.00000
 H          0.59079        2.03435        0.00000
@@ -329,7 +346,9 @@ def test_bsse_opt_hf_trimer(optimizer, bsse_type, sio):
     # from qcmanybody.models.generalized_optimization import GeneralizedOptimizationInput
     # opt_data = GeneralizedOptimizationInput(**opt_data)
 
-    ret = qcng.compute_procedure(opt_data, optimizer, raise_error=True)
+    # opt_data = checkver_and_convert(opt_data, request.node.name, "pre")
+    ret = qcng.compute(opt_data, optimizer, raise_error=True, return_version=retver)
+    # ret = checkver_and_convert(ret, request.node.name, "post")
 
     print("FFFFFFFFFF")
     pprint.pprint(ret.dict(), width=200)

qcengine/util.py

@@ -237,6 +237,13 @@ def handle_output_metadata(
             inp_ret = output_fusion
         ret = model(success=success_ret, error=error_ret, input_data=inp_ret)
 
+    # temp while ManyBody has no v2. empty string for FailedOp
+    if getattr(ret, "schema_name", "") == "qcschema_manybodyresult":
+        if return_dict:
+            return json.loads(ret.json())
+        else:
+            return ret
+
     if convert_version > 0:
         ret = ret.convert_v(convert_version)