Merge pull request #43 from dgasmith/procedures

Procedures Base Models

qcengine/__init__.py

@@ -3,14 +3,14 @@
 """
 
 from . import config
-
-from .programs import get_program, list_all_programs, list_available_programs
 from .compute import compute, compute_procedure
 from .config import get_config
-from .stock_mols import get_molecule
-
 # Handle versioneer
 from .extras import get_information
+from .procedures import get_procedure, list_all_procedures, list_available_procedures
+from .programs import get_program, list_all_programs, list_available_programs
+from .stock_mols import get_molecule
+
 __version__ = get_information('version')
 __git_revision__ = get_information('git_revision')
 del get_information

qcengine/compute.py

@@ -6,17 +6,30 @@
 from qcelemental.models import ComputeError, FailedOperation, Optimization, OptimizationInput, ResultInput
 
 from .config import get_config
-from .programs import get_program
+from .procedures import get_procedure, list_all_procedures, list_available_procedures
+from .programs import get_program, list_all_programs, list_available_programs
 from .util import compute_wrapper, get_module_function, handle_output_metadata, model_wrapper
 
 __all__ = ["compute", "compute_procedure"]
 
 
+def _process_failure_and_return(model, return_dict, raise_error):
+    if isinstance(model, FailedOperation):
+        if raise_error:
+            raise ValueError(model.error.error_message)
+        elif return_dict:
+            return model.dict()
+        else:
+            return model
+    else:
+        return False
+
+
 def compute(input_data: Union[Dict[str, Any], 'ResultInput'],
             program: str,
-            raise_error: bool = False,
-            local_options: Optional[Dict[str, str]] = None,
-            return_dict: bool = False) -> 'Result':
+            raise_error: bool=False,
+            local_options: Optional[Dict[str, str]]=None,
+            return_dict: bool=False) -> 'Result':
     """Executes a single quantum chemistry program given a QC Schema input.
 
     The full specification can be found at:
