import copy import os import torch import time from argparse import ArgumentParser, Namespace, FileType from rdkit.Chem import RemoveHs from functools import partial import numpy as np import pandas as pd from rdkit import RDLogger from rdkit.Chem import MolFromSmiles, AddHs from torch_geometric.loader import DataLoader from datasets.process_mols import read_molecule, generate_conformer, write_mol_with_coords from datasets.pdbbind import PDBBind from utils.diffusion_utils import t_to_sigma as t_to_sigma_compl, get_t_schedule from utils.sampling import randomize_position, sampling from utils.utils import get_model from utils.visualise import PDBFile from tqdm import tqdm RDLogger.DisableLog('rdApp.*') import yaml parser = ArgumentParser() parser.add_argument('--config', type=FileType(mode='r'), default=None) parser.add_argument('--protein_ligand_csv', type=str, default=None, help='Path to a .csv file specifying the input as described in the README. If this is not None, it will be used instead of the --protein_path and --ligand parameters') parser.add_argument('--protein_path', type=str, default='data/dummy_data/1a0q_protein.pdb', help='Path to the protein .pdb file') parser.add_argument('--ligand', type=str, default='COc(cc1)ccc1C#N', help='Either a SMILES string or the path to a molecule file that rdkit can read') parser.add_argument('--out_dir', type=str, default='results/user_inference', help='Directory where the outputs will be written to') parser.add_argument('--esm_embeddings_path', type=str, default='data/esm2_output', help='If this is set then the LM embeddings at that path will be used for the receptor features') parser.add_argument('--save_visualisation', action='store_true', default=False, help='Save a pdb file with all of the steps of the reverse diffusion') parser.add_argument('--samples_per_complex', type=int, default=10, help='Number of samples to generate') parser.add_argument('--model_dir', type=str, default='workdir/paper_score_model', help='Path to folder with trained score model and hyperparameters') parser.add_argument('--ckpt', type=str, default='best_ema_inference_epoch_model.pt', help='Checkpoint to use for the score model') parser.add_argument('--confidence_model_dir', type=str, default='workdir/paper_confidence_model', help='Path to folder with trained confidence model and hyperparameters') parser.add_argument('--confidence_ckpt', type=str, default='best_model_epoch75.pt', help='Checkpoint to use for the confidence model') parser.add_argument('--batch_size', type=int, default=32, help='') parser.add_argument('--cache_path', type=str, default='data/cache', help='Folder from where to load/restore cached dataset') parser.add_argument('--no_random', action='store_true', default=False, help='Use no randomness in reverse diffusion') parser.add_argument('--no_final_step_noise', action='store_true', default=False, help='Use no noise in the final step of the reverse diffusion') parser.add_argument('--ode', action='store_true', default=False, help='Use ODE formulation for inference') parser.add_argument('--inference_steps', type=int, default=20, help='Number of denoising steps') parser.add_argument('--num_workers', type=int, default=1, help='Number of workers for creating the dataset') parser.add_argument('--sigma_schedule', type=str, default='expbeta', help='') parser.add_argument('--actual_steps', type=int, default=None, help='Number of denoising steps that are actually performed') parser.add_argument('--keep_local_structures', action='store_true', default=False, help='Keeps the local structure when specifying an input with 3D coordinates instead of generating them with RDKit') args = parser.parse_args() if args.config: config_dict = yaml.load(args.config, Loader=yaml.FullLoader) arg_dict = args.__dict__ for key, value in config_dict.items(): if isinstance(value, list): for v in value: arg_dict[key].append(v) else: arg_dict[key] = value os.makedirs(args.out_dir, exist_ok=True) with open(f'{args.model_dir}/model_parameters.yml') as f: score_model_args = Namespace(**yaml.full_load(f)) if args.confidence_model_dir is not None: with open(f'{args.confidence_model_dir}/model_parameters.yml') as f: confidence_args = Namespace(**yaml.full_load(f)) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if args.protein_ligand_csv is not None: df = pd.read_csv(args.protein_ligand_csv) protein_path_list = df['protein_path'].tolist() ligand_descriptions = df['ligand'].tolist() else: protein_path_list = [args.protein_path] ligand_descriptions = [args.ligand] test_dataset = PDBBind(transform=None, root='', protein_path_list=protein_path_list, ligand_descriptions=ligand_descriptions, receptor_radius=score_model_args.receptor_radius, cache_path=args.cache_path, remove_hs=score_model_args.remove_hs, max_lig_size=None, c_alpha_max_neighbors=score_model_args.c_alpha_max_neighbors, matching=False, keep_original=False, popsize=score_model_args.matching_popsize, maxiter=score_model_args.matching_maxiter, all_atoms=score_model_args.all_atoms, atom_radius=score_model_args.atom_radius, atom_max_neighbors=score_model_args.atom_max_neighbors, esm_embeddings_path= args.esm_embeddings_path if score_model_args.esm_embeddings_path is not None else None, require_ligand=True, num_workers=args.num_workers, keep_local_structures=args.keep_local_structures) test_loader = DataLoader(dataset=test_dataset, batch_size=1, shuffle=False) if args.confidence_model_dir is not None: if not (confidence_args.use_original_model_cache or confidence_args.transfer_weights): # if the confidence model uses the same type of data as the original model then we do not need this dataset and can just use the complexes print('HAPPENING | confidence model uses different type of graphs than the score model. Loading (or creating if not existing) the data for the confidence model now.') confidence_test_dataset = PDBBind(transform=None, root='', protein_path_list=protein_path_list, ligand_descriptions=ligand_descriptions, receptor_radius=confidence_args.receptor_radius, cache_path=args.cache_path, remove_hs=confidence_args.remove_hs, max_lig_size=None, c_alpha_max_neighbors=confidence_args.c_alpha_max_neighbors, matching=False, keep_original=False, popsize=confidence_args.matching_popsize, maxiter=confidence_args.matching_maxiter, all_atoms=confidence_args.all_atoms, atom_radius=confidence_args.atom_radius, atom_max_neighbors=confidence_args.atom_max_neighbors, esm_embeddings_path=args.esm_embeddings_path if confidence_args.esm_embeddings_path is not None else None, require_ligand=True, num_workers=args.num_workers) confidence_complex_dict = {d.name: d for d in confidence_test_dataset} t_to_sigma = partial(t_to_sigma_compl, args=score_model_args) model = get_model(score_model_args, device, t_to_sigma=t_to_sigma, no_parallel=True) state_dict = torch.load(f'{args.model_dir}/{args.ckpt}', map_location=torch.device('cpu')) model.load_state_dict(state_dict, strict=True) model = model.to(device) model.eval() if args.confidence_model_dir is not None: if confidence_args.transfer_weights: with open(f'{confidence_args.original_model_dir}/model_parameters.yml') as f: confidence_model_args = Namespace(**yaml.full_load(f)) else: confidence_model_args = confidence_args confidence_model = get_model(confidence_model_args, device, t_to_sigma=t_to_sigma, no_parallel=True, confidence_mode=True) state_dict = torch.load(f'{args.confidence_model_dir}/{args.confidence_ckpt}', map_location=torch.device('cpu')) confidence_model.load_state_dict(state_dict, strict=True) confidence_model = confidence_model.to(device) confidence_model.eval() else: confidence_model = None confidence_args = None confidence_model_args = None tr_schedule = get_t_schedule(inference_steps=args.inference_steps) rot_schedule = tr_schedule tor_schedule = tr_schedule print('common t schedule', tr_schedule) failures, skipped, confidences_list, names_list, run_times, min_self_distances_list = 0, 0, [], [], [], [] N = args.samples_per_complex print('Size of test dataset: ', len(test_dataset)) for idx, orig_complex_graph in tqdm(enumerate(test_loader)): if confidence_model is not None and not (confidence_args.use_original_model_cache or confidence_args.transfer_weights) and orig_complex_graph.name[0] not in confidence_complex_dict.keys(): skipped += 1 print(f"HAPPENING | The confidence dataset did not contain {orig_complex_graph.name[0]}. We are skipping this complex.") continue try: data_list = [copy.deepcopy(orig_complex_graph) for _ in range(N)] randomize_position(data_list, score_model_args.no_torsion, args.no_random,score_model_args.tr_sigma_max) pdb = None lig = orig_complex_graph.mol[0] if args.save_visualisation: visualization_list = [] for graph in data_list: pdb = PDBFile(lig) pdb.add(lig, 0, 0) pdb.add((orig_complex_graph['ligand'].pos + orig_complex_graph.original_center).detach().cpu(), 1, 0) pdb.add((graph['ligand'].pos + graph.original_center).detach().cpu(), part=1, order=1) visualization_list.append(pdb) else: visualization_list = None start_time = time.time() if confidence_model is not None and not (confidence_args.use_original_model_cache or confidence_args.transfer_weights): confidence_data_list = [copy.deepcopy(confidence_complex_dict[orig_complex_graph.name[0]]) for _ in range(N)] else: confidence_data_list = None data_list, confidence = sampling(data_list=data_list, model=model, inference_steps=args.actual_steps if args.actual_steps is not None else args.inference_steps, tr_schedule=tr_schedule, rot_schedule=rot_schedule, tor_schedule=tor_schedule, device=device, t_to_sigma=t_to_sigma, model_args=score_model_args, no_random=args.no_random, ode=args.ode, visualization_list=visualization_list, confidence_model=confidence_model, confidence_data_list=confidence_data_list, confidence_model_args=confidence_model_args, batch_size=args.batch_size, no_final_step_noise=args.no_final_step_noise) ligand_pos = np.asarray([complex_graph['ligand'].pos.cpu().numpy() + orig_complex_graph.original_center.cpu().numpy() for complex_graph in data_list]) run_times.append(time.time() - start_time) if confidence is not None and isinstance(confidence_args.rmsd_classification_cutoff, list): confidence = confidence[:,0] if confidence is not None: confidence = confidence.cpu().numpy() re_order = np.argsort(confidence)[::-1] confidence = confidence[re_order] confidences_list.append(confidence) ligand_pos = ligand_pos[re_order] write_dir = f'{args.out_dir}/index{idx}_{data_list[0]["name"][0].replace("/","-")}' os.makedirs(write_dir, exist_ok=True) for rank, pos in enumerate(ligand_pos): mol_pred = copy.deepcopy(lig) if score_model_args.remove_hs: mol_pred = RemoveHs(mol_pred) if rank == 0: write_mol_with_coords(mol_pred, pos, os.path.join(write_dir, f'rank{rank+1}.sdf')) write_mol_with_coords(mol_pred, pos, os.path.join(write_dir, f'rank{rank+1}_confidence{confidence[rank]:.2f}.sdf')) self_distances = np.linalg.norm(ligand_pos[:, :, None, :] - ligand_pos[:, None, :, :], axis=-1) self_distances = np.where(np.eye(self_distances.shape[2]), np.inf, self_distances) min_self_distances_list.append(np.min(self_distances, axis=(1, 2))) if args.save_visualisation: if confidence is not None: for rank, batch_idx in enumerate(re_order): visualization_list[batch_idx].write(os.path.join(write_dir, f'rank{rank+1}_reverseprocess.pdb')) else: for rank, batch_idx in enumerate(ligand_pos): visualization_list[batch_idx].write(os.path.join(write_dir, f'rank{rank+1}_reverseprocess.pdb')) names_list.append(orig_complex_graph.name[0]) except Exception as e: print("Failed on", orig_complex_graph["name"], e) failures += 1 raise e print(f'Failed for {failures} complexes') print(f'Skipped {skipped} complexes') min_self_distances = np.array(min_self_distances_list) confidences = np.array(confidences_list) names = np.array(names_list) run_times = np.array(run_times) np.save(f'{args.out_dir}/min_self_distances.npy', min_self_distances) np.save(f'{args.out_dir}/confidences.npy', confidences) np.save(f'{args.out_dir}/run_times.npy', run_times) np.save(f'{args.out_dir}/complex_names.npy', np.array(names)) print(f'Results are in {args.out_dir}')