File size: 9,026 Bytes
6e7d4ba |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 |
import argparse
import sys
import os
import warnings
import tempfile
import pandas as pd
from Bio.PDB import PDBParser
from pathlib import Path
from rdkit import Chem
from torch.utils.data import DataLoader
from functools import partial
basedir = Path(__file__).resolve().parent.parent
sys.path.append(str(basedir))
warnings.filterwarnings("ignore")
from src import utils
from src.data.dataset import ProcessedLigandPocketDataset
from src.data.data_utils import TensorDict, process_raw_pair
from src.model.lightning import DrugFlow
from src.sbdd_metrics.metrics import FullEvaluator
from tqdm import tqdm
from pdb import set_trace
def aggregate_metrics(table):
agg_col = 'posebusters'
total = 0
table[agg_col] = 0
for column in table.columns:
if column.startswith(agg_col) and column != agg_col:
table[agg_col] += table[column].fillna(0).astype(float)
total += 1
table[agg_col] = table[agg_col] / total
agg_col = 'reos'
total = 0
table[agg_col] = 0
for column in table.columns:
if column.startswith(agg_col) and column != agg_col:
table[agg_col] += table[column].fillna(0).astype(float)
total += 1
table[agg_col] = table[agg_col] / total
agg_col = 'chembl_ring_systems'
total = 0
table[agg_col] = 0
for column in table.columns:
if column.startswith(agg_col) and column != agg_col and not column.endswith('smi'):
table[agg_col] += table[column].fillna(0).astype(float)
total += 1
table[agg_col] = table[agg_col] / total
return table
if __name__ == "__main__":
p = argparse.ArgumentParser()
p.add_argument('--protein', type=str, required=True, help="Input PDB file.")
p.add_argument('--ref_ligand', type=str, required=True, help="SDF file with reference ligand used to define the pocket.")
p.add_argument('--checkpoint', type=str, required=True, help="Model checkpoint file.")
p.add_argument('--molecule_size', type=str, required=False, default=None, help="Maximum number of atoms in the sampled molecules. Can be a single number or a range, e.g. '15,20'. If None, size will be sampled.")
p.add_argument('--output', type=str, required=False, default='samples.sdf', help="Output file.")
p.add_argument('--n_samples', type=int, required=False, default=10, help="Number of sampled molecules.")
p.add_argument('--batch_size', type=int, required=False, default=32, help="Batch size.")
p.add_argument('--pocket_distance_cutoff', type=float, required=False, default=8.0, help="Distance cutoff to define the pocket around the reference ligand.")
p.add_argument('--n_steps', type=int, required=False, default=None, help="Number of denoising steps.")
p.add_argument('--device', type=str, required=False, default='cuda:0', help="Device to use.")
p.add_argument('--datadir', type=Path, required=False, default=Path(basedir, 'src', 'default'), help="Needs to be specified to sample molecule sizes.")
p.add_argument('--seed', type=int, required=False, default=42, help="Random seed.")
p.add_argument('--filter', action='store_true', required=False, default=False, help="Apply basic filters and keep sampling until `n_samples` molecules passing these filters are found.")
p.add_argument('--metrics_output', type=str, required=False, default=None, help="If provided, metrics will be computed and saved in csv format at this location.")
p.add_argument('--gnina', type=str, required=False, default=None, help="Path to a gnina executable. Required for computing docking scores.")
p.add_argument('--reduce', type=str, required=False, default=None, help="Path to a reduce executable. Required for computing interactions.")
args = p.parse_args()
utils.set_deterministic(seed=args.seed)
utils.disable_rdkit_logging()
if args.molecule_size is None and (args.datadir is None or not args.datadir.exists()):
raise NotImplementedError(
"Please provide a path to the processed dataset (using `--datadir`) "\
"to infer the number of nodes. It contains the size distribution histogram."
)
if not args.filter:
args.batch_size = min(args.batch_size, args.n_samples)
# Loading model
chkpt_path = Path(args.checkpoint)
chkpt_name = chkpt_path.parts[-1].split('.')[0]
model = DrugFlow.load_from_checkpoint(args.checkpoint, map_location=args.device, strict=False)
if args.datadir is not None:
model.datadir = args.datadir
model.setup(stage='generation')
model.batch_size = model.eval_batch_size = args.batch_size
model.eval().to(args.device)
if args.n_steps is not None:
model.T = args.n_steps
# Loading size model
size_model = None
molecule_size = None
molecule_size_boundaries = None
if args.molecule_size is not None:
if args.molecule_size.isdigit():
molecule_size = int(args.molecule_size)
print(f'Will generate molecules of size {molecule_size}')
else:
boundaries = [x.strip() for x in args.molecule_size.split(',')]
assert len(boundaries) == 2 and boundaries[0].isdigit() and boundaries[1].isdigit()
left = int(boundaries[0])
right = int(boundaries[1])
molecule_size = f"uniform_{left}_{right}"
print(f'Will generate molecules with numbers of atoms sampled from U({left}, {right})')
# Preparing input
pdb_model = PDBParser(QUIET=True).get_structure('', args.protein)[0]
rdmol = Chem.SDMolSupplier(str(args.ref_ligand))[0]
ligand, pocket = process_raw_pair(
pdb_model, rdmol,
dist_cutoff=args.pocket_distance_cutoff,
pocket_representation=model.pocket_representation,
compute_nerf_params=True,
nma_input=args.protein if model.dynamics.add_nma_feat else None
)
ligand['name'] = 'ligand'
dataset = [{'ligand': ligand, 'pocket': pocket} for _ in range(args.batch_size)]
dataloader = DataLoader(
dataset=dataset,
batch_size=args.batch_size,
collate_fn=partial(ProcessedLigandPocketDataset.collate_fn, ligand_transform=None),
pin_memory=True
)
# Start sampling
smiles = set()
sampled_molecules = []
metrics = []
Path(args.output).parent.absolute().mkdir(parents=True, exist_ok=True)
print(f'Will generate {args.n_samples} samples')
evaluator = FullEvaluator(gnina=args.gnina, reduce=args.reduce)
with tqdm(total=args.n_samples) as pbar:
while len(sampled_molecules) < args.n_samples:
for i, data in enumerate(dataloader):
new_data = {
'ligand': TensorDict(**data['ligand']).to(args.device),
'pocket': TensorDict(**data['pocket']).to(args.device),
}
rdmols, rdpockets, _ = model.sample(
new_data,
n_samples=1,
timesteps=args.n_steps,
num_nodes=molecule_size,
)
if args.filter or (args.metrics_output is not None):
results = []
with tempfile.TemporaryDirectory() as tmpdir:
for mol, receptor in zip(rdmols, rdpockets):
receptor_path = Path(tmpdir, 'receptor.pdb')
Chem.MolToPDBFile(receptor, str(receptor_path))
results.append(evaluator(mol, receptor_path))
table = pd.DataFrame(results)
table['novel'] = ~table['representation.smiles'].isin(smiles)
table = aggregate_metrics(table)
added_molecules = 0
if args.filter:
table['passed_filters'] = (
(table['posebusters'] == 1) &
# (table['reos'] == 1) &
(table['chembl_ring_systems'] == 1) &
(table['novel'] == 1)
)
for i, (passed, smi) in enumerate(table[['passed_filters', 'representation.smiles']].values):
if passed:
sampled_molecules.append(rdmols[i])
smiles.add(smi)
added_molecules += 1
if args.metrics_output is not None:
metrics.append(table[table['passed_filters']])
else:
sampled_molecules.extend(rdmols)
added_molecules = len(rdmols)
if args.metrics_output is not None:
metrics.append(table)
pbar.update(added_molecules)
# Write results
utils.write_sdf_file(args.output, sampled_molecules)
if args.metrics_output is not None:
metrics = pd.concat(metrics)
metrics.to_csv(args.metrics_output, index=False)
|