sample_MOAD_GlintDM.py

import argparse
import os
import shutil
import time

import numpy as np
import torch
from torch_geometric.data import Batch
from torch_geometric.transforms import Compose
from torch_scatter import scatter_sum, scatter_mean
from tqdm.auto import tqdm

import utils.misc as misc
import utils.transforms as trans
from datasets import get_dataset
from datasets.pl_data import FOLLOW_BATCH
from models.GlintDM import *
from utils.evaluation import atom_num

from utils.evaluation.docking_qvina import QVinaDockingTask
from utils.evaluation.docking_vina import VinaDockingTask
from utils.evaluation.docking_vina import get_vina_moad

from utils.evaluation import eval_atom_type, scoring_func, analyze, eval_bond_length
from utils import misc, reconstruct, transforms
from rdkit import Chem

from rdkit.Chem.QED import qed
from utils.evaluation.sascorer import compute_sa_score

#%matplotlib inline
import matplotlib.pyplot as plt

def unbatch_v_traj(ligand_v_traj, n_data, ligand_cum_atoms):
    all_step_v = [[] for _ in range(n_data)]
    for v in ligand_v_traj:  # step_i
        v_array = v.cpu().numpy()
        for k in range(n_data):
            all_step_v[k].append(v_array[ligand_cum_atoms[k]:ligand_cum_atoms[k + 1]])
    all_step_v = [np.stack(step_v) for step_v in all_step_v]  # num_samples * [num_steps, num_atoms_i]
    return all_step_v

def sampling(model, data, device='cuda:0', n_data = 0, 
             num_steps=None, Nsampling = 5, ligand_num_atoms = 0,
             init_ligand_pos=None, init_log_ligand_v = None, step_size= 50, starting_T = 1000):
    
    #n_data = batch_size if i < num_batch - 1 else num_samples - batch_size * (num_batch - 1)
    batch = Batch.from_data_list([data.clone() for _ in range(n_data)], follow_batch=FOLLOW_BATCH).to(device)
    
    r_list = []; ligand_num_atoms_list = []
    
    for n in range(Nsampling):
        with torch.no_grad():
            batch_protein = batch.protein_element_batch
            batch_ligand = torch.repeat_interleave(torch.arange(n_data), torch.tensor(ligand_num_atoms)).to(device)
            ligand_num_atoms_list += ligand_num_atoms
            
            r = model.sample_diffusion_refinement(
                protein_pos=batch.protein_pos,
                protein_v=batch.protein_atom_feature.float(),
                batch_protein=batch_protein,
                init_ligand_pos=init_ligand_pos,
                init_log_ligand_v=init_log_ligand_v,
                batch_ligand=batch_ligand,
                num_steps=num_steps,
                pos_only=False,
                center_pos_mode='protein',
                step_size = step_size,
                t = starting_T,#5*(Nrepeat-re) #100,#999,#*(Nrepeat-re),
            )
           
            r_list.append(r)
        
    N_r = len(r_list)
    
    r['pos'] = torch.concat([r_list[i]['pos'] for i in range(N_r)])
    r['v'] = torch.concat([r_list[i]['v'] for i in range(N_r)])
    r['pos_traj'] =[r['pos'].detach().cpu()]
    r['v_traj'] =[r['v'].argmax(-1).detach().cpu()]
    r['v0_traj'] =r['pos_traj']
    r['vt_traj'] = r['v_traj']
    
    return r, ligand_num_atoms_list
    
