Source code for qcengine.compute

"""
Integrates the computes together
"""
from typing import TYPE_CHECKING, Any, Dict, Optional, Union

from qcelemental.models import AtomicInput, AtomicResult, FailedOperation

from .config import get_config
from .exceptions import InputError, RandomError
from .procedures import get_procedure
from .programs import get_program
from .util import compute_wrapper, environ_context, handle_output_metadata, model_wrapper

if TYPE_CHECKING:
    from pydantic.main import BaseModel
    from qcelemental.models import AtomicResult


__all__ = ["compute", "compute_procedure"]


def _process_failure_and_return(model, return_dict, raise_error):
    if isinstance(model, FailedOperation):
        if raise_error:
            raise InputError(model.error.error_message)
        elif return_dict:
            return model.dict()
        else:
            return model
    else:
        return False


[docs]def compute( input_data: Union[Dict[str, Any], "AtomicInput"], program: str, raise_error: bool = False, local_options: Optional[Dict[str, Any]] = None, return_dict: bool = False, ) -> "AtomicResult": """Executes a single CMS program given a QCSchema input. The full specification can be found at: http://molssi-qc-schema.readthedocs.io/en/latest/index.html# Parameters ---------- input_data A QCSchema input specification in dictionary or model from QCElemental.models program The CMS program with which to execute the input. raise_error Determines if compute should raise an error or not. retries : int, optional The number of random tries to retry for. local_options A dictionary of local configuration options return_dict Returns a dict instead of qcelemental.models.AtomicResult Returns ------- result A computed AtomicResult object. """ output_data = input_data.copy() # lgtm [py/multiple-definition] with compute_wrapper(capture_output=False, raise_error=raise_error) as metadata: # Grab the executor and build the input model executor = get_program(program) # Build the model and validate input_data = model_wrapper(input_data, AtomicInput) # Build out local options if local_options is None: local_options = {} input_engine_options = input_data.extras.pop("_qcengine_local_config", {}) local_options = {**local_options, **input_engine_options} config = get_config(local_options=local_options) # Set environment parameters and execute with environ_context(config=config): # Handle optional retries for x in range(config.retries + 1): try: output_data = executor.compute(input_data, config) break except RandomError as e: if x == config.retries: raise e else: metadata["retries"] += 1 except: raise return handle_output_metadata(output_data, metadata, raise_error=raise_error, return_dict=return_dict)
[docs]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": """Runs a procedure (a collection of the quantum chemistry executions) Parameters ---------- input_data : dict or qcelemental.models.OptimizationInput A JSON input specific to the procedure executed in dictionary or model from QCElemental.models procedure : {"geometric", "berny"} The name of the procedure to run raise_error : bool, option Determines if compute should raise an error or not. local_options : dict, optional A dictionary of local configuration options return_dict : bool, optional, default True Returns a dict instead of qcelemental.models.AtomicInput Returns ------ dict, OptimizationResult, FailedOperation A QC Schema representation of the requested output, type depends on return_dict key. """ output_data = input_data.copy() # lgtm [py/multiple-definition] with compute_wrapper(capture_output=False, raise_error=raise_error) as metadata: # 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) # Create a base output data in case of errors output_data = input_data.copy() # lgtm [py/multiple-definition] # Set environment parameters and execute with environ_context(config=config): output_data = executor.compute(input_data, config) return handle_output_metadata(output_data, metadata, raise_error=raise_error, return_dict=return_dict)