mirror of
https://github.com/gcorso/DiffDock.git
synced 2026-06-05 02:14:21 +08:00
222 lines
13 KiB
Python
222 lines
13 KiB
Python
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}')
|
|
|
|
|