def sample_diffusion_ligand(model, data, num_samples, batch_size=16, device='cuda:0', TopN = 20,
                            num_steps=None, pos_only=False, center_pos_mode='protein', starting_T = 1000,
                            sample_num_atoms='prior',  step_size= 50, Nrepeat = 5, Nsampling= 5):
    all_pred_pos, all_pred_v = [], []
    all_pred_pos_traj, all_pred_v_traj = [], []
    all_pred_v0_traj, all_pred_vt_traj = [], []
    time_list = []; pred_score_list = []
    num_batch = int(np.ceil(num_samples / batch_size))
    current_i = 0
    
    for i in tqdm(range(num_batch)):
        
        init_ligand_pos, init_log_ligand_v = None, None
        n_data = batch_size if i < num_batch - 1 else num_samples - batch_size * (num_batch - 1)
        batch = Batch.from_data_list([data.clone() for _ in range(n_data)], follow_batch=FOLLOW_BATCH).to(device)
            
        best_indices = []
                    
        t1 = time.time()
        while len(best_indices) == 0:
            with torch.no_grad():
                batch_protein = batch.protein_element_batch
                
                if sample_num_atoms == 'prior':
                    pocket_size = atom_num.get_space_size(data.protein_pos.detach().cpu().numpy())
                    ligand_num_atoms = [atom_num.sample_atom_num(pocket_size).astype(int) for _ in range(n_data)]
                    batch_ligand = torch.repeat_interleave(torch.arange(n_data), torch.tensor(ligand_num_atoms)).to(device)
                    
                elif sample_num_atoms == 'range':
                    ligand_num_atoms = list(range(current_i + 1, current_i + n_data + 1))
                    batch_ligand = torch.repeat_interleave(torch.arange(n_data), torch.tensor(ligand_num_atoms)).to(device)
                elif sample_num_atoms == 'ref':
                    batch_ligand = batch.ligand_element_batch
                    ligand_num_atoms = scatter_sum(torch.ones_like(batch_ligand), batch_ligand, dim=0).tolist()
                else:
                    raise ValueError
                
                # init ligand pos
                center_pos = scatter_mean(batch.protein_pos, batch_protein, dim=0)
                
                batch_center_pos = center_pos[batch_ligand]
                init_ligand_pos = batch_center_pos + torch.randn_like(batch_center_pos)
                # init ligand v
                if pos_only:
                    init_ligand_pos = batch.ligand_pos #+ torch.randn_like(batch.ligand_pos)
                    init_ligand_v = batch.ligand_atom_feature_full
                    init_log_ligand_v = index_to_log_onehot(init_ligand_v, model.num_classes)
                    
                else:
                    uniform_logits = torch.zeros(len(batch_ligand), model.num_classes).to(device)
                    init_ligand_v = log_sample_categorical(uniform_logits)
                    init_log_ligand_v = index_to_log_onehot(init_ligand_v, model.num_classes)
                    
                    uniform = torch.rand_like(init_log_ligand_v)
                    gumbel_noise = -torch.log(-torch.log(uniform + 1e-30) + 1e-30)
                    
                    init_log_ligand_v = log_add_exp(init_log_ligand_v, gumbel_noise)
                    init_log_ligand_v = F.log_softmax(init_log_ligand_v,dim=-1)
                        
                for re in range(Nrepeat):
                    r = model.sample_diffusion_large_step(
                        protein_pos=batch.protein_pos,
                        protein_v=batch.protein_atom_feature.float(),
                        batch_protein=batch_protein,
                        init_ligand_pos=init_ligand_pos,
                        init_log_ligand_v=init_log_ligand_v,
                        batch_ligand=batch_ligand,
                        num_steps=num_steps,
                        pos_only=pos_only,
                        center_pos_mode=center_pos_mode,
                        step_size = refine_step_size
                    )
                    
                    init_ligand_pos, init_log_ligand_v = r['pos'], r['v'] 
                    gumbel_noise = -0.1*torch.log(-torch.log(uniform + 1e-30) + 1e-30)
                    init_log_ligand_v = F.log_softmax(gumbel_noise,dim=-1)
                    
                ligand_pos, ligand_log_v, ligand_pos_traj, ligand_v_traj = r['pos'], r['v'], r['pos_traj'], r['v_traj']
                
                ligand_v = ligand_log_v.argmax(-1)
                # unbatch pos
                ligand_cum_atoms = np.cumsum([0] + ligand_num_atoms)
                ligand_pos_array = ligand_pos.cpu().numpy().astype(np.float64)
                pred_pos = [ligand_pos_array[ligand_cum_atoms[k]:ligand_cum_atoms[k + 1]] for k in
                                 range(n_data)]  # num_samples * [num_atoms_i, 3]
                
                ligand_v_array = ligand_v.cpu().numpy()
                pred_v = [ligand_v_array[ligand_cum_atoms[k]:ligand_cum_atoms[k + 1]] for k in range(n_data)]
    
                score_list = []; vina_score_list = []; smiles_list = []
                for idx, (pos, v) in enumerate(zip(pred_pos, pred_v)):
                    try:
                        pred_atom_type = transforms.get_atomic_number_from_index(v, mode='add_aromatic')
                        pred_aromatic = transforms.is_aromatic_from_index(v, mode='add_aromatic')
                        
                        mol = reconstruct.reconstruct_from_generated(pos, pred_atom_type, pred_aromatic)
                        smiles = Chem.MolToSmiles(mol)
                        
                        if '.' in smiles:
                            vina_score = 1000
                            score = vina_score
                            smiles = None
                            
                        else:
                            QED = qed(mol); SA = compute_sa_score(mol)
                            pdb_file = pdb_root+data['protein_filename'].split('_')[0]+'.pdb'
                            vina_score = get_vina_moad(mol, pdb_file)
                            
                            if MAX_QED == 0 or MAX_SA ==0:
                                score = vina_score
                            else:
                                score = vina_score*min(QED,MAX_QED)*min(SA,MAX_SA)
                                
                    except:
                        vina_score = 1000
                        score = vina_score
                        smiles = None
                        
                    score_list.append(score)
                    vina_score_list.append(vina_score)
                    smiles_list.append(smiles)
                    
                pred_score_list.append(vina_score_list)
                score_list = np.array(score_list)
                vina_score_list = np.array(vina_score_list)
                
                best_indices = np.argsort(score_list)[:TopN]
                best_indices = best_indices[vina_score_list[best_indices] < 100]
                
        Nsampling = int(num_samples/len(best_indices))
        ligand_cum_atoms = torch.tensor(ligand_cum_atoms).to(device)
        
        ligand_pos_list = []; ligand_log_v_list = []; ligand_num_atoms_list = []; 
        
        for k in range(len(best_indices)):
            index = best_indices[k]
            
            pred_log_v = ligand_log_v[ligand_cum_atoms[index]:ligand_cum_atoms[index+1]]
            tensor_pos = torch.tensor(pred_pos[index]).to(device,torch.float32)
            
            ligand_pos_list.append(tensor_pos); ligand_log_v_list.append(pred_log_v); 
            ligand_num_atoms_list.append(pred_v[index].shape[0])
            
        n_data = len(ligand_num_atoms_list); ligand_num_atoms = ligand_num_atoms_list
        
        ligand_pos = torch.concat(ligand_pos_list); ligand_log_v = torch.concat(ligand_log_v_list)
        
        r, ligand_num_atoms_list = sampling(model, data, device = device, n_data = n_data, ligand_num_atoms = ligand_num_atoms,
                                           num_steps = num_steps, Nsampling = Nsampling, init_ligand_pos=ligand_pos, 
                                           init_log_ligand_v = ligand_log_v, step_size= step_size, starting_T = starting_T)
                            
        
        ligand_num_atoms = ligand_num_atoms_list; n_data = len(ligand_num_atoms_list)
        ligand_pos, ligand_log_v, ligand_pos_traj, ligand_v_traj = r['pos'], r['v'], r['pos_traj'], r['v_traj']
        
        ligand_v = ligand_log_v.argmax(-1)
        ligand_v0_traj, ligand_vt_traj = r['v0_traj'], r['vt_traj']
            
        # unbatch pos
        ligand_cum_atoms = np.cumsum([0] + ligand_num_atoms)
        ligand_pos_array = ligand_pos.cpu().numpy().astype(np.float64)
        all_pred_pos += [ligand_pos_array[ligand_cum_atoms[k]:ligand_cum_atoms[k + 1]] for k in
                         range(n_data)]  # num_samples * [num_atoms_i, 3]

        all_step_pos = [[] for _ in range(n_data)]
        for p in ligand_pos_traj:  # step_i
            p_array = p.cpu().numpy().astype(np.float64)
            for k in range(n_data):
                all_step_pos[k].append(p_array[ligand_cum_atoms[k]:ligand_cum_atoms[k + 1]])
        all_step_pos = [np.stack(step_pos) for step_pos in
                        all_step_pos]  # num_samples * [num_steps, num_atoms_i, 3]
        
        all_pred_pos_traj += [p for p in all_step_pos]

        # unbatch v
        ligand_v_array = ligand_v.cpu().numpy()
        all_pred_v += [ligand_v_array[ligand_cum_atoms[k]:ligand_cum_atoms[k + 1]] for k in range(n_data)]

        all_step_v = unbatch_v_traj(ligand_v_traj, n_data, ligand_cum_atoms)
        all_pred_v_traj += [v for v in all_step_v]

        if not pos_only:
            all_step_v0 = unbatch_v_traj(ligand_v0_traj, n_data, ligand_cum_atoms)
            all_pred_v0_traj += [v for v in all_step_v0]
            all_step_vt = unbatch_v_traj(ligand_vt_traj, n_data, ligand_cum_atoms)
            all_pred_vt_traj += [v for v in all_step_vt]
            
            #print(n_data)
                
        t2 = time.time()
        time_list.append(t2 - t1)
        current_i += n_data
        
    return all_pred_pos, all_pred_v, all_pred_pos_traj, all_pred_v_traj, all_pred_v0_traj, all_pred_vt_traj, time_list, pred_score_list

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('config', type=str, nargs='?', 
                        default = '/NAS_Storage4/leo8544/3D_DrugGeneration/targetdiff/configs/sampling.yml')
    parser.add_argument('-i', '--data_id', type=int, default=98)
    parser.add_argument('--device', type=str, default='cuda:0')
    parser.add_argument('--batch_size', type=int, default=5)
    parser.add_argument('--result_path', type=str, default='./outputs/original')
    args = parser.parse_args()

    #args.device = 'cpu'
    args.device = 'cuda:1'
    
    logger = misc.get_logger('sampling')

    # Load config
    config = misc.load_config(args.config)
    logger.info(config)
    misc.seed_all(config.sample.seed)
 
    config.model.checkpoint = './checkpoints/pretrained_GlintDM.pt'
    
    ckpt = torch.load(config.model.checkpoint, map_location=args.device,weights_only=False)
    logger.info(f"Training Config: {ckpt['config']}")

    # follow_batch = ['protein_element', 'ligand_element']
    collate_exclude_keys = ['ligand_nbh_list']
    # Transforms
    protein_featurizer = trans.FeaturizeProteinAtom()
    ligand_atom_mode = ckpt['config'].data.transform.ligand_atom_mode
    ligand_featurizer = trans.FeaturizeLigandAtom(ligand_atom_mode)
    transform = Compose([
        protein_featurizer,
        ligand_featurizer,
        trans.FeaturizeLigandBond(),
    ])
    
    from datasets.pl_moad import ProteinLigandDataset
    pdb_root = '/NAS_Storage4/leo8544/3D_DrugGeneration/targetdiff/data/processed_noH_full/test/'
    data_path = "/NAS_Storage4/leo8544/3D_DrugGeneration/targetdiff/data/processed_noH_full/"
    
    try:
        print("Test proteins: ", len(test_set))
    except:
        test_set = ProteinLigandDataset(data_path, index_file='moad_test_index.pkl', transform= transform)
    
    # Load model
    model = GlintDM(
        ckpt['config'].model,
        protein_atom_feature_dim=protein_featurizer.feature_dim,
        ligand_atom_feature_dim=ligand_featurizer.feature_dim,
        device=args.device,
        prop_func=None,
    ).to(args.device)
    
    model.load_state_dict(ckpt['model'], strict = False)
    logger.info(f'Successfully load the model! {config.model.checkpoint}')
    
    Nrepeat = 10; Nsampling = 10; config.sample.num_samples = 100; starting_T = 1000
    args.batch_size = config.sample.num_samples; step_size = 100; refine_step_size = 200;
    TopN = int(config.sample.num_samples/Nsampling); MAX_QED = 0.6; MAX_SA = 0.6
    
    QEDs = []; SAs = []; atom_stabilities=[]; smiles_list=[]
    vinas = []; vina_mins = []; vina_docks = []; hits = []
    
    mol_stable = [];
    for i in range(0, 130):
        args.data_id= i; 
        data = test_set[args.data_id]
        aromatic = transforms.is_aromatic_from_index(data.ligand_atom_feature_full.detach().cpu().numpy(),
                                                          mode='add_aromatic')
        mol = reconstruct.reconstruct_from_generated(data.ligand_pos.detach().to('cpu',torch.double).numpy(), 
                                                     data.ligand_element.detach().cpu().numpy().tolist(), aromatic)
        
        pdb_file = pdb_root+data['protein_filename'].split('_')[0]+'.pdb'
        vina_score = get_vina_moad(mol, pdb_file, exhaustiveness=16)
        
        pred_pos, pred_v, pred_pos_traj, pred_v_traj, pred_v0_traj, pred_vt_traj, time_list, pred_score_list = sample_diffusion_ligand(
            model, data, config.sample.num_samples,
            batch_size=args.batch_size, device=args.device,
            num_steps= 1000,
            pos_only=config.sample.pos_only, starting_T = starting_T,
            center_pos_mode=config.sample.center_pos_mode,
            sample_num_atoms=config.sample.sample_num_atoms, TopN = TopN, 
            step_size= step_size, Nrepeat = Nrepeat, Nsampling= Nsampling,
        )
        
        result = {
            'data': data,
            'pred_ligand_pos': pred_pos,
            'pred_ligand_v': pred_v,
            'pred_ligand_pos_traj': pred_pos_traj,
            'pred_ligand_v_traj': pred_v_traj,
            'time': time_list,
            'vina': vina_score,
            'ligand_pos': data.ligand_pos.numpy(),
            'ligand_v': data.ligand_atom_feature_full.numpy(),
        }
        logger.info('Sample done!')
        
        result_path = args.result_path
        os.makedirs(result_path, exist_ok=True)
        shutil.copyfile(args.config, os.path.join(result_path, 'sample.yml'))
        torch.save(result, os.path.join(result_path, 'GlintDM_MOAD_'+\
                                                    '_QED'+str(MAX_QED)+'_sa_'+str(MAX_SA)+\
                                                    '_Nmol_'+str(config.sample.num_samples)+'_repeat_'+str(Nrepeat)+\
                                                    '_TopN_'+str(TopN)+'_ST_'+str(starting_T)+\
                                                   '_step_'+str(step_size)+f'_N1000_{args.data_id}.pt'))