|
|
import argparse |
|
|
import sys |
|
|
import yaml |
|
|
import torch |
|
|
import numpy as np |
|
|
import pickle |
|
|
from argparse import Namespace |
|
|
|
|
|
from pathlib import Path |
|
|
|
|
|
basedir = Path(__file__).resolve().parent.parent |
|
|
sys.path.append(str(basedir)) |
|
|
|
|
|
from src import utils |
|
|
from src.utils import dict_to_namespace, namespace_to_dict |
|
|
from src.analysis.visualization_utils import mols_to_pdbfile, mol_as_pdb |
|
|
from src.data.data_utils import TensorDict, Residues |
|
|
from src.data.postprocessing import process_all |
|
|
from src.model.lightning import DrugFlow |
|
|
from src.sbdd_metrics.evaluation import compute_all_metrics_drugflow |
|
|
|
|
|
from tqdm import tqdm |
|
|
from pdb import set_trace |
|
|
|
|
|
|
|
|
def combine(base_args, override_args): |
|
|
assert not isinstance(base_args, dict) |
|
|
assert not isinstance(override_args, dict) |
|
|
|
|
|
arg_dict = base_args.__dict__ |
|
|
for key, value in override_args.__dict__.items(): |
|
|
if key not in arg_dict or arg_dict[key] is None: |
|
|
print(f"Add parameter {key}: {value}") |
|
|
arg_dict[key] = value |
|
|
elif isinstance(value, Namespace): |
|
|
arg_dict[key] = combine(arg_dict[key], value) |
|
|
else: |
|
|
print(f"Replace parameter {key}: {arg_dict[key]} -> {value}") |
|
|
arg_dict[key] = value |
|
|
return base_args |
|
|
|
|
|
|
|
|
def path_to_str(input_dict): |
|
|
for key, value in input_dict.items(): |
|
|
if isinstance(value, dict): |
|
|
input_dict[key] = path_to_str(value) |
|
|
else: |
|
|
input_dict[key] = str(value) if isinstance(value, Path) else value |
|
|
return input_dict |
|
|
|
|
|
|
|
|
def sample(cfg, model_params, samples_dir, job_id=0, n_jobs=1): |
|
|
print('Sampling...') |
|
|
model = DrugFlow.load_from_checkpoint(cfg.checkpoint, map_location=cfg.device, strict=False, |
|
|
**model_params) |
|
|
model.setup(stage='fit' if cfg.set == 'train' else cfg.set) |
|
|
model.eval().to(cfg.device) |
|
|
|
|
|
dataloader = getattr(model, f'{cfg.set}_dataloader')() |
|
|
print(f'Real batch size is {dataloader.batch_size * cfg.n_samples}') |
|
|
|
|
|
name2count = {} |
|
|
for i, data in enumerate(tqdm(dataloader)): |
|
|
if i % n_jobs != job_id: |
|
|
print(f'Skipping batch {i}') |
|
|
continue |
|
|
|
|
|
new_data = { |
|
|
'ligand': TensorDict(**data['ligand']).to(cfg.device), |
|
|
'pocket': Residues(**data['pocket']).to(cfg.device), |
|
|
} |
|
|
try: |
|
|
rdmols, rdpockets, names = model.sample( |
|
|
data=new_data, |
|
|
n_samples=cfg.n_samples, |
|
|
num_nodes=("ground_truth" if cfg.sample_with_ground_truth_size else None) |
|
|
) |
|
|
except Exception as e: |
|
|
if cfg.set == 'train': |
|
|
names = data['ligand']['name'] |
|
|
print(f'Failed to sample for {names}: {e}') |
|
|
continue |
|
|
else: |
|
|
raise e |
|
|
|
|
|
for mol, pocket, name in zip(rdmols, rdpockets, names): |
|
|
name = name.replace('.sdf', '') |
|
|
idx = name2count.setdefault(name, 0) |
|
|
output_dir = Path(samples_dir, name) |
|
|
output_dir.mkdir(parents=True, exist_ok=True) |
|
|
if cfg.postprocess: |
|
|
mol = process_all(mol, largest_frag=True, adjust_aromatic_Ns=True, relax_iter=0) |
|
|
|
|
|
for prop in mol.GetAtoms()[0].GetPropsAsDict().keys(): |
|
|
|
|
|
mol.SetDoubleProp(prop, np.mean([a.GetDoubleProp(prop) for a in mol.GetAtoms()])) |
|
|
|
|
|
|
|
|
out_pdb_path = Path(output_dir, f'{idx}_ligand_{prop}.pdb') |
|
|
mol_as_pdb(mol, out_pdb_path, bfactor=prop) |
|
|
|
|
|
out_sdf_path = Path(output_dir, f'{idx}_ligand.sdf') |
|
|
out_pdb_path = Path(output_dir, f'{idx}_pocket.pdb') |
|
|
utils.write_sdf_file(out_sdf_path, [mol]) |
|
|
mols_to_pdbfile([pocket], out_pdb_path) |
|
|
|
|
|
name2count[name] += 1 |
|
|
|
|
|
|
|
|
def evaluate(cfg, model_params, samples_dir): |
|
|
print('Evaluation...') |
|
|
data, table_detailed, table_aggregated = compute_all_metrics_drugflow( |
|
|
in_dir=samples_dir, |
|
|
gnina_path=model_params['train_params'].gnina, |
|
|
reduce_path=cfg.reduce, |
|
|
reference_smiles_path=Path(model_params['train_params'].datadir, 'train_smiles.npy'), |
|
|
n_samples=cfg.n_samples, |
|
|
exclude_evaluators=[] if cfg.exclude_evaluators is None else cfg.exclude_evaluators, |
|
|
) |
|
|
with open(Path(samples_dir, 'metrics_data.pkl'), 'wb') as f: |
|
|
pickle.dump(data, f) |
|
|
table_detailed.to_csv(Path(samples_dir, 'metrics_detailed.csv'), index=False) |
|
|
table_aggregated.to_csv(Path(samples_dir, 'metrics_aggregated.csv'), index=False) |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
p = argparse.ArgumentParser() |
|
|
p.add_argument('--config', type=str) |
|
|
p.add_argument('--job_id', type=int, default=0, help='Job ID') |
|
|
p.add_argument('--n_jobs', type=int, default=1, help='Number of jobs') |
|
|
args = p.parse_args() |
|
|
|
|
|
with open(args.config, 'r') as f: |
|
|
cfg = yaml.safe_load(f) |
|
|
cfg = dict_to_namespace(cfg) |
|
|
|
|
|
utils.set_deterministic(seed=cfg.seed) |
|
|
utils.disable_rdkit_logging() |
|
|
|
|
|
model_params = torch.load(cfg.checkpoint, map_location=cfg.device)['hyper_parameters'] |
|
|
if 'model_args' in cfg: |
|
|
ckpt_args = dict_to_namespace(model_params) |
|
|
model_params = combine(ckpt_args, cfg.model_args).__dict__ |
|
|
|
|
|
ckpt_path = Path(cfg.checkpoint) |
|
|
ckpt_name = ckpt_path.parts[-1].split('.')[0] |
|
|
n_steps = model_params['simulation_params'].n_steps |
|
|
samples_dir = Path(cfg.sample_outdir, cfg.set, f'{ckpt_name}_T={n_steps}') or \ |
|
|
Path(ckpt_path.parent.parent, 'samples', cfg.set, f'{ckpt_name}_T={n_steps}') |
|
|
assert cfg.set in {'val', 'test', 'train'} |
|
|
samples_dir.mkdir(parents=True, exist_ok=True) |
|
|
|
|
|
|
|
|
with open(Path(samples_dir, 'model_params.yaml'), 'w') as f: |
|
|
yaml.dump(path_to_str(namespace_to_dict(model_params)), f) |
|
|
with open(Path(samples_dir, 'sampling_params.yaml'), 'w') as f: |
|
|
yaml.dump(path_to_str(namespace_to_dict(cfg)), f) |
|
|
|
|
|
if cfg.sample: |
|
|
sample(cfg, model_params, samples_dir, job_id=args.job_id, n_jobs=args.n_jobs) |
|
|
|
|
|
if cfg.evaluate: |
|
|
assert args.job_id == 0 and args.n_jobs == 1, 'Evaluation is not parallelised on GPU machines' |
|
|
evaluate(cfg, model_params, samples_dir) |
