mirror of
https://github.com/gcorso/DiffDock.git
synced 2026-06-04 18:04:23 +08:00
9d0bf3d884
With some low (but non-trivial) frequency, processing through the convolutional layers diverge and node attributes become a mixture of nan and inf (which seem to all turn to nan). This later throws an exception during the Kabsch transform, which ruins results for the whole complex. Setting these to 0 basically skips an iteration, at worst it ruins one of the sampled complexes, but leaves the others. Note this is only applied to the main model, *not* the confidence model.
269 lines
15 KiB
Python
269 lines
15 KiB
Python
import copy
|
|
import random
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torch_geometric.data import Batch
|
|
from torch_geometric.loader import DataLoader
|
|
|
|
from utils.diffusion_utils import modify_conformer, set_time, modify_conformer_batch
|
|
from utils.torsion import modify_conformer_torsion_angles
|
|
from scipy.spatial.transform import Rotation as R
|
|
from utils.utils import crop_beyond
|
|
from utils.logging_utils import get_logger
|
|
|
|
|
|
def randomize_position(data_list, no_torsion, no_random, tr_sigma_max, pocket_knowledge=False, pocket_cutoff=7,
|
|
initial_noise_std_proportion=-1.0, choose_residue=False):
|
|
# in place modification of the list
|
|
center_pocket = data_list[0]['receptor'].pos.mean(dim=0)
|
|
if pocket_knowledge:
|
|
complex = data_list[0]
|
|
d = torch.cdist(complex['receptor'].pos, torch.from_numpy(complex['ligand'].orig_pos[0]).float() - complex.original_center)
|
|
label = torch.any(d < pocket_cutoff, dim=1)
|
|
|
|
if torch.any(label):
|
|
center_pocket = complex['receptor'].pos[label].mean(dim=0)
|
|
else:
|
|
print("No pocket residue below minimum distance ", pocket_cutoff, "taking closest at", torch.min(d))
|
|
center_pocket = complex['receptor'].pos[torch.argmin(torch.min(d, dim=1)[0])]
|
|
|
|
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[0], 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 + center_pocket
|
|
# base_rmsd = np.sqrt(np.sum((complex_graph['ligand'].pos.cpu().numpy() - orig_complex_graph['ligand'].pos.numpy()) ** 2, axis=1).mean())
|
|
|
|
if not no_random: # note for now the torsion angles are still randomised
|
|
if choose_residue:
|
|
idx = random.randint(0, len(complex_graph['receptor'].pos)-1)
|
|
tr_update = torch.normal(mean=complex_graph['receptor'].pos[idx:idx+1], std=0.01)
|
|
elif initial_noise_std_proportion >= 0.0:
|
|
std_rec = torch.sqrt(torch.mean(torch.sum(complex_graph['receptor'].pos ** 2, dim=1)))
|
|
tr_update = torch.normal(mean=0, std=std_rec * initial_noise_std_proportion / 1.73, size=(1, 3))
|
|
else:
|
|
# if initial_noise_std_proportion < 0.0, we use the tr_sigma_max multiplied by -initial_noise_std_proportion
|
|
tr_update = torch.normal(mean=0, std=-initial_noise_std_proportion * tr_sigma_max, size=(1, 3))
|
|
complex_graph['ligand'].pos += tr_update
|
|
|
|
|
|
def is_iterable(arr):
|
|
try:
|
|
some_object_iterator = iter(arr)
|
|
return True
|
|
except TypeError as te:
|
|
return False
|
|
|
|
|
|
def sampling(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,
|
|
t_schedule=None, batch_size=32, no_final_step_noise=False, pivot=None, return_full_trajectory=False,
|
|
temp_sampling=1.0, temp_psi=0.0, temp_sigma_data=0.5, return_features=False):
|
|
N = len(data_list)
|
|
trajectory = []
|
|
logger = get_logger()
|
|
if return_features:
|
|
lig_features, rec_features = [], []
|
|
assert batch_size >= N, "Not implemented yet"
|
|
|
|
loader = DataLoader(data_list, batch_size=batch_size)
|
|
assert not (return_full_trajectory or return_features or pivot), "Not implemented yet in new inference version"
|
|
|
|
mask_rotate = torch.from_numpy(data_list[0]['ligand'].mask_rotate[0]).to(device)
|
|
|
|
confidence = None
|
|
if confidence_model is not None:
|
|
confidence_loader = iter(DataLoader(confidence_data_list, batch_size=batch_size))
|
|
confidence = []
|
|
|
|
with torch.no_grad():
|
|
for batch_id, complex_graph_batch in enumerate(loader):
|
|
b = complex_graph_batch.num_graphs
|
|
n = len(complex_graph_batch['ligand'].pos) // b
|
|
complex_graph_batch = complex_graph_batch.to(device)
|
|
|
|
for t_idx in range(inference_steps):
|
|
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]
|
|
|
|
tr_sigma, rot_sigma, tor_sigma = t_to_sigma(t_tr, t_rot, t_tor)
|
|
|
|
if hasattr(model_args, 'crop_beyond') and model_args.crop_beyond is not None:
|
|
#print('Cropping beyond', tr_sigma * 3 + model_args.crop_beyond, 'for score model')
|
|
mod_complex_graph_batch = copy.deepcopy(complex_graph_batch).to_data_list()
|
|
for batch in mod_complex_graph_batch:
|
|
crop_beyond(batch, tr_sigma * 3 + model_args.crop_beyond, model_args.all_atoms)
|
|
mod_complex_graph_batch = Batch.from_data_list(mod_complex_graph_batch)
|
|
else:
|
|
mod_complex_graph_batch = complex_graph_batch
|
|
|
|
set_time(mod_complex_graph_batch, t_schedule[t_idx] if t_schedule is not None else None, t_tr, t_rot, t_tor, b,
|
|
'all_atoms' in model_args and model_args.all_atoms, device)
|
|
|
|
tr_score, rot_score, tor_score = model(mod_complex_graph_batch)[:3]
|
|
mean_scores = torch.mean(tr_score, dim=-1)
|
|
num_nans = torch.sum(torch.isnan(mean_scores))
|
|
if num_nans > 0:
|
|
name = complex_graph_batch['name']
|
|
if isinstance(name, list):
|
|
name = name[0]
|
|
logger.warning(f"Complex {name} Batch {batch_id+1} Inference Iteration {t_idx}: "
|
|
f"{num_nans} / {mean_scores.numel()} samples failed")
|
|
|
|
# Set the nan values to a small value, just want to disturb slightly
|
|
# Hopefully won't get nan the next iteration
|
|
tr_score.nan_to_num_(nan=(eps := 0.01*torch.nanmean(tr_score.abs())), posinf=eps, neginf=-eps)
|
|
rot_score.nan_to_num_(nan=(eps := 0.01*torch.nanmean(rot_score.abs())), posinf=eps, neginf=-eps)
|
|
tor_score.nan_to_num_(nan=(eps := 0.01*torch.nanmean(tor_score.abs())), posinf=eps, neginf=-eps)
|
|
del eps
|
|
|
|
tr_g = tr_sigma * torch.sqrt(torch.tensor(2 * np.log(model_args.tr_sigma_max / model_args.tr_sigma_min)))
|
|
rot_g = rot_sigma * torch.sqrt(torch.tensor(2 * 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)
|
|
rot_perturb = (0.5 * rot_score * dt_rot * rot_g ** 2)
|
|
else:
|
|
tr_z = torch.zeros((min(batch_size, N), 3), device=device) if no_random or (no_final_step_noise and t_idx == inference_steps - 1) \
|
|
else torch.normal(mean=0, std=1, size=(min(batch_size, N), 3), device=device)
|
|
tr_perturb = (tr_g ** 2 * dt_tr * tr_score + tr_g * np.sqrt(dt_tr) * tr_z)
|
|
|
|
rot_z = torch.zeros((min(batch_size, N), 3), device=device) if no_random or (no_final_step_noise and t_idx == inference_steps - 1) \
|
|
else torch.normal(mean=0, std=1, size=(min(batch_size, N), 3), device=device)
|
|
rot_perturb = (rot_score * dt_rot * rot_g ** 2 + rot_g * np.sqrt(dt_rot) * rot_z)
|
|
|
|
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)
|
|
else:
|
|
tor_z = torch.zeros(tor_score.shape, device=device) 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, device=device)
|
|
tor_perturb = (tor_g ** 2 * dt_tor * tor_score + tor_g * np.sqrt(dt_tor) * tor_z)
|
|
torsions_per_molecule = tor_perturb.shape[0] // b
|
|
else:
|
|
tor_perturb = None
|
|
|
|
if not is_iterable(temp_sampling):
|
|
temp_sampling = [temp_sampling] * 3
|
|
if not is_iterable(temp_psi):
|
|
temp_psi = [temp_psi] * 3
|
|
|
|
if not is_iterable(temp_sampling): temp_sampling = [temp_sampling] * 3
|
|
if not is_iterable(temp_psi): temp_psi = [temp_psi] * 3
|
|
if not is_iterable(temp_sigma_data): temp_sigma_data = [temp_sigma_data] * 3
|
|
|
|
assert len(temp_sampling) == 3
|
|
assert len(temp_psi) == 3
|
|
assert len(temp_sigma_data) == 3
|
|
|
|
if temp_sampling[0] != 1.0:
|
|
tr_sigma_data = np.exp(temp_sigma_data[0] * np.log(model_args.tr_sigma_max) + (1 - temp_sigma_data[0]) * np.log(model_args.tr_sigma_min))
|
|
lambda_tr = (tr_sigma_data + tr_sigma) / (tr_sigma_data + tr_sigma / temp_sampling[0])
|
|
tr_perturb = (tr_g ** 2 * dt_tr * (lambda_tr + temp_sampling[0] * temp_psi[0] / 2) * tr_score + tr_g * np.sqrt(dt_tr * (1 + temp_psi[0])) * tr_z)
|
|
|
|
if temp_sampling[1] != 1.0:
|
|
rot_sigma_data = np.exp(temp_sigma_data[1] * np.log(model_args.rot_sigma_max) + (1 - temp_sigma_data[1]) * np.log(model_args.rot_sigma_min))
|
|
lambda_rot = (rot_sigma_data + rot_sigma) / (rot_sigma_data + rot_sigma / temp_sampling[1])
|
|
rot_perturb = (rot_g ** 2 * dt_rot * (lambda_rot + temp_sampling[1] * temp_psi[1] / 2) * rot_score + rot_g * np.sqrt(dt_rot * (1 + temp_psi[1])) * rot_z)
|
|
|
|
if temp_sampling[2] != 1.0:
|
|
tor_sigma_data = np.exp(temp_sigma_data[2] * np.log(model_args.tor_sigma_max) + (1 - temp_sigma_data[2]) * np.log(model_args.tor_sigma_min))
|
|
lambda_tor = (tor_sigma_data + tor_sigma) / (tor_sigma_data + tor_sigma / temp_sampling[2])
|
|
tor_perturb = (tor_g ** 2 * dt_tor * (lambda_tor + temp_sampling[2] * temp_psi[2] / 2) * tor_score + tor_g * np.sqrt(dt_tor * (1 + temp_psi[2])) * tor_z)
|
|
|
|
# Apply noise
|
|
complex_graph_batch['ligand'].pos = \
|
|
modify_conformer_batch(complex_graph_batch['ligand'].pos, complex_graph_batch, tr_perturb, rot_perturb,
|
|
tor_perturb if not model_args.no_torsion else None, mask_rotate)
|
|
|
|
if visualization_list is not None:
|
|
for idx_b in range(b):
|
|
visualization_list[batch_id * batch_size + idx_b].add((
|
|
complex_graph_batch['ligand'].pos[idx_b*n:n*(idx_b+1)].detach().cpu() +
|
|
data_list[batch_id * batch_size + idx_b].original_center.detach().cpu()),
|
|
part=1, order=t_idx + 2)
|
|
|
|
for i in range(b):
|
|
data_list[batch_id * batch_size + i]['ligand'].pos = complex_graph_batch['ligand'].pos[i*n:n*(i+1)]
|
|
|
|
if visualization_list is not None:
|
|
for idx, visualization in enumerate(visualization_list):
|
|
visualization.add((data_list[idx]['ligand'].pos.detach().cpu() + data_list[idx].original_center.detach().cpu()),
|
|
part=1, order=2)
|
|
|
|
if confidence_model is not None:
|
|
if confidence_data_list is not None:
|
|
confidence_complex_graph_batch = next(confidence_loader)
|
|
confidence_complex_graph_batch['ligand'].pos = complex_graph_batch['ligand'].pos.cpu()
|
|
|
|
if hasattr(confidence_model_args, 'crop_beyond') and confidence_model_args.crop_beyond is not None:
|
|
confidence_complex_graph_batch = confidence_complex_graph_batch.to_data_list()
|
|
for batch in confidence_complex_graph_batch:
|
|
crop_beyond(batch, confidence_model_args.crop_beyond, confidence_model_args.all_atoms)
|
|
confidence_complex_graph_batch = Batch.from_data_list(confidence_complex_graph_batch)
|
|
|
|
confidence_complex_graph_batch = confidence_complex_graph_batch.to(device)
|
|
set_time(confidence_complex_graph_batch, 0, 0, 0, 0, b, confidence_model_args.all_atoms, device)
|
|
out = confidence_model(confidence_complex_graph_batch)
|
|
else:
|
|
out = confidence_model(complex_graph_batch)
|
|
|
|
if type(out) is tuple:
|
|
out = out[0]
|
|
confidence.append(out)
|
|
|
|
if confidence_model is not None:
|
|
confidence = torch.cat(confidence, dim=0)
|
|
confidence = torch.nan_to_num(confidence, nan=-1000)
|
|
|
|
if return_full_trajectory:
|
|
return data_list, confidence, trajectory
|
|
elif return_features:
|
|
lig_features = torch.cat(lig_features, dim=0)
|
|
rec_features = torch.cat(rec_features, dim=0)
|
|
return data_list, confidence, lig_features, rec_features
|
|
|
|
return data_list, confidence
|
|
|
|
|
|
def compute_affinity(data_list, affinity_model, affinity_data_list, device, parallel, all_atoms, include_miscellaneous_atoms):
|
|
|
|
with torch.no_grad():
|
|
if affinity_model is not None:
|
|
assert parallel <= len(data_list)
|
|
loader = DataLoader(data_list, batch_size=parallel)
|
|
complex_graph_batch = next(iter(loader)).to(device)
|
|
positions = complex_graph_batch['ligand'].pos
|
|
|
|
assert affinity_data_list is not None
|
|
complex_graph = affinity_data_list[0]
|
|
N = complex_graph['ligand'].num_nodes
|
|
complex_graph['ligand'].x = complex_graph['ligand'].x.repeat(parallel, 1)
|
|
complex_graph['ligand'].edge_mask = complex_graph['ligand'].edge_mask.repeat(parallel)
|
|
complex_graph['ligand', 'ligand'].edge_index = torch.cat(
|
|
[N * i + complex_graph['ligand', 'ligand'].edge_index for i in range(parallel)], dim=1)
|
|
complex_graph['ligand', 'ligand'].edge_attr = complex_graph['ligand', 'ligand'].edge_attr.repeat(parallel, 1)
|
|
complex_graph['ligand'].pos = positions
|
|
|
|
affinity_loader = DataLoader([complex_graph], batch_size=1)
|
|
affinity_batch = next(iter(affinity_loader)).to(device)
|
|
set_time(affinity_batch, 0, 0, 0, 0, 1, all_atoms, device, include_miscellaneous_atoms=include_miscellaneous_atoms)
|
|
_, affinity = affinity_model(affinity_batch)
|
|
else:
|
|
affinity = None
|
|
|
|
return affinity |