import numpy as np import torch,os from torch_geometric.loader import DataLoader import traceback from utils.diffusion_utils import modify_conformer, set_time from utils.torsion import modify_conformer_torsion_angles from scipy.spatial.transform import Rotation as R import warnings # from datasets.process_mols import write_mol_with_coords from force_optimize.minimize_utils import UpdateGrpah,GetfixedPDB,GetFFGenerator from openmm.app import Modeller from joblib import Parallel,delayed from tqdm import tqdm from loguru import logger def randomize_position(data_list, no_torsion, no_random, tr_sigma_max,ligand_to_pocket_center = False): # in place modification of the list if not no_torsion: # randomize torsion angles for complex_graph in data_list: torsion_updates = np.random.uniform(low=-np.pi, high=np.pi, size=complex_graph['ligand'].edge_mask.sum()) complex_graph['ligand'].pos = \ modify_conformer_torsion_angles(complex_graph['ligand'].pos, complex_graph['ligand', 'ligand'].edge_index.T[ complex_graph['ligand'].edge_mask], complex_graph['ligand'].mask_rotate, torsion_updates) for complex_graph in data_list: # randomize position molecule_center = torch.mean(complex_graph['ligand'].pos, dim=0, keepdim=True) random_rotation = torch.from_numpy(R.random().as_matrix()).float() complex_graph['ligand'].pos = (complex_graph['ligand'].pos - molecule_center) @ random_rotation.T # base_rmsd = np.sqrt(np.sum((complex_graph['ligand'].pos.cpu().numpy() - orig_complex_graph['ligand'].pos.numpy()) ** 2, axis=1).mean()) # put the molecule in the center of the pocket by caoduanhua if ligand_to_pocket_center and complex_graph['receptor'].pocket_center is not None: # 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)) 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) logger.info('Use predict pocket center to put ligand in the center of pocket!') else: if not no_random: # note for now the torsion angles are still randomised tr_update = torch.normal(mean=0, std=tr_sigma_max, size=(1, 3)) complex_graph['ligand'].pos += tr_update def inferenceFFOptimize(data_list,args,receptor_path,N=40): # loaded ligand docking pose and add Hs fixer = GetfixedPDB(receptor_path) modeller = Modeller(fixer.topology, fixer.positions) protein_atoms = list(fixer.topology.atoms()) system_generator = GetFFGenerator() with Parallel(n_jobs=max(args.num_process,N)) as parallel: logger.info('Use force field to do energy minimized!') 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) result = np.array([i if type(i) == int else 0 for i in new_data_list]) new_data_list = list(filter(lambda x:type(x)!=int,new_data_list)) if result.sum() == 0: return new_data_list,[] else: indices = np.where(result == 1) failed_graphs = [data_list[i] for i in indices[0]] logger.info(f'Minimized not Completed:{receptor_path}, {len(failed_graphs)} sdf not be minimized by default forcefield , try use gaff-2.11 forcefield!') with Parallel(n_jobs=max(args.num_process,len(failed_graphs))) as parallel: new_data_list_add = parallel(delayed(UpdateGrpah)(graph,system_generator,modeller,protein_atoms) for graph in failed_graphs) result = np.array([i if type(i) == int else 0 for i in new_data_list_add ]) new_data_list_add = list(filter(lambda x:type(x)!=int,new_data_list_add)) new_data_list = new_data_list + new_data_list_add if result.sum() != 0: indices = np.where(result == 1) failed_graphs = [failed_graphs[i] for i in indices[0]] logger.info(f'Minimized not Completed:{receptor_path}, {len(failed_graphs)} sdf not be minimized by default forcefield , try use gaff-2.11 forcefield!') else: failed_graphs = [] return new_data_list,failed_graphs @logger.catch def sampling(input_data_list, model, inference_steps, tr_schedule, rot_schedule, tor_schedule, device, t_to_sigma, model_args, no_random=False, ode=False, visualization_list=None, confidence_model=None, confidence_data_list=None, confidence_model_args=None, batch_size=32, no_final_step_noise=False,args = None): data_list = input_data_list N = len(data_list) pred_score = [] for t_idx in range(inference_steps): # use prediction score as a ranking metric , implemented by caoduanhua # pred_score = [] t_tr, t_rot, t_tor = tr_schedule[t_idx], rot_schedule[t_idx], tor_schedule[t_idx] dt_tr = tr_schedule[t_idx] - tr_schedule[t_idx + 1] if t_idx < inference_steps - 1 else tr_schedule[t_idx] dt_rot = rot_schedule[t_idx] - rot_schedule[t_idx + 1] if t_idx < inference_steps - 1 else rot_schedule[t_idx] dt_tor = tor_schedule[t_idx] - tor_schedule[t_idx + 1] if t_idx < inference_steps - 1 else tor_schedule[t_idx] loader = DataLoader(data_list, batch_size=batch_size) new_data_list = [] for complex_graph_batch in loader: b = complex_graph_batch.num_graphs complex_graph_batch = complex_graph_batch.to(device) tr_sigma, rot_sigma, tor_sigma = t_to_sigma(t_tr, t_rot, t_tor) set_time(complex_graph_batch, t_tr, t_rot, t_tor, b, model_args.all_atoms, device) with torch.no_grad(): tr_score, rot_score, tor_score = model(complex_graph_batch) # tr_g = tr_sigma * torch.sqrt(torch.tensor(2 * np.log(model_args.tr_sigma_max / model_args.tr_sigma_min))) rot_g = 2 * rot_sigma * torch.sqrt(torch.tensor(np.log(model_args.rot_sigma_max / model_args.rot_sigma_min))) if ode: tr_perturb = (0.5 * tr_g ** 2 * dt_tr * tr_score.cpu()).cpu() rot_perturb = (0.5 * rot_score.cpu() * dt_rot * rot_g ** 2).cpu() else: tr_z = torch.zeros((b, 3)) if no_random or (no_final_step_noise and t_idx == inference_steps - 1) \ else torch.normal(mean=0, std=1, size=(b, 3)) tr_perturb = (tr_g ** 2 * dt_tr * tr_score.cpu() + tr_g * np.sqrt(dt_tr) * tr_z).cpu() rot_z = torch.zeros((b, 3)) if no_random or (no_final_step_noise and t_idx == inference_steps - 1) \ else torch.normal(mean=0, std=1, size=(b, 3)) rot_perturb = (rot_score.cpu() * dt_rot * rot_g ** 2 + rot_g * np.sqrt(dt_rot) * rot_z).cpu() if not model_args.no_torsion: tor_g = tor_sigma * torch.sqrt(torch.tensor(2 * np.log(model_args.tor_sigma_max / model_args.tor_sigma_min))) if ode: tor_perturb = (0.5 * tor_g ** 2 * dt_tor * tor_score.cpu()).numpy() else: tor_z = torch.zeros(tor_score.shape) if no_random or (no_final_step_noise and t_idx == inference_steps - 1) \ else torch.normal(mean=0, std=1, size=tor_score.shape) tor_perturb = (tor_g ** 2 * dt_tor * tor_score.cpu() + tor_g * np.sqrt(dt_tor) * tor_z).numpy() else: tor_perturb = None # Apply noise tor_count_head = 0 tor_count_tail = 0 # node_head = 0 # node_tail = 0 for i, complex_graph in enumerate(complex_graph_batch.to('cpu').to_data_list()): # node_tail += complex_graph['ligand'].pos.shape[0] # complex_graph['ligand']['final_ligand'] = lig_node_attr[node_head:node_tail] # node_head+=complex_graph['ligand'].pos.shape[0] # if i==0: if type(complex_graph['ligand'].mask_rotate) is list: complex_graph['ligand'].mask_rotate = complex_graph['ligand'].mask_rotate[0] tor_count_tail += complex_graph['ligand'].mask_rotate.shape[0] try: new_data_list.append(modify_conformer(complex_graph, tr_perturb[i:i + 1], rot_perturb[i:i + 1].squeeze(0), tor_perturb[tor_count_head :tor_count_tail] if not model_args.no_torsion else None)) except: new_data_list.append(complex_graph) tor_count_head += complex_graph['ligand'].mask_rotate.shape[0] data_list = new_data_list if visualization_list is not None: for idx, visualization in enumerate(visualization_list): visualization.add((data_list[idx]['ligand'].pos + data_list[idx].original_center).detach().cpu(), part=1, order=t_idx + 2) #Before scoring the final conformers, we need to use force field to do energy minimized if args is not None and args.force_optimize: """ step 1 : load mol to make energy minimize step 2 : load fixed protein or pocket """ try: receptor_path = data_list[0]["protein_path"] logger.info('recptor path: {}',receptor_path) new_data_list,failed_graphs = inferenceFFOptimize(data_list,args,receptor_path,N) if len(failed_graphs) != 0: receptor_path = data_list[0]["pocket_path"] logger.info('Some minimized failed ! Use pocket file to do energy minimized!') new_data_list_pocket,failed_graphs = inferenceFFOptimize(failed_graphs,args,receptor_path,len(failed_graphs)) new_data_list += new_data_list_pocket logger.info('Return sucessed examples: {}',len(new_data_list)) logger.info('Return failed examples: {}',len(failed_graphs)) data_list = new_data_list + failed_graphs except Exception as e: error_info = traceback.format_exc() logger.info(error_info) warnings.warn(f'Complex {data_list[0]["name"]} will scoring without energy minimized!') pass with torch.no_grad(): if confidence_model is not None: loader = DataLoader(data_list, batch_size=batch_size) # try use forcefields to do energy minimized! if confidence_data_list is not None: confidence_loader = iter(DataLoader(confidence_data_list, batch_size=batch_size)) confidence = [] for complex_graph_batch in loader: complex_graph_batch = complex_graph_batch.to(device) if confidence_data_list is not None: confidence_complex_graph_batch = next(confidence_loader).to(device) confidence_complex_graph_batch['ligand'].pos = complex_graph_batch['ligand'].pos # confidence need all_atoms or not set_time(confidence_complex_graph_batch, 0, 0, 0, N, confidence_model_args.all_atoms, device) confidence.append(confidence_model(confidence_complex_graph_batch)[-1]) else: b = complex_graph_batch.num_graphs set_time(complex_graph_batch, 0, 0, 0, b, confidence_model_args.all_atoms, device) confidence.append(confidence_model(complex_graph_batch)[-1]) confidence = torch.cat(confidence, dim=0) else: confidence = None if confidence is not None: pred_score = confidence return data_list, pred_score