@@ -44,12 +57,33 @@ def compute(input_data: Union[Dict[str, Any], 'ResultInput'],
 
     """
 
+    program = program.lower()
+    if program not in list_all_programs():
+        input_data = FailedOperation(
+            input_data=input_data,
+            error=ComputeError(
+                error_type="not_registered",
+                error_message="QCEngine Call Error:\n"
+                "Program {} is not registered with QCEngine".format(program)))
+    elif program not in list_available_programs():
+        input_data = FailedOperation(
+            input_data=input_data,
+            error=ComputeError(
+                error_type="not_available",
+                error_message="QCEngine Call Error:\n"
+                "Program {} is registered with QCEngine, but cannot be found".format(program)))
+    error = _process_failure_and_return(input_data, return_dict, raise_error)
+    if error:
+        return error
+
     # Build the model and validate
-    input_data = model_wrapper(input_data, ResultInput, raise_error)
-    if isinstance(input_data, FailedOperation):
-        if return_dict:
-            return input_data.dict()
-        return input_data
+    input_data = model_wrapper(input_data, ResultInput)
+    error = _process_failure_and_return(input_data, return_dict, raise_error)
+    if error:
+        return error
+
+    # Grab the executor and build the input model
+    executor = get_program(program)
 
     # Build out local options
     if local_options is None:
@@ -63,26 +97,17 @@ def compute(input_data: Union[Dict[str, Any], 'ResultInput'],
     # Run the program
     with compute_wrapper(capture_output=False) as metadata:
 
-        output_data = input_data.copy()  # Initial in case of error handling
-        try:
-            output_data = get_program(program).compute(input_data, config)
-        except KeyError as e:
-            output_data = FailedOperation(
-                input_data=output_data.dict(),
-                success=False,
-                error=ComputeError(
-                    error_type='program_error',
-                    error_message="QCEngine Call Error:\nProgram {} not understood."
-                    "\nError Message: {}".format(program, str(e))))
+        output_data = input_data.copy()  # lgtm [py/multiple-definition]
+        output_data = executor.compute(input_data, config)
 
     return handle_output_metadata(output_data, metadata, raise_error=raise_error, return_dict=return_dict)
 
 
-def compute_procedure(input_data: Dict[str, Any],
+def compute_procedure(input_data: Union[Dict[str, Any], 'BaseModel'],
                       procedure: str,
-                      raise_error: bool = False,
-                      local_options: Optional[Dict[str, str]] = None,
-                      return_dict: bool = False) -> 'BaseModel':
+                      raise_error: bool=False,
+                      local_options: Optional[Dict[str, str]]=None,
+                      return_dict: bool=False) -> 'BaseModel':
     """Runs a procedure (a collection of the quantum chemistry executions)
 
     Parameters
@@ -104,44 +129,39 @@ def compute_procedure(input_data: Dict[str, Any],
         A QC Schema representation of the requested output, type depends on return_dict key.
     """
 
-    # Build the model and validate
-    input_data = model_wrapper(input_data, OptimizationInput, raise_error)
-    if isinstance(input_data, FailedOperation):
-        if return_dict:
-            return input_data.dict()
-        return input_data
+    procedure = procedure.lower()
+    if procedure not in list_all_procedures():
+        input_data = FailedOperation(
+            input_data=input_data,
+            error=ComputeError(
+                error_type="not_registered",
+                error_message="QCEngine Call Error:\n"
+                "Procedure {} is not registered with QCEngine".format(procedure)))
+    elif procedure not in list_available_procedures():
+        input_data = FailedOperation(
+            input_data=input_data,
+            error=ComputeError(
+                error_type="not_available",
+                error_message="QCEngine Call Error:\n"
+                "Procedure {} is registered with QCEngine, but cannot be found".format(procedure)))
+    error = _process_failure_and_return(input_data, return_dict, raise_error)
+    if error:
+        return error
+
+    # Grab the executor and build the input model
+    executor = get_procedure(procedure)
 
     config = get_config(local_options=local_options)
+    input_data = executor.build_input_model(input_data)
+    error = _process_failure_and_return(input_data, return_dict, raise_error)
+    if error:
+        return error
 
     # Run the procedure
     with compute_wrapper(capture_output=False) as metadata:
+
         # Create a base output data in case of errors
         output_data = input_data.copy()  # lgtm [py/multiple-definition]
-        if procedure == "geometric":
-            # Augment the input
-            geometric_input = input_data.dict()
-
-            # Older QCElemental compat, can be removed in v0.6
-            if "extras" not in geometric_input["input_specification"]:
-                geometric_input["input_specification"]["extras"] = {}
-
-            geometric_input["input_specification"]["extras"]["_qcengine_local_config"] = config.dict()
-
-            # Run the program
-            output_data = get_module_function(procedure, "run_json.geometric_run_json")(geometric_input)
-
-            output_data["schema_name"] = "qcschema_optimization_output"
-            output_data["input_specification"]["extras"].pop("_qcengine_local_config", None)
-            if output_data["success"]:
-                output_data = Optimization(**output_data)
-
-        else:
-            output_data = FailedOperation(
-                input_data=input_data.dict(),
-                success=False,
-                error=ComputeError(
-                    error_type="program_error",
-                    error_message="QCEngine Call Error:"
-                    "\nProcedure {} not understood".format(procedure)))
+        output_data = executor.compute(input_data, config)
 
     return handle_output_metadata(output_data, metadata, raise_error=raise_error, return_dict=return_dict)

qcengine/procedures/__init__.py

@@ -0,0 +1 @@
+from .base import get_procedure, list_all_procedures, list_available_procedures, register_procedure

qcengine/procedures/base.py

@@ -0,0 +1,53 @@
+"""
+Imports the various procedure backends
+"""
+
+from typing import List, Set
+
+from .geometric import GeometricProcedure
+
+__all__ = ["register_procedure", "get_procedure", "list_all_procedures", "list_available_procedures"]
+
+procedures = {}
+
+
+def register_procedure(entry_point: 'BaseProcedure') -> None:
+    """
+    Register a new BaseProcedure with QCEngine
+    """
+
+    name = entry_point.name
+    if name.lower() in procedures.keys():
+        raise ValueError('{} is already a registered procedure.'.format(name))
+
+    procedures[name.lower()] = entry_point
+
+
+def get_procedure(name: str) -> 'BaseProcedure':
+    """
+    Returns a procedures executor class
+    """
+    return procedures[name.lower()]
+
+
+def list_all_procedures() -> Set[str]:
+    """
+    List all procedures registered by QCEngine.
+    """
+    return set(procedures.keys())
+
+
+def list_available_procedures() -> Set[str]:
+    """
+    List all procedures that can be exectued (found) by QCEngine.
+    """
+
+    ret = set()
+    for k, p in procedures.items():
+        if p.found():
+            ret.add(k)
+
+    return ret
+
+
+register_procedure(GeometricProcedure())

qcengine/procedures/geometric.py

@@ -0,0 +1,56 @@
+from typing import Any, Dict, Union
+
+from qcelemental.models import ComputeError, FailedOperation, Optimization, OptimizationInput
+
+from .procedure_model import BaseProcedure
+
+
+class GeometricProcedure(BaseProcedure):
+
+    _defaults = {"name": "geomeTRIC", "procedure": "optimization"}
+
+    class Config(BaseProcedure.Config):
+        pass
+
+    def __init__(self, **kwargs):
+        super().__init__(**{**self._defaults, **kwargs})
+
+    def build_input_model(self, data: Union[Dict[str, Any], 'OptimizationInput']) -> 'OptimizationInput':
+        return self._build_model(data, OptimizationInput)
+
+    def compute(self, input_data: 'OptimizationInput', config: 'JobConfig') -> 'Optimization':
+        try:
+            import geometric
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError("Could not find geomeTRIC in the Python path.")
+
+        geometric_input = input_data.dict()
+
+        # Older QCElemental compat, can be removed in v0.6
+        if "extras" not in geometric_input["input_specification"]:
+            geometric_input["input_specification"]["extras"] = {}
+
+        geometric_input["input_specification"]["extras"]["_qcengine_local_config"] = config.dict()
+
+        # Run the program
+        output_data = geometric.run_json.geometric_run_json(geometric_input)
+
+        output_data["provenance"] = {
+            "creator": "geomeTRIC",
+            "routine": "geometric.run_json.geometric_run_json",
+            "version": geometric.__version__
+        }
+
+        output_data["schema_name"] = "qcschema_optimization_output"
+        output_data["input_specification"]["extras"].pop("_qcengine_local_config", None)
+        if output_data["success"]:
+            output_data = Optimization(**output_data)
+
+        return output_data
+
+    def found(self) -> bool:
+        try:
+            import geometric
+            return True
+        except ModuleNotFoundError:
+            return False

qcengine/procedures/procedure_model.py

@@ -0,0 +1,57 @@
+import abc
+from typing import Any, Dict, Union
+
+from pydantic import BaseModel
+
+from ..util import model_wrapper
+
+
+class BaseProcedure(BaseModel, abc.ABC):
+
+    name: str
+    procedure: str
+
+    class Config:
+        allow_mutation: False
+        extra: "forbid"
+
+    @abc.abstractmethod
+    def build_input_model(self, data: Union[Dict[str, Any], 'BaseModel'], raise_error: bool=True) -> 'BaseModel':
+        """
+        Build and validate the input model, passes if the data was a normal BaseModel input.
+
+        Parameters
+        ----------
+        data : Union[Dict[str, Any], 'BaseModel']
+            A data blob to construct the model from or the input model itself
+        raise_error : bool, optional
+            Raise an error or not if the operation failed.
+
+        Returns
+        -------
+        BaseModel
+            The input model for the procedure.
+        """
+
+    @abc.abstractmethod
+    def compute(self, input_data: 'BaseModel', config: 'JobConfig') -> 'BaseModel':
+        pass
+
+    @abc.abstractmethod
+    def found(self) -> bool:
+        """
+        Checks if the program can be found.
+
+        Returns
+        -------
+        bool
+            If the proceudre was found or not.
+        """
+        pass
+
+    def _build_model(self, data: Dict[str, Any], model: 'BaseModel') -> 'BaseModel':
+        """
+        Quick wrapper around util.model_wrapper for inherited classes
+        """
+
+        return model_wrapper(data, model)

qcengine/programs/tests/test_programs.py

@@ -20,7 +20,7 @@ def test_missing_key():
 
 
 def test_missing_key_raises():
-    with pytest.raises(TypeError):
+    with pytest.raises(ValueError):
         ret = qcng.compute({"hello": "hi"}, "bleh", raise_error=True)