File size: 12,284 Bytes
96f168d | 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 206 207 208 209 210 211 212 213 | 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
|