def test_get_initial_state(tmpdir, starting_conformations):
"""
Make sure we can correctly build a starting state using the torsiondrive api.
"""
with tmpdir.as_cwd():
mol = Ligand.from_file(get_data("ethane.sdf"))
bond = mol.find_rotatable_bonds()[0]
dihedral = mol.dihedrals[bond.indices][0]
tdriver = TorsionDriver(starting_conformations=starting_conformations)
# make the scan data
dihedral_data = TorsionScan(torsion=dihedral, scan_range=(-165, 180))
td_state = tdriver._create_initial_state(molecule=mol,
dihedral_data=dihedral_data,
qc_spec=QCOptions())
assert td_state["dihedrals"] == [
dihedral,
]
assert td_state["elements"] == [
atom.atomic_symbol for atom in mol.atoms
]
assert td_state["dihedral_ranges"] == [
(-165, 180),
]
assert np.allclose((mol.coordinates * constants.ANGS_TO_BOHR),
td_state["init_coords"][0])
# make sure we have tried to generate conformers
assert len(td_state["init_coords"]) <= tdriver.starting_conformations
def test_full_tdrive(tmpdir, workers, capsys):
"""
Try and run a full torsiondrive for ethane with a cheap rdkit method.
"""
with tmpdir.as_cwd():
ethane = Ligand.from_file(get_data("ethane.sdf"))
# make the scan data
bond = ethane.find_rotatable_bonds()[0]
dihedral = ethane.dihedrals[bond.indices][0]
dihedral_data = TorsionScan(torsion=dihedral, scan_range=(-165, 180))
qc_spec = QCOptions(program="rdkit", basis=None, method="uff")
local_ops = LocalResource(cores=workers, memory=2)
tdriver = TorsionDriver(
n_workers=workers,
grid_spacing=60,
)
_ = tdriver.run_torsiondrive(
molecule=ethane,
dihedral_data=dihedral_data,
qc_spec=qc_spec,
local_options=local_ops,
)
captured = capsys.readouterr()
# make sure a fresh torsiondrive is run
assert "Starting new torsiondrive" in captured.out
def test_get_new_jobs(ethane_state):
"""
Make sure that for a given initial state we can get the next jobs to be done.
"""
tdriver = TorsionDriver()
new_jobs = tdriver._get_new_jobs(td_state=ethane_state)
assert "-60" in new_jobs
assert new_jobs["-60"][0] == pytest.approx([
-1.44942051524959,
0.015117815022160003,
-0.030235630044320005,
1.44942051524959,
-0.015117815022160003,
0.030235630044320005,
-2.2431058039129903,
-0.18897268777700002,
1.8708296089923,
-2.16562700192442,
1.78768162637042,
-0.82203119182995,
-2.1920831782132,
-1.5325684978714702,
-1.18863820611733,
2.1920831782132,
1.5306787709937002,
1.18863820611733,
2.2431058039129903,
0.18897268777700002,
-1.8708296089923,
2.16562700192442,
-1.78957135324819,
0.82014146495218,
])
def ethane_state() -> Dict[str, Any]:
"""
build an initial state for a ethane scan.
"""
mol = Ligand.from_file(get_data("ethane.sdf"))
bond = mol.find_rotatable_bonds()[0]
dihedral = mol.dihedrals[bond.indices][0]
tdriver = TorsionDriver(grid_spacing=15)
# make the scan data
dihedral_data = TorsionScan(torsion=dihedral, scan_range=(-165, 180))
td_state = tdriver._create_initial_state(molecule=mol,
dihedral_data=dihedral_data)
return td_state
def ethane_state(tmpdir) -> Dict[str, Any]:
"""
build an initial state for a ethane scan.
"""
with tmpdir.as_cwd():
mol = Ligand.from_file(get_data("ethane.sdf"))
bond = mol.find_rotatable_bonds()[0]
dihedral = mol.dihedrals[bond.indices][0]
tdriver = TorsionDriver(grid_spacing=15)
# make the scan data
dihedral_data = TorsionScan(torsion=dihedral, scan_range=(-165, 180))
qc_spec = QCOptions(program="rdkit", basis=None, method="uff")
td_state = tdriver._create_initial_state(molecule=mol,
dihedral_data=dihedral_data,
qc_spec=qc_spec)
return td_state
def test_load_old_state(tmpdir, ethane_state, qc_options, scan_range,
compatible):
"""
Make sure we can load and cross-check torsiondrive state files.
"""
with tmpdir.as_cwd():
# dump the basic ethane result to file
td_api.current_state_json_dump(current_state=ethane_state,
jsonfilename="torsiondrive_state.json")
td = TorsionDriver()
state = td._load_state(
qc_spec=qc_options,
torsion_scan=TorsionScan(ethane_state["dihedrals"][0],
scan_range=scan_range),
)
if compatible:
assert state is not None
else:
assert state is None
def test_optimise_grid_point_and_update(tmpdir, ethane_state):
"""
Try and perform a single grid point optimisation.
"""
with tmpdir.as_cwd():
mol = Ligand.from_file(get_data("ethane.sdf"))
tdriver = TorsionDriver(n_workers=1)
qc_spec = QCOptions(program="rdkit", basis=None, method="uff")
local_ops = LocalResource(cores=1, memory=1)
geo_opt = tdriver._build_geometry_optimiser()
# get the job inputs
new_jobs = tdriver._get_new_jobs(td_state=ethane_state)
coords = new_jobs["-60"][0]
result = optimise_grid_point(
geometry_optimiser=geo_opt,
qc_spec=qc_spec,
local_options=local_ops,
molecule=mol,
coordinates=coords,
dihedral=ethane_state["dihedrals"][0],
dihedral_angle=-60,
job_id=0,
)
new_state = tdriver._update_state(
td_state=ethane_state,
result_data=[
result,
],
)
next_jobs = tdriver._get_new_jobs(td_state=new_state)
assert "-75" in next_jobs
assert "-45" in next_jobs
def test_full_tdrive(tmpdir):
"""
Try and run a full torsiondrive for ethane with a cheap rdkit method.
"""
with tmpdir.as_cwd():
ethane = Ligand.from_file(get_data("ethane.sdf"))
# make the scan data
bond = ethane.find_rotatable_bonds()[0]
dihedral = ethane.dihedrals[bond.indices][0]
dihedral_data = TorsionScan(torsion=dihedral, scan_range=(-165, 180))
tdriver = TorsionDriver(
program="rdkit",
method="uff",
basis=None,
memory=2,
cores=1,
n_workers=1,
grid_spacing=60,
)
_ = tdriver.run_torsiondrive(molecule=ethane,
dihedral_data=dihedral_data)
def test_initial_state_coords_passed(tmpdir):
"""
Make sure any seed conformations are used in the initial state
"""
with tmpdir.as_cwd():
mol = Ligand.from_file(get_data("ethane.sdf"))
bond = mol.find_rotatable_bonds()[0]
dihedral = mol.dihedrals[bond.indices][0]
tdriver = TorsionDriver()
# make the scan data
dihedral_data = TorsionScan(torsion=dihedral, scan_range=(-165, 180))
# make some mock coords
coords = [np.random.random(size=(mol.n_atoms, 3)) for _ in range(4)]
td_state = tdriver._create_initial_state(
molecule=mol,
dihedral_data=dihedral_data,
qc_spec=QCOptions(),
seed_coordinates=coords,
)
assert len(td_state["init_coords"]) == 4
# make sure they are the same random coords
for i in range(4):
assert np.allclose((coords[i] * constants.ANGS_TO_BOHR),
td_state["init_coords"][i])
def test_double_dihedral(tmpdir):
"""Test running a molecule with two rotatable bonds."""
with tmpdir.as_cwd():
mol = Ligand.from_smiles("CCO", "ethanol")
# build a scanner with grid spacing 60 and clear out avoided methyl
tdrive = TorsionDriver(
program="rdkit",
method="uff",
basis=None,
cores=1,
memory=1,
n_workers=1,
grid_spacing=60,
)
t_scan = TorsionScan1D(torsion_driver=tdrive)
t_scan.clear_avoided_torsions()
result_mol = t_scan.run(molecule=mol)
assert len(result_mol.qm_scans) == 2
# make sure input molecule coords were not changed
assert np.allclose(mol.coordinates, result_mol.coordinates)
def test_tdrive_restarts(capsys, ethane_state, tmpdir):
"""
Make sure that an old torsiondrive is continued when possible from the current state file.
"""
with tmpdir.as_cwd():
ethane_state["grid_spacing"] = [
60,
]
mol = Ligand.from_file(get_data("ethane.sdf"))
tdriver = TorsionDriver(n_workers=1, grid_spacing=60)
qc_spec = QCOptions(program="rdkit", basis=None, method="uff")
local_ops = LocalResource(cores=1, memory=1)
geo_opt = tdriver._build_geometry_optimiser()
# get the job inputs
new_jobs = tdriver._get_new_jobs(td_state=ethane_state)
coords = new_jobs["-60"][0]
result = optimise_grid_point(
geometry_optimiser=geo_opt,
qc_spec=qc_spec,
local_options=local_ops,
molecule=mol,
coordinates=coords,
dihedral=ethane_state["dihedrals"][0],
dihedral_angle=-60,
job_id=0,
)
_ = tdriver._update_state(
td_state=ethane_state,
result_data=[
result,
],
)
# now start a run and make sure it continues
_ = tdriver.run_torsiondrive(
molecule=mol,
dihedral_data=TorsionScan(torsion=ethane_state["dihedrals"][0],
scan_range=(-165, 180)),
qc_spec=qc_spec,
local_options=local_ops,
)
capture = capsys.readouterr()
assert ("Compatible TorsionDrive state found restarting torsiondrive!"
in capture.out)
class TorsionScan1D(StageBase):
"""
A 1D torsion scanner.
Note:
By default this will scan all rotatable bonds not involving methyl or amine terminal groups.
"""
special_torsions: List[TargetTorsion] = Field(
[],
description=
"A list of special target torsions to be scanned and their scan range.",
)
default_scan_range: Tuple[int, int] = Field(
(-165, 180),
description=
"The default scan range for any torsions not covered by a special rule.",
)
avoided_torsions: List[AvoidedTorsion] = Field(
[
AvoidedTorsion(smirks="[*:1]-[CH3:2]"),
AvoidedTorsion(smirks="[*:1]-[NH2:2]"),
],
description="The list of torsion patterns that should be avoided.",
)
torsion_driver: TorsionDriver = Field(
TorsionDriver(),
description=
"The torsion drive engine used to compute the reference data.",
)
type: Literal["TorsionScan1D"] = "TorsionScan1D"
def start_message(self, **kwargs) -> str:
return f"Performing QM-constrained optimisation with Torsiondrive and {kwargs['qc_spec'].program}"
def finish_message(self, **kwargs) -> str:
return "Torsiondrive finished and QM results saved."
@classmethod
def is_available(cls) -> bool:
"""
Make sure geometric and torsiondrive are installed.
"""
geo = which_import(
"geometric",
return_bool=True,
raise_error=True,
raise_msg=
"Please install via `conda install geometric -c conda-forge`.",
)
tdrive = which_import(
"torsiondrive",
return_bool=True,
raise_error=True,
raise_msg=
"Please install via `conda install torsiondrive -c conda-forge`.",
)
engine = which_import(
"qcengine",
return_bool=True,
raise_error=True,
raise_msg=
"Please install via `conda install qcengine -c conda-forge`.",
)
return geo and tdrive and engine
def run(self, molecule: "Ligand", **kwargs) -> "Ligand":
"""
Run any possible torsiondrives for this molecule given the list of allowed and disallowed torsion patterns.
Note:
This function just validates the molecule and builds a list of torsions to scan before passing to the main method.
Important:
We work with a copy of the input molecule as we change the coordinates throughout the calculation.
"""
molecule.qm_scans = []
# work with a copy as we change coordinates from the qm a lot!
drive_mol = copy.deepcopy(molecule)
# first find all rotatable bonds, while removing the unwanted scans
bonds = drive_mol.find_rotatable_bonds(smirks_to_remove=[
torsion.smirks for torsion in self.avoided_torsions
])
if bonds is None:
print("No rotatable bonds found to scan!")
return molecule
# remove symmetry duplicates
bonds = self._get_symmetry_unique_bonds(molecule=drive_mol,
bonds=bonds)
torsion_scans = []
for bond in bonds:
# get the scan range and a torsion for the bond
torsion = find_heavy_torsion(molecule=drive_mol, bond=bond)
scan_range = self._get_scan_range(molecule=drive_mol, bond=bond)
torsion_scans.append(
TorsionScan(torsion=torsion, scan_range=scan_range))
result_mol = self._run_torsion_drives(
molecule=drive_mol,
torsion_scans=torsion_scans,
qc_spec=kwargs["qc_spec"],
local_options=kwargs["local_options"],
)
# make sure we preserve the input coords
result_mol.coordinates = copy.deepcopy(molecule.coordinates)
# make sure we have all of the scans we expect
assert len(result_mol.qm_scans) == len(bonds)
return result_mol
def _get_symmetry_unique_bonds(self, molecule: "Ligand",
bonds: List["Bond"]) -> List["Bond"]:
"""
For a list of central torsion bonds deduplicate the list by bond symmetry types.
"""
atom_types = molecule.atom_types
unique_bonds = {}
for bond in bonds:
bond_type = f"{atom_types[bond.atom1_index]}-{atom_types[bond.atom2_index]}"
if bond_type not in unique_bonds and bond_type[::
-1] not in unique_bonds:
unique_bonds[bond_type] = bond
return list(unique_bonds.values())
def _get_scan_range(self, molecule: "Ligand",
bond: "Bond") -> Tuple[int, int]:
"""
Get the scan range for the target bond based on the allowed substructure list.
Note:
We loop over the list of targets checking each meaning that the last match in the list will be applied to substructure.
So generic matches should be placed at the start of the list with more specific ones at the end.
"""
scan_range = self.default_scan_range
for target_torsion in self.special_torsions:
matches = molecule.get_smarts_matches(smirks=target_torsion.smirks)
for match in matches:
if len(match) == 4:
atoms = match[1:3]
else:
atoms = match
if set(atoms) == set(bond.indices):
scan_range = target_torsion.scan_range
return scan_range
def _run_torsion_drives(
self,
molecule: "Ligand",
torsion_scans: List[TorsionScan],
qc_spec: QCOptions,
local_options: LocalResource,
) -> "Ligand":
"""
Run the list of validated torsion drives.
Note:
We do not change the initial coordinates passed at this point.
Args:
molecule: The molecule to be scanned.
torsion_scans: A list of TorsionScan jobs to perform detailing the dihedral and the scan range.
Returns:
The updated molecule object with the scan results.
"""
for torsion_scan in torsion_scans:
# make a folder and move into to run the calculation
folder = "SCAN_"
folder += "_".join([str(t) for t in torsion_scan.torsion])
with folder_setup(folder):
print(
f"Running scan for dihedral: {torsion_scan.torsion} with range: {torsion_scan.scan_range}"
)
result_mol = self.torsion_driver.run_torsiondrive(
molecule=molecule,
dihedral_data=torsion_scan,
qc_spec=qc_spec,
local_options=local_options,
)
return result_mol
def add_special_torsion(
self, smirks: str, scan_range: Tuple[int, int] = (-165, 180)) -> None:
"""
Add a new allowed torsion scan.
Args:
smirks: The smirks pattern that should be used to identify the torsion.
scan_range: The scan range for this type of dihedral.
"""
target = TargetTorsion(smirks=smirks, scan_range=scan_range)
self.special_torsions.append(target)
def clear_special_torsions(self) -> None:
"""Remove all allowed torsion scans."""
self.special_torsions = []
def add_avoided_torsion(self, smirks: str) -> None:
"""
Add a new torsion pattern to avoid scanning.
Args:
smirks: The valid smirks pattern which describes a torsion not to be scanned.
"""
torsion = AvoidedTorsion(smirks=smirks)
self.avoided_torsions.append(torsion)
def clear_avoided_torsions(self) -> None:
"""Remove all avoided torsions."""
self.avoided_torsions = []