def test_adding_params_parameterize_flag():
"""
Test adding new smirks patterns with cosmetic attributes.
"""
ff = ForceFieldEditor(forcefield_name="openff-1.0.0.offxml")
# add an atom smirks for boron
boron = AtomSmirks(smirks="[#5:1]", parameterize={"epsilon"}, atoms={(0,)}, epsilon=0.04, rmin_half=3)
# add boron with the flag
ff.add_smirks(smirks=[boron, ], parameterize=True)
with temp_directory():
ff.forcefield.to_file(filename="boron.offxml")
# this should fail if the flag was added
with pytest.raises(SMIRNOFFSpecError):
_ = ForceField("boron.offxml", allow_cosmetic_attributes=False)
boron_ff = ForceField("boron.offxml", allow_cosmetic_attributes=True)
# now look for the parameter we added
boron_param = boron_ff.get_parameter_handler(
"vdW"
).parameters["[#5:1]"]
# now make sure it has the param flag
param_dict = boron_param.__dict__
assert param_dict["_cosmetic_attribs"] == ["parameterize"]
assert param_dict["_parameterize"] == "epsilon"
def collect_results(self, schema: OptimizationSchema) -> OptimizationSchema:
"""
Collect the results of a forcebalance optimization.
Check the exit state of the optimization before attempting to update the final smirks parameters.
Parameters
----------
schema: OptimizationSchema
The workflow schema that should be updated with the results of the current optimization.
Returns
-------
OptimizationSchema
The updated workflow schema.
"""
import copy
# look for the result
result = self.read_output()
schema.status = result["status"]
ff = ForceFieldEditor(result["forcefield"])
# make a new list as smirks are updated in place
final_smirks = copy.deepcopy(schema.target_smirks)
ff.update_smirks_parameters(smirks=final_smirks)
# put the new smirks back in the schema
schema.final_smirks = final_smirks
return schema
def test_label_molecule():
"""
Test that labeling a molecule with the editor works.
"""
ff = ForceFieldEditor(forcefield_name="openff-1.0.0.offxml")
ethane = Molecule.from_file(file_path=get_data("ethane.sdf"), file_format="sdf")
labels = ff.label_molecule(molecule=ethane)
for param_type in ["Bonds", "Angles", "ProperTorsions", "ImproperTorsions", "vdW"]:
assert param_type in labels
def get_fitting_forcefield(self) -> ForceField:
"""
Take the initial forcefield and edit it to add the new terms and return the OpenFF FF object.
Parameters:
initial_forcefield: The name of the initial Forcefield we will be starting at.
"""
from openff.bespokefit.forcefield_tools import ForceFieldEditor
# get all of the new target smirks
target_smirks = self.parameterize_smirks
ff = ForceFieldEditor(self.initial_forcefield)
ff.add_smirks(target_smirks, parameterize=True)
# if there are any parameters from a different optimization stage add them here without parameterize tags
return ff.forcefield
def _get_all_bespoke_smirks(
self,
molecule: Molecule,
forcefield_editor: ForceFieldEditor,
central_bonds: Optional[List[Tuple[int, int]]] = None,
) -> Smirks:
"""
The main worker method for generating new bespoke smirks, this will check which parameters are wanted and call each method.
The new smirks will then have any dummy values set by the initial forcefield values.
"""
bespoke_smirks = []
if SmirksType.Vdw in self.target_smirks:
atom_smirks = self._get_bespoke_atom_smirks(molecule=molecule)
bespoke_smirks.extend(atom_smirks)
if SmirksType.Bonds in self.target_smirks:
bond_smirks = self._get_bespoke_bond_smirks(molecule=molecule)
bespoke_smirks.extend(bond_smirks)
if SmirksType.Angles in self.target_smirks:
angle_smirks = self._get_bespoke_angle_smirks(molecule=molecule)
bespoke_smirks.extend(angle_smirks)
if SmirksType.ProperTorsions in self.target_smirks:
torsion_smirks = self._get_bespoke_torsion_smirks(
molecule=molecule, central_bonds=central_bonds)
bespoke_smirks.extend(torsion_smirks)
# now we need to update all smirks
updated_smirks = forcefield_editor.get_initial_parameters(
molecule=molecule, smirks=bespoke_smirks, clear_existing=True)
return updated_smirks
def _get_all_smirks(
self,
molecule: Molecule,
forcefield_editor: ForceFieldEditor,
central_bonds: Optional[List[Tuple[int, int]]] = None,
) -> Smirks:
"""
The main worker method for extracting current smirks patterns for the molecule, this will only extract parameters for the requested groups.
"""
# generate a list of all of the required parameter types
requested_smirks = []
if SmirksType.Vdw in self.target_smirks:
atom_indices = [(i, ) for i in range(molecule.n_atoms)]
requested_smirks.extend(atom_indices)
if SmirksType.Bonds in self.target_smirks:
bond_indices = [(bond.atom1_index, bond.atom2_index)
for bond in molecule.bonds]
requested_smirks.extend(bond_indices)
if SmirksType.Angles in self.target_smirks:
angle_indices = [
tuple([atom.molecule_atom_index for atom in angle])
for angle in molecule.angles
]
requested_smirks.extend(angle_indices)
if SmirksType.ProperTorsions in self.target_smirks:
torsions = self._get_torsion_indices(molecule=molecule,
central_bonds=central_bonds)
requested_smirks.extend(torsions)
# now request all of these smirks from the forcefield
new_smirks = forcefield_editor.get_smirks_parameters(
molecule=molecule, atoms=requested_smirks)
return new_smirks
def test_get_all_bespoke_smirks():
"""
Generate bespoke smirks for all parameters in a molecule, make sure they all hit the intended atoms, and every term is now bespoke.
"""
gen = SmirksGenerator()
gen.target_smirks = [
SmirksType.Vdw, SmirksType.Bonds, SmirksType.Angles,
SmirksType.ProperTorsions
]
mol = Molecule.from_smiles("CO")
all_bespoke_smirks = gen._get_all_bespoke_smirks(
molecule=mol,
forcefield_editor=ForceFieldEditor(
"openff_unconstrained-1.3.0.offxml"))
# this is a list of all bespoke smirks with real initial values
all_matches = []
for smirk in all_bespoke_smirks:
atoms = condense_matches(mol.chemical_environment_matches(
smirk.smirks))
assert compare_matches(atoms, smirk.atoms) is True
all_matches.extend(atoms)
assert all_covered(all_matches, mol) is True
def test_updating_smirks(smirks_data):
"""
Test that each smirks type can be updated from a given force field file.
"""
smirks, updated_params = smirks_data
ff = ForceFieldEditor(forcefield_name="openff-1.0.0.offxml")
ff.update_smirks_parameters(smirks=[smirks, ])
if smirks.type.value != "ProperTorsions":
for param, value in updated_params.items():
assert getattr(smirks, param) == value
else:
term = smirks.terms["3"]
for param, value in updated_params.items():
assert getattr(term, param) == value
def test_adding_new_smirks_types(smirks):
"""
Test adding new smirks to a forcefield with and without the parameterize flag.
"""
param, param_id = smirks
ff = ForceFieldEditor("openff-1.0.0.offxml")
# now make some new smirks pattern
ff.add_smirks(smirks=[param, ], parameterize=True)
# now make sure it was added under the correct parameter handler
with temp_directory():
ff.forcefield.to_file(filename="bespoke.offxml")
new_ff = ForceField("bespoke.offxml", allow_cosmetic_attributes=True)
parameter = new_ff.get_parameter_handler(param.type.value).parameters[param.smirks]
param_dict = parameter.__dict__
assert param_dict["_cosmetic_attribs"] == ["parameterize"]
assert set(param_dict["_parameterize"].split()) == param.parameterize
# make sure the id is correct
assert param_id in parameter.id
def generate_smirks(
self,
molecule: Molecule,
central_bonds: Optional[List[Tuple[int, int]]] = None) -> Smirks:
"""
The main method of the class which takes an input molecule and can generate new smirks patterns for it corresponding to the types set in the target smirks list.
Parameters:
molecule: The openforcefield molecule for which we should make the smirks patterns
central_bonds: An optional list of central bonds which are used with the TargetTorsions option to specify which torsions need new terms.
Returns:
A list of new bespoke smirks parameters for the molecule.
"""
if not self.target_smirks:
raise SMIRKSTypeError(
"No smirks targets were provided so no new patterns were made, set a target and run again."
)
# now we need to set the forcefield
if isinstance(self.initial_forcefield, ForceFieldEditor):
ff = self.initial_forcefield
else:
ff = ForceFieldEditor(forcefield_name=self.initial_forcefield)
# for each requested smirks type generate the parameters
if self.generate_bespoke_terms:
new_smirks = self._get_all_bespoke_smirks(
molecule=molecule,
forcefield_editor=ff,
central_bonds=central_bonds)
else:
new_smirks = self._get_all_smirks(molecule=molecule,
forcefield_editor=ff,
central_bonds=central_bonds)
# now sort the smirks into a dict
# all_smirks = dict()
# for smirk in new_smirks:
# all_smirks.setdefault(smirk.type.value, []).append(smirk)
# now we need to check if we need to expand any torsion smirks
if self.expand_torsion_terms:
for smirk in new_smirks:
if smirk.type == SmirksType.ProperTorsions:
for i in range(1, 5):
if str(i) not in smirk.terms:
smirk.add_torsion_term(f"k{i}")
return new_smirks
def test_loading_forcefields():
"""
Test that loading the forcefield always strips out any constraints.
"""
# load in the initial FF with constraints
ff = ForceFieldEditor(forcefield_name="openff-1.0.0.offxml")
with temp_directory():
# write out the ff
ff.forcefield.to_file(filename="bespoke.offxml")
# read the file and look for the constraints tag
new_ff = ForceField("bespoke.offxml")
assert "Constraints" not in new_ff._parameter_handlers
def test_get_all_smirks():
"""
Get the full list of smirks which cover this molecule from the forcefield, no new smirks should be generated here.
"""
gen = SmirksGenerator()
gen.target_smirks = [
SmirksType.Vdw, SmirksType.Bonds, SmirksType.Angles,
SmirksType.ProperTorsions
]
mol = Molecule.from_smiles("CO")
all_smirks = gen._get_all_smirks(molecule=mol,
forcefield_editor=ForceFieldEditor(
"openff_unconstrained-1.3.0.offxml"))
# this is a list of all of the smirks from the forcefield
all_matches = []
for smirk in all_smirks:
atoms = condense_matches(mol.chemical_environment_matches(
smirk.smirks))
assert compare_matches(atoms, smirk.atoms) is True
all_matches.extend(atoms)
assert all_covered(all_matches, mol) is True
def get_final_forcefield(
self,
generate_bespoke_terms: bool = True,
drop_out_value: Optional[float] = None,
) -> ForceField:
"""
Generate the final bespoke forcefield for this molecule by collecting together all optimized smirks.
Note:
It is know that when creating fitting smirks for the fragments that they can hit unintended dihedrals in other fragments if they are similar during fitting. To ensure the correct parameters are used in their intended positions
on the parent molecule each fragment is typed with the fitting forcefield and parameters again and a new bespoke term for the parent is made which uses the same parameters.
Parameters:
generate_bespoke_terms: If molecule specific bespoke terms should be made, this is recommended as some fragment smirks patterns may not transfer back to the parent correctly due to fragmentation.
drop_out_value: Any torsion force constants below this value will be dropped as they are probably negligible.
"""
# TODO change this to a util function remove from target base class
# check that all optimizations are finished
if self.status != Status.Complete:
raise OptimizerError(
f"The molecule has not completed all optimization stages which are required to generate the final forcefield."
)
if self.final_smirks is None:
raise OptimizerError(
f"The optimization status is complete but no optimized smirks were found."
)
from openff.bespokefit.forcefield_tools import ForceFieldEditor
# get all of the target smirks
target_smirks = self.final_smirks
# build the parent molecule
parent_molecule = self.target_molecule.molecule
if drop_out_value is not None:
# loop over all of the target smirks and drop and torsion k values lower than the drop out
for smirk in target_smirks:
if smirk.type == SmirksType.ProperTorsions:
# keep a list of terms to remove
to_remove = []
for p, term in smirk.terms.items():
if abs(float(term.k.split("*")[0])) < drop_out_value:
to_remove.append(p)
# now remove the low values
for p in to_remove:
del smirk.terms[p]
# the final fitting force field should have all final smirks propagated through
fitting_ff = ForceFieldEditor(forcefield_name=self.initial_forcefield)
fitting_ff.add_smirks(smirks=target_smirks, parameterize=False)
# here we type the fragment with the final forcefield and then build a bespoke dihedral term for the parent
# to hit the atoms that map from the fragment to the parent
# we do not modify any other smirks types but this needs testing to make sure that they do transfer.
if generate_bespoke_terms:
bespoke_smirks = []
# get a list of unique fragments and all of the torsions that have been targeted
for target in self.targets:
if target.collection_workflow == "torsion1d":
for task in target.tasks:
# check if the smirks are from a fragment
if task.fragment:
new_smirks = self._generate_bespoke_torsions(
forcefield=fitting_ff,
parent_molecule=parent_molecule,
task_data=task,
)
bespoke_smirks.extend(new_smirks)
# make a new ff object with the new terms
bespoke_ff = ForceFieldEditor(
forcefield_name=self.initial_forcefield)
# get a list of non torsion smirks
new_smirks = [
smirk for smirk in target_smirks
if smirk.type != SmirksType.ProperTorsions
]
new_smirks.extend(bespoke_smirks)
bespoke_ff.add_smirks(smirks=new_smirks, parameterize=False)
return bespoke_ff.forcefield
else:
return fitting_ff.forcefield