def test_prep_for_fitting_no_ref():
"""
Make sure an error is raised if we try to fit with missing reference data.
"""
torsion_target = AbInitio_SMIRNOFF()
target_schema = biphenyl_target(target=torsion_target)
with pytest.raises(MissingReferenceError):
torsion_target.prep_for_fitting(fitting_target=target_schema)
def test_abinitio_fitting_prep_no_gradient():
"""
Test preparing for fitting using the abinitio target.
"""
torsion_target = AbInitio_SMIRNOFF()
torsion_target.fit_gradient = False
target_schema = biphenyl_target(target=torsion_target)
biphenyl = Molecule.from_file(file_path=get_data("biphenyl.sdf"),
file_format="sdf")
# now load in a scan result we have saved
result_data = TorsionDriveCollectionResult.parse_file(
get_data("biphenyl.json.xz"))
# now try and update the results
target_schema.update_with_results(results=result_data)
assert target_schema.ready_for_fitting is True
# now try and prep for fitting
with temp_directory():
torsion_target.prep_for_fitting(fitting_target=target_schema)
# we should only have one torsion drive to do here
folders = os.listdir(".")
assert len(folders) == 1
target_files = os.listdir(folders[0])
assert "molecule.pdb" in target_files
assert "scan.xyz" in target_files
assert "molecule.mol2" in target_files
assert "qdata.txt" in target_files
# now we need to make sure the pdb order was not changed
mol = Molecule.from_file(os.path.join(folders[0], "molecule.pdb"),
file_format="pdb")
isomorphic, atom_map = Molecule.are_isomorphic(biphenyl,
mol,
return_atom_map=True)
assert isomorphic is True
assert atom_map == dict((i, i) for i in range(biphenyl.n_atoms))
# also make sure charges are in the mol2 file
mol = Molecule.from_file(os.path.join(folders[0], "molecule.mol2"),
"mol2")
assert mol.partial_charges is not None
# make sure the scan coords and energies match
qdata_file = os.path.join(folders[0], "qdata.txt")
coords, energies, gradients = read_qdata(qdata_file=qdata_file)
# make sure no gradients were written
assert not gradients
reference_data = target_schema.tasks[0].reference_data()
for i, (coord, energy) in enumerate(zip(coords, energies)):
# find the reference data
data = reference_data[i]
assert data.energy == energy
assert coord == data.molecule.geometry.flatten().tolist()
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 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_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_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_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()
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_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
"""
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):
-----
This function can be used to generate many optimize in files so many force balance jobs can be ran simultaneously.
"""
# check that all of the fitting targets have been set
target_names = [target.name.lower() for target in self.optimization_targets]
for target_name in fitting_targets.keys():
if target_name.lower() not in target_names:
raise TargetNotSetError(
f"The target {target_name} is not setup for this optimizer and is required, please add it with runtime options using `set_optimization_target`."
)
# grab the template file
template_file = get_data(os.path.join("templates", "optimize.txt"))
with open(template_file) as file:
template = Template(file.read())
data = self.dict()
# function to collect the priors from the targets.
data["priors"] = priors
# now we need to collect the fitting target data from the schema
data["fitting_targets"] = fitting_targets
rendered_template = template.render(**data)
with open("optimize.in", "w") as opt_in:
opt_in.write(rendered_template)
# register all of the available targets.
ForceBalanceOptimizer.register_target(AbInitio_SMIRNOFF())
ForceBalanceOptimizer.register_target(TorsionProfile_SMIRNOFF())