def test_get_registered_targets():
"""
Test listing all of the registered targets for this optimizer.
"""
fb = ForceBalanceOptimizer()
targets = fb.get_registered_targets()
assert targets == [AbInitio_SMIRNOFF(), TorsionProfile_SMIRNOFF()]
def test_task_from_results():
"""
Test making an individual task from a set of results
"""
# load a client and pull some results
client = FractalClient()
# grab a dataset with small fragments in it
result = TorsionDriveCollectionResult.from_server(
client=client,
spec_name="default",
dataset_name="OpenFF-benchmark-ligand-fragments-v1.0",
final_molecule_only=True,
subset=bace_entries[:1])
# grab the only result
result = list(result.collection.values())[0]
# set up the workflow
workflow = WorkflowFactory()
fb = ForceBalanceOptimizer()
fb.set_optimization_target(target=AbInitio_SMIRNOFF())
workflow.set_optimizer(optimizer=fb)
# this should be a simple biphenyl molecule
opt_schema = workflow._task_from_results(results=[
result,
], index=1)
assert opt_schema.initial_forcefield == workflow.initial_forcefield
assert opt_schema.optimizer_name == fb.optimizer_name
assert opt_schema.job_id == "bespoke_task_1"
assert bool(opt_schema.target_smirks) is True
assert opt_schema.target_parameters == workflow.target_parameters
assert result.molecule == opt_schema.target_molecule.molecule
assert opt_schema.n_tasks == 1
assert opt_schema.n_targets == 1
assert opt_schema.ready_for_fitting is True
def test_forcebalance_collect_results():
"""
Test trying to collect results that have been successful and updated the parameters.
"""
workflow = biphenyl_workflow(target=AbInitio_SMIRNOFF)
# first make sure the target smirks are set to the default value
target_smirks = workflow.target_smirks
for smirk in target_smirks:
for param in smirk.terms.values():
# starting value
assert param.k == "1.048715180139 * mole**-1 * kilocalorie"
# set up the dummy output folder
with temp_directory():
# copy the file over
shutil.copy(get_data("complete.out"), "optimize.out")
results_folder = os.path.join("result", "optimize")
os.makedirs(results_folder, exist_ok=True)
ff_path = os.path.join(results_folder, "bespoke.offxml")
shutil.copy(get_data("bespoke.offxml"), ff_path)
fb = ForceBalanceOptimizer()
result_workflow = fb.collect_results(schema=workflow)
# make sure the smirks have been updated
new_smirks = result_workflow.final_smirks
for smirk in new_smirks:
for param in smirk.terms.values():
assert param.k != "1.048715180139 * mole**-1 * kilocalorie"
def test_make_fitting_schema_from_results():
"""
Test that new fitting schemas can be made from results and that all results are full
"""
# build the workflow
workflow = WorkflowFactory()
fb = ForceBalanceOptimizer()
fb.set_optimization_target(target=AbInitio_SMIRNOFF())
workflow.set_optimizer(optimizer=fb)
# set up the client and load the results
# load a client and pull some results
client = FractalClient()
# grab a dataset with small fragments in it
result = TorsionDriveCollectionResult.from_server(
client=client,
spec_name="default",
dataset_name="OpenFF-benchmark-ligand-fragments-v1.0",
final_molecule_only=True,
subset=bace_entries)
schema = workflow.fitting_schema_from_results(results=result, combine=True)
# there should be 2 total molecules as we have combined two results
assert schema.n_molecules == 2
# there are a total of 3 torsiondrives
assert schema.n_tasks == 3
# make sure each task has results and is ready to fit
for task in schema.tasks:
assert task.ready_for_fitting is True
def test_forcebalance_name_change():
"""
Check that the optimizer name can not be changed.
"""
fb = ForceBalanceOptimizer()
# try a random name
with pytest.raises(ValidationError):
fb.optimizer_name = "fancy_opt"
def test_forcebalance_provenance():
"""
Make sure the correct forcebalance version is returned.
"""
import forcebalance
fb = ForceBalanceOptimizer()
provenance = fb.provenance()
assert provenance["forcebalance"] == forcebalance.__version__
def test_get_registered_target_names():
"""
Make sure the names of the targets are returned.
"""
fb = ForceBalanceOptimizer()
names = fb.get_registered_target_names()
assert AbInitio_SMIRNOFF().name.lower() in names
assert TorsionProfile_SMIRNOFF().name.lower() in names
def test_pre_run_check_no_smirks():
"""
Make sure that the pre run check checks that some target smirks have been supplied.
"""
workflow = WorkflowFactory()
ethane = Molecule.from_file(file_path=get_data("ethane.sdf"),
file_format="sdf")
fb = ForceBalanceOptimizer()
fb.set_optimization_target(target=AbInitio_SMIRNOFF())
workflow.set_optimizer(optimizer=fb)
workflow.target_smirks = []
with pytest.raises(TargetNotSetError):
_ = workflow.fitting_schema_from_molecules(molecules=ethane)
def test_pre_run_check_no_params():
"""
Make sure that the pre run check catches if we have not set any parameters to optimise, like bond length.
"""
workflow = WorkflowFactory()
ethane = Molecule.from_file(file_path=get_data("ethane.sdf"),
file_format="sdf")
fb = ForceBalanceOptimizer()
fb.set_optimization_target(target=AbInitio_SMIRNOFF())
workflow.set_optimizer(optimizer=fb)
workflow.target_parameters = []
with pytest.raises(TargetNotSetError):
_ = workflow.fitting_schema_from_molecules(molecules=ethane)
def test_pre_run_check_no_frag():
"""
Make sure the pre run check catches if there is no fragmentation method set.
"""
workflow = WorkflowFactory()
ethane = Molecule.from_file(file_path=get_data("ethane.sdf"),
file_format="sdf")
fb = ForceBalanceOptimizer()
fb.set_optimization_target(target=AbInitio_SMIRNOFF())
workflow.set_optimizer(optimizer=fb)
workflow.fragmentation_engine = None
with pytest.raises(FragmenterError):
_ = workflow.fitting_schema_from_molecules(molecules=ethane)
def test_generate_opt_in():
"""
Test generating the optimize in file with various input settings.
"""
fb = ForceBalanceOptimizer(penalty_type="L1", max_iterations=150)
# make sure an error is raised if the targets were not set
with temp_directory():
with pytest.raises(TargetNotSetError):
fb.generate_optimize_in(
priors={"test": 1.23},
fitting_targets={"AbInitio_SMIRNOFF": ["job1", "job2"]})
# now set them and run again
fb.set_optimization_target(AbInitio_SMIRNOFF())
fb.generate_optimize_in(
priors={"test": 1.23},
fitting_targets={"AbInitio_SMIRNOFF": ["job1", "job2"]})
# now load in the file and check the attributes
with open("optimize.in") as opt_in:
data = opt_in.readlines()
assert " test : 1.23\n" in data
assert "penalty_type L1\n" in data
assert "maxstep 150\n" in data
assert "type AbInitio_SMIRNOFF\n" in data
assert "name job1\n" in data
assert "name job2\n" in data
def biphenyl_workflow(target) -> OptimizationSchema:
"""
Create a workflow schema which targets the rotatable bond in ethane.
"""
mol = Molecule.from_file(get_data("biphenyl.sdf"), "sdf")
workflow = WorkflowFactory()
# turn off bespoke terms we want fast fitting
workflow.generate_bespoke_terms = False
workflow.expand_torsion_terms = False
fb = ForceBalanceOptimizer()
target = target()
fb.set_optimization_target(target=target)
workflow.set_optimizer(optimizer=fb)
schema = workflow.fitting_schema_from_molecules(molecules=mol)
return schema.tasks[0]
def test_forcebalance_collect_result_error():
"""
Test trying to collect the result when the workflow has an error.
"""
workflow = biphenyl_workflow(target=AbInitio_SMIRNOFF)
# we need to set up a dummy folder with the error
with temp_directory():
# copy the file over
shutil.copy(get_data("error.out"), "optimize.out")
results_folder = os.path.join("result", "optimize")
os.makedirs(results_folder, exist_ok=True)
ff_path = os.path.join(results_folder, "bespoke.offxml")
shutil.copy(get_data("bespoke.offxml"), ff_path)
fb = ForceBalanceOptimizer()
result_workflow = fb.collect_results(schema=workflow)
assert result_workflow.status == Status.ConvergenceError
def test_get_optimizer_no_settings():
"""
Test getting an optimizer with no extra settings.
"""
optimizer = get_optimizer(optimizer_name="ForceBalanceOptimizer")
# make sure the settings are the default
assert optimizer.dict() == ForceBalanceOptimizer().dict()
def test_remove_optimization_stages():
"""
Test removing optimization stages from the workflow.
"""
workflow = WorkflowFactory(optimization_workflow=[ForceBalanceOptimizer()])
workflow.clear_optimizer()
assert workflow.optimizer is None
def test_forcebalance_readoutput(output):
"""
Test reading the output of a forcebalance run.
"""
file_name, status = output
with temp_directory():
# copy the file over
shutil.copy(get_data(file_name), "optimize.out")
# now we have to make sum dummy folders
results_folder = os.path.join("result", "optimize")
os.makedirs(results_folder, exist_ok=True)
with open(os.path.join(results_folder, "bespoke.offxml"), "w") as xml:
xml.write("test")
fb = ForceBalanceOptimizer()
result = fb.read_output()
assert result["status"] == status
assert "bespoke.offxml" in result["forcefield"]
def test_pre_run_check_no_target():
"""
Make sure that the pre run check catches if there are no targets set up
"""
workflow = WorkflowFactory()
ethane = Molecule.from_file(file_path=get_data("ethane.sdf"),
file_format="sdf")
fb = ForceBalanceOptimizer()
workflow.set_optimizer(optimizer=fb)
with pytest.raises(OptimizerError):
_ = workflow.fitting_schema_from_molecules(molecules=ethane)
def test_task_from_molecule():
"""
Test the workflow function which makes the optimization schema from a molecule
"""
bace = Molecule.from_file(file_path=get_data("bace.sdf"),
file_format="sdf")
workflow = WorkflowFactory()
fb = ForceBalanceOptimizer()
fb.set_optimization_target(target=AbInitio_SMIRNOFF())
workflow.set_optimizer(optimizer=fb)
opt_schema = workflow._task_from_molecule(molecule=bace, index=1)
assert opt_schema.initial_forcefield == workflow.initial_forcefield
assert opt_schema.optimizer_name == fb.optimizer_name
assert opt_schema.job_id == "bespoke_task_1"
assert bool(opt_schema.target_smirks) is True
assert opt_schema.target_parameters == workflow.target_parameters
assert opt_schema.target_molecule.molecule == bace
assert opt_schema.n_tasks == 3
assert opt_schema.n_targets == 1
def test_forcebalance_optimize(optimization_target):
"""
Test running the full optimization stage for a simple biphenyl system using different targets.
The data has been extracted from qcarchive.
"""
from openff.qcsubmit.results import TorsionDriveCollectionResult
workflow = biphenyl_workflow(target=optimization_target)
with temp_directory():
# load the computed results and add them to the workflow
torsiondrive_result = TorsionDriveCollectionResult.parse_file(
get_data("biphenyl.json.xz"))
workflow.update_with_results(results=torsiondrive_result)
# setup the optimizer
fb = ForceBalanceOptimizer()
result = fb.optimize(schema=workflow)
assert result.status == Status.Complete
new_smirks = result.target_smirks
for smirk in new_smirks:
for param in smirk.terms.values():
assert param.k != "1e-05 * mole**-1 * kilocalorie"
def test_make_fitting_schema_from_molecules():
"""
Test making a fitting schema for a simple molecule using the default settings.
Bace is a small molecule that should split into 2 fragments for a total of 3 torsiondrives.
"""
bace = Molecule.from_file(file_path=get_data("bace.sdf"),
file_format="sdf")
workflow = WorkflowFactory()
fb = ForceBalanceOptimizer()
fb.set_optimization_target(target=AbInitio_SMIRNOFF())
workflow.set_optimizer(optimizer=fb)
schema = workflow.fitting_schema_from_molecules(molecules=bace)
# make sure one ethane torsion dirve is made
assert schema.n_molecules == 1
assert schema.n_tasks == 3
assert bace in schema.molecules
assert bace not in schema.entry_molecules
# get the qcsubmit dataset
datasets = schema.generate_qcsubmit_datasets()
assert len(datasets) == 1
assert datasets[0].dataset_type == "TorsiondriveDataset"
assert datasets[0].n_records == 3
def test_deregister_optimizers(optimizers):
"""
Test removing optimizers from the registered list.
"""
optimizer, error = optimizers
if error is None:
deregister_optimizer(optimizer=optimizer)
# make sure it is missing
assert "forcebalanceoptimizer" not in list_optimizers()
# now add it back
register_optimizer(ForceBalanceOptimizer())
else:
with pytest.raises(error):
deregister_optimizer(optimizer=optimizer)
def test_adding_optimization_stages_missing():
"""
Test adding an optimization stage with an optimizer which is not registered.
"""
# first we need to remove the forcebalance optimizer
fb = ForceBalanceOptimizer()
deregister_optimizer(optimizer=fb)
# now try and add it to the workflow
workflow = WorkflowFactory()
with pytest.raises(OptimizerError):
workflow.set_optimizer(fb)
# register it again
register_optimizer(optimizer=fb)
def test_workflow_export_import():
"""
Test exporting and importing a workflow.
"""
workflow = WorkflowFactory()
# add fb and a target with non standard settings
fb = ForceBalanceOptimizer(
penalty_type="L1",
optimization_targets=[AbInitio_SMIRNOFF(fragmentation=False)])
workflow.set_optimizer(optimizer=fb)
with temp_directory():
workflow.export_workflow(file_name="test.json")
# now read it back in
workflow2 = WorkflowFactory.parse_file("test.json")
assert workflow.dict() == workflow2.dict()
with pytest.raises(error):
_ = get_optimizer(optimizer_name=optimzer_name, **settings)
def test_get_optimizer_no_settings():
"""
Test getting an optimizer with no extra settings.
"""
optimizer = get_optimizer(optimizer_name="ForceBalanceOptimizer")
# make sure the settings are the default
assert optimizer.dict() == ForceBalanceOptimizer().dict()
@pytest.mark.parametrize("optimizers", [
pytest.param(("ForceBalanceOptimizer", None), id="Forcebalance by name"),
pytest.param((ForceBalanceOptimizer(), None), id="Forcebalance by class"),
pytest.param(("FancyOptimizer", OptimizerError), id="Missing optimizer")
])
def test_deregister_optimizers(optimizers):
"""
Test removing optimizers from the registered list.
"""
optimizer, error = optimizers
if error is None:
deregister_optimizer(optimizer=optimizer)
# make sure it is missing
assert "forcebalanceoptimizer" not in list_optimizers()
# now add it back
register_optimizer(ForceBalanceOptimizer())
else:
def test_forcebalance_available():
"""
Make sure forcebalance is correctly found when installed.
"""
assert ForceBalanceOptimizer.is_available() is True
def test_register_optimization_target():
"""
Test adding new optimization targets via the api.
"""
fb = ForceBalanceOptimizer()
# first try and add a target again with different settings
target = TorsionProfile_SMIRNOFF(keywords={"test": True})
with pytest.raises(TargetRegisterError):
# this will not allow the target to be registered
fb.register_target(target=target, replace=False)
opt_target = fb.get_optimization_target(target_name=target.name)
assert opt_target.keywords != {"test": True}
# now use the replace flag
fb.register_target(target=target, replace=True)
opt_target = fb.get_optimization_target(target_name=target.name)
assert opt_target.keywords["test"] is True
# now add a totally new target
target = TorsionDrive1D()
fb.register_target(target=target, replace=False)
assert target in fb.get_registered_targets()
"""
force_file, error = forcefield
factory = WorkflowFactory()
if error is None:
factory.initial_forcefield = force_file
else:
with pytest.raises(ForceFieldError):
factory.initial_forcefield = force_file
@pytest.mark.parametrize("optimization_data", [
pytest.param(
("ForceBalanceOptimizer", None), id="Forcebalance string pass"),
pytest.param(("BadOptimizer", OptimizerError), id="Missing optimizer"),
pytest.param(
(ForceBalanceOptimizer(optimization_targets=[AbInitio_SMIRNOFF()]),
None),
id="Forcebalance class with target.")
])
def test_adding_optimization_stages(optimization_data):
"""
Test adding optimization stages to the workflow.
"""
stage, error = optimization_data
workflow = WorkflowFactory()
if error is None:
workflow.set_optimizer(optimizer=stage)
assert workflow.optimizer is not None
else:
with pytest.raises(error):
