Spaces:
Sleeping
Sleeping
File size: 4,599 Bytes
4a3f787 |
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 |
import math
import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
from scipy.stats import beta
from utils.geometry import axis_angle_to_matrix, rigid_transform_Kabsch_3D_torch
from utils.torsion import modify_conformer_torsion_angles
def t_to_sigma(t_tr, t_rot, t_tor, args):
tr_sigma = args.tr_sigma_min ** (1-t_tr) * args.tr_sigma_max ** t_tr
rot_sigma = args.rot_sigma_min ** (1-t_rot) * args.rot_sigma_max ** t_rot
tor_sigma = args.tor_sigma_min ** (1-t_tor) * args.tor_sigma_max ** t_tor
return tr_sigma, rot_sigma, tor_sigma
def modify_conformer(data, tr_update, rot_update, torsion_updates):
lig_center = torch.mean(data['ligand'].pos, dim=0, keepdim=True)
rot_mat = axis_angle_to_matrix(rot_update.squeeze())
rigid_new_pos = (data['ligand'].pos - lig_center) @ rot_mat.T + tr_update + lig_center
if torsion_updates is not None:
flexible_new_pos = modify_conformer_torsion_angles(rigid_new_pos,
data['ligand', 'ligand'].edge_index.T[data['ligand'].edge_mask],
data['ligand'].mask_rotate if isinstance(data['ligand'].mask_rotate, np.ndarray) else data['ligand'].mask_rotate[0],
torsion_updates).to(rigid_new_pos.device)
R, t = rigid_transform_Kabsch_3D_torch(flexible_new_pos.T, rigid_new_pos.T)
aligned_flexible_pos = flexible_new_pos @ R.T + t.T
data['ligand'].pos = aligned_flexible_pos
else:
data['ligand'].pos = rigid_new_pos
return data
def sinusoidal_embedding(timesteps, embedding_dim, max_positions=10000):
""" from https://github.com/hojonathanho/diffusion/blob/master/diffusion_tf/nn.py """
assert len(timesteps.shape) == 1
half_dim = embedding_dim // 2
emb = math.log(max_positions) / (half_dim - 1)
emb = torch.exp(torch.arange(half_dim, dtype=torch.float32, device=timesteps.device) * -emb)
emb = timesteps.float()[:, None] * emb[None, :]
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
if embedding_dim % 2 == 1: # zero pad
emb = F.pad(emb, (0, 1), mode='constant')
assert emb.shape == (timesteps.shape[0], embedding_dim)
return emb
class GaussianFourierProjection(nn.Module):
"""Gaussian Fourier embeddings for noise levels.
from https://github.com/yang-song/score_sde_pytorch/blob/1618ddea340f3e4a2ed7852a0694a809775cf8d0/models/layerspp.py#L32
"""
def __init__(self, embedding_size=256, scale=1.0):
super().__init__()
self.W = nn.Parameter(torch.randn(embedding_size//2) * scale, requires_grad=False)
def forward(self, x):
x_proj = x[:, None] * self.W[None, :] * 2 * np.pi
emb = torch.cat([torch.sin(x_proj), torch.cos(x_proj)], dim=-1)
return emb
def get_timestep_embedding(embedding_type, embedding_dim, embedding_scale=10000):
if embedding_type == 'sinusoidal':
emb_func = (lambda x : sinusoidal_embedding(embedding_scale * x, embedding_dim))
elif embedding_type == 'fourier':
emb_func = GaussianFourierProjection(embedding_size=embedding_dim, scale=embedding_scale)
else:
raise NotImplemented
return emb_func
def get_t_schedule(inference_steps):
return np.linspace(1, 0, inference_steps + 1)[:-1]
def set_time(complex_graphs, t_tr, t_rot, t_tor, batchsize, all_atoms, device):
complex_graphs['ligand'].node_t = {
'tr': t_tr * torch.ones(complex_graphs['ligand'].num_nodes).to(device),
'rot': t_rot * torch.ones(complex_graphs['ligand'].num_nodes).to(device),
'tor': t_tor * torch.ones(complex_graphs['ligand'].num_nodes).to(device)}
complex_graphs['receptor'].node_t = {
'tr': t_tr * torch.ones(complex_graphs['receptor'].num_nodes).to(device),
'rot': t_rot * torch.ones(complex_graphs['receptor'].num_nodes).to(device),
'tor': t_tor * torch.ones(complex_graphs['receptor'].num_nodes).to(device)}
complex_graphs.complex_t = {'tr': t_tr * torch.ones(batchsize).to(device),
'rot': t_rot * torch.ones(batchsize).to(device),
'tor': t_tor * torch.ones(batchsize).to(device)}
if all_atoms:
complex_graphs['atom'].node_t = {
'tr': t_tr * torch.ones(complex_graphs['atom'].num_nodes).to(device),
'rot': t_rot * torch.ones(complex_graphs['atom'].num_nodes).to(device),
'tor': t_tor * torch.ones(complex_graphs['atom'].num_nodes).to(device)} |