Merge pull request #981 from MolSSI/vendor_qcel

Vendor in some models from qcelemental

Author: bennybp

docs/source/user_guide/records/optimization.rst

@@ -101,9 +101,8 @@ when adding to datasets).
 
       .. code-block:: py3
 
-        from qcportal.optimization import OptimizationSpecification
+        from qcportal.optimization import OptimizationSpecification, OptimizationProtocols
         from qcportal.singlepoint import QCSpecification
-        from qcelemental.models.procedures import OptimizationProtocols
 
         opt_spec = OptimizationSpecification(
             program="geometric",

qcfractal/qcfractal/components/gridoptimization/testing_helpers.py

@@ -8,13 +8,12 @@
 except ImportError:
     import pydantic
 from qcelemental.models import Molecule, FailedOperation, ComputeError, OptimizationResult as QCEl_OptimizationResult
-from qcelemental.models.procedures import OptimizationProtocols
 
 from qcarchivetesting.helpers import read_procedure_data, read_record_data
 from qcfractal.components.gridoptimization.record_db_models import GridoptimizationRecordORM
 from qcfractal.testing_helpers import run_service
 from qcportal.gridoptimization import GridoptimizationSpecification, GridoptimizationKeywords, GridoptimizationRecord
-from qcportal.optimization import OptimizationSpecification
+from qcportal.optimization import OptimizationSpecification, OptimizationProtocols
 from qcportal.record_models import PriorityEnum, RecordStatusEnum, RecordTask
 from qcportal.singlepoint import SinglepointProtocols, QCSpecification
 from qcportal.utils import recursive_normalizer

qcfractal/qcfractal/components/optimization/record_socket.py

@@ -308,7 +308,7 @@ def add_specifications(
         to_add = []
 
         for opt_spec in opt_specs:
-            protocols_dict = opt_spec.protocols.dict(exclude_defaults=True)
+            protocols_dict = opt_spec.protocols.dict(exclude_defaults=True, exclude_unset=True)
 
             # Don't include lower specifications in the hash
             opt_spec_dict = opt_spec.dict(exclude={"protocols", "qc_specification"})

qcfractal/qcfractal/components/optimization/test_record_socket.py

@@ -108,7 +108,7 @@ def test_optimization_socket_task_spec(
 
         task_input = t.function_kwargs["input_data"]
         assert task_input["keywords"] == kw_with_prog
-        assert task_input["protocols"] == spec.protocols.dict(exclude_defaults=True)
+        assert task_input["protocols"] == spec.protocols.dict(exclude_defaults=True, exclude_unset=True)
 
         # Forced to gradient in the qcschema input
         assert task_input["input_specification"]["driver"] == SinglepointDriver.gradient

qcfractal/qcfractal/components/singlepoint/record_socket.py

@@ -234,7 +234,7 @@ def add_specifications(
         to_add = []
 
         for qc_spec in qc_specs:
-            protocols_dict = qc_spec.protocols.dict(exclude_defaults=True)
+            protocols_dict = qc_spec.protocols.dict(exclude_defaults=True, exclude_unset=True)
 
             # TODO - if error_correction is manually specified as the default, then it will be an empty dict
             if "error_correction" in protocols_dict:

qcfractal/qcfractal/components/singlepoint/test_record_socket.py

@@ -108,7 +108,9 @@ def test_singlepoint_socket_task_spec(
     for t in tasks:
         function_kwargs = t.function_kwargs
         assert function_kwargs["input_data"]["model"] == {"method": spec.method, "basis": spec.basis}
-        assert function_kwargs["input_data"]["protocols"] == spec.protocols.dict(exclude_defaults=True)
+        assert function_kwargs["input_data"]["protocols"] == spec.protocols.dict(
+            exclude_defaults=True, exclude_unset=True
+        )
         assert function_kwargs["input_data"]["keywords"] == spec.keywords
         assert function_kwargs["program"] == spec.program
         assert t.compute_tag == "tag1"

qcfractal/qcfractal/components/torsiondrive/testing_helpers.py

@@ -8,12 +8,11 @@
 except ImportError:
     import pydantic
 from qcelemental.models import Molecule, FailedOperation, ComputeError, OptimizationResult as QCEl_OptimizationResult
-from qcelemental.models.procedures import OptimizationProtocols
 
 from qcarchivetesting.helpers import read_procedure_data, read_record_data
 from qcfractal.components.torsiondrive.record_db_models import TorsiondriveRecordORM
 from qcfractal.testing_helpers import run_service
-from qcportal.optimization import OptimizationSpecification
+from qcportal.optimization import OptimizationSpecification, OptimizationProtocols
 from qcportal.record_models import PriorityEnum, RecordStatusEnum, RecordTask
 from qcportal.singlepoint import SinglepointProtocols, QCSpecification
 from qcportal.torsiondrive import TorsiondriveSpecification, TorsiondriveKeywords, TorsiondriveRecord

qcportal/qcportal/optimization/record_models.py

@@ -1,40 +1,56 @@
 from __future__ import annotations
 
 from copy import deepcopy
+from enum import Enum
+from typing import Iterable
 from typing import Optional, Union, Any, List, Dict
 
-try:
-    from pydantic.v1 import BaseModel, Field, constr, validator, Extra
-except ImportError:
-    from pydantic import BaseModel, Field, constr, validator, Extra
+from pydantic.v1 import BaseModel, Field, constr, validator, Extra
 from qcelemental.models import Molecule
 from qcelemental.models.procedures import (
     OptimizationResult,
-    OptimizationProtocols,
     QCInputSpecification,
-    Model as AtomicResultModel,
 )
 from typing_extensions import Literal
-from typing import Iterable
 
 from qcportal.base_models import RestModelBase
+from qcportal.cache import get_records_with_cache
 from qcportal.record_models import (
     BaseRecord,
     RecordAddBodyBase,
     RecordQueryFilters,
     RecordStatusEnum,
     compare_base_records,
 )
-from qcportal.utils import is_included
-from qcportal.cache import get_records_with_cache
 from qcportal.singlepoint import (
     SinglepointProtocols,
     SinglepointRecord,
     QCSpecification,
     SinglepointDriver,
     compare_singlepoint_records,
 )
+from qcportal.utils import is_included
+
+
+class TrajectoryProtocolEnum(str, Enum):
+    """
+    Which gradient evaluations to keep in an optimization trajectory.
+    """
+
+    all = "all"
+    initial_and_final = "initial_and_final"
+    final = "final"
+    none = "none"
+
+
+class OptimizationProtocols(BaseModel):
+    """
+    Protocols regarding the manipulation of a Optimization output data.
+    """
 
+    trajectory: TrajectoryProtocolEnum = Field(
+        TrajectoryProtocolEnum.all, description=str(TrajectoryProtocolEnum.__doc__)
+    )
 
 class OptimizationSpecification(BaseModel):
     """
@@ -216,7 +232,7 @@ def to_qcschema_result(self) -> OptimizationResult:
             keywords=new_keywords,
             input_specification=QCInputSpecification(
                 driver=SinglepointDriver.gradient,  # forced
-                model=AtomicResultModel(
+                model=dict(
                     method=self.specification.qc_specification.method,
                     basis=self.specification.qc_specification.basis,
                 ),

qcportal/qcportal/record_models.py

@@ -8,25 +8,30 @@
 from typing import Optional, Dict, Any, List, Union, Iterable, Tuple, Type, Sequence, ClassVar, TypeVar
 
 from dateutil.parser import parse as date_parser
-
-try:
-    from pydantic.v1 import BaseModel, Extra, constr, validator, PrivateAttr, Field, parse_obj_as, root_validator
-except ImportError:
-    from pydantic import BaseModel, Extra, constr, validator, PrivateAttr, Field, parse_obj_as, root_validator
-from qcelemental.models.results import Provenance
+from pydantic.v1 import BaseModel, Extra, constr, validator, PrivateAttr, Field, root_validator
 
 from qcportal.base_models import (
     RestModelBase,
     QueryModelBase,
     QueryIteratorBase,
 )
-
 from qcportal.cache import RecordCache, get_records_with_cache
 from qcportal.compression import CompressionEnum, decompress, get_compressed_ext
 
 _T = TypeVar("_T")
 
 
+class Provenance(BaseModel):
+    """Provenance information."""
+
+    creator: str = Field(..., description="The name of the program, library, or person who created the object.")
+    version: str = Field("", description="The version of the creator, blank otherwise")
+    routine: str = Field("", description="The name of the routine or function within the creator, blank otherwise.")
+
+    class Config(BaseModel.Config):
+        extra: str = "allow"
+
+
 class PriorityEnum(int, Enum):
     """
     The priority of a Task. Higher priority will be pulled first.

qcportal/qcportal/singlepoint/record_models.py

@@ -4,22 +4,17 @@
 from enum import Enum
 from typing import Optional, Union, Any, List, Dict, Tuple
 
-try:
-    from pydantic.v1 import BaseModel, Field, constr, validator, Extra, PrivateAttr
-except ImportError:
-    from pydantic import BaseModel, Field, constr, validator, Extra, PrivateAttr
+from pydantic.v1 import BaseModel, Field, constr, validator, Extra, PrivateAttr
 from qcelemental.models import Molecule
 from qcelemental.models.results import (
     AtomicResult,
-    Model as AtomicResultModel,
-    AtomicResultProtocols as SinglepointProtocols,
     AtomicResultProperties,
     WavefunctionProperties,
 )
 from typing_extensions import Literal
 
-from qcportal.compression import CompressionEnum, decompress
 from qcportal.base_models import RestModelBase
+from qcportal.compression import CompressionEnum, decompress
 from qcportal.record_models import (
     RecordStatusEnum,
     BaseRecord,
@@ -29,6 +24,24 @@
 )
 
 
+class Model(BaseModel):
+    """The computational molecular sciences model to run."""
+
+    method: str = Field(  # type: ignore
+        ...,
+        description="The quantum chemistry method to evaluate (e.g., B3LYP, PBE, ...). "
+        "For MM, name of the force field.",
+    )
+    basis: Optional[Union[str, BasisSet]] = Field(  # type: ignore
+        None,
+        description="The quantum chemistry basis set to evaluate (e.g., 6-31g, cc-pVDZ, ...). Can be ``None`` for "
+        "methods without basis sets. For molecular mechanics, name of the atom-typer.",
+    )
+
+    class Config(BaseModel.Config):
+        extra: str = "allow"
+
+
 class SinglepointDriver(str, Enum):
     # Copied from qcelemental to add "deferred"
     energy = "energy"
@@ -38,6 +51,59 @@ class SinglepointDriver(str, Enum):
     deferred = "deferred"
 
 
+class WavefunctionProtocolEnum(str, Enum):
+    r"""Wavefunction to keep from a computation."""
+
+    all = "all"
+    orbitals_and_eigenvalues = "orbitals_and_eigenvalues"
+    occupations_and_eigenvalues = "occupations_and_eigenvalues"
+    return_results = "return_results"
+    none = "none"
+
+
+class ErrorCorrectionProtocol(BaseModel):
+    r"""Configuration for how computationaal chemistry programs handle error correction
+    """
+
+    default_policy: bool = Field(
+        True, description="Whether to allow error corrections to be used " "if not directly specified in `policies`"
+    )
+    policies: Optional[Dict[str, bool]] = Field(
+        None,
+        description="Settings that define whether specific error corrections are allowed. "
+        "Keys are the name of a known error and values are whether it is allowed to be used.",
+    )
+
+    def allows(self, policy: str):
+        if self.policies is None:
+            return self.default_policy
+        return self.policies.get(policy, self.default_policy)
+
+
+class NativeFilesProtocolEnum(str, Enum):
+    r"""Any program-specific files to keep from a computation."""
+
+    all = "all"
+    input = "input"
+    none = "none"
+
+
+class SinglepointProtocols(BaseModel):
+    r"""Protocols regarding the manipulation of computational result data."""
+
+    wavefunction: WavefunctionProtocolEnum = Field(
+        WavefunctionProtocolEnum.none, description=str(WavefunctionProtocolEnum.__doc__)
+    )
+    stdout: bool = Field(True, description="Primary output file to keep from the computation")
+    error_correction: ErrorCorrectionProtocol = Field(
+        default_factory=ErrorCorrectionProtocol, description="Policies for error correction"
+    )
+    native_files: NativeFilesProtocolEnum = Field(
+        NativeFilesProtocolEnum.none,
+        description="Policies for keeping processed files from the computation",
+    )
+
+
 class QCSpecification(BaseModel):
     class Config:
         extra = Extra.forbid
@@ -57,7 +123,7 @@ class Config:
         "methods without basis sets.",
     )
     keywords: Dict[str, Any] = Field({}, description="Program-specific keywords to use for the computation")
-    protocols: SinglepointProtocols = Field(SinglepointProtocols(), description=str(SinglepointProtocols.__base_doc__))
+    protocols: SinglepointProtocols = Field(SinglepointProtocols())
 
     @validator("basis", pre=True)
     def _convert_basis(cls, v):
@@ -178,7 +244,7 @@ def to_qcschema_result(self) -> AtomicResult:
 
         return AtomicResult(
             driver=self.specification.driver,
-            model=AtomicResultModel(
+            model=dict(
                 method=self.specification.method,
                 basis=self.specification.basis,
             ),