From 5cd96661d179a1eda1a0be1a877bec5d1f87f6b9 Mon Sep 17 00:00:00 2001 From: SimonBoothroyd Date: Mon, 23 May 2022 07:13:07 -0400 Subject: [PATCH 1/2] Generalize QC task models to allow pre-optimization --- .../executor/services/qcgenerator/cache.py | 53 +++++-- .../executor/services/qcgenerator/worker.py | 94 ++++++------ openff/bespokefit/schema/tasks.py | 144 +++++++++++++++--- .../tests/cli/executor/test_submit.py | 12 +- .../executor/services/qcgenerator/test_app.py | 25 ++- .../services/qcgenerator/test_cache.py | 31 ++-- .../services/qcgenerator/test_worker.py | 29 ++-- .../bespokefit/tests/schema/test_targets.py | 14 +- openff/bespokefit/workflows/bespoke.py | 23 +-- 9 files changed, 290 insertions(+), 135 deletions(-) diff --git a/openff/bespokefit/executor/services/qcgenerator/cache.py b/openff/bespokefit/executor/services/qcgenerator/cache.py index 08df620b..89051f4a 100644 --- a/openff/bespokefit/executor/services/qcgenerator/cache.py +++ b/openff/bespokefit/executor/services/qcgenerator/cache.py @@ -6,9 +6,9 @@ from openff.bespokefit.executor.services.qcgenerator import worker from openff.bespokefit.schema.tasks import ( + BaseTaskSpec, HessianTask, OptimizationTask, - QCGenerationTask, Torsion1DTask, ) from openff.bespokefit.utilities.molecule import canonical_order_atoms @@ -48,18 +48,21 @@ def _canonicalize_task(task: _T) -> _T: task.central_bond = (1, 2) - else: + elif isinstance(task, (HessianTask, OptimizationTask)): canonical_smiles = canonical_molecule.to_smiles( isomeric=True, explicit_hydrogens=True, mapped=False ) + else: + raise NotImplementedError() + task.smiles = canonical_smiles return task -def _hash_task(task: QCGenerationTask) -> str: +def _hash_task(task: BaseTaskSpec) -> str: """Returns a hashed representation of a QC task""" return hashlib.sha512(task.json().encode()).hexdigest() @@ -85,6 +88,24 @@ def _cache_task_id( redis_connection.hset("qcgenerator:task-ids", task_hash, task_id) +def _compute_hessian_task() -> str: + raise NotImplementedError() + + +def _compute_optimization_task(task: OptimizationTask): + + if task.pre_optimization_spec is not None or task.evaluation_spec is not None: + raise NotImplementedError() + + task_id = worker.compute_optimization.delay( + smiles=task.smiles, + optimization_spec_json=task.optimization_spec.json(), + n_conformers=task.n_conformers, + ).id + + return task_id + + def _compute_torsion_drive_task( task: Torsion1DTask, redis_connection: redis.Redis ) -> str: @@ -94,7 +115,7 @@ def _compute_torsion_drive_task( task_id = None torsion_drive_task = task.copy(deep=True) - torsion_drive_task.sp_specification = None + torsion_drive_task.evaluation_spec = None torsion_drive_hash = _hash_task(torsion_drive_task) torsion_drive_id = _retrieve_cached_task_id(torsion_drive_hash, redis_connection) @@ -103,19 +124,22 @@ def _compute_torsion_drive_task( # There are no cached torsion drives at the 'pre-optimise' level of theory # we need to run a torsion drive and then optionally a single point - if task.sp_specification is None: + if task.evaluation_spec is None: torsion_drive_id = worker.compute_torsion_drive.delay( - task_json=task.json() + smiles=task.smiles, + central_bond=task.central_bond, + grid_spacing=task.grid_spacing, + scan_range=task.scan_range, + optimization_spec_json=task.optimization_spec, ).id else: task_future: AsyncResult = ( worker.compute_torsion_drive.s(task_json=task.json()) - | worker.evaluate_torsion_drive.s( - model_json=task.sp_specification.model.json(), - program=task.sp_specification.program, + | worker.re_evaluate_torsion_drive.s( + evaluation_spec_json=task.evaluation_spec.json(), ) ).delay() @@ -126,7 +150,7 @@ def _compute_torsion_drive_task( torsion_drive_id, task.type, torsion_drive_hash, redis_connection ) - if task.sp_specification is None: + if task.evaluation_spec is None: return torsion_drive_id if task_id is None: @@ -136,9 +160,8 @@ def _compute_torsion_drive_task( task_id = ( ( worker.wait_for_task.s(torsion_drive_id) - | worker.evaluate_torsion_drive.s( - model_json=task.sp_specification.model.json(), - program=task.sp_specification.program, + | worker.re_evaluate_torsion_drive.s( + evaluation_spec_json=task.evaluation_spec.json(), ) ) .delay() @@ -168,9 +191,9 @@ def cached_compute_task( if isinstance(task, Torsion1DTask): task_id = _compute_torsion_drive_task(task, redis_connection) elif isinstance(task, OptimizationTask): - task_id = worker.compute_optimization.delay(task_json=task.json()).id + task_id = _compute_optimization_task(task) elif isinstance(task, HessianTask): - task_id = worker.compute_hessian.delay(task_json=task.json()).id + task_id = _compute_hessian_task() else: raise NotImplementedError() diff --git a/openff/bespokefit/executor/services/qcgenerator/worker.py b/openff/bespokefit/executor/services/qcgenerator/worker.py index 41c24306..3d5e4526 100644 --- a/openff/bespokefit/executor/services/qcgenerator/worker.py +++ b/openff/bespokefit/executor/services/qcgenerator/worker.py @@ -1,5 +1,5 @@ import logging -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional, Tuple import psutil import qcelemental @@ -9,7 +9,7 @@ from celery.utils.log import get_task_logger from openff.toolkit.topology import Atom, Molecule from qcelemental.models import AtomicInput, AtomicResult -from qcelemental.models.common_models import DriverEnum, Model, Provenance +from qcelemental.models.common_models import DriverEnum, Provenance from qcelemental.models.procedures import ( OptimizationInput, OptimizationResult, @@ -25,7 +25,7 @@ from openff.bespokefit.executor.services import current_settings from openff.bespokefit.executor.utilities.celery import configure_celery_app from openff.bespokefit.executor.utilities.redis import connect_to_default_redis -from openff.bespokefit.schema.tasks import OptimizationTask, Torsion1DTask +from openff.bespokefit.schema.tasks import OptimizationSpec, SinglePointSpec celery_app = configure_celery_app( "qcgenerator", connect_to_default_redis(validate=False) @@ -66,23 +66,30 @@ def _select_atom(atoms: List[Atom]) -> int: @celery_app.task(acks_late=True) -def compute_torsion_drive(task_json: str) -> str: +def compute_torsion_drive( + smiles: str, + central_bond: Tuple[int, int], + grid_spacing: int, + scan_range: Optional[Tuple[int, int]], + optimization_spec_json: str, + n_conformers: int, +) -> str: """Runs a torsion drive using QCEngine.""" - task = Torsion1DTask.parse_raw(task_json) - _task_logger.info(f"running 1D scan with {_task_config()}") - molecule: Molecule = Molecule.from_smiles(task.smiles) - molecule.generate_conformers(n_conformers=task.n_conformers) + optimization_spec = OptimizationSpec.parse_raw(optimization_spec_json) + + molecule: Molecule = Molecule.from_smiles(smiles) + molecule.generate_conformers(n_conformers=n_conformers) map_to_atom_index = { map_index: atom_index for atom_index, map_index in molecule.properties["atom_map"].items() } - index_2 = map_to_atom_index[task.central_bond[0]] - index_3 = map_to_atom_index[task.central_bond[1]] + index_2 = map_to_atom_index[central_bond[0]] + index_3 = map_to_atom_index[central_bond[1]] index_1_atoms = [ atom @@ -107,8 +114,8 @@ def compute_torsion_drive(task_json: str) -> str: _select_atom(index_4_atoms), ) ], - grid_spacing=[task.grid_spacing], - dihedral_ranges=[task.scan_range] if task.scan_range is not None else None, + grid_spacing=[grid_spacing], + dihedral_ranges=[scan_range] if scan_range is not None else None, ), extras={ "canonical_isomeric_explicit_hydrogen_mapped_smiles": molecule.to_smiles( @@ -119,14 +126,14 @@ def compute_torsion_drive(task_json: str) -> str: molecule.to_qcschema(conformer=i) for i in range(molecule.n_conformers) ], input_specification=QCInputSpecification( - model=task.model, + model=optimization_spec.model, driver=DriverEnum.gradient, ), optimization_spec=OptimizationSpecification( - procedure=task.optimization_spec.program, + procedure=optimization_spec.procedure.program, keywords={ - **task.optimization_spec.dict(exclude={"program", "constraints"}), - "program": task.program, + **optimization_spec.procedure.dict(exclude={"program", "constraints"}), + "program": optimization_spec.program, }, ), ) @@ -146,21 +153,18 @@ def compute_torsion_drive(task_json: str) -> str: @celery_app.task(acks_late=True) -def evaluate_torsion_drive( +def re_evaluate_torsion_drive( result_json: str, - model_json: str, - program: str, + evaluation_spec_json: str, ) -> str: """ Re-evaluates the energies at each optimised geometry along a torsion drive at a new level of theory. """ - model = Model.parse_raw(model_json) + evaluation_spec = SinglePointSpec.parse_raw(evaluation_spec_json) - _task_logger.info( - f"performing single point evaluations using {model} and {program}" - ) + _task_logger.info(f"performing single point evaluations using {evaluation_spec}") original_result = TorsionDriveResult.parse_raw(result_json) @@ -170,8 +174,7 @@ def evaluate_torsion_drive( energies = { grid_point: _compute_single_point( molecule=molecule, - model=model, - program=program, + spec=evaluation_spec, config=qcengine_config, ).return_result for grid_point, molecule in original_result.final_molecules.items() @@ -181,7 +184,7 @@ def evaluate_torsion_drive( keywords=original_result.keywords, extras=original_result.extras, input_specification=QCInputSpecification( - driver=DriverEnum.gradient, model=model + driver=DriverEnum.gradient, model=evaluation_spec.model ), initial_molecule=original_result.initial_molecule, optimization_spec=original_result.optimization_spec, @@ -199,24 +202,26 @@ def evaluate_torsion_drive( @celery_app.task(acks_late=True) def compute_optimization( - task_json: str, + smiles: str, + optimization_spec_json: str, + n_conformers: int, ) -> 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) + # or the first optimisation to work? _task_logger.info(f"running opt with {_task_config()}") - molecule: Molecule = Molecule.from_smiles(task.smiles) - molecule.generate_conformers(n_conformers=task.n_conformers) + optimization_spec = OptimizationSpec.parse_raw(optimization_spec_json) + + molecule: Molecule = Molecule.from_smiles(smiles) + molecule.generate_conformers(n_conformers=n_conformers) input_schemas = [ OptimizationInput( keywords={ - **task.optimization_spec.dict(exclude={"program", "constraints"}), - "program": task.program, + **optimization_spec.procedure.dict(exclude={"program", "constraints"}), + "program": optimization_spec.program, }, extras={ "canonical_isomeric_explicit_hydrogen_mapped_smiles": molecule.to_smiles( @@ -224,7 +229,7 @@ def compute_optimization( ) }, input_specification=QCInputSpecification( - model=task.model, + model=optimization_spec.model, driver=DriverEnum.gradient, ), initial_molecule=molecule.to_qcschema(conformer=i), @@ -238,7 +243,7 @@ def compute_optimization( return_value = qcengine.compute_procedure( input_schema, - task.optimization_spec.program, + optimization_spec.procedure.program, raise_error=True, local_options=_task_config(), ) @@ -257,26 +262,27 @@ def compute_optimization( return serialize(return_values, "json") -@celery_app.task(acks_late=True) -def compute_hessian(task_json: str) -> AtomicResult: - """Runs a set of hessian evaluations using QCEngine.""" - raise NotImplementedError() +# @celery_app.task(acks_late=True) +# def compute_hessian() -> AtomicResult: +# """Runs a set of hessian evaluations using QCEngine.""" +# raise NotImplementedError() def _compute_single_point( molecule: qcelemental.models.Molecule, - model: Model, - program: str, + spec: SinglePointSpec, config: Dict[str, Any], ) -> AtomicResult: """ Perform a single point calculation on the input ``qcelemental`` molecule. """ - qc_input = AtomicInput(molecule=molecule, driver=DriverEnum.energy, model=model) + qc_input = AtomicInput( + molecule=molecule, driver=DriverEnum.energy, model=spec.model + ) return qcengine.compute( input_data=qc_input, - program=program, + program=spec.program, raise_error=True, local_options=config, ) diff --git a/openff/bespokefit/schema/tasks.py b/openff/bespokefit/schema/tasks.py index 2efb9d36..f906b938 100644 --- a/openff/bespokefit/schema/tasks.py +++ b/openff/bespokefit/schema/tasks.py @@ -12,15 +12,60 @@ from openff.bespokefit.utilities.pydantic import BaseModel -class QCGenerationTask(BaseModel, abc.ABC): +class SinglePointSpec(BaseModel): + """Defines the specification for performing a single point calculation.""" - type: Literal["base-task"] = "base-task" + program: str = Field(description="The program to use to evaluate the model.") + model: Model = Field(description="The QC model to perform the single point using.") - program: str = Field(..., description="The program to use to evaluate the model.") - model: Model = Field(..., description=str(Model.__doc__)) +class OptimizationSpec(SinglePointSpec): + """Defines the specification for performing a conformer optimization.""" + + procedure: GeometricProcedure = Field( + GeometricProcedure(), + description="The procedure to follow when optimizing the conformer.", + ) + + +class BaseTaskSpec(BaseModel, abc.ABC): + """The base for tasks that will generate QC reference data that force field + parameters will be trained to. + """ + + type: Literal["base-task"] -class HessianTaskSpec(QCGenerationTask): + +class HessianTaskSpec(BaseTaskSpec): + """ + + Examples: + + Define the specification for a task that will optimize an initial set of *up to* + 10 diverse conformers at the B3LYP-D3BJ/DZVP level of theory before finally + evaluting the hessian of the conformer with the lowest energy: + + >>> HessianTaskSpec( + >>> optimization_spec=OptimizationSpec( + >>> model=Model(method="B3LYP-D3BJ", basis="DZVP"), program="psi4", + >>> ) + >>> n_conformers=10, + >>> ) + + Create a specification for a task that will pre-optimize a set of conformers + using GFN2XTB, followed by a second optimization using B3LYP-D3BJ/DZVP, before + finally evaluting the hessian of the conformer with the lowest energy using + B3LYP-D3BJ/DZVP: + + >>> HessianTaskSpec( + >>> pre_optimization_spec=OptimizationSpec( + >>> model=Model(method="GFN2XTB", basis=None), program="xtb", + >>> ), + >>> optimization_spec=OptimizationSpec( + >>> model=Model(method="B3LYP-D3BJ", basis="DZVP"), program="psi4", + >>> ) + >>> ) + """ type: Literal["hessian"] = "hessian" @@ -30,35 +75,86 @@ class HessianTaskSpec(QCGenerationTask): "hessian. Each conformer will be minimized and the one with the lowest energy " "will have its hessian computed.", ) - optimization_spec: GeometricProcedure = Field( - GeometricProcedure(), + + pre_optimization_spec: Optional[OptimizationSpec] = Field( + None, + description="The (optional) specification to follow when pre-optimizing each " + "conformer using a 'cheaper' level of theory. If no value is provided, a " + "pre-optimization will not be performed.", + ) + optimization_spec: OptimizationSpec = Field( description="The specification for how to optimize each conformer before " - "computing the hessian.", + "computing the hessian of the lowest energy conformer.", ) class HessianTask(HessianTaskSpec): smiles: str = Field( - ..., description="A fully indexed SMILES representation of the molecule to compute " "the hessian for.", ) class OptimizationTaskSpec(HessianTaskSpec): + """ + Examples: + + Optimize multiple conformers of a molecule at the B3LYP-D3BJ/DZVP level of + theory + + >>> OptimizationTaskSpec( + >>> optimization_spec=OptimizationSpec( + >>> model=Model(method="B3LYP-D3BJ", basis="DZVP"), program="psi4", + >>> ) + >>> ) + + Optimize multiple conformers of a molecule first using GFN2XTB, followed by + a second optimization using B3LYP-D3BJ/DZVP: + + >>> OptimizationTaskSpec( + >>> pre_optimization_spec=OptimizationSpec( + >>> model=Model(method="GFN2XTB", basis=None), program="xtb", + >>> ), + >>> optimization_spec=OptimizationSpec( + >>> model=Model(method="B3LYP-D3BJ", basis="DZVP"), program="psi4", + >>> ) + >>> ) + + Optimize multiple conformers of a molecule using GFN2XTB before evaluating + the final energy of each conformer using B3LYP-D3BJ/DZVP: + + >>> OptimizationTaskSpec( + >>> optimization_spec=OptimizationSpec( + >>> model=Model(method="GFN2XTB", basis=None), program="xtb", + >>> ), + >>> evaluation_spec=SinglePointSpec( + >>> model=Model(method="B3LYP-D3BJ", basis="DZVP"), program="psi4", + >>> ) + >>> ) + """ + type: Literal["optimization"] = "optimization" + # Redefine base field to give a more specific description. + optimization_spec: OptimizationSpec = Field( + description="The specification for how to optimize each conformer.", + ) + evaluation_spec: Optional[SinglePointSpec] = Field( + None, + description="The (optional) specification to follow when evaluating properties " + "of the final conformer such as the energy and gradient.", + ) + class OptimizationTask(OptimizationTaskSpec): smiles: str = Field( - ..., description="A fully indexed SMILES representation of the molecule to optimize.", ) -class Torsion1DTaskSpec(QCGenerationTask): +class Torsion1DTaskSpec(BaseTaskSpec): type: Literal["torsion1d"] = "torsion1d" @@ -67,26 +163,26 @@ class Torsion1DTaskSpec(QCGenerationTask): None, description="The range of grid angles to scan." ) - 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.", + optimization_spec: OptimizationSpec = Field( + description="The specification for how to optimize each conformer at each " + "grid angle.", + ) + evaluation_spec: Optional[SinglePointSpec] = Field( + None, + description="The (optional) specification to follow when evaluating evaluating " + "the energy at each grid angle. If no value is provided the level of theory " + "used to optimize the conformer at each grid angle will be used.", ) n_conformers: conint(gt=0) = Field( 10, description="The number of initial conformers to seed the torsion drive with.", ) - sp_specification: Optional[QCGenerationTask] = Field( - None, - description="An extra optional specification used to compute the reference energy surface on the optimised geometries.", - ) class Torsion1DTask(Torsion1DTaskSpec): smiles: str = Field( - ..., description="An indexed SMILES representation of the molecule to drive.", ) central_bond: Tuple[int, int] = Field( @@ -122,13 +218,13 @@ def task_from_result(result): smiles=off_mol.to_smiles( isomeric=True, explicit_hydrogens=True, mapped=True ), - program=result.extras["program"], - model=result.input_specification.model, central_bond=(dihedral[1] + 1, dihedral[2] + 1), grid_spacing=result.keywords.grid_spacing[0], scan_range=result.keywords.dihedral_ranges, - optimization_spec=GeometricProcedure.from_opt_spec( - result.optimization_spec + optimization_spec=OptimizationSpec( + program=result.extras["program"], + model=result.input_specification.model, + procedure=GeometricProcedure.from_opt_spec(result.optimization_spec), ), ) else: diff --git a/openff/bespokefit/tests/cli/executor/test_submit.py b/openff/bespokefit/tests/cli/executor/test_submit.py index e9a31697..f449a107 100644 --- a/openff/bespokefit/tests/cli/executor/test_submit.py +++ b/openff/bespokefit/tests/cli/executor/test_submit.py @@ -183,11 +183,13 @@ def test_to_input_schema_overwrite_spec(tmpdir): workflow_file_name=None, ) - qc_spec = input_schema.stages[0].targets[0].calculation_specification - eval_spec = input_schema.stages[0].targets[0].reference_data.spec.sp_specification - assert qc_spec.program == "xtb" - assert qc_spec.model.method == "gfn2xtb" - assert qc_spec.model.basis is None + opt_spec = ( + input_schema.stages[0].targets[0].calculation_specification.optimization_spec + ) + eval_spec = input_schema.stages[0].targets[0].reference_data.spec.evaluation_spec + assert opt_spec.program == "xtb" + assert opt_spec.model.method == "gfn2xtb" + assert opt_spec.model.basis is None assert eval_spec.program == "torchani" assert eval_spec.model.basis is None assert eval_spec.model.method == "ani2x" diff --git a/openff/bespokefit/tests/executor/services/qcgenerator/test_app.py b/openff/bespokefit/tests/executor/services/qcgenerator/test_app.py index e9382e43..ce7a4260 100644 --- a/openff/bespokefit/tests/executor/services/qcgenerator/test_app.py +++ b/openff/bespokefit/tests/executor/services/qcgenerator/test_app.py @@ -21,7 +21,12 @@ QCGeneratorPOSTResponse, ) from openff.bespokefit.executor.utilities.depiction import IMAGE_UNAVAILABLE_SVG -from openff.bespokefit.schema.tasks import HessianTask, OptimizationTask, Torsion1DTask +from openff.bespokefit.schema.tasks import ( + HessianTask, + OptimizationSpec, + OptimizationTask, + Torsion1DTask, +) from openff.bespokefit.tests.executor.mocking.celery import mock_celery_task @@ -147,8 +152,10 @@ def test_get_qc_result( Torsion1DTask( smiles="[CH2:1][CH2:2]", central_bond=(1, 2), - program="rdkit", - model=Model(method="uff", basis=None), + optimization_spec=OptimizationSpec( + program="rdkit", + model=Model(method="uff", basis=None), + ), ), "compute_torsion_drive", ), @@ -156,16 +163,20 @@ def test_get_qc_result( OptimizationTask( smiles="[CH2:1][CH2:2]", n_conformers=1, - program="rdkit", - model=Model(method="uff", basis=None), + optimization_spec=OptimizationSpec( + program="rdkit", + model=Model(method="uff", basis=None), + ), ), "compute_optimization", ), ( HessianTask( smiles="[CH2:1][CH2:2]", - program="rdkit", - model=Model(method="uff", basis=None), + optimization_spec=OptimizationSpec( + program="rdkit", + model=Model(method="uff", basis=None), + ), ), "compute_hessian", ), diff --git a/openff/bespokefit/tests/executor/services/qcgenerator/test_cache.py b/openff/bespokefit/tests/executor/services/qcgenerator/test_cache.py index 0b183079..78b29f2f 100644 --- a/openff/bespokefit/tests/executor/services/qcgenerator/test_cache.py +++ b/openff/bespokefit/tests/executor/services/qcgenerator/test_cache.py @@ -6,7 +6,12 @@ _canonicalize_task, cached_compute_task, ) -from openff.bespokefit.schema.tasks import HessianTask, OptimizationTask, Torsion1DTask +from openff.bespokefit.schema.tasks import ( + HessianTask, + OptimizationSpec, + OptimizationTask, + Torsion1DTask, +) from openff.bespokefit.tests.executor.mocking.celery import mock_celery_task @@ -15,8 +20,10 @@ def test_canonicalize_torsion_task(): original_task = Torsion1DTask( smiles="[H:1][C:2]([H:3])([H:4])[O:5][H:6]", central_bond=(2, 5), - program="rdkit", - model=Model(method="uff", basis=None), + optimization_spec=OptimizationSpec( + program="rdkit", + model=Model(method="uff", basis=None), + ), ) canonical_task = _canonicalize_task(original_task) @@ -31,8 +38,10 @@ def test_canonicalize_torsion_task(): Torsion1DTask( smiles="[CH2:1][CH2:2]", central_bond=(1, 2), - program="rdkit", - model=Model(method="uff", basis=None), + optimization_spec=OptimizationSpec( + program="rdkit", + model=Model(method="uff", basis=None), + ), ), "compute_torsion_drive", ), @@ -40,16 +49,20 @@ def test_canonicalize_torsion_task(): OptimizationTask( smiles="[CH2:1][CH2:2]", n_conformers=1, - program="rdkit", - model=Model(method="uff", basis=None), + optimization_spec=OptimizationSpec( + program="rdkit", + model=Model(method="uff", basis=None), + ), ), "compute_optimization", ), ( HessianTask( smiles="[CH2:1][CH2:2]", - program="rdkit", - model=Model(method="uff", basis=None), + optimization_spec=OptimizationSpec( + program="rdkit", + model=Model(method="uff", basis=None), + ), ), "compute_hessian", ), diff --git a/openff/bespokefit/tests/executor/services/qcgenerator/test_worker.py b/openff/bespokefit/tests/executor/services/qcgenerator/test_worker.py index d4dceace..66d800cd 100644 --- a/openff/bespokefit/tests/executor/services/qcgenerator/test_worker.py +++ b/openff/bespokefit/tests/executor/services/qcgenerator/test_worker.py @@ -5,28 +5,22 @@ from qcelemental.models.procedures import OptimizationResult, TorsionDriveResult from openff.bespokefit.executor.services.qcgenerator import worker -from openff.bespokefit.schema.tasks import ( - OptimizationTask, - QCGenerationTask, - Torsion1DTask, -) +from openff.bespokefit.schema.tasks import OptimizationSpec def test_compute_torsion_drive(): - task = Torsion1DTask( + result_json = worker.compute_torsion_drive( smiles="[F][CH2:1][CH2:2][F]", central_bond=(1, 2), grid_spacing=180, scan_range=(-180, 180), - program="rdkit", - model=Model(method="uff", basis=None), - sp_specification=QCGenerationTask( - program="rdkit", model=Model(method="mmff94", basis=None) - ), + optimization_spec_json=OptimizationSpec( + program="rdkit", + model=Model(method="uff", basis=None), + ).json(), + n_conformers=1, ) - - result_json = worker.compute_torsion_drive(task.json()) assert isinstance(result_json, str) result_dict = json.loads(result_json) @@ -50,14 +44,13 @@ def test_compute_torsion_drive(): def test_compute_optimization(): - task = OptimizationTask( + result_json = worker.compute_optimization( smiles="CCCCC", + optimization_spec_json=OptimizationSpec( + program="rdkit", model=Model(method="uff", basis=None) + ).json(), n_conformers=2, - program="rdkit", - model=Model(method="uff", basis=None), ) - - result_json = worker.compute_optimization(task.json()) assert isinstance(result_json, str) result_dicts = json.loads(result_json) diff --git a/openff/bespokefit/tests/schema/test_targets.py b/openff/bespokefit/tests/schema/test_targets.py index e7de18b4..a2ed6afb 100644 --- a/openff/bespokefit/tests/schema/test_targets.py +++ b/openff/bespokefit/tests/schema/test_targets.py @@ -4,7 +4,11 @@ from openff.bespokefit.schema.data import BespokeQCData from openff.bespokefit.schema.targets import TorsionProfileTargetSchema -from openff.bespokefit.schema.tasks import HessianTaskSpec, Torsion1DTaskSpec +from openff.bespokefit.schema.tasks import ( + HessianTaskSpec, + OptimizationSpec, + Torsion1DTaskSpec, +) def test_check_reference_data(qc_torsion_drive_results): @@ -16,7 +20,9 @@ def test_check_reference_data(qc_torsion_drive_results): TorsionProfileTargetSchema( reference_data=BespokeQCData( spec=Torsion1DTaskSpec( - program="rdkit", model=Model(method="uff", basis=None) + optimization_spec=OptimizationSpec( + program="rdkit", model=Model(method="uff", basis=None) + ) ) ) ) @@ -27,7 +33,9 @@ def test_check_reference_data(qc_torsion_drive_results): TorsionProfileTargetSchema( reference_data=BespokeQCData( spec=HessianTaskSpec( - program="rdkit", model=Model(method="uff", basis=None) + optimization_spec=OptimizationSpec( + program="rdkit", model=Model(method="uff", basis=None) + ) ) ) ) diff --git a/openff/bespokefit/workflows/bespoke.py b/openff/bespokefit/workflows/bespoke.py index 85807502..9022a95f 100644 --- a/openff/bespokefit/workflows/bespoke.py +++ b/openff/bespokefit/workflows/bespoke.py @@ -48,8 +48,9 @@ ) from openff.bespokefit.schema.tasks import ( HessianTaskSpec, + OptimizationSpec, OptimizationTaskSpec, - QCGenerationTask, + SinglePointSpec, Torsion1DTaskSpec, ) from openff.bespokefit.utilities import parallel @@ -485,17 +486,19 @@ def _build_optimization_schema( # set the calculation specification for provenance and caching task_type = target_schema.bespoke_task_type() target_specification = task_type_to_spec[task_type]( - program=default_qc_spec.program.lower(), - # lower to hit the cache more often - model=Model( - method=default_qc_spec.method.lower(), - basis=default_qc_spec.basis.lower() - if default_qc_spec.basis is not None - else default_qc_spec.basis, - ), + optimization_spec=OptimizationSpec( + program=default_qc_spec.program.lower(), + # lower to hit the cache more often + model=Model( + method=default_qc_spec.method.lower(), + basis=default_qc_spec.basis.lower() + if default_qc_spec.basis is not None + else default_qc_spec.basis, + ), + ) ) if task_type == "torsion1d" and self.evaluation_qc_spec is not None: - target_specification.sp_specification = QCGenerationTask( + target_specification.evaluation_spec = SinglePointSpec( program=self.evaluation_qc_spec.program.lower(), model=Model( method=self.evaluation_qc_spec.method.lower(), From a37cf2538d3179919e9aff8330c79bd96846fb4d Mon Sep 17 00:00:00 2001 From: SimonBoothroyd Date: Mon, 23 May 2022 07:33:15 -0400 Subject: [PATCH 2/2] Fix tests --- .../executor/services/qcgenerator/cache.py | 21 ++++++++------- openff/bespokefit/tests/cli/test_cache.py | 2 +- .../executor/services/qcgenerator/test_app.py | 26 +++++++++---------- .../services/qcgenerator/test_cache.py | 21 +++++++-------- 4 files changed, 35 insertions(+), 35 deletions(-) diff --git a/openff/bespokefit/executor/services/qcgenerator/cache.py b/openff/bespokefit/executor/services/qcgenerator/cache.py index 89051f4a..5a5a232b 100644 --- a/openff/bespokefit/executor/services/qcgenerator/cache.py +++ b/openff/bespokefit/executor/services/qcgenerator/cache.py @@ -124,20 +124,21 @@ def _compute_torsion_drive_task( # There are no cached torsion drives at the 'pre-optimise' level of theory # we need to run a torsion drive and then optionally a single point - if task.evaluation_spec is None: - - torsion_drive_id = worker.compute_torsion_drive.delay( - smiles=task.smiles, - central_bond=task.central_bond, - grid_spacing=task.grid_spacing, - scan_range=task.scan_range, - optimization_spec_json=task.optimization_spec, - ).id + compute_torsion_drive_func = worker.compute_torsion_drive.s( + smiles=task.smiles, + central_bond=task.central_bond, + grid_spacing=task.grid_spacing, + scan_range=task.scan_range, + optimization_spec_json=task.optimization_spec.json(), + n_conformers=task.n_conformers, + ) + if task.evaluation_spec is None: + torsion_drive_id = compute_torsion_drive_func.delay().id else: task_future: AsyncResult = ( - worker.compute_torsion_drive.s(task_json=task.json()) + compute_torsion_drive_func | worker.re_evaluate_torsion_drive.s( evaluation_spec_json=task.evaluation_spec.json(), ) diff --git a/openff/bespokefit/tests/cli/test_cache.py b/openff/bespokefit/tests/cli/test_cache.py index 54f6c55b..5d045c33 100644 --- a/openff/bespokefit/tests/cli/test_cache.py +++ b/openff/bespokefit/tests/cli/test_cache.py @@ -131,7 +131,7 @@ def test_update_from_qcsubmit(redis_connection): # find the result in redis task_id = redis_connection.hget( name="qcgenerator:task-ids", - key="21212802afd507cae91fa9b8af76d7fa76174cc1ba4ab04b7c251aa5263598fbb79d9d85a7d0654257afe684312db399bbdbeb174366425a981ab20a31eb6938", + key="d68837c929ff61fb0d8eeda94648ad0da29146214452c5ad4e4c68b911099e05638a984e6101b3bf102ee639281893da85937457863d598b0931fe20253009c1", ).decode() assert redis_connection.hget("qcgenerator:types", task_id).decode() == "torsion1d" diff --git a/openff/bespokefit/tests/executor/services/qcgenerator/test_app.py b/openff/bespokefit/tests/executor/services/qcgenerator/test_app.py index ce7a4260..e2df7070 100644 --- a/openff/bespokefit/tests/executor/services/qcgenerator/test_app.py +++ b/openff/bespokefit/tests/executor/services/qcgenerator/test_app.py @@ -13,7 +13,6 @@ from openff.bespokefit.executor.services.qcgenerator import worker from openff.bespokefit.executor.services.qcgenerator.app import _retrieve_qc_result -from openff.bespokefit.executor.services.qcgenerator.cache import _canonicalize_task from openff.bespokefit.executor.services.qcgenerator.models import ( QCGeneratorGETPageResponse, QCGeneratorGETResponse, @@ -22,7 +21,6 @@ ) from openff.bespokefit.executor.utilities.depiction import IMAGE_UNAVAILABLE_SVG from openff.bespokefit.schema.tasks import ( - HessianTask, OptimizationSpec, OptimizationTask, Torsion1DTask, @@ -170,16 +168,16 @@ def test_get_qc_result( ), "compute_optimization", ), - ( - HessianTask( - smiles="[CH2:1][CH2:2]", - optimization_spec=OptimizationSpec( - program="rdkit", - model=Model(method="uff", basis=None), - ), - ), - "compute_hessian", - ), + # ( + # HessianTask( + # smiles="[CH2:1][CH2:2]", + # optimization_spec=OptimizationSpec( + # program="rdkit", + # model=Model(method="uff", basis=None), + # ), + # ), + # "compute_hessian", + # ), ], ) def test_post_qc_result( @@ -193,7 +191,9 @@ def test_post_qc_result( ) request.raise_for_status() - assert submitted_task_kwargs["task_json"] == _canonicalize_task(task).json() + assert ( + submitted_task_kwargs["optimization_spec_json"] == task.optimization_spec.json() + ) assert redis_connection.hget("qcgenerator:types", "1").decode() == task.type result = QCGeneratorPOSTResponse.parse_raw(request.text) diff --git a/openff/bespokefit/tests/executor/services/qcgenerator/test_cache.py b/openff/bespokefit/tests/executor/services/qcgenerator/test_cache.py index 78b29f2f..83047d0f 100644 --- a/openff/bespokefit/tests/executor/services/qcgenerator/test_cache.py +++ b/openff/bespokefit/tests/executor/services/qcgenerator/test_cache.py @@ -7,7 +7,6 @@ cached_compute_task, ) from openff.bespokefit.schema.tasks import ( - HessianTask, OptimizationSpec, OptimizationTask, Torsion1DTask, @@ -56,16 +55,16 @@ def test_canonicalize_torsion_task(): ), "compute_optimization", ), - ( - HessianTask( - smiles="[CH2:1][CH2:2]", - optimization_spec=OptimizationSpec( - program="rdkit", - model=Model(method="uff", basis=None), - ), - ), - "compute_hessian", - ), + # ( + # HessianTask( + # smiles="[CH2:1][CH2:2]", + # optimization_spec=OptimizationSpec( + # program="rdkit", + # model=Model(method="uff", basis=None), + # ), + # ), + # "compute_hessian", + # ), ], ) def test_cached_compute_task(