OneScience-Group/SurfDock

1import numpy as np2import torch,os3from torch_geometric.loader import DataLoader4import traceback5from utils.diffusion_utils import modify_conformer, set_time6from utils.torsion import modify_conformer_torsion_angles7from scipy.spatial.transform import Rotation as R8import warnings9# from datasets.process_mols import write_mol_with_coords10from force_optimize.minimize_utils import UpdateGrpah,GetfixedPDB,GetFFGenerator11from openmm.app import Modeller12from joblib import Parallel,delayed13from tqdm import tqdm14from loguru import logger15def randomize_position(data_list, no_torsion, no_random, tr_sigma_max,ligand_to_pocket_center = False):16    # in place modification of the list17    if not no_torsion:18        # randomize torsion angles19        for complex_graph in data_list:20            torsion_updates = np.random.uniform(low=-np.pi, high=np.pi, size=complex_graph['ligand'].edge_mask.sum())21            complex_graph['ligand'].pos = \22                modify_conformer_torsion_angles(complex_graph['ligand'].pos,23                                                complex_graph['ligand', 'ligand'].edge_index.T[24                                                    complex_graph['ligand'].edge_mask],25                                                complex_graph['ligand'].mask_rotate, torsion_updates)26 27    for complex_graph in data_list:28        # randomize position29        molecule_center = torch.mean(complex_graph['ligand'].pos, dim=0, keepdim=True)30        random_rotation = torch.from_numpy(R.random().as_matrix()).float()31        complex_graph['ligand'].pos = (complex_graph['ligand'].pos - molecule_center) @ random_rotation.T32        # base_rmsd = np.sqrt(np.sum((complex_graph['ligand'].pos.cpu().numpy() - orig_complex_graph['ligand'].pos.numpy()) ** 2, axis=1).mean())33        # put the molecule in the center of the pocket by caoduanhua34        if ligand_to_pocket_center and complex_graph['receptor'].pocket_center is not None:35            # logger.info('Use predict pocket center to put ligand in the center of pocket! {},{},{}'.format(complex_graph['ligand'].pos.shape,complex_graph['receptor'].pocket_center.shape,torch.mean(complex_graph['ligand'].pos, dim=0, keepdim=True).shape))36            complex_graph['ligand'].pos  = complex_graph['ligand'].pos - torch.mean(complex_graph['ligand'].pos, dim=0, keepdim=True) + complex_graph['receptor'].pocket_center.to(complex_graph['ligand'].pos.device)37            logger.info('Use predict pocket center to put ligand in the center of pocket!')38        else:39            if not no_random:  # note for now the torsion angles are still randomised40                tr_update = torch.normal(mean=0, std=tr_sigma_max, size=(1, 3))41                complex_graph['ligand'].pos += tr_update42 43def inferenceFFOptimize(data_list,args,receptor_path,N=40):44    # loaded ligand docking pose and add Hs45    fixer = GetfixedPDB(receptor_path)46    modeller = Modeller(fixer.topology, fixer.positions)47    protein_atoms = list(fixer.topology.atoms())48    system_generator = GetFFGenerator()49    with Parallel(n_jobs=max(args.num_process,N)) as parallel:50        logger.info('Use force field to do energy minimized!')51        new_data_list = parallel(delayed(UpdateGrpah)(graph,system_generator,modeller,protein_atoms) for graph in data_list) # succssed return graph object ,error return int(1)52        53    result = np.array([i if type(i) == int else 0 for i in new_data_list])54 55    new_data_list = list(filter(lambda x:type(x)!=int,new_data_list))56    57    if result.sum() == 0:58        return new_data_list,[]59    else:60        indices = np.where(result == 1)61        failed_graphs = [data_list[i] for i in indices[0]]62        logger.info(f'Minimized not Completed:{receptor_path}, {len(failed_graphs)} sdf not be minimized by default forcefield , try use gaff-2.11 forcefield!')63        with Parallel(n_jobs=max(args.num_process,len(failed_graphs))) as parallel:64            new_data_list_add =  parallel(delayed(UpdateGrpah)(graph,system_generator,modeller,protein_atoms) for graph in failed_graphs)65            66        result = np.array([i if type(i) == int else 0 for i in new_data_list_add ])67        new_data_list_add = list(filter(lambda x:type(x)!=int,new_data_list_add))68        new_data_list = new_data_list + new_data_list_add69        if result.sum() != 0:70            indices = np.where(result == 1)71            failed_graphs = [failed_graphs[i] for i in indices[0]]72            logger.info(f'Minimized not Completed:{receptor_path}, {len(failed_graphs)} sdf not be minimized by default forcefield , try use gaff-2.11 forcefield!')73        else:74            failed_graphs = []75        return new_data_list,failed_graphs76@logger.catch77def sampling(input_data_list, model, inference_steps, tr_schedule, rot_schedule, tor_schedule, device, t_to_sigma, model_args,78             no_random=False, ode=False, visualization_list=None, confidence_model=None, confidence_data_list=None,79             confidence_model_args=None, batch_size=32, no_final_step_noise=False,args = None):80    data_list = input_data_list81    N = len(data_list)82    pred_score = []83    for t_idx in range(inference_steps):84        # use prediction score as a ranking metric , implemented by caoduanhua85        # pred_score = []86        t_tr, t_rot, t_tor = tr_schedule[t_idx], rot_schedule[t_idx], tor_schedule[t_idx]87        dt_tr = tr_schedule[t_idx] - tr_schedule[t_idx + 1] if t_idx < inference_steps - 1 else tr_schedule[t_idx]88        dt_rot = rot_schedule[t_idx] - rot_schedule[t_idx + 1] if t_idx < inference_steps - 1 else rot_schedule[t_idx]89        dt_tor = tor_schedule[t_idx] - tor_schedule[t_idx + 1] if t_idx < inference_steps - 1 else tor_schedule[t_idx]90 91        loader = DataLoader(data_list, batch_size=batch_size)92        new_data_list = []93 94        for complex_graph_batch in loader:95            b = complex_graph_batch.num_graphs96            complex_graph_batch = complex_graph_batch.to(device)97            tr_sigma, rot_sigma, tor_sigma = t_to_sigma(t_tr, t_rot, t_tor)98            set_time(complex_graph_batch, t_tr, t_rot, t_tor, b, model_args.all_atoms, device)99            with torch.no_grad():100                tr_score, rot_score, tor_score = model(complex_graph_batch) 101                102                #103                104            tr_g = tr_sigma * torch.sqrt(torch.tensor(2 * np.log(model_args.tr_sigma_max / model_args.tr_sigma_min)))105            rot_g = 2 * rot_sigma * torch.sqrt(torch.tensor(np.log(model_args.rot_sigma_max / model_args.rot_sigma_min)))106 107            if ode:108                tr_perturb = (0.5 * tr_g ** 2 * dt_tr * tr_score.cpu()).cpu()109                rot_perturb = (0.5 * rot_score.cpu() * dt_rot * rot_g ** 2).cpu()110            else:111                tr_z = torch.zeros((b, 3)) if no_random or (no_final_step_noise and t_idx == inference_steps - 1) \112                    else torch.normal(mean=0, std=1, size=(b, 3))113                tr_perturb = (tr_g ** 2 * dt_tr * tr_score.cpu() + tr_g * np.sqrt(dt_tr) * tr_z).cpu()114 115                rot_z = torch.zeros((b, 3)) if no_random or (no_final_step_noise and t_idx == inference_steps - 1) \116                    else torch.normal(mean=0, std=1, size=(b, 3))117                rot_perturb = (rot_score.cpu() * dt_rot * rot_g ** 2 + rot_g * np.sqrt(dt_rot) * rot_z).cpu()118 119            if not model_args.no_torsion:120                tor_g = tor_sigma * torch.sqrt(torch.tensor(2 * np.log(model_args.tor_sigma_max / model_args.tor_sigma_min)))121                if ode:122                    tor_perturb = (0.5 * tor_g ** 2 * dt_tor * tor_score.cpu()).numpy()123                else:124                    tor_z = torch.zeros(tor_score.shape) if no_random or (no_final_step_noise and t_idx == inference_steps - 1) \125                        else torch.normal(mean=0, std=1, size=tor_score.shape)126                    tor_perturb = (tor_g ** 2 * dt_tor * tor_score.cpu() + tor_g * np.sqrt(dt_tor) * tor_z).numpy()127            else:128                tor_perturb = None129 130            # Apply noise131            tor_count_head = 0132            tor_count_tail = 0133            # node_head = 0134            # node_tail = 0135            for i, complex_graph in enumerate(complex_graph_batch.to('cpu').to_data_list()):136                # node_tail += complex_graph['ligand'].pos.shape[0]137                # complex_graph['ligand']['final_ligand'] = lig_node_attr[node_head:node_tail]138                # node_head+=complex_graph['ligand'].pos.shape[0]139                # if i==0:140                if type(complex_graph['ligand'].mask_rotate) is list:141                    complex_graph['ligand'].mask_rotate = complex_graph['ligand'].mask_rotate[0]142                tor_count_tail += complex_graph['ligand'].mask_rotate.shape[0]143                try:144                    new_data_list.append(modify_conformer(complex_graph, tr_perturb[i:i + 1], rot_perturb[i:i + 1].squeeze(0),145                                            tor_perturb[tor_count_head :tor_count_tail] if not model_args.no_torsion else None))146                except:147                    new_data_list.append(complex_graph)148                tor_count_head += complex_graph['ligand'].mask_rotate.shape[0]149                150                151        data_list = new_data_list152        153 154        if visualization_list is not None:155            for idx, visualization in enumerate(visualization_list):156                visualization.add((data_list[idx]['ligand'].pos + data_list[idx].original_center).detach().cpu(),157                                  part=1, order=t_idx + 2)158    #Before scoring the final conformers, we need to use force field to do energy minimized159    if args is not None and args.force_optimize:160        161        """162        step 1 : load mol to make energy minimize163        step 2 : load fixed protein or pocket164        """165        try:166            receptor_path = data_list[0]["protein_path"]167            logger.info('recptor path: {}',receptor_path)168            new_data_list,failed_graphs = inferenceFFOptimize(data_list,args,receptor_path,N)169            if len(failed_graphs) != 0:170                receptor_path = data_list[0]["pocket_path"]171                logger.info('Some minimized failed ! Use pocket file to do energy minimized!')172                new_data_list_pocket,failed_graphs = inferenceFFOptimize(failed_graphs,args,receptor_path,len(failed_graphs))173                new_data_list += new_data_list_pocket174            logger.info('Return sucessed examples: {}',len(new_data_list))175            logger.info('Return failed examples: {}',len(failed_graphs))176            data_list = new_data_list + failed_graphs177            178        except Exception as e:179            error_info = traceback.format_exc()180            logger.info(error_info)181 182            warnings.warn(f'Complex {data_list[0]["name"]} will scoring without energy minimized!')183            pass184        185 186 187    with torch.no_grad():188        if confidence_model is not None:189            loader = DataLoader(data_list, batch_size=batch_size)190            # try use forcefields to do energy minimized!191            if confidence_data_list is not None:192                confidence_loader = iter(DataLoader(confidence_data_list, batch_size=batch_size))193            confidence = []194            for complex_graph_batch in loader:195                complex_graph_batch = complex_graph_batch.to(device)196                if confidence_data_list is not None:197                    confidence_complex_graph_batch = next(confidence_loader).to(device)198                    confidence_complex_graph_batch['ligand'].pos = complex_graph_batch['ligand'].pos199                    # confidence need all_atoms or not200                    set_time(confidence_complex_graph_batch, 0, 0, 0, N, confidence_model_args.all_atoms, device)201                    confidence.append(confidence_model(confidence_complex_graph_batch)[-1])202                else:203                    b = complex_graph_batch.num_graphs204                    set_time(complex_graph_batch, 0, 0, 0, b, confidence_model_args.all_atoms, device)205                    confidence.append(confidence_model(complex_graph_batch)[-1])206            confidence = torch.cat(confidence, dim=0)207        else:208            confidence = None209    if confidence is not None:210        pred_score = confidence211 212    return data_list, pred_score213