diff --git a/fegrow/cli/cli.py b/fegrow/cli/cli.py index 359da357..447d833c 100644 --- a/fegrow/cli/cli.py +++ b/fegrow/cli/cli.py @@ -3,6 +3,7 @@ import click from fegrow.cli.scoring import score from fegrow.cli.utils import Settings +from fegrow.cli.report import report @click.group() @@ -11,13 +12,14 @@ def cli(): cli.add_command(score) +cli.add_command(report) @cli.command() @click.option( "-g", "--gnina-path", - type=click.Path(exists=True, file_okay=True, dir_okay=False, executable=True), + type=click.Path(exists=True, file_okay=True, dir_okay=False, executable=True, resolve_path=True), help="The path to the gnina executable which will override the settings.", ) def settings(gnina_path: pathlib.Path): @@ -25,6 +27,6 @@ def settings(gnina_path: pathlib.Path): Create a runtime settings object for scoring runs which can be configured. """ config = Settings(gnina_path=gnina_path) - with open("settings.json") as output: + with open("settings.json", "w") as output: output.write(config.json(indent=2)) diff --git a/fegrow/cli/report.py b/fegrow/cli/report.py new file mode 100644 index 00000000..6e47f642 --- /dev/null +++ b/fegrow/cli/report.py @@ -0,0 +1,105 @@ +import click +import pathlib +from typing import Optional + + +@click.command() +@click.option( + '-l', + '--ligands', + help='The sdf file which contains the scored ligands to create the report for.', + type=click.Path(file_okay=True, dir_okay=False, path_type=pathlib.Path, exists=True), + multiple=True +) +@click.option( + "-o", + "--output", + type=click.Path(dir_okay=True, path_type=pathlib.Path), + help="The name of the output folder.", +) +@click.option( + '-t', + '--title', + type=click.STRING, + help='The title of the dataset which will be used for the report and sdf file names' +) +@click.option( + '-tn', + '--top-n', + help='The top N molecules to keep after sorting by predicted affinity', + type=click.INT +) +def report( + ligands: list[pathlib.Path], + output: pathlib.Path, + title: str, + top_n: Optional[int] = None +): + """ + Generate an interactive HTML report of the molecules from the scored SDF files and an SDF file ordered by the top + performing molecules. Use top_n to restrict the report and sdf file to only include n molecules. + """ + from rdkit import Chem + from fegrow.cli.utils import draw_mol, has_bad_geometry + import pandas as pd + import tqdm + import panel + import bokeh.models.widgets.tables + + output.mkdir(exist_ok=True, parents=True) + # load all best scoring ligands + molecules_and_affinities = [] + for file in ligands: + # load the low energy ligands and the predicted affinity + supplier = Chem.SDMolSupplier(file.as_posix(), removeHs=True) + for molecule in supplier: + if not has_bad_geometry(molecule): + continue + try: + molecules_and_affinities.append((molecule, float(molecule.GetProp('cnnaffinity')))) + except KeyError: + # catch cases where the sd tags were not writen + continue + + click.echo(f'Total number of molecules and predictions: {len(molecules_and_affinities)}') + + # sort all molecules by predicted affinity + ranked_all_mols = sorted(molecules_and_affinities, key=lambda x: x[1]) + # write out all molecules to a single sdf sorted by affinity + output_file = output.joinpath(f'{title}.sdf') + with Chem.SDWriter(output_file.as_posix()) as sdf_out: + rows = [] + for mol, affinity in tqdm.tqdm(ranked_all_mols[:top_n], desc="Creating report...", ncols=80): + try: + smiles = mol.GetProp('smiles') + except KeyError: + smiles = Chem.MolToSmiles(mol) + mol.SetProp('smiles', smiles) + rows.append( + { + "Smiles": smiles, + "Molecule": draw_mol(smiles), + "CNNaffinity": affinity, + "IC50 (nM)": mol.GetProp('cnnaffinityIC50').split()[0], + } + ) + sdf_out.write(mol) + # create the report + df = pd.DataFrame(rows) + + number_format = bokeh.models.widgets.tables.NumberFormatter(format="0.0000") + layout = panel.Column( + panel.widgets.Tabulator( + df, + show_index=False, + selectable=False, + disabled=True, + formatters={"Smiles": "html", "Molecule": "html", "CNNaffinity": number_format, + "IC50 (nM)": number_format}, + configuration={"rowHeight": 400}, + ), + sizing_mode="stretch_width", + scroll=True + ) + + layout.save(output.joinpath(f'{title}.html').as_posix(), title=title, embed=True) diff --git a/fegrow/cli/scoring.py b/fegrow/cli/scoring.py index 70b6879e..9ea64ea3 100644 --- a/fegrow/cli/scoring.py +++ b/fegrow/cli/scoring.py @@ -1,5 +1,4 @@ import pathlib -import time import click from typing import Optional @@ -9,38 +8,56 @@ @click.option( "-c", "--core-ligand", - type=click.Path(exists=True, file_okay=True, dir_okay=False), + type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=pathlib.Path), help="The path to the SDF file of the core ligand which will be used to restrain the geometries of the scored ligands this should be a substructure of the ligands to be scored.", ) @click.option( - "-l" "--ligands", - type=click.Path(exists=True, file_okay=True, dir_okay=False), + "-l", + "--ligands", + type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=pathlib.Path), help="The path to the ligands to be scored can be in any supported format by RDKit such as csv. smiles or SDF.", ) @click.option( "-r", "--receptor", - type=click.Path(exists=True, file_okay=True, dir_okay=False), + type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=pathlib.Path), help="The path of the receptor PDB file, this should only contain the receptor and not the reference ligand.", ) @click.option( "-s", "--settings", - type=click.Path(exists=True, file_okay=True, dir_okay=False), + type=click.Path(exists=True, file_okay=True, dir_okay=False, path_type=pathlib.Path), help="The path of the settings file which configures the scoring run.", ) @click.option( "-o", "--output", - type=click.Path(dir_okay=True), + type=click.Path(dir_okay=True, path_type=pathlib.Path), help="The name of the output folder.", ) @click.option( "-g", "--gnina-path", - type=click.Path(exists=True, file_okay=True, dir_okay=False, executable=True), + type=click.Path(exists=True, file_okay=True, dir_okay=False, executable=True, resolve_path=True, path_type=pathlib.Path), help="The path to the gnina executable which will override the settings.", ) +@click.option( + '-t', + '--threads', + default=1, + show_default=True, + type=click.INT, + help='The number of threads per worker.' +) +@click.option( + '-w', + '--workers', + default=4, + show_default=True, + type=click.INT, + help='The number of workers to use.' + +) def score( core_ligand: pathlib.Path, ligands: pathlib.Path, @@ -48,86 +65,100 @@ def score( output: pathlib.Path, settings: Optional[pathlib.Path] = None, gnina_path: Optional[pathlib.Path] = None, + threads: int = 1, + workers: int = 4, ): """ - Score the list of input ligands using Gnina after optimising in the receptor. + Score the list of input ligands using Gnina after optimising in the receptor. Tasks are distributed using Dask. """ from dask.distributed import Client from rdkit import Chem import traceback import dask import tqdm - + import warnings + import os + import logging from fegrow.cli.utils import score_ligand, Settings, load_target_ligands - - try: - from mycluster import create_cluster - except ImportError: - - def create_cluster(): - from dask.distributed import LocalCluster - - return LocalCluster() - - client = Client(create_cluster()) - # create the cluster - click.echo(f"Client created {client}") - - if settings is not None: - config = Settings.parse_file(settings) - if gnina_path is not None: - config.gnina_path = gnina_path - else: - # build the base settings object this needs the gnina path - config = Settings(gnina_path=gnina_path) - - click.echo(f"Loading core ligand from {core_ligand}") - # we remove all Hs rather than specific ones at the attachment point - core = Chem.MolFromMolFile(core_ligand, removeHs=True) - core_dask = dask.delayed(core) - - # load the target ligands - target_smiles = load_target_ligands(ligand_file=ligands) - - # build a list of tasks to submit - for_submission = [ - score_ligand( - core_ligand=core_dask, - target_smiles=smiles, - receptor=receptor, - settings=config, - ) - for smiles in target_smiles - ] - - submitted = client.compute(for_submission) - jobs = dict((job, i) for i, job in enumerate(submitted)) - - output_path = pathlib.Path(output) - output_path.mkdir(exist_ok=True) - - molecule_output = Chem.SDWriter(output_path.joinpath("scored_molecules.sdf")) - with tqdm.tqdm(total=len(submitted), desc="Scoring molecules...", ncols=80) as pbar: - while len(jobs) > 0: - for job, index in jobs.items(): - if not job.done(): - continue - - # remove the job - del jobs[job] - pbar.update(1) - - try: - mol_data = job.result() - rmol = mol_data.pop("molecule") - # recover the properties (they are not passed with serialisation) - [rmol.SetProp(k, str(v)) for k, v in mol_data.items()] - # write the molecule out when they complete incase we crash - molecule_output.write(rmol) - except Exception as E: - print("error for index, ", index) - traceback.print_exc() - - time.sleep(5) - - click.echo("All molecules scored") + from fegrow.package import RMol + from dask.distributed import LocalCluster + import time + + + # hide warnings and logs from openff-toolkit + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + os.environ["PYTHONWARNINGS"] = "ignore" + logging.getLogger('fegrow').setLevel(logging.ERROR) + + client = Client(LocalCluster(threads_per_worker=threads, n_workers=workers)) + # create the cluster + click.echo(f"Client created {client}") + + if settings is not None: + config = Settings.parse_file(settings) + if gnina_path is not None: + config.gnina_path = gnina_path.as_posix() + else: + # build the base settings object this needs the gnina path + config = Settings(gnina_path=gnina_path.as_posix()) + # set the gnina directory + RMol.set_gnina(loc=config.gnina_path) + click.echo(f'Setting Gnina path to: {config.gnina_path}') + + click.echo(f"Loading core ligand from {core_ligand}") + # we remove all Hs rather than specific ones at the attachment point + core = Chem.MolFromMolFile(core_ligand.as_posix(), removeHs=True) + core_dask = dask.delayed(core) + + # load the target ligands + target_smiles = load_target_ligands(ligand_file=ligands) + + # build a list of tasks to submit + for_submission = [ + score_ligand( + core_ligand=core_dask, + target_smiles=smiles, + receptor=receptor, + settings=config, + ) + for smiles in target_smiles + ] + + submitted = client.compute(for_submission) + + output.mkdir(exist_ok=True) + + molecule_output = Chem.SDWriter(output.joinpath("scored_molecules.sdf").as_posix()) + with tqdm.tqdm(total=len(submitted), desc="Scoring molecules...", ncols=80) as pbar: + while len(submitted) > 0: + to_remove = [] + for job in submitted: + if not job.done(): + continue + + # remove the job + to_remove.append(job) + pbar.update(1) + + try: + mol_data = job.result() + # if the molecule could not be scored as it had no conformers do not save the geometry + if mol_data['cnnaffinity'] == 0: + continue + + rmol = mol_data.pop("molecule") + # recover the properties (they are not passed with serialisation) + [rmol.SetProp(k, str(v)) for k, v in mol_data.items()] + # write the molecule out when they complete incase we crash + molecule_output.write(rmol) + except Exception as E: + print("error for molecule") + traceback.print_exc() + + # remove jobs before next round + for job in to_remove: + submitted.remove(job) + time.sleep(5) + + click.echo("All molecules scored") diff --git a/fegrow/cli/utils.py b/fegrow/cli/utils.py index 6e2eb62d..b2e60b85 100644 --- a/fegrow/cli/utils.py +++ b/fegrow/cli/utils.py @@ -6,7 +6,7 @@ from fegrow import RMol import pathlib import pandas as pd -from rdkit.Chem import PandasTools +import warnings class Settings(BaseModel): @@ -67,51 +67,57 @@ def score_ligand( Note: We assume the core does not need to be altered and is a substructure of the target ligand """ - # create the target ligand with Hs - candidate_mol = Chem.MolFromSmiles(target_smiles) - candidate_mol = Chem.AddHs(candidate_mol) - rmol = RMol(candidate_mol) - # set up the core as the template - rmol._save_template(core_ligand) - - # conformer gen - rmol.generate_conformers( - num_conf=settings.num_confs, minimum_conf_rms=settings.conf_rms - ) - # remove missing - rmol.remove_clashing_confs(protein=receptor.as_posix()) - - # optimise - rmol.optimise_in_receptor( - receptor_file=receptor, - ligand_force_field=settings.ligand_force_field, - use_ani=settings.use_ani, - sigma_scale_factor=settings.sigma_scale_factor, - relative_permittivity=settings.relative_permittivity, - water_model=settings.water_model, - platform_name=settings.platform_name, - ) - - if rmol.GetNumConformers() == 0: - # set a pentalty - cnnaffinity = 0 - cnnaffinityIC50 = 0 - else: - # score only the lowest energy conformer - rmol.sort_conformers(energy_range=settings.energy_filter) # kcal/mol - # purge all but the lowest energy conformers - rmol = Rmol(rmol, confId=0) - affinities = rmol.gnina(receptor_file=receptor.as_posix()) - cnnaffinity = -affinities.CNNaffinity.values[0] - cnnaffinityIC50 = affinities["CNNaffinity->IC50s"].values[0] - - data = { - "cnnaffinity": cnnaffinity, - "cnnaffinityIC50": cnnaffinityIC50, - "molecule": rmol, - } - - return data + # hide the openff warnings as we catch errors later + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + + # create the target ligand with Hs + candidate_mol = Chem.MolFromSmiles(target_smiles) + candidate_mol = Chem.AddHs(candidate_mol) + rmol = RMol(candidate_mol) + # set up the core as the template + rmol._save_template(core_ligand) + + # conformer gen + rmol.generate_conformers( + num_conf=settings.num_confs, minimum_conf_rms=settings.conf_rms + ) + # remove missing + rmol.remove_clashing_confs(protein=receptor.as_posix()) + + # optimise + rmol.optimise_in_receptor( + receptor_file=receptor.as_posix(), + ligand_force_field=settings.ligand_force_field, + use_ani=settings.use_ani, + sigma_scale_factor=settings.sigma_scale_factor, + relative_permittivity=settings.relative_permittivity, + water_model=settings.water_model, + platform_name=settings.platform_name, + show_progress=False + ) + + if rmol.GetNumConformers() == 0: + # set a pentalty + cnnaffinity = 0 + cnnaffinityIC50 = 0 + else: + # score only the lowest energy conformer + rmol.sort_conformers(energy_range=settings.energy_filter) # kcal/mol + # purge all but the lowest energy conformers + rmol = RMol(rmol, confId=0) + affinities = rmol.gnina(receptor_file=receptor.as_posix()) + cnnaffinity = -affinities.CNNaffinity.values[0] + cnnaffinityIC50 = affinities["Kd"].values[0] + + data = { + "cnnaffinity": cnnaffinity, + "cnnaffinityIC50": cnnaffinityIC50, + "molecule": rmol, + 'smiles': target_smiles + } + + return data def load_target_ligands(ligand_file: pathlib.Path) -> list[str]: @@ -121,16 +127,55 @@ def load_target_ligands(ligand_file: pathlib.Path) -> list[str]: Note: For CSV we assume that the smiles have the column name "Smiles" """ - if ligand_file.stem.lower == "csv": - target_molecules = pd.read_csv(ligand_file) - return list(target_molecules.Smiles.values) - - if ligand_file.stem.lower() in ["sdf", "mol"]: - ligands = list(Chem.SDMolSupplier(ligand_file, removeHs=False)) - elif ligand_file.stem.lower() == "smi": - ligands = list(Chem.SmilesMolSupplier(ligand_file, remoeHs=False)) + if ligand_file.suffix.lower() == ".csv": + target_molecules = pd.read_csv(ligand_file.as_posix()) + return target_molecules.Smiles.values.tolist() + + if ligand_file.suffix.lower() in [".sdf", ".mol"]: + ligands = list(Chem.SDMolSupplier(ligand_file.as_posix(), removeHs=False)) + elif ligand_file.suffix.lower() == ".smi": + ligands = list(Chem.SmilesMolSupplier(ligand_file.as_posix(), remoeHs=False)) else: - raise RuntimeError(f"Can extract smiles from input file {ligand_file}") + raise RuntimeError(f"Can not extract smiles from input file {ligand_file}") smiles = [Chem.MolToSmiles(mol) for mol in ligands] return smiles + + +def draw_mol(smiles: str) -> str: + import base64 + from rdkit.Chem import Draw + + rdkit_mol = Chem.RemoveHs(Chem.MolFromSmiles(smiles)) + rdkit_mol = Draw.PrepareMolForDrawing(rdkit_mol, forceCoords=True) + drawer = Draw.rdMolDraw2D.MolDraw2DSVG(400, 200) + drawer.DrawMolecule(rdkit_mol) + drawer.FinishDrawing() + + data = base64.b64encode(drawer.GetDrawingText().encode()).decode() + return f'' + + +def has_bad_geometry(molecule: Chem.Mol) -> bool: + """ + Try to detect if a molecule has as bad geometry from the ML/MM optimisation. + We find that rings buckle and protons can dissociate causing inaccurate geometries. + + Notes: + We try to detect the bonding pattern using qcelemental based on the overlap of atomic radii. + """ + from qcelemental.molutil import guess_connectivity + from qcelemental.exceptions import ValidationError + from qcelemental import constants + from qcelemental import periodictable + # convert the rdkit molecule to qcelemental + try: + symbols = [periodictable.to_symbol(atom.GetAtomicNum()) for atom in molecule.GetAtoms()] + geometry = molecule.GetConformer().GetPositions() * constants.conversion_factor('angstrom', 'bohr') + # guess the connectivity from the geometry + connectivity_guess = {tuple(sorted([a + 1, b + 1])) for a, b in guess_connectivity(symbols=symbols, geometry=geometry)} + expected_connectivity = {tuple(sorted([bond.GetBeginAtomIdx() + 1, bond.GetEndAtomIdx() + 1])) for bond in molecule.GetBonds()} + return expected_connectivity == connectivity_guess + except ValidationError: + # some atoms get too close and raise a validation error + return False diff --git a/fegrow/conformers.py b/fegrow/conformers.py index 15594922..b75c3f79 100644 --- a/fegrow/conformers.py +++ b/fegrow/conformers.py @@ -94,7 +94,7 @@ def to_ties_atoms(rdkit_mol): coordinates_map[matched_index.id] = core_atom_coordinate manmap.append((matched_index.id, core_index.id)) # - print("Used the TIES (Bieniek et al) package to get the mapping") + logger.info("Used the TIES (Bieniek et al) package to get the mapping") # use a reproducible random seed randomseed = 194715 @@ -121,7 +121,7 @@ def to_ties_atoms(rdkit_mol): if dup_count: logger.info(f"Removed {dup_count} duplicated conformations, leaving {rmol.GetNumConformers()} in total. ") - print(f"Generated {rmol.GetNumConformers()} conformers. ") + logger.info(f"Generated {rmol.GetNumConformers()} conformers. ") return rmol diff --git a/fegrow/package.py b/fegrow/package.py index 83ab232f..fc2e7ade 100644 --- a/fegrow/package.py +++ b/fegrow/package.py @@ -150,7 +150,7 @@ def optimise_in_receptor(self, *args, **kwargs): :param water_model: can be used to set the force field for any water molecules present in the binding site. """ if self.GetNumConformers() == 0: - print("Warning: no conformers so cannot optimise_in_receptor. Ignoring.") + logger.info("No conformers so cannot optimise_in_receptor. Ignoring.") return from .receptor import optimise_in_receptor @@ -187,7 +187,7 @@ def sort_conformers(self, energy_range=5): above the minimum, for which conformers should be kept. """ if self.GetNumConformers() == 0: - print("An rmol doesn't have any conformers. Ignoring.") + logger.info("An rmol doesn't have any conformers. Ignoring.") return None elif self.opt_energies is None: raise AttributeError( @@ -323,7 +323,7 @@ def remove_clashing_confs(self, ) rm_counter += 1 break - print(f"Removed {rm_counter} conformers. ") + logger.info(f"Removed {rm_counter} conformers. ") @staticmethod @@ -366,7 +366,7 @@ def _check_download_gnina(): pass # gnina is not found, try downloading it - print(f"Gnina not found or set. Download gnina (~500MB) into {RMol.gnina_dir}") + logger.info(f"Gnina not found or set. Download gnina (~500MB) into {RMol.gnina_dir}") gnina = os.path.join(RMol.gnina_dir, "gnina") # fixme - currently download to the working directory (Home could be more applicable). urllib.request.urlretrieve( @@ -551,7 +551,7 @@ def generate_conformers( self, num_conf: int, minimum_conf_rms: Optional[float] = [], **kwargs ): for i, rmol in enumerate(self): - print(f"RMol index {i}") + logger.info(f"RMol index {i}") rmol.generate_conformers(num_conf, minimum_conf_rms, **kwargs) def GetNumConformers(self): @@ -559,7 +559,7 @@ def GetNumConformers(self): def remove_clashing_confs(self, prot, min_dst_allowed=1): for i, rmol in enumerate(self): - print(f"RMol index {i}") + logger.info(f"RMol index {i}") rmol.remove_clashing_confs(prot, min_dst_allowed=min_dst_allowed) def optimise_in_receptor(self, *args, **kwargs): @@ -570,7 +570,7 @@ def optimise_in_receptor(self, *args, **kwargs): dfs = [] for i, rmol in enumerate(self): - print(f"RMol index {i}") + logger.info(f"RMol index {i}") dfs.append(rmol.optimise_in_receptor(*args, **kwargs)) df = pandas.concat(dfs) @@ -580,7 +580,7 @@ def optimise_in_receptor(self, *args, **kwargs): def sort_conformers(self, energy_range=5): dfs = [] for i, rmol in enumerate(self): - print(f"RMol index {i}") + logger.info(f"RMol index {i}") dfs.append(rmol.sort_conformers(energy_range)) df = pandas.concat(dfs) @@ -590,7 +590,7 @@ def sort_conformers(self, energy_range=5): def gnina(self, receptor_file): dfs = [] for i, rmol in enumerate(self): - print(f"RMol index {i}") + logger.info(f"RMol index {i}") dfs.append(rmol.gnina(receptor_file)) df = pandas.concat(dfs) @@ -605,7 +605,7 @@ def discard_missing(self): for rmol in self[::-1]: if rmol.GetNumConformers() == 0: rmindex = self.index(rmol) - print( + logger.info( f"Discarding a molecule (id {rmindex}) due to the lack of conformers. " ) self.remove(rmol) diff --git a/fegrow/receptor.py b/fegrow/receptor.py index 9697bd15..2fd26802 100644 --- a/fegrow/receptor.py +++ b/fegrow/receptor.py @@ -1,8 +1,8 @@ import logging +import pathlib import tempfile from copy import deepcopy from typing import List, Tuple, Union -import warnings import parmed from openmmforcefields.generators import SystemGenerator @@ -18,11 +18,8 @@ import numpy # fix for new openmm versions -try: - from openmm import app, openmm, unit, Platform -except (ImportError, ModuleNotFoundError): - from simtk.openmm import app, openmm - from simtk import unit +from openmm import app, openmm, unit, Platform + from openff.toolkit.topology import Molecule as OFFMolecule @@ -92,7 +89,8 @@ def optimise_in_receptor( sigma_scale_factor: float = 0.8, relative_permittivity: float = 4, water_model: str = "tip3p.xml", - platform_name: str = "CPU" + platform_name: str = "CPU", + show_progress: bool = True ) -> Tuple[RMol, List[float]]: """ For each of the input molecule conformers optimise the system using the chosen force field with the receptor held fixed. @@ -116,6 +114,8 @@ def optimise_in_receptor( platform_name: The OpenMM platform name, 'cuda' if available, with the 'cpu' used by default. See the OpenMM documentation of Platform. + show_progress: + If a progress bar should be displayed. Returns: A copy of the input molecule with the optimised positions. @@ -169,17 +169,18 @@ def optimise_in_receptor( system.setParticleMass(i, 0) # if we want to use ani2x check we can and adapt the system if use_ani and _can_use_ani2x(openff_mol): - print("using ani2x") + logger.info("using ani2x") + # use the torchani standard NNpops should not make much difference here potential = MLPotential("ani2x", platform_name=platform_name) # save the torch model animodel.pt to a temporary file to ensure this is thread safe _, tmpfile = tempfile.mkstemp() - + logger.info('writing ani jit model to ', tmpfile) complex_system = potential.createMixedSystem( - complex_structure.topology, system, ligand_idx, filename=tmpfile + complex_structure.topology, system, ligand_idx, filename=tmpfile, implementation='torchani' ) else: - print("Using force field") + logger.info("Using force field") complex_system = system # scale the charges and sigma values _scale_system( @@ -209,7 +210,7 @@ def optimise_in_receptor( final_mol.RemoveAllConformers() energies = [] for i, conformer in enumerate( - tqdm(openff_mol.conformers, desc="Optimising conformer: ", ncols=80) + tqdm(openff_mol.conformers, desc="Optimising conformer: ", ncols=80, disable=not show_progress) ): # make the ligand coords lig_vec = unit.Quantity([c.m.tolist() for c in conformer], unit=unit.angstrom) @@ -245,6 +246,13 @@ def optimise_in_receptor( ) final_mol.AddConformer(final_conformer, assignId=True) + # clean up the ani tmp file if used + if use_ani and _can_use_ani2x(openff_mol): + logger.info('cleaning up ani jit model ', tmpfile) + tmp_file = pathlib.Path(tmpfile) + if tmp_file.exists(): + tmp_file.unlink() + return final_mol, energies