Bespoke clean up (#77)

* update `from_file`, fix incorrect initial parameter values for bespoke torsions. Correct TD base settings in line with torsiondrive, use qcsubmit optimisation spec.

* PR feedback, switch to use QCSpec from qcsubmit

* remove QCSpec for now

* moved keep files to be a setting

openff/bespokefit/cli/executor/submit.py

@@ -82,7 +82,7 @@ def _to_input_schema(
             "openff.bespokefit",
         )
 
-    workflow_factory = BespokeWorkflowFactory.parse_file(spec_file_name)
+    workflow_factory = BespokeWorkflowFactory.from_file(spec_file_name)
     workflow_factory.initial_force_field = force_field_path
 
     return workflow_factory.optimization_schema_from_molecule(molecule)

openff/bespokefit/executor/services/optimizer/worker.py

@@ -26,6 +26,6 @@ def optimize(self, optimization_input_json: str) -> str:
     input_schema.id = self.request.id
 
     optimizer = get_optimizer(input_schema.optimizer.type)
-    result = optimizer.optimize(input_schema, keep_files=False)
+    result = optimizer.optimize(input_schema, keep_files=settings.BEFLOW_KEEP_FILES)
 
     return serialize(result, encoding="json")

openff/bespokefit/executor/services/qcgenerator/worker.py

@@ -2,7 +2,7 @@
 
 import qcengine
 import redis
-from openff.toolkit.topology import Molecule
+from openff.toolkit.topology import Atom, Molecule
 from qcelemental.models import AtomicResult
 from qcelemental.models.common_models import DriverEnum
 from qcelemental.models.procedures import (
@@ -28,6 +28,17 @@
 celery_app = configure_celery_app("qcgenerator", redis_connection)
 
 
+def _select_atom(atoms: List[Atom]) -> int:
+    """
+    For a list of atoms chose the heaviest atom.
+    """
+    candidate = atoms[0]
+    for atom in atoms:
+        if atom.atomic_number > candidate.atomic_number:
+            candidate = atom
+    return candidate.molecule_atom_index
+
+
 @celery_app.task(acks_late=True)
 def compute_torsion_drive(task_json: str) -> TorsionDriveResult:
     """Runs a torsion drive using QCEngine."""
@@ -46,24 +57,31 @@ def compute_torsion_drive(task_json: str) -> TorsionDriveResult:
     index_2 = map_to_atom_index[task.central_bond[0]]
     index_3 = map_to_atom_index[task.central_bond[1]]
 
-    index_1 = [
-        atom.molecule_atom_index
+    index_1_atoms = [
+        atom
         for atom in molecule.atoms[index_2].bonded_atoms
         if atom.molecule_atom_index != index_3
-    ][0]
-    index_4 = [
-        atom.molecule_atom_index
+    ]
+    index_4_atoms = [
+        atom
         for atom in molecule.atoms[index_3].bonded_atoms
         if atom.molecule_atom_index != index_2
-    ][0]
+    ]
 
     del molecule.properties["atom_map"]
 
     input_schema = TorsionDriveInput(
         keywords=TDKeywords(
-            dihedrals=[(index_1, index_2, index_3, index_4)],
+            dihedrals=[
+                (
+                    _select_atom(index_1_atoms),
+                    index_2,
+                    index_3,
+                    _select_atom(index_4_atoms),
+                )
+            ],
             grid_spacing=[task.grid_spacing],
-            dihedral_ranges=[task.scan_range],
+            dihedral_ranges=[task.scan_range] if task.scan_range is not None else None,
         ),
         extras={
             "canonical_isomeric_explicit_hydrogen_mapped_smiles": molecule.to_smiles(
@@ -72,13 +90,13 @@ def compute_torsion_drive(task_json: str) -> TorsionDriveResult:
         },
         initial_molecule=molecule.to_qcschema(),
         input_specification=QCInputSpecification(
-            model=task.model, driver=DriverEnum.gradient
+            model=task.model,
+            driver=DriverEnum.gradient,
         ),
         optimization_spec=OptimizationSpecification(
-            procedure=task.optimization_spec.procedure,
+            procedure=task.optimization_spec.program,
             keywords={
-                "coordsys": "dlc",
-                "maxiter": task.optimization_spec.max_iterations,
+                **task.optimization_spec.dict(exclude={"program", "constraints"}),
                 "program": task.program,
             },
         ),
@@ -104,6 +122,8 @@ def compute_optimization(
     task_json: str,
 ) -> List[OptimizationResult]:
     """Runs a set of geometry optimizations using QCEngine."""
+    # TODO: should we only return the lowest energy optimization?
+    # or the first optimisation to work?
 
     task = OptimizationTask.parse_raw(task_json)
 
@@ -113,8 +133,7 @@ def compute_optimization(
     input_schemas = [
         OptimizationInput(
             keywords={
-                "coordsys": "dlc",
-                "maxiter": task.optimization_spec.max_iterations,
+                **task.optimization_spec.dict(exclude={"program", "constraints"}),
                 "program": task.program,
             },
             extras={
@@ -123,7 +142,8 @@ def compute_optimization(
                 )
             },
             input_specification=QCInputSpecification(
-                model=task.model, driver=DriverEnum.gradient
+                model=task.model,
+                driver=DriverEnum.gradient,
             ),
             initial_molecule=molecule.to_qcschema(conformer=i),
         )
@@ -135,7 +155,7 @@ def compute_optimization(
     for input_schema in input_schemas:
 
         return_value = qcengine.compute_procedure(
-            input_schema, task.optimization_spec.procedure, raise_error=True
+            input_schema, task.optimization_spec.program, raise_error=True
         )
 
         if isinstance(return_value, OptimizationResult):

openff/bespokefit/executor/services/settings.py

@@ -35,3 +35,4 @@ class Settings(BaseSettings):
     BEFLOW_OPTIMIZER_PREFIX = "optimizations"
     BEFLOW_OPTIMIZER_ROUTER = "openff.bespokefit.executor.services.optimizer.app:router"
     BEFLOW_OPTIMIZER_WORKER = "openff.bespokefit.executor.services.optimizer.worker"
+    BEFLOW_KEEP_FILES = False

openff/bespokefit/optimizers/forcebalance/forcebalance.py

@@ -103,14 +103,11 @@ def _optimize(cls, schema: BaseOptimizationSchema) -> OptimizerResultsType:
 
             _logger.debug("Launching Forcebalance")
 
-            current_env = os.environ
-            current_env["ENABLE_FB_SMIRNOFF_CACHING"] = "false"
             subprocess.run(
                 "ForceBalance optimize.in",
                 shell=True,
                 stdout=log,
                 stderr=log,
-                env=current_env,
             )
 
             results = cls._collect_results("", schema=schema)

openff/bespokefit/schema/fitting.py

@@ -66,18 +66,27 @@ def n_targets(self) -> int:
         return len(self.targets)
 
     @property
-    def initial_parameter_values(self) -> Dict[BaseSMIRKSParameter, unit.Quantity]:
+    def initial_parameter_values(
+        self,
+    ) -> Dict[BaseSMIRKSParameter, Dict[str, unit.Quantity]]:
         """A list of the refit force field parameters."""
 
         initial_force_field = ForceField(self.initial_force_field)
 
         return {
-            parameter: getattr(
-                initial_force_field[parameter.type].parameters[parameter.smirks],
-                attribute,
+            parameter: dict(
+                (
+                    attribute,
+                    getattr(
+                        initial_force_field[parameter.type].parameters[
+                            parameter.smirks
+                        ],
+                        attribute,
+                    ),
+                )
+                for attribute in parameter.attributes
             )
             for parameter in self.parameters
-            for attribute in parameter.attributes
         }
 
     def get_fitting_force_field(self) -> ForceField:

openff/bespokefit/schema/results.py

@@ -44,15 +44,17 @@ class BaseOptimizationResults(SchemaBase, abc.ABC):
     # )
 
     @property
-    def initial_parameter_values(self) -> Dict[BaseSMIRKSParameter, unit.Quantity]:
+    def initial_parameter_values(
+        self,
+    ) -> Dict[BaseSMIRKSParameter, Dict[str, unit.Quantity]]:
         """A list of the refit force field parameters."""
 
         return self.input_schema.initial_parameter_values
 
     @property
     def refit_parameter_values(
         self,
-    ) -> Optional[Dict[BaseSMIRKSParameter, unit.Quantity]]:
+    ) -> Optional[Dict[BaseSMIRKSParameter, Dict[str, unit.Quantity]]]:
         """A list of the refit force field parameters."""
 
         if self.refit_force_field is None:
@@ -61,12 +63,17 @@ def refit_parameter_values(
         refit_force_field = ForceField(self.refit_force_field)
 
         return {
-            parameter: getattr(
-                refit_force_field[parameter.type].parameters[parameter.smirks],
-                attribute,
+            parameter: dict(
+                (
+                    attribute,
+                    getattr(
+                        refit_force_field[parameter.type].parameters[parameter.smirks],
+                        attribute,
+                    ),
+                )
+                for attribute in parameter.attributes
             )
             for parameter in self.input_schema.parameters
-            for attribute in parameter.attributes
         }
 
 

openff/bespokefit/schema/tasks.py

@@ -1,28 +1,21 @@
 import abc
-from typing import Tuple
+from typing import Optional, Tuple
 
+from openff.qcsubmit.procedures import GeometricProcedure
 from pydantic import Field, conint
 from qcelemental.models.common_models import Model
 from typing_extensions import Literal
 
 from openff.bespokefit.utilities.pydantic import BaseModel
 
 
-class OptimizationSpec(BaseModel):
-
-    procedure: Literal["geometric"] = "geometric"
-
-    max_iterations: conint(gt=0) = Field(
-        300,
-        description="The maximum number of iterations to perform before raising a "
-        "convergence failure exception.",
-    )
-
-
 class QCGenerationTask(BaseModel, abc.ABC):
 
     type: Literal["base-task"]
 
+    program: str = Field(..., description="The program to use to evaluate the model.")
+    model: Model = Field(..., description=str(Model.__doc__))
+
 
 class HessianTaskSpec(QCGenerationTask):
 
@@ -34,12 +27,8 @@ class HessianTaskSpec(QCGenerationTask):
         "hessian. Each conformer will be minimized and the one with the lowest energy "
         "will have its hessian computed.",
     )
-
-    program: str = Field(..., description="The program to use to evaluate the model.")
-    model: Model = Field(..., description=str(Model.__doc__))
-
-    optimization_spec: OptimizationSpec = Field(
-        OptimizationSpec(),
+    optimization_spec: GeometricProcedure = Field(
+        GeometricProcedure(),
         description="The specification for how to optimize each conformer before "
         "computing the hessian.",
     )
@@ -54,22 +43,9 @@ class HessianTask(HessianTaskSpec):
     )
 
 
-class OptimizationTaskSpec(QCGenerationTask):
+class OptimizationTaskSpec(HessianTaskSpec):
     type: Literal["optimization"] = "optimization"
 
-    n_conformers: conint(gt=0) = Field(
-        ...,
-        description="The maximum number of conformers to begin the optimization from.",
-    )
-
-    program: str = Field(..., description="The program to use to evaluate the model.")
-    model: Model = Field(..., description=str(Model.__doc__))
-
-    optimization_spec: OptimizationSpec = Field(
-        OptimizationSpec(),
-        description="The specification for how to optimize each conformer.",
-    )
-
 
 class OptimizationTask(OptimizationTaskSpec):
 
@@ -84,15 +60,12 @@ class Torsion1DTaskSpec(QCGenerationTask):
     type: Literal["torsion1d"] = "torsion1d"
 
     grid_spacing: int = Field(15, description="The spacing between grid angles.")
-    scan_range: Tuple[int, int] = Field(
-        (-180, 165), description="The range of grid angles to scan."
+    scan_range: Optional[Tuple[int, int]] = Field(
+        None, description="The range of grid angles to scan."
     )
 
-    program: str = Field(..., description="The program to use to evaluate the model.")
-    model: Model = Field(..., description=str(Model.__doc__))
-
-    optimization_spec: OptimizationSpec = Field(
-        OptimizationSpec(),
+    optimization_spec: GeometricProcedure = Field(
+        GeometricProcedure(enforce=0.1, reset=True, qccnv=True, epsilon=0.0),
         description="The specification for how to optimize the structure at each angle "
         "in the scan.",
     )

openff/bespokefit/tests/executor/services/qcgenerator/test_app.py

@@ -34,7 +34,7 @@ def mock_atomic_result() -> AtomicResult:
     return AtomicResult(
         molecule=molecule.to_qcschema(),
         driver=DriverEnum.hessian,
-        model=Model(method="rdkit", basis=None),
+        model=Model(method="uff", basis=None),
         return_result=5.2,
         success=True,
         provenance=Provenance(creator="pytest"),
@@ -51,7 +51,7 @@ def mock_torsion_drive_result() -> TorsionDriveResult:
     return TorsionDriveResult(
         keywords=TDKeywords(dihedrals=[(0, 1, 2, 3)], grid_spacing=[15]),
         input_specification=QCInputSpecification(
-            model=Model(method="rdkit", basis=None), driver=DriverEnum.gradient
+            model=Model(method="uff", basis=None), driver=DriverEnum.gradient
         ),
         initial_molecule=molecule.to_qcschema(),
         optimization_spec=OptimizationSpecification(procedure="geometric